Add scene callbacks to viewer/renderer.py and viewer/viewer.py

PiperOrigin-RevId: 472485161
Change-Id: Ic5fa30c97e1a33b6e0c8878818d0f51539795acb

dm_control/viewer/renderer.py

@@ -378,7 +378,8 @@ def __init__(self,
                data,
                options,
                settings=None,
-               zoom_factor=_FULL_SCENE_ZOOM_FACTOR):
+               zoom_factor=_FULL_SCENE_ZOOM_FACTOR,
+               scene_callback=None):
     """Instance initializer.
 
     Args:
@@ -388,6 +389,9 @@ def __init__(self,
       settings: Optional, internal camera settings obtained from another
         SceneCamera instance using 'settings' property.
       zoom_factor: The initial zoom factor for zooming into the scene.
+      scene_callback: Scene callback.
+        This is a callable of the form: `my_callable(MjModel, MjData, MjvScene)`
+        that gets applied to every rendered scene.
     """
     # Design notes:
     # We need to recreate the camera for each new model, because each model
@@ -403,6 +407,7 @@ def __init__(self,
     self._camera.fixedcamid = _FREE_CAMERA_INDEX
     self._camera.type_ = mujoco.mjtCamera.mjCAMERA_FREE
     self._zoom_factor = zoom_factor
+    self._scene_callback = scene_callback
 
     if settings is not None:
       self._settings = settings
@@ -499,6 +504,10 @@ def render(self, perturbation=None):
                            self._options.visualization.ptr, perturb_to_render,
                            self._camera.ptr, mujoco.mjtCatBit.mjCAT_ALL,
                            self._scene.ptr)
+
+    # Apply callback if defined.
+    if self._scene_callback is not None:
+      self._scene_callback(self._model, self._data, self._scene)
     return self._scene
 
   def zoom_to_scene(self):

dm_control/viewer/viewer.py

@@ -95,7 +95,7 @@ class Viewer:
   """Viewport displaying the contents of a physics world."""
 
   def __init__(self, viewport, mouse, keyboard, camera_settings=None,
-               zoom_factor=_FULL_SCENE_ZOOM_FACTOR):
+               zoom_factor=_FULL_SCENE_ZOOM_FACTOR, scene_callback=None):
     """Instance initializer.
 
     Args:
@@ -104,6 +104,9 @@ def __init__(self, viewport, mouse, keyboard, camera_settings=None,
       keyboard: A keyboard device.
       camera_settings: Properties of the scene MjvCamera.
       zoom_factor: Initial scale factor for zooming into the scene.
+      scene_callback: Scene callback.
+        This is a callable of the form: `my_callable(MjModel, MjData, MjvScene)`
+        that gets applied to every rendered scene.
     """
     self._viewport = viewport
     self._mouse = mouse
@@ -119,6 +122,7 @@ def __init__(self, viewport, mouse, keyboard, camera_settings=None,
     self._free_camera = None
     self._camera_select = None
     self._zoom_factor = zoom_factor
+    self._scene_callback = scene_callback
 
   def __del__(self):
     del self._camera
@@ -136,8 +140,12 @@ def initialize(self, physics, renderer_instance, touchpad):
       touchpad: A boolean, use input dedicated to touchpad.
     """
     self._camera = renderer.SceneCamera(
-        physics.model, physics.data, self._render_settings,
-        settings=self._camera_settings, zoom_factor=self._zoom_factor)
+        physics.model,
+        physics.data,
+        self._render_settings,
+        settings=self._camera_settings,
+        zoom_factor=self._zoom_factor,
+        scene_callback=self._scene_callback)
 
     self._manipulator = ManipulationController(
         self._viewport, self._camera, self._mouse)