Enabling rasterizer camera sensor to work with n_envs > 1.

Author: Narsil

genesis/engine/sensors/camera.py

@@ -24,7 +24,12 @@
 from genesis.options.renderers import BatchRenderer as BatchRendererOptions
 from genesis.options.vis import VisOptions
 from genesis.vis.rasterizer_context import RasterizerContext
-from .base_sensor import Sensor, SharedSensorMetadata, RigidSensorMixin, RigidSensorMetadataMixin
+from .base_sensor import (
+    Sensor,
+    SharedSensorMetadata,
+    RigidSensorMixin,
+    RigidSensorMetadataMixin,
+)
 from .sensor_manager import register_sensor
 
 
@@ -195,7 +200,13 @@ class BaseCameraSensor(RigidSensorMixin, Sensor[SharedSensorMetadata]):
     - Shared read() method returning torch tensors
     """
 
-    def __init__(self, options: "SensorOptions", idx: int, data_cls: Type[CameraData], manager: "gs.SensorManager"):
+    def __init__(
+        self,
+        options: "SensorOptions",
+        idx: int,
+        data_cls: Type[CameraData],
+        manager: "gs.SensorManager",
+    ):
         super().__init__(options, idx, data_cls, manager)
         self._stale: bool = True
 
@@ -211,7 +222,9 @@ def _get_cache_dtype(cls) -> torch.dtype:
 
     @classmethod
     def _update_shared_ground_truth_cache(
-        cls, shared_metadata: SharedSensorMetadata, shared_ground_truth_cache: torch.Tensor
+        cls,
+        shared_metadata: SharedSensorMetadata,
+        shared_ground_truth_cache: torch.Tensor,
     ):
         pass
 
@@ -246,7 +259,10 @@ def move_to_attach(self):
             offset_T = torch.tensor(self._options.offset_T, dtype=gs.tc_float, device=gs.device)
         else:
             pos = torch.tensor(self._options.pos, dtype=gs.tc_float, device=gs.device)
-            offset_T = trans_quat_to_T(pos, torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=gs.tc_float, device=gs.device))
+            offset_T = trans_quat_to_T(
+                pos,
+                torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=gs.tc_float, device=gs.device),
+            )
 
         link_pos = self._link.get_pos()
         link_quat = self._link.get_quat()
@@ -423,11 +439,29 @@ def build(self):
 
     def _create_standalone_context(self, scene):
         """Create a simplified RasterizerContext for camera sensors."""
+        if not scene.sim._rigid_only and scene.n_envs > 1:
+            gs.raise_exception("Rasterizer with n_envs > 1, does not work when using non rigid simulation")
+        if sys.platform == "darwin":
+            if scene.n_envs > 1:
+                gs.raise_exception(
+                    "Rasterizer with n_envs > 1, does not work on Metal because it doesn't support OpenGL 4.2"
+                )
+            env_separate_rigid = False
+        else:
+            if self._link is not None:
+                gs.raise_exception("Rasterizer with n_envs > 1, does not work with attached cameras yet.")
+
+            if scene.n_envs > 1:
+                gs.logger.warning(
+                    "Rasterizer with n_envs > 1 is slow as it doesn't do batched rendering consider using BatchRenderer instead."
+                )
+            env_separate_rigid = True
         vis_options = VisOptions(
             show_world_frame=False,
             show_link_frame=False,
             show_cameras=False,
             rendered_envs_idx=range(max(self._manager._sim._B, 1)),
+            env_separate_rigid=env_separate_rigid,
         )
 
         context = RasterizerContext(vis_options)
@@ -501,22 +535,19 @@ def _render_current_state(self):
         self._shared_metadata.context.update(force_render=True)
 
         rgb_arr, _, _, _ = self._shared_metadata.renderer.render_camera(
-            self._get_camera_wrapper(), rgb=True, depth=False, segmentation=False, normal=False
+            self._get_camera_wrapper(),
+            rgb=True,
+            depth=False,
+            segmentation=False,
+            normal=False,
         )
 
         rgb_tensor = torch.from_numpy(rgb_arr.copy()).to(dtype=torch.uint8, device=gs.device)
 
-        # Store in cache
-        n_envs = self._manager._sim._B
-        if n_envs <= 1:
-            # Single environment case - add batch dimension
-            self._shared_metadata.image_cache[self._idx][0] = rgb_tensor
-        else:
-            # Multi-environment rendering is not yet supported for Rasterizer cameras
-            gs.raise_exception(
-                f"Rasterizer camera sensors do not support multi-environment rendering (n_envs={n_envs}). "
-                "Use BatchRenderer camera sensors for batched rendering."
-            )
+        if len(rgb_tensor.shape) == 3:
+            # Single environment rendered - add batch dimension.
+            rgb_tensor = rgb_tensor.unsqueeze(0)
+        self._shared_metadata.image_cache[self._idx][:] = rgb_tensor
 
 
 # ========================== Raytracer Camera Sensor ==========================
@@ -625,7 +656,10 @@ def build(self):
             from genesis.utils.geom import trans_quat_to_T
 
             pos = torch.tensor(opts.pos, dtype=gs.tc_float, device=gs.device)
-            offset_T = trans_quat_to_T(pos, torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=gs.tc_float, device=gs.device))
+            offset_T = trans_quat_to_T(
+                pos,
+                torch.tensor([1.0, 0.0, 0.0, 0.0], dtype=gs.tc_float, device=gs.device),
+            )
             self._camera_obj.attach(self._link, offset_T)
 
         h, w = self._options.res[1], self._options.res[0]
@@ -774,7 +808,12 @@ def _render_current_state(self):
         self._shared_metadata.renderer.update_scene(force_render=True)
 
         rgb_arr, *_ = self._shared_metadata.renderer.render(
-            rgb=True, depth=False, segmentation=False, normal=False, antialiasing=False, force_render=True
+            rgb=True,
+            depth=False,
+            segmentation=False,
+            normal=False,
+            antialiasing=False,
+            force_render=True,
         )
 
         # rgb_arr might be a tuple of arrays (one per camera) or a single array

genesis/ext/pyrender/jit_render.py

@@ -374,7 +374,7 @@ def gen_func_ptr(self):
         self.gl = GLWrapper()
 
         IS_OPENGL_42_AVAILABLE = hasattr(self.gl.wrapper_instance, "glDrawElementsInstancedBaseInstance")
-        OPENGL_42_ERROR_MSG = "Seperated env rendering not supported because OpenGL 4.2 not available on this machine."
+        OPENGL_42_ERROR_MSG = "Separated env rendering not supported because OpenGL 4.2 not available on this machine."
 
         @nb.jit(
             nb.none(

genesis/ext/pyrender/renderer.py

@@ -152,7 +152,7 @@ def render(self, scene, flags, seg_node_map=None, *, is_first_pass=True, force_s
 
         if flags & RenderFlags.ENV_SEPARATE and flags & RenderFlags.OFFSCREEN:
             n_envs = scene.n_envs
-            use_env_idx = True
+            use_env_idx = True and scene.n_envs > 1
         else:
             n_envs = 1
             use_env_idx = False

genesis/vis/camera.py

@@ -156,7 +156,7 @@ def build(self):
                 self._raytracer.add_camera(self)
             else:
                 self._is_batched = False
-                if self._visualizer.scene.n_envs > 0 and self._visualizer._context.env_separate_rigid:
+                if self._visualizer.scene.n_envs > 1 and self._visualizer._context.env_separate_rigid:
                     gs.logger.warning(
                         "Batched rendering via 'VisOptions.env_separate_rigid=True' is only partially supported by "
                         "Rasterizer for now. The same camera transform will be used for all the environments."
@@ -165,7 +165,7 @@ def build(self):
                 if self._env_idx is None:
                     if not self._is_batched:
                         self._env_idx = int(self._visualizer._context.rendered_envs_idx[0])
-                        if self._visualizer.scene.n_envs > 0:
+                        if self._visualizer.scene.n_envs > 1:
                             gs.logger.info(
                                 "Raytracer and Rasterizer requires binding to the camera with a specific environment "
                                 "index. Defaulting to 'rendered_envs_idx[0]'. Please specify 'env_idx' if necessary."

tests/test_sensor_camera.py

@@ -1,9 +1,13 @@
+import sys
+
 import numpy as np
+from numpy.__config__ import show
 import pytest
+import torch
 
 import genesis as gs
 from genesis.utils.misc import tensor_to_array
-from .utils import assert_allclose
+from .utils import assert_allclose, rgb_array_to_png_bytes
 
 
 @pytest.mark.required
@@ -130,3 +134,92 @@ def _get_camera_world_pos(sensor):
     cam_move_dist_offset_T = np.linalg.norm(cam_pos_final_offset_T - cam_pos_initial_offset_T)
     assert cam_move_dist_offset_T > 1e-2
     assert_allclose(cam_move_dist_offset_T, cam_move_dist, atol=1e-2)
+
+
+# ========================== Multi-environment tests ==========================
+
+
+@pytest.mark.required
+def test_rasterizer_camera_sensor_n_envs(show_viewer, png_snapshot):
+    if sys.platform == "darwin":
+        pytest.skip(
+            "Batched rendering with Rasterizer backend is not supported on this machine because it requires OpenGL 4.2."
+        )
+
+    scene = gs.Scene(
+        rigid_options=gs.options.RigidOptions(),
+        renderer=gs.renderers.Rasterizer(),
+        show_viewer=show_viewer,
+    )
+
+    # Add a sphere
+    sphere = scene.add_entity(
+        morph=gs.morphs.Sphere(pos=(0.0, 0.0, 1.0), radius=0.3),
+        surface=gs.surfaces.Smooth(color=(1.0, 0.5, 0.5)),
+    )
+
+    options = gs.sensors.RasterizerCameraOptions(
+        res=(64, 64),
+        pos=(3.0, 0.0, 1.5),
+        lookat=(0.0, 0.0, 0.5),
+        fov=60.0,
+        draw_debug=show_viewer,
+    )
+    camera = scene.add_sensor(options)
+    n_envs = 2
+
+    scene.build(n_envs=n_envs)
+    sphere.set_pos([[0.0, 0.0, 1.0], [0.2, 0.0, 0.5]])
+
+    scene.step()
+
+    data = camera.read()
+
+    assert data.rgb.shape == (2, 64, 64, 3)
+    assert data.rgb.dtype == torch.uint8
+    assert (data.rgb[0] != data.rgb[1]).any(), "We should have different frames"
+
+    for i in range(n_envs):
+        assert rgb_array_to_png_bytes(data.rgb[i]) == png_snapshot
+
+
+@pytest.mark.required
+def test_rasterizer_camera_sensor_n_envs_attached_camera():
+    if sys.platform == "darwin":
+        pytest.skip(
+            "Batched rendering with Rasterizer backend is not supported on this machine because it requires OpenGL 4.2."
+        )
+
+    scene = gs.Scene(
+        rigid_options=gs.options.RigidOptions(
+            enable_collision=True,
+            gravity=(0, 0, -9.8),
+        ),
+        renderer=gs.renderers.Rasterizer(),
+        show_viewer=False,
+    )
+
+    # Add a plane
+    scene.add_entity(
+        morph=gs.morphs.Plane(),
+        surface=gs.surfaces.Rough(color=(0.4, 0.4, 0.4)),
+    )
+
+    # Add a sphere
+    sphere = scene.add_entity(
+        morph=gs.morphs.Sphere(pos=(0.0, 0.0, 1.0), radius=0.3),
+        surface=gs.surfaces.Smooth(color=(1.0, 0.5, 0.5)),
+    )
+
+    options = gs.sensors.RasterizerCameraOptions(
+        res=(64, 64),
+        pos=(3.0, 0.0, 1.5),
+        lookat=(0.0, 0.0, 0.5),
+        fov=60.0,
+        entity_idx=sphere.idx,
+    )
+    scene.add_sensor(options)
+    n_envs = 2
+
+    with pytest.raises(gs.GenesisException, match="does not work with attached cameras yet."):
+        scene.build(n_envs=n_envs)

tests/utils.py

@@ -36,7 +36,7 @@
 DEFAULT_BRANCH_NAME = "main"
 
 HUGGINGFACE_ASSETS_REVISION = "701f78c1465f0a98f6540bae6c9daacaa551b7bf"
-HUGGINGFACE_SNAPSHOT_REVISION = "ea6ae70386c2b2fbae1387f93ba0e4de1ed7abf7"
+HUGGINGFACE_SNAPSHOT_REVISION = "d190d8f651266c56fc54ec9a16b468731211b315"
 
 MESH_EXTENSIONS = (".mtl", *MESH_FORMATS, *GLTF_FORMATS, *USD_FORMATS)
 IMAGE_EXTENSIONS = (".png", ".jpg")