BiobotGrammar/view.py at main · benquick123/BiobotGrammar · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
import numpy as np

import pyrobotdesign as rd
import time


class CameraTracker(object):
    def __init__(self, viewer, sim, robot_idx):
        self.viewer = viewer
        self.sim = sim
        self.robot_idx = robot_idx

        self.reset()

    def update(self, time_step):
        lower = np.zeros(3)
        upper = np.zeros(3)
        self.sim.get_robot_world_aabb(self.robot_idx, lower, upper)

        # Update camera position to track the robot smoothly
        target_pos = 0.5 * (lower + upper)
        camera_pos = self.viewer.camera_params.position.copy()
        camera_pos += 5.0 * time_step * (target_pos - camera_pos)
        self.viewer.camera_params.position = camera_pos

    def reset(self):
        lower = np.zeros(3)
        upper = np.zeros(3)
        self.sim.get_robot_world_aabb(self.robot_idx, lower, upper)

        self.viewer.camera_params.position = 0.5 * (lower + upper)


def prepare_viewer(sim):
    viewer = rd.GLFWViewer()

    # Get robot bounds
    lower = np.zeros(3)
    upper = np.zeros(3)
    sim.get_robot_world_aabb(0, lower, upper)

    # Set initial camera parameters
    viewer.camera_params.yaw = -np.pi / 4
    viewer.camera_params.pitch = -np.pi / 6
    viewer.camera_params.distance = 1.5 * np.linalg.norm(upper - lower)

    tracker = CameraTracker(viewer, sim, 0)

    return viewer, tracker


def viewer_step(sim, task, actions, viewer, tracker, step=0, torques=None):
    if torques is not None:
        sim.update_torques(0, torques)

    for i in range(task.interval):
        step_idx = step * task.interval + i

        sim.set_joint_targets(0, actions)

        task.add_noise(sim, step_idx)
        sim.step()
        tracker.update(task.time_step)
        viewer.update(task.time_step)

    viewer.render(sim)