From cb3c83fadef4eecf2d4a689652d93b26f5dcf8bf Mon Sep 17 00:00:00 2001
From: yuecideng <dengyueci@qq.com>
Date: Sat, 31 Jan 2026 23:03:59 +0800
Subject: [PATCH] wip

---
 embodichain/lab/gym/envs/base_env.py        | 20 ++++++++++++++++++--
 embodichain/lab/gym/envs/managers/record.py |  7 +++++--
 embodichain/lab/sim/sensors/camera.py       |  9 +++++++++
 embodichain/lab/sim/sim_manager.py          |  7 +++++--
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py
index fd0f873c..c5bdd411 100644
--- a/embodichain/lab/gym/envs/base_env.py
+++ b/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
@@ -219,6 +219,16 @@ 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] = []
+        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.
@@ -245,6 +255,12 @@ def _setup_scene(self, **kwargs):
 
         self.sensors = self._setup_sensors(**kwargs)
 
+        # Setup camera groups for rendering.
+        self._camera_group_ids: List[int] = []
+        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.
 
@@ -337,7 +353,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)
diff --git a/embodichain/lab/gym/envs/managers/record.py b/embodichain/lab/gym/envs/managers/record.py
index efaadd39..71875b45 100644
--- a/embodichain/lab/gym/envs/managers/record.py
+++ b/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,9 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
             )
         )
 
+        # Add this camera's group ID to the environment for batch rendering.
+        env.add_camera_group_id(self.camera.group_id)
+
         self._current_episode = 0
         self._frames: List[np.ndarray] = []
 
diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py
index a4aeb749..bbdcf6c8 100644
--- a/embodichain/lab/sim/sensors/camera.py
+++ b/embodichain/lab/sim/sensors/camera.py
@@ -248,6 +248,15 @@ 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.
+        """
+        return self._frame_buffer.get_group_id()
+
     def update(self, **kwargs) -> None:
         """Update the sensor data.
 
diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py
index 5e6b47ce..5116d5ad 100644
--- a/embodichain/lab/sim/sim_manager.py
+++ b/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."