-
Notifications
You must be signed in to change notification settings - Fork 9
/
simple_replay.py
52 lines (42 loc) · 1.99 KB
/
simple_replay.py
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
from pathlib import Path
import numpy as np
from hand_teleop.env.sim_env.constructor import add_default_scene_light
from hand_teleop.env.sim_env.mug_flip_env import MugFlipEnv
from hand_teleop.env.sim_env.relocate_env import RelocateEnv
from hand_teleop.env.sim_env.table_door_env import TableDoorEnv
from hand_teleop.gui.teleop_gui import GUIBase, DEFAULT_TABLE_TOP_CAMERAS
from hand_teleop.kinematics.mano_robot_hand import MANORobotHand
def replay(path_to_pickle: str):
path = Path(path_to_pickle)
all_data = np.load(str(path.resolve()), allow_pickle=True)
meta_data = all_data["meta_data"]
data = all_data["data"]
env_class = meta_data["env_class"].lower()
if "relocate" in env_class:
env = RelocateEnv(**meta_data["env_kwargs"])
elif "door" in env_class:
env = TableDoorEnv(**meta_data["env_kwargs"])
elif "flip" in env_class:
env = MugFlipEnv(**meta_data["env_kwargs"])
else:
raise ValueError(env_class)
# Setup viewer and camera
add_default_scene_light(env.scene, env.renderer)
gui = GUIBase(env.scene, env.renderer)
for name, params in DEFAULT_TABLE_TOP_CAMERAS.items():
gui.create_camera(**params)
gui.viewer.set_camera_rpy(0, -0.7, 0.01)
gui.viewer.set_camera_xyz(-0.4, 0, 0.45)
scene = env.scene
# Create the mano robot hand during teleoperation
mano_robot = MANORobotHand(env.scene, env.renderer, init_joint_pos=meta_data["zero_joint_pos"],
control_interval=env.frame_skip * scene.get_timestep(), scale=1)
for i in range(meta_data["data_len"] - 10):
scene.unpack(data[i]["simulation"])
for _ in range(2):
gui.render(render_all_views=False)
if __name__ == '__main__':
# pkl_path = Path(__file__).parent / "example_teleop_data/relocate_tomato_soup_can.pkl"
# pkl_path = Path(__file__).parent / "example_teleop_data/flip_mug.pkl"
pkl_path = Path(__file__).parent / "example_teleop_data/open_door.pkl"
replay(str(pkl_path))