Skip to content

Commit 9e1980d

Browse files
[FEATURE] allow mounted cameras to be run in parallel simulations (#1323)
1 parent bd1fd48 commit 9e1980d

File tree

2 files changed

+117
-7
lines changed

2 files changed

+117
-7
lines changed

genesis/vis/camera.py

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import genesis as gs
99
import genesis.utils.geom as gu
1010
from genesis.repr_base import RBC
11+
from genesis.utils.misc import tensor_to_array
1112

1213

1314
class Camera(RBC):
@@ -91,6 +92,7 @@ def __init__(
9192
self._is_built = False
9293
self._attached_link = None
9394
self._attached_offset_T = None
95+
self._attached_env_idx = None
9496

9597
self._in_recording = False
9698
self._recorded_imgs = []
@@ -124,7 +126,7 @@ def _build(self):
124126
self._is_built = True
125127
self.set_pose(self._transform, self._pos, self._lookat, self._up)
126128

127-
def attach(self, rigid_link, offset_T):
129+
def attach(self, rigid_link, offset_T, env_idx: int | None = None):
128130
"""
129131
Attach the camera to a rigid link in the scene.
130132
@@ -136,9 +138,26 @@ def attach(self, rigid_link, offset_T):
136138
The rigid link to which the camera should be attached.
137139
offset_T : np.ndarray, shape (4, 4)
138140
The transformation matrix specifying the camera's pose relative to the rigid link.
141+
env_idx : int
142+
The environment index this camera should be tied to. Offsets the `offset_T` accordingly. Must be specified
143+
if running parallel environments
144+
145+
Raises
146+
------
147+
Exception
148+
If running parallel simulations but env_idx is not specified.
149+
Exception
150+
If invalid env_idx is specified (env_idx >= n_envs)
139151
"""
140152
self._attached_link = rigid_link
141153
self._attached_offset_T = offset_T
154+
if self._visualizer._scene.n_envs > 0 and env_idx is None:
155+
gs.raise_exception("Must specify env_idx when running parallel simulations")
156+
if env_idx is not None:
157+
n_envs = self._visualizer._scene.n_envs
158+
if env_idx >= n_envs:
159+
gs.raise_exception(f"Invalid env_idx {env_idx} for camera, configured for {n_envs} environments")
160+
self._attached_env_idx = env_idx
142161

143162
def detach(self):
144163
"""
@@ -148,6 +167,7 @@ def detach(self):
148167
"""
149168
self._attached_link = None
150169
self._attached_offset_T = None
170+
self._attached_env_idx = None
151171

152172
@gs.assert_built
153173
def move_to_attach(self):
@@ -160,16 +180,15 @@ def move_to_attach(self):
160180
------
161181
Exception
162182
If the camera has not been mounted using `attach()`.
163-
Exception
164-
If the simulation is running in parallel (`n_envs > 0`), which is currently unsupported for mounted cameras.
165183
"""
166184
if self._attached_link is None:
167185
gs.raise_exception(f"The camera hasn't been mounted!")
168-
if self._visualizer._scene.n_envs > 0:
169-
gs.raise_exception(f"Mounted camera not supported in parallel simulation!")
170186

171-
link_pos = self._attached_link.get_pos().cpu().numpy()
172-
link_quat = self._attached_link.get_quat().cpu().numpy()
187+
link_pos = tensor_to_array(self._attached_link.get_pos(envs_idx=self._attached_env_idx))
188+
link_quat = tensor_to_array(self._attached_link.get_quat(envs_idx=self._attached_env_idx))
189+
if self._attached_env_idx is not None:
190+
link_pos = link_pos[0] + self._visualizer._scene.envs_offset[self._attached_env_idx]
191+
link_quat = link_quat[0]
173192
link_T = gu.trans_quat_to_T(link_pos, link_quat)
174193
transform = link_T @ self._attached_offset_T
175194
self.set_pose(transform=transform)

tests/test_rigid_physics.py

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import math
22
import os
3+
import queue
34
import sys
45
import tempfile
56
import xml.etree.ElementTree as ET
@@ -1038,6 +1039,96 @@ def test_batched_offscreen_rendering(show_viewer, tol):
10381039
assert_allclose(steps_rgb_arrays[0][i], steps_rgb_arrays[1][i], tol=tol)
10391040

10401041

1042+
@pytest.mark.required
1043+
@pytest.mark.skipif(sys.platform == "darwin", reason="Segfault inside 'shadow_mapping_pass' on MacOS VM.")
1044+
@pytest.mark.xfail(reason="This test is not passing on all platforms for now.")
1045+
def test_batched_mounted_camera_rendering(show_viewer, tol):
1046+
scene = gs.Scene(
1047+
vis_options=gs.options.VisOptions(
1048+
env_separate_rigid=False,
1049+
),
1050+
show_viewer=show_viewer,
1051+
show_FPS=False,
1052+
)
1053+
plane = scene.add_entity(
1054+
morph=gs.morphs.Plane(),
1055+
surface=gs.surfaces.Aluminium(
1056+
ior=10.0,
1057+
),
1058+
)
1059+
scene.add_entity(
1060+
morph=gs.morphs.Mesh(
1061+
file="meshes/sphere.obj",
1062+
scale=0.1,
1063+
pos=(-0.2, -0.5, 0.2),
1064+
fixed=True,
1065+
),
1066+
surface=gs.surfaces.Rough(
1067+
diffuse_texture=gs.textures.ColorTexture(
1068+
color=(1.0, 0.5, 0.5),
1069+
),
1070+
),
1071+
)
1072+
robot = scene.add_entity(
1073+
gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
1074+
)
1075+
n_envs = 3
1076+
env_spacing = (2.0, 2.0)
1077+
cams = [scene.add_camera(GUI=show_viewer, fov=70) for _ in range(n_envs)]
1078+
scene.build(n_envs=n_envs, env_spacing=env_spacing)
1079+
1080+
T = np.eye(4)
1081+
T[:3, :3] = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
1082+
T[:3, 3] = np.array([0.1, 0.0, 0.1])
1083+
for nenv, cam in enumerate(cams):
1084+
cam.attach(robot.get_link("hand"), T, nenv)
1085+
1086+
target_quat = np.tile(np.array([0, 1, 0, 0]), [n_envs, 1]) # pointing downwards
1087+
center = np.tile(np.array([-0.25, -0.25, 0.5]), [n_envs, 1])
1088+
rng = np.random.default_rng(42)
1089+
angular_speed = rng.uniform(-10, 10, n_envs)
1090+
r = 0.25
1091+
1092+
ee_link = robot.get_link("hand")
1093+
1094+
steps_rgb_queue: queue.Queue[list[np.ndarray]] = queue.Queue(maxsize=2)
1095+
1096+
for i in range(50):
1097+
target_pos = np.zeros([n_envs, 3])
1098+
target_pos[:, 0] = center[:, 0] + np.cos(i / 360 * np.pi * angular_speed) * r
1099+
target_pos[:, 1] = center[:, 1] + np.sin(i / 360 * np.pi * angular_speed) * r
1100+
target_pos[:, 2] = center[:, 2]
1101+
1102+
q = robot.inverse_kinematics(
1103+
link=ee_link,
1104+
pos=target_pos,
1105+
quat=target_quat,
1106+
rot_mask=[False, False, True], # for demo purpose: only restrict direction of z-axis
1107+
)
1108+
1109+
robot.set_qpos(q)
1110+
scene.step()
1111+
if i < 10:
1112+
# skip the first few frames because the robots start off with the same state
1113+
for cam in cams:
1114+
cam.render()
1115+
continue
1116+
robots_rgb_arrays = []
1117+
for cam in cams:
1118+
rgb_array, *_ = cam.render()
1119+
assert np.std(rgb_array) > 10.0
1120+
robots_rgb_arrays.append(rgb_array)
1121+
steps_rgb_queue.put(robots_rgb_arrays)
1122+
1123+
if steps_rgb_queue.full(): # we have a set of 2 consecutive frames
1124+
diff_tol = 0.02 # expect atlest 2% difference between each frame
1125+
frames_t_minus_1 = steps_rgb_queue.get()
1126+
frames_t = steps_rgb_queue.get()
1127+
for i in range(n_envs):
1128+
diff = frames_t[i] - frames_t_minus_1[i]
1129+
assert np.count_nonzero(diff) > diff_tol * np.prod(diff.shape)
1130+
1131+
10411132
@pytest.mark.required
10421133
@pytest.mark.parametrize("backend", [gs.cpu])
10431134
def test_pd_control(show_viewer):

0 commit comments

Comments
 (0)