allow mounted cameras to be run in parallel simulations; added unit test

Author: abhijitmajumdar

genesis/vis/camera.py

@@ -91,6 +91,7 @@ def __init__(
         self._is_built = False
         self._attached_link = None
         self._attached_offset_T = None
+        self._attached_env_idx = None
 
         self._in_recording = False
         self._recorded_imgs = []
@@ -124,7 +125,7 @@ def _build(self):
         self._is_built = True
         self.set_pose(self._transform, self._pos, self._lookat, self._up)
 
-    def attach(self, rigid_link, offset_T):
+    def attach(self, rigid_link, offset_T, env_idx: int | None = None):
         """
         Attach the camera to a rigid link in the scene.
 
@@ -136,9 +137,25 @@ def attach(self, rigid_link, offset_T):
             The rigid link to which the camera should be attached.
         offset_T : np.ndarray, shape (4, 4)
             The transformation matrix specifying the camera's pose relative to the rigid link.
+        env_idx : int
+            The environment index this camera should be tied to. Offsets the `offset_T` accordingly. Must be specified if running parallel environments
+
+        Raises
+        ------
+        Exception
+            If running parallel simulations but env_idx is not specified.
+        Exception
+            If invalid env_idx is specified (env_idx >= n_envs)
         """
         self._attached_link = rigid_link
         self._attached_offset_T = offset_T
+        if self._visualizer._scene.n_envs > 0 and env_idx is None:
+            gs.raise_exception("Must specify env_idx when running parallel simulations")
+        if env_idx is not None:
+            n_envs = self._visualizer._scene.n_envs
+            if env_idx >= n_envs:
+                gs.raise_exception(f"Invalid env_idx {env_idx} for camera, configured for {n_envs} environments")
+            self._attached_env_idx = env_idx
 
     def detach(self):
         """
@@ -148,6 +165,7 @@ def detach(self):
         """
         self._attached_link = None
         self._attached_offset_T = None
+        self._attached_env_idx = None
 
     @gs.assert_built
     def move_to_attach(self):
@@ -160,16 +178,15 @@ def move_to_attach(self):
         ------
         Exception
             If the camera has not been mounted using `attach()`.
-        Exception
-            If the simulation is running in parallel (`n_envs > 0`), which is currently unsupported for mounted cameras.
         """
         if self._attached_link is None:
             gs.raise_exception(f"The camera hasn't been mounted!")
-        if self._visualizer._scene.n_envs > 0:
-            gs.raise_exception(f"Mounted camera not supported in parallel simulation!")
 
-        link_pos = self._attached_link.get_pos().cpu().numpy()
-        link_quat = self._attached_link.get_quat().cpu().numpy()
+        link_pos = self._attached_link.get_pos(envs_idx=self._attached_env_idx).cpu().numpy()
+        link_quat = self._attached_link.get_quat(envs_idx=self._attached_env_idx).cpu().numpy()
+        if self._attached_env_idx is not None:
+            link_pos = link_pos[0] + self._visualizer._scene.envs_offset[self._attached_env_idx]
+            link_quat = link_quat[0]
         link_T = gu.trans_quat_to_T(link_pos, link_quat)
         transform = link_T @ self._attached_offset_T
         self.set_pose(transform=transform)

tests/test_rigid_physics.py

@@ -1,5 +1,6 @@
 import math
 import os
+import queue
 import sys
 import xml.etree.ElementTree as ET
 
@@ -1034,6 +1035,183 @@ def test_batched_offscreen_rendering(show_viewer, tol):
             assert_allclose(steps_rgb_arrays[0][i], steps_rgb_arrays[1][i], tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.skipif(sys.platform == "darwin", reason="Segfault inside 'shadow_mapping_pass' on MacOS VM.")
+@pytest.mark.xfail(reason="This test is not passing on all platforms for now.")
+def test_batched_mounted_camera_rendering(show_viewer, tol):
+    scene = gs.Scene(
+        vis_options=gs.options.VisOptions(
+            env_separate_rigid=False,
+        ),
+        show_viewer=show_viewer,
+        show_FPS=False,
+    )
+    plane = scene.add_entity(
+        morph=gs.morphs.Plane(),
+        surface=gs.surfaces.Aluminium(
+            ior=10.0,
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(-0.2, -0.8, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Rough(
+            diffuse_texture=gs.textures.ColorTexture(
+                color=(1.0, 0.5, 0.5),
+            ),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(-0.2, -0.5, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Rough(
+            color=(1.0, 1.0, 1.0),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(-0.2, -0.2, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Smooth(
+            color=(0.6, 0.8, 1.0),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(-0.2, 0.2, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Iron(
+            color=(1.0, 1.0, 1.0),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(-0.2, 0.5, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Gold(
+            color=(1.0, 1.0, 1.0),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(-0.2, 0.8, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Glass(
+            color=(1.0, 1.0, 1.0),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/sphere.obj",
+            scale=0.1,
+            pos=(0.2, -0.8, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Smooth(
+            color=(1.0, 1.0, 1.0, 0.5),
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/wooden_sphere_OBJ/wooden_sphere.obj",
+            scale=0.025,
+            pos=(0.2, -0.5, 0.2),
+            fixed=True,
+        ),
+    )
+    scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/wooden_sphere_OBJ/wooden_sphere.obj",
+            scale=0.025,
+            pos=(0.2, -0.2, 0.2),
+            fixed=True,
+        ),
+        surface=gs.surfaces.Rough(
+            diffuse_texture=gs.textures.ImageTexture(
+                image_path="textures/checker.png",
+            )
+        ),
+    )
+    robot = scene.add_entity(
+        gs.morphs.MJCF(file="xml/franka_emika_panda/panda.xml"),
+    )
+    n_envs = 3
+    env_spacing = (2.0, 2.0)
+    cams = [scene.add_camera(GUI=show_viewer, fov=70) for _ in range(n_envs)]
+    scene.build(n_envs=n_envs, env_spacing=env_spacing)
+
+    T = np.eye(4)
+    T[:3, :3] = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
+    T[:3, 3] = np.array([0.1, 0.0, 0.1])
+    for nenv, cam in enumerate(cams):
+        cam.attach(robot.get_link("hand"), T, nenv)
+
+    target_quat = np.tile(np.array([0, 1, 0, 0]), [n_envs, 1])  # pointing downwards
+    center = np.tile(np.array([-0.25, -0.25, 0.5]), [n_envs, 1])
+    rng = np.random.default_rng(42)
+    angular_speed = rng.uniform(-10, 10, n_envs)
+    r = 0.25
+
+    ee_link = robot.get_link("hand")
+
+    steps_rgb_queue: queue.Queue[list[np.ndarray]] = queue.Queue(maxsize=2)
+
+    for i in range(50):
+        target_pos = np.zeros([n_envs, 3])
+        target_pos[:, 0] = center[:, 0] + np.cos(i / 360 * np.pi * angular_speed) * r
+        target_pos[:, 1] = center[:, 1] + np.sin(i / 360 * np.pi * angular_speed) * r
+        target_pos[:, 2] = center[:, 2]
+
+        q = robot.inverse_kinematics(
+            link=ee_link,
+            pos=target_pos,
+            quat=target_quat,
+            rot_mask=[False, False, True],  # for demo purpose: only restrict direction of z-axis
+        )
+
+        robot.set_qpos(q)
+        scene.step()
+        if i < 10:
+            # skip the first few frames because the robots start off with the same state
+            for cam in cams:
+                cam.render()
+            continue
+        robots_rgb_arrays = []
+        for cam in cams:
+            rgb_array, *_ = cam.render()
+            assert np.std(rgb_array) > 10.0
+            robots_rgb_arrays.append(rgb_array)
+        steps_rgb_queue.put(robots_rgb_arrays)
+
+        if steps_rgb_queue.full():  # we have a set of 2 consecutive frames
+            diff_tol = 0.05  # expect atlest 5% difference between each frame
+            frames_t_minus_1 = steps_rgb_queue.get()
+            frames_t = steps_rgb_queue.get()
+            for i in range(n_envs):
+                diff = frames_t[i] - frames_t_minus_1[i]
+                assert np.count_nonzero(diff) > diff_tol * np.prod(diff.shape)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_pd_control(show_viewer):