Suport camera group ids for batch rendering on RT backend (#108)

Author: yuecideng

embodichain/lab/gym/envs/base_env.py

@@ -24,7 +24,7 @@
 from embodichain.lab.sim.types import EnvObs, EnvAction
 from embodichain.lab.sim import SimulationManagerCfg, SimulationManager
 from embodichain.lab.sim.objects import Robot
-from embodichain.lab.sim.sensors import BaseSensor
+from embodichain.lab.sim.sensors import BaseSensor, Camera
 from embodichain.lab.gym.utils import gym_utils
 from embodichain.utils import configclass
 from embodichain.utils import logger, set_seed
@@ -225,6 +225,17 @@ def get_sensor(self, name: str, **kwargs) -> BaseSensor:
 
         return self.sensors[name]
 
+    def add_camera_group_id(self, group_id: int) -> None:
+        """Add a camera group ID for rendering.
+
+        Args:
+            group_id: The camera group ID to be added.
+        """
+        if not hasattr(self, "_camera_group_ids"):
+            self._camera_group_ids: List[int] = []
+        if self.sim.is_rt_enabled:
+            self._camera_group_ids.append(group_id)
+
     def _setup_scene(self, **kwargs):
         # Init sim manager.
         # we want to open gui window when the scene is setup, so init sim manager in headless mode first.
@@ -251,6 +262,13 @@ def _setup_scene(self, **kwargs):
 
         self.sensors = self._setup_sensors(**kwargs)
 
+        # Setup camera groups for rendering.
+        self._camera_group_ids: List[int] = []
+        if self.sim.is_rt_enabled:
+            for sensor in self.sensors.values():
+                if isinstance(sensor, Camera):
+                    self._camera_group_ids.append(sensor.group_id)
+
     def _setup_robot(self, **kwargs) -> Robot:
         """Load the robot agent, setup the controller and action space.
 
@@ -343,7 +361,7 @@ def _get_sensor_obs(self, **kwargs) -> Dict[str, any]:
         fetch_only = False
         if self.sim.is_rt_enabled:
             fetch_only = True
-            self.sim.render_camera_group()
+            self.sim.render_camera_group(self._camera_group_ids)
 
         for sensor_name, sensor in self.sensors.items():
             sensor.update(fetch_only=fetch_only)

embodichain/lab/gym/envs/managers/record.py

@@ -24,7 +24,7 @@
 
 from dexsim.utility import images_to_video
 from embodichain.lab.gym.envs.managers import Functor, FunctorCfg
-from embodichain.lab.sim.sensors.camera import CameraCfg
+from embodichain.lab.sim.sensors.camera import CameraCfg, Camera
 
 if TYPE_CHECKING:
     from embodichain.lab.gym.envs import EmbodiedEnv
@@ -69,7 +69,7 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
             "intrinsics", (600, 600, int(resolution[0] / 2), int(resolution[1] / 2))
         )
 
-        self.camera = env.sim.add_sensor(
+        self.camera: Camera = env.sim.add_sensor(
             sensor_cfg=CameraCfg(
                 uid=self._name,
                 width=resolution[0],
@@ -79,6 +79,10 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
             )
         )
 
+        # Add this camera's group ID to the environment for batch rendering when RT is enabled.
+        if getattr(env.sim, "is_rt_enabled", False):
+            env.add_camera_group_id(self.camera.group_id)
+
         self._current_episode = 0
         self._frames: List[np.ndarray] = []
 

embodichain/lab/sim/sensors/camera.py

@@ -248,6 +248,21 @@ def is_rt_enabled(self) -> bool:
         """
         return is_rt_enabled()
 
+    @cached_property
+    def group_id(self) -> int:
+        """Get the camera group ID in the dexsim world.
+
+        Returns:
+            int: The camera group ID.
+        """
+        if self.is_rt_enabled:
+            return self._frame_buffer.get_group_id()
+        else:
+            logger.log_warning(
+                "Camera group ID is only available for Ray Tracing renderer. Returning -1 for non-RT renderer."
+            )
+            return -1
+
     def update(self, **kwargs) -> None:
         """Update the sensor data.
 

embodichain/lab/sim/sim_manager.py

@@ -485,14 +485,17 @@ def init_gpu_physics(self) -> None:
 
         self._is_initialized_gpu_physics = True
 
-    def render_camera_group(self) -> None:
+    def render_camera_group(self, group_ids: list[int]) -> None:
         """Render all camera group in the simulation.
 
+        Args:
+            group_ids (list[int]): The list of camera group ids to render.
+
         Note: This interface is only valid when Ray Tracing rendering backend is enabled.
         """
 
         if self.is_rt_enabled:
-            self._world.render_camera_group()
+            self._world.render_camera_group(group_ids)
         else:
             logger.log_warning(
                 "This interface is only valid when Ray Tracing rendering backend is enabled."

pyproject.toml

@@ -26,7 +26,7 @@ dynamic = ["version"]
 # Core install dependencies (kept from requirements.txt). Some VCS links are
 # specified using PEP 508 direct references where present.
 dependencies = [
-  "dexsim_engine==0.3.9",
+  "dexsim_engine==0.3.10",
   "setuptools>=78.1.1",
   "gymnasium>=0.29.1",
   "langchain",

tests/gym/envs/test_embodied_env.py

@@ -119,12 +119,14 @@
 class EmbodiedEnvTest:
     """Shared test logic for CPU and CUDA."""
 
-    def setup_simulation(self, sim_device):
+    def setup_simulation(self, sim_device, enable_rt):
         cfg: EmbodiedEnvCfg = config_to_cfg(
             METADATA, manager_modules=DEFAULT_MANAGER_MODULES
         )
         cfg.num_envs = NUM_ENVS
-        cfg.sim_cfg = SimulationManagerCfg(headless=True, sim_device=sim_device)
+        cfg.sim_cfg = SimulationManagerCfg(
+            headless=True, sim_device=sim_device, enable_rt=enable_rt
+        )
 
         self.env = gym.make(id=METADATA["id"], cfg=cfg)
 
@@ -162,10 +164,15 @@ def teardown_method(self):
 
 class TestCPU(EmbodiedEnvTest):
     def setup_method(self):
-        self.setup_simulation("cpu")
+        self.setup_simulation("cpu", enable_rt=False)
+
+
+class TestCPURT(EmbodiedEnvTest):
+    def setup_method(self):
+        self.setup_simulation("cpu", enable_rt=True)
 
 
 @pytest.mark.skip(reason="Skipping CUDA tests temporarily")
 class TestCUDA(EmbodiedEnvTest):
     def setup_method(self):
-        self.setup_simulation("cuda")
+        self.setup_simulation("cuda", enable_rt=False)