diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 7ae36271fc..a866b07aa5 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -8,6 +8,8 @@ Upcoming version (not yet released) Added ^^^^^ +- Added ``CameraSensorCfg.render_update_period`` for reducing camera render + frequency while keeping reset observations fresh. - Added ``BuiltinDcMotorActuator``, a native MuJoCo ```` wrapper. Supports voltage / position / velocity input modes with back-EMF, configurable motor constants, and optional integral, slew, inductance, diff --git a/docs/source/sensors/rgbd_camera.rst b/docs/source/sensors/rgbd_camera.rst index cd7d8c845d..fb6bf36bab 100644 --- a/docs/source/sensors/rgbd_camera.rst +++ b/docs/source/sensors/rgbd_camera.rst @@ -160,7 +160,8 @@ Render settings --------------- All camera sensors in a scene must share identical values for -``use_textures``, ``use_shadows``, and ``enabled_geom_groups``. This +``use_textures``, ``use_shadows``, ``enabled_geom_groups``, and +``render_update_period``. This is a constraint of the underlying MuJoCo Warp rendering system, which uses a single ``RenderContext`` for all cameras. Mismatched settings raise a ``ValueError`` at scene construction. @@ -186,6 +187,28 @@ raise a ``ValueError`` at scene construction. ) +Render cadence +-------------- + +Set ``render_update_period`` to reuse camera images across multiple +``sim.sense()`` calls. This is useful when visual observations are +more expensive than the policy or control rate requires. + +.. code-block:: python + + CameraSensorCfg( + name="wrist_cam", + camera_name="robot/wrist_camera", + data_types=("rgb", "depth"), + render_update_period=5, # render every fifth sense call + ) + +Skipped render calls leave the previous camera buffers unchanged. +Environment resets force a fresh camera render so post-reset +observations are current, even when the reset happens between scheduled +camera updates. + + Output ------ diff --git a/src/mjlab/envs/manager_based_rl_env.py b/src/mjlab/envs/manager_based_rl_env.py index b123ef872c..244fe01af9 100644 --- a/src/mjlab/envs/manager_based_rl_env.py +++ b/src/mjlab/envs/manager_based_rl_env.py @@ -370,7 +370,7 @@ def reset( self.scene.write_data_to_sim() self.sim.forward() self.command_manager.compute(dt=0.0) - self.sim.sense() + self.sim.sense(force_camera_render=True) self.obs_buf = self.observation_manager.compute(update_history=True) self.recorder_manager.record_post_reset(env_ids) return self.obs_buf, self.extras @@ -460,7 +460,7 @@ def step(self, action: torch.Tensor) -> types.VecEnvStepReturn: if "interval" in self.event_manager.available_modes: self.event_manager.apply(mode="interval", dt=self.step_dt) - self.sim.sense() + self.sim.sense(force_camera_render=self.cfg.auto_reset and len(reset_env_ids) > 0) self.obs_buf = self.observation_manager.compute(update_history=True) if self.cfg.auto_reset and len(reset_env_ids) > 0: diff --git a/src/mjlab/sensor/camera_sensor.py b/src/mjlab/sensor/camera_sensor.py index 93b9f35c02..07828c063c 100644 --- a/src/mjlab/sensor/camera_sensor.py +++ b/src/mjlab/sensor/camera_sensor.py @@ -93,6 +93,14 @@ class CameraSensorCfg(SensorCfg): Set to True if you modify the returned data in-place. """ + render_update_period: int = 1 + """Render camera images every N ``sim.sense()`` calls. + + Values greater than 1 keep the previous camera buffer on skipped + calls. Environment resets still force a fresh render so reset + observations are current. + """ + def __post_init__(self) -> None: valid = {"rgb", "depth", "segmentation"} invalid = {dt for dt in self.data_types if dt not in valid} @@ -100,6 +108,8 @@ def __post_init__(self) -> None: raise ValueError(f"Invalid camera data types: {invalid}. Valid types: {valid}") if not self.data_types: raise ValueError("At least one data type must be specified.") + if self.render_update_period < 1: + raise ValueError("render_update_period must be at least 1.") def build(self) -> CameraSensor: return CameraSensor(self) diff --git a/src/mjlab/sensor/sensor_context.py b/src/mjlab/sensor/sensor_context.py index adccd69e77..bb0dbfb06a 100644 --- a/src/mjlab/sensor/sensor_context.py +++ b/src/mjlab/sensor/sensor_context.py @@ -69,6 +69,11 @@ def __init__( if self.camera_sensors: self._validate_sensor_settings() + self.camera_render_update_period = ( + self.camera_sensors[0].cfg.render_update_period if self.camera_sensors else 1 + ) + self._camera_render_step = 0 + self._rgb_unpacked: wp.array | None = None self._rgb_torch: torch.Tensor | None = None self._depth_torch: torch.Tensor | None = None @@ -120,6 +125,17 @@ def finalize(self) -> None: for sensor in self.raycast_sensors: sensor.postprocess_rays() + def should_render_cameras(self, force: bool = False) -> bool: + """Whether this sense call should update camera render buffers.""" + return self.has_cameras and ( + force or self._camera_render_step % self.camera_render_update_period == 0 + ) + + def advance_sense_counter(self) -> None: + """Advance render scheduling after a completed sense call.""" + if self.has_cameras: + self._camera_render_step += 1 + def get_rgb(self, cam_idx: int) -> torch.Tensor: """Get unpacked RGB data for a camera. @@ -250,6 +266,12 @@ def _validate_sensor_settings(self) -> None: f"enabled_geom_groups. '{s.cfg.name}' differs from " f"'{ref.name}'." ) + if cfg.render_update_period != ref.render_update_period: + raise ValueError( + "All camera sensors must share the same " + f"render_update_period. '{s.cfg.name}' differs from " + f"'{ref.name}'." + ) def _raycast_geom_groups(self) -> set[int]: """Compute the union of geom groups across all raycast sensors.""" diff --git a/src/mjlab/sim/sim.py b/src/mjlab/sim/sim.py index 0ead49cd65..1b6a98b3f6 100644 --- a/src/mjlab/sim/sim.py +++ b/src/mjlab/sim/sim.py @@ -326,6 +326,7 @@ def create_graph(self) -> None: self.forward_graph = None self.reset_graph = None self.sense_graph = None + self.sense_no_camera_render_graph = None if self.use_cuda_graph: with _suspend_gc(), wp.ScopedDevice(self.wp_device): with wp.ScopedCapture() as capture: @@ -339,8 +340,15 @@ def create_graph(self) -> None: self.reset_graph = capture.graph if self._sensor_context is not None: with wp.ScopedCapture() as capture: - self._sense_kernel() + self._sense_kernel(render_cameras=True) self.sense_graph = capture.graph + if ( + self._sensor_context.has_cameras + and self._sensor_context.camera_render_update_period > 1 + ): + with wp.ScopedCapture() as capture: + self._sense_kernel(render_cameras=False) + self.sense_no_camera_render_graph = capture.graph # Properties. @@ -483,30 +491,41 @@ def set_sensor_context(self, ctx: SensorContext) -> None: self._sensor_context = ctx self.create_graph() - def sense(self) -> None: + def sense(self, *, force_camera_render: bool = False) -> None: """Execute the sense pipeline: prepare -> graph -> finalize. Runs BVH refit, camera rendering, and raycasting in a single CUDA graph launch. Should be called once per env step, right before observation computation. + + Args: + force_camera_render: Refresh camera buffers even when their configured + render period would normally skip this sense call. Reset paths use this + so post-reset visual observations are current. """ if self._sensor_context is None: return ctx = self._sensor_context ctx.prepare() + render_cameras = ctx.should_render_cameras(force=force_camera_render) with wp.ScopedDevice(self.wp_device): if self.use_cuda_graph and self.sense_graph is not None: - wp.capture_launch(self.sense_graph) + if render_cameras or not ctx.has_cameras: + wp.capture_launch(self.sense_graph) + else: + assert self.sense_no_camera_render_graph is not None + wp.capture_launch(self.sense_no_camera_render_graph) else: - self._sense_kernel() + self._sense_kernel(render_cameras=render_cameras) ctx.finalize() + ctx.advance_sense_counter() # Private methods. - def _sense_kernel(self) -> None: + def _sense_kernel(self, *, render_cameras: bool = True) -> None: """GPU kernel sequence for sensing (captured in sense_graph).""" assert self._sensor_context is not None ctx = self._sensor_context @@ -514,7 +533,7 @@ def _sense_kernel(self) -> None: mjwarp.refit_bvh(self.wp_model, self.wp_data, rc) - if ctx.has_cameras: + if render_cameras and ctx.has_cameras: mjwarp.render(self.wp_model, self.wp_data, rc) ctx.unpack_rgb() diff --git a/src/mjlab/viewer/viser/viewer.py b/src/mjlab/viewer/viser/viewer.py index 56957d32b5..fd3df8176c 100644 --- a/src/mjlab/viewer/viser/viewer.py +++ b/src/mjlab/viewer/viser/viewer.py @@ -332,7 +332,7 @@ def _handle_gui_reset(self, all_envs: bool) -> None: if env.command_manager.apply_gui_reset(env_ids): env.scene.write_data_to_sim() env.sim.forward() - env.sim.sense() + env.sim.sense(force_camera_render=True) self._pending_update_reasons.add(UpdateReason.ACTION) self._sync_ui_state() diff --git a/tests/test_camera_sensor.py b/tests/test_camera_sensor.py index 1d10e49967..98ed4e3531 100644 --- a/tests/test_camera_sensor.py +++ b/tests/test_camera_sensor.py @@ -328,6 +328,62 @@ def test_cam_intrinsic_dr_changes_rendered_image(device): assert torch.equal(after[1], baseline[1]), "Env 1 image changed unexpectedly" +def test_camera_render_period_skip_force(device): + """Camera buffers are held on skipped sense calls and update when forced.""" + cam_cfg = CameraSensorCfg( + name="test_cam", + camera_name="world/intrinsic_cam", + width=32, + height=24, + data_types=("rgb",), + render_update_period=3, + ) + scene, sim = _make_scene_and_sim( + device, sensors=(cam_cfg,), xml=INTRINSIC_CAM_XML, num_envs=2 + ) + sim.expand_model_fields(("cam_intrinsic",)) + + sim.forward() + sim.sense() + baseline = scene["test_cam"].data.rgb.clone() + + cam_id = sim.mj_model.camera("world/intrinsic_cam").id + sim.model.cam_intrinsic[0, cam_id, :2] *= 5.0 + + sim.forward() + sim.sense() + skipped = scene["test_cam"].data.rgb + assert torch.equal(skipped, baseline) + + sim.sense(force_camera_render=True) + forced = scene["test_cam"].data.rgb + assert not torch.equal(forced[0], baseline[0]) + assert torch.equal(forced[1], baseline[1]) + + +def test_camera_render_period_mismatch(device): + """Render cadence is shared because cameras use one RenderContext.""" + cam_a = CameraSensorCfg( + name="cam_a", + camera_name="world/overhead_cam", + width=16, + height=12, + data_types=("rgb",), + render_update_period=1, + ) + cam_b = CameraSensorCfg( + name="cam_b", + camera_name="world/overhead_cam", + width=16, + height=12, + data_types=("depth",), + render_update_period=2, + ) + + with pytest.raises(ValueError, match="render_update_period"): + _make_scene_and_sim(device, sensors=(cam_a, cam_b)) + + def test_camera_segmentation_shape(device): """Segmentation data has correct shape [num_envs, height, width, 2].""" cam_cfg = CameraSensorCfg(