Add a callback to Camera.render for modifying the scene.

Camera.render calls mjv_updateScene before rendering, which makes it hard for users to add geoms to the scene.
This callback is called after mjv_updateScene and before rendering.

PiperOrigin-RevId: 455351943
Change-Id: I58de48218ad2012b23eaeffe3396b2321b905e3f

Author: nimrod-gileadi

dm_control/mujoco/engine.py

@@ -34,7 +34,7 @@
 import collections
 import contextlib
 import threading
-from typing import NamedTuple
+from typing import Callable, NamedTuple, Optional, Union
 
 from absl import logging
 
@@ -177,6 +177,8 @@ def render(
       segmentation=False,
       scene_option=None,
       render_flag_overrides=None,
+      scene_callback: Optional[Callable[['Physics', mujoco.MjvScene],
+                                        None]] = None,
   ):
     """Returns a camera view as a NumPy array of pixel values.
 
@@ -204,12 +206,18 @@ def render(
         `{'wireframe': True}` or `{mujoco.mjtRndFlag.mjRND_WIREFRAME: True}`.
         See `mujoco.mjtRndFlag` for the set of valid flags. Must be None if
         either `depth` or `segmentation` is True.
+      scene_callback: Called after the scene has been created and before
+        it is rendered. Can be used to add more geoms to the scene.
 
     Returns:
       The rendered RGB, depth or segmentation image.
     """
     camera = Camera(
-        physics=self, height=height, width=width, camera_id=camera_id)
+        physics=self,
+        height=height,
+        width=width,
+        camera_id=camera_id,
+        scene_callback=scene_callback)
     image = camera.render(
         overlays=overlays, depth=depth, segmentation=segmentation,
         scene_option=scene_option, render_flag_overrides=render_flag_overrides)
@@ -602,12 +610,16 @@ class Camera:
   `camera_id`, for example to render the same view at different resolutions.
   """
 
-  def __init__(self,
-               physics,
-               height=240,
-               width=320,
-               camera_id=-1,
-               max_geom=None):
+  def __init__(
+      self,
+      physics,
+      height: int = 240,
+      width: int = 320,
+      camera_id: Union[int, str] = -1,
+      max_geom: Optional[int] = None,
+      scene_callback: Optional[Callable[[Physics, mujoco.MjvScene],
+                                        None]] = None,
+  ):
     """Initializes a new `Camera`.
 
     Args:
@@ -621,6 +633,8 @@ def __init__(self,
       max_geom: Optional integer specifying the maximum number of geoms that can
         be rendered in the same scene. If None this will be chosen automatically
         based on the estimated maximum number of renderable geoms in the model.
+      scene_callback: Called after the scene has been created and before
+        it is rendered. Can be used to add more geoms to the scene.
     Raises:
       ValueError: If `camera_id` is outside the valid range, or if `width` or
         `height` exceed the dimensions of MuJoCo's offscreen framebuffer.
@@ -652,6 +666,7 @@ def __init__(self,
     self._width = width
     self._height = height
     self._physics = physics
+    self._scene_callback = scene_callback
 
     # Variables corresponding to structs needed by Mujoco's rendering functions.
     self._scene = wrapper.MjvScene(model=physics.model, max_geom=max_geom)
@@ -844,6 +859,9 @@ def render(
     # Update scene geometry.
     self.update(scene_option=scene_option)
 
+    if self._scene_callback:
+      self._scene_callback(self._physics, self._scene)
+
     # Enable flags to compute segmentation labels
     if segmentation:
       render_flag_overrides.update({

dm_control/mujoco/engine_test.py

@@ -132,6 +132,41 @@ def testSegmentationRender(self, enable_geom_frame_rendering):
       contains_decor = np.any(pixels[:, :, 1] == obj_type_decor)
       self.assertEqual(contains_decor, enable_geom_frame_rendering)
 
+  def testSceneCallback(self):
+    empty_world = """
+    <mujoco>
+      <worldbody>
+        <camera name="cam" pos="0 0 3"/>
+      </worldbody>
+    </mujoco>
+    """
+
+    def callback(_, scn: mujoco.MjvScene):
+      # Add a red box to the scene
+      scn.ngeom += 1
+      mujoco.mjv_initGeom(
+          scn.geoms[scn.ngeom - 1],
+          mujoco.mjtGeom.mjGEOM_BOX.value,
+          size=np.array([0.2, 0.2, 0.2]),
+          pos=np.zeros(3),
+          mat=np.eye(3).flatten(),
+          rgba=np.array([1, 0, 0, 1], dtype=np.float32))
+
+    physics = engine.Physics.from_xml_string(empty_world)
+
+    # Without the callback, render should return a black image.
+    empty_image = physics.render(
+        height=8, width=8, camera_id='cam', scene_callback=None)
+    np.testing.assert_array_equal(
+        np.zeros((8, 8, 3), dtype=np.uint8), empty_image)
+
+    # With the callback, there should be a red box.
+    pixels = physics.render(
+        height=8, width=8, camera_id='cam', scene_callback=callback)
+    # Are there any pixels where red component is bigger than green and blue?
+    any_red_pixels = np.any(pixels[:, :, 0] > np.max(pixels[:, :, 1:3], axis=2))
+    self.assertTrue(any_red_pixels, 'Expecting some red pixels.')
+
   def testTextOverlay(self):
     height, width = 480, 640
     overlay = engine.TextOverlay(title='Title', body='Body', style='big',