Skip to content

Commit dee7799

Browse files
yzqinyzqin
committed
[update] update SAPIEN maniskill2 example
1 parent 0353387 commit dee7799

File tree

5 files changed

+191
-6
lines changed

5 files changed

+191
-6
lines changed

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ Web Visualizer on [IsaacGymEnvs](https://github.com/NVIDIA-Omniverse/IsaacGymEnv
4747
Check the [SAPIEN Example](example/sapien/README.md) for information on running the
4848
Web Visualizer on [ManiSkill2](https://github.com/haosulab/ManiSkill2).
4949

50+
![sapien](doc/sapien_side_by_side.gif)
51+
5052
### URDF Web Visualizer
5153

5254
Check the [URDF Example](example/visualize_urdf/README.md) for information to visualize a static URDF file on Web

‎doc/sapien_side_by_side.gif‎

770 KB
Loading

‎example/sapien/README.md‎

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,36 @@
11
## Using the web visualizer on ManiSkill2
2+
import sapien.core as sapien
3+
import sapien.core as sapien
24

35
ManiSkill2 is a comprehensive environment suite for robot learning. For more information, please visit
46
the [ManiSkill2 website](https://maniskill2.github.io/).
57

68
This example provides a minimal code modification to port all visualization to the web browser.
79

8-
![sapien](../../doc/sapien.gif)
910

1011
## Examples
1112

13+
#### Replay ManiSkill2 Demonstration in Web Visualizer
14+
![sapien](../../doc/sapien_side_by_side.gif)
15+
16+
17+
```bash
18+
# Install the dependencies for ManiSkill2 if you have not
19+
pip install mani-skill2
20+
# or pip install ."[sapien]" in the project root to install all required dep for the SAPIEN ManiSkill example
21+
cd example/sapien
22+
23+
# Download ManiSkill2 demonstrations, you can also replace the `PegInsertionSide` with other task name
24+
python -m mani_skill2.utils.download_demo "PegInsertionSide-v0"
25+
python replay_maniskill_demo.py --traj-path demos/v0/rigid_body/PegInsertionSide-v0/trajectory.h5
26+
27+
```
28+
Open a browser window and navigate to `http://127.0.0.1:7000/static/` to visualize the demonstration trajectory
29+
30+
31+
#### Visualize ManiSkill2 environment with random robot action
32+
![sapien](../../doc/sapien.gif)
33+
1234
1. Start the Meshcat server in the terminal. The default URL is `http://127.0.0.1:7000/static/` for localhost.
1335

1436
```shell
@@ -21,10 +43,7 @@ meshcat-server
2143
the [ManiSkill2 Quickstart documentation](https://haosulab.github.io/ManiSkill2/getting_started/quickstart.html).
2244

2345
```bash
24-
# Install the dependencies for ManiSkill2 if you have not
25-
pip install mani-skill2
26-
# or pip install ."[sapien]" in the project root
27-
cd example/sapien
46+
2847

2948
# You may need to download the ManiSkill2 assets before running
3049
python -m mani_skill2.utils.download_asset all

‎example/sapien/replay_maniskill_demo.py‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
# Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# Licensed under the MIT License [see LICENSE for details].
4+
5+
6+
"""
7+
Port ManiSkill2 demonstration replay visualization to web visualizer
8+
# Reference: https://github.com/haosulab/ManiSkill2/tree/main/examples/tutorials/imitation-learning
9+
10+
"""
11+
12+
13+
import argparse
14+
import os
15+
from typing import Optional
16+
17+
import gymnasium as gym
18+
import h5py
19+
import numpy as np
20+
import sapien.core as sapien
21+
from mani_skill2.envs import sapien_env
22+
from mani_skill2.utils.io_utils import load_json
23+
from tqdm.auto import tqdm
24+
25+
################################################################################################
26+
# The following code is only used for visualizer and not presented in the original ManiSkill
27+
################################################################################################
28+
from meshcat.servers.zmqserver import start_zmq_server_as_subprocess
29+
from sim_web_visualizer import bind_visualizer_to_sapien_scene, create_sapien_visualizer
30+
31+
32+
def wrapped_setup_scene(self: sapien_env.BaseEnv, scene_config: Optional[sapien.SceneConfig] = None):
33+
if scene_config is None:
34+
scene_config = self._get_default_scene_config()
35+
self._scene = self._engine.create_scene(scene_config)
36+
self._scene.set_timestep(1.0 / self._sim_freq)
37+
self._scene = bind_visualizer_to_sapien_scene(self._scene, self._engine, self._renderer)
38+
39+
40+
def wrapped_setup_viewer(self):
41+
self._viewer.set_scene(self._scene._scene)
42+
self._viewer.scene = self._scene
43+
self._viewer.toggle_axes(False)
44+
self._viewer.toggle_camera_lines(False)
45+
46+
47+
start_zmq_server_as_subprocess()
48+
# Set to True if you want to keep both the original viewer and the web visualizer. A display is needed for True
49+
keep_on_screen_renderer = True
50+
51+
create_sapien_visualizer(port=6000, host="localhost", keep_default_viewer=keep_on_screen_renderer)
52+
sapien_env.BaseEnv._setup_scene = wrapped_setup_scene
53+
sapien_env.BaseEnv._setup_viewer = wrapped_setup_viewer
54+
################################################################################################
55+
# End of visualizer code
56+
################################################################################################
57+
58+
59+
def parse_args(args=None):
60+
parser = argparse.ArgumentParser()
61+
parser.add_argument("--traj-path", type=str, required=True)
62+
parser.add_argument("--verbose", action="store_true")
63+
parser.add_argument(
64+
"--count",
65+
type=int,
66+
default=None,
67+
help="number of demonstrations to replay before exiting. By default will replay all demonstrations",
68+
)
69+
return parser.parse_args(args)
70+
71+
72+
def main(args):
73+
pbar = tqdm(position=0, leave=None, unit="step", dynamic_ncols=True)
74+
75+
# Load HDF5 containing trajectories
76+
traj_path = args.traj_path
77+
ori_h5_file = h5py.File(traj_path, "r")
78+
79+
# Load associated json
80+
json_path = traj_path.replace(".h5", ".json")
81+
json_data = load_json(json_path)
82+
83+
env_info = json_data["env_info"]
84+
env_id = env_info["env_id"]
85+
ori_env_kwargs = env_info["env_kwargs"]
86+
87+
# Create a main env for replay
88+
env_kwargs = ori_env_kwargs.copy()
89+
env_kwargs["obs_mode"] = "state"
90+
env_kwargs[
91+
"render_mode"
92+
] = "rgb_array" # note this only affects the videos saved as RecordEpisode wrapper calls env.render
93+
env = gym.make(env_id, **env_kwargs).unwrapped
94+
if pbar is not None:
95+
pbar.set_postfix(
96+
{
97+
"control_mode": env_kwargs.get("control_mode"),
98+
"obs_mode": env_kwargs.get("obs_mode"),
99+
}
100+
)
101+
102+
# Prepare for recording
103+
output_dir = os.path.dirname(traj_path)
104+
ori_traj_name = os.path.splitext(os.path.basename(traj_path))[0]
105+
suffix = "{}.{}".format(env.obs_mode, env.control_mode)
106+
new_traj_name = ori_traj_name + "." + suffix
107+
108+
episodes = json_data["episodes"][: args.count]
109+
n_ep = len(episodes)
110+
inds = np.arange(n_ep)
111+
inds = np.array_split(inds, 1)[0]
112+
113+
# Replay
114+
for ind in inds:
115+
ep = episodes[ind]
116+
episode_id = ep["episode_id"]
117+
traj_id = f"traj_{episode_id}"
118+
if pbar is not None:
119+
pbar.set_description(f"Replaying {traj_id}")
120+
121+
if traj_id not in ori_h5_file:
122+
tqdm.write(f"{traj_id} does not exist in {traj_path}")
123+
continue
124+
125+
reset_kwargs = ep["reset_kwargs"].copy()
126+
if "seed" in reset_kwargs:
127+
assert reset_kwargs["seed"] == ep["episode_seed"]
128+
else:
129+
reset_kwargs["seed"] = ep["episode_seed"]
130+
seed = reset_kwargs.pop("seed")
131+
132+
env.reset(seed=seed, options=reset_kwargs)
133+
env.render_human()
134+
135+
# Original actions to replay
136+
ori_actions = ori_h5_file[traj_id]["actions"][:]
137+
138+
# Original env states to replay
139+
ori_env_states = ori_h5_file[traj_id]["env_states"][1:]
140+
141+
info = {}
142+
143+
# Without conversion between control modes
144+
n = len(ori_actions)
145+
if pbar is not None:
146+
pbar.reset(total=n)
147+
for t, a in enumerate(ori_actions):
148+
if pbar is not None:
149+
pbar.update()
150+
env.render_human()
151+
env.set_state(ori_env_states[t])
152+
153+
success = info.get("success", False)
154+
155+
# Cleanup
156+
env.close()
157+
ori_h5_file.close()
158+
159+
if pbar is not None:
160+
pbar.close()
161+
162+
163+
if __name__ == "__main__":
164+
main(parse_args())

‎example/sapien/run_sapien_viz.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def wrapped_setup_viewer(self):
8585
try:
8686
obs = env.reset()
8787
done = False
88-
for _ in range(100):
88+
for _ in range(500):
8989
action = env.action_space.sample()
9090
obs, reward, terminated, truncated, info = env.step(action)
9191
if keep_on_screen_renderer:

0 commit comments

Comments
 (0)