Default to setting max_geom to the estimated max number of renderable geoms in the model

DeepMind · alimuldal · commit bd4cd0eb16bf · 2020-07-03T12:41:02.000+01:00
The previous fixed default of 1000 meant that geoms were sometimes omitted for complex scenes.

PiperOrigin-RevId: 319497192
Change-Id: I852430de4660838557ed51707fbd1ecf2509c03a
diff --git a/dm_control/mujoco/engine.py b/dm_control/mujoco/engine.py
@@ -590,7 +590,7 @@ def __init__(self,
                height=240,
                width=320,
                camera_id=-1,
-               max_geom=1000):
+               max_geom=None):
     """Initializes a new `Camera`.
 
     Args:
@@ -601,8 +601,9 @@ def __init__(self,
         camera, which is always defined. A nonnegative integer or string
         corresponds to a fixed camera, which must be defined in the model XML.
         If `camera_id` is a string then the camera must also be named.
-      max_geom: (optional) An integer specifying the maximum number of geoms
-        that can be represented in the scene.
+      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.
     Raises:
       ValueError: If `camera_id` is outside the valid range, or if `width` or
         `height` exceed the dimensions of MuJoCo's offscreen framebuffer.
diff --git a/dm_control/mujoco/wrapper/core.py b/dm_control/mujoco/wrapper/core.py
@@ -850,19 +850,48 @@ def free(self):
 }
 
 
+def _estimate_max_renderable_geoms(model):
+  """Estimates the maximum number of renderable geoms for a given model."""
+  # Only one type of object frame can be rendered at once.
+  max_nframes = max(
+      [model.nbody, model.ngeom, model.nsite, model.ncam, model.nlight])
+  # This is probably an underestimate, but it is unlikely that all possible
+  # rendering options will be enabled simultaneously, or that all renderable
+  # geoms will be present within the viewing frustum at the same time.
+  return (
+      3 * max_nframes +  # 1 geom per axis for each frame.
+      4 * model.ngeom +  # geom itself + contacts + 2 * split contact forces.
+      3 * model.nbody +  # COM + inertia box + perturbation force.
+      model.nsite +
+      model.ntendon +
+      model.njnt +
+      model.nu +
+      model.nskin +
+      model.ncam +
+      model.nlight)
+
+
 class MjvScene(wrappers.MjvSceneWrapper):  # pylint: disable=missing-docstring
 
-  def __init__(self, model=None, max_geom=1000):
+  def __init__(self, model=None, max_geom=None):
     """Initializes a new `MjvScene` instance.
 
     Args:
       model: (optional) An `MjModel` instance.
       max_geom: (optional) An integer specifying the maximum number of geoms
-        that can be represented in the scene.
+        that can be represented in the scene. If None, this will be chosen
+        automatically based on `model`.
     """
     model_ptr = model.ptr if model is not None else None
     scene_ptr = ctypes.pointer(types.MJVSCENE())
 
+    if max_geom is None:
+      if model is None:
+        max_renderable_geoms = 0
+      else:
+        max_renderable_geoms = _estimate_max_renderable_geoms(model)
+      max_geom = max(1000, max_renderable_geoms)
+
     # Allocate and initialize resources for the abstract scene.
     mjlib.mjv_makeScene(model_ptr, scene_ptr, max_geom)
 
diff --git a/dm_control/viewer/renderer_test.py b/dm_control/viewer/renderer_test.py
@@ -295,14 +295,16 @@ def test_cycling_backward_through_labeling_modes(self):
 
 class SceneCameraTest(parameterized.TestCase):
 
-  def setUp(self):
+  @mock.patch.object(renderer.wrapper.core,
+                     '_estimate_max_renderable_geoms',
+                     return_value=1000)
+  @mock.patch.object(renderer.wrapper.core.mjlib, 'mjv_makeScene')
+  def setUp(self, mock_make_scene, _):
     super(SceneCameraTest, self).setUp()
     self.model = mock.MagicMock()
     self.data = mock.MagicMock()
     self.options = mock.MagicMock()
-    with mock.patch.object(renderer.wrapper.core.mjlib,
-                           'mjv_makeScene') as mock_make_scene:
-      self.camera = renderer.SceneCamera(self.model, self.data, self.options)
+    self.camera = renderer.SceneCamera(self.model, self.data, self.options)
     mock_make_scene.assert_called_once()
 
   def test_freelook_mode(self):
@@ -394,16 +396,18 @@ def test_is_camera_initialized(self, frustum_near, frustum_far, result):
 
 class RaycastsTest(absltest.TestCase):
 
-  def setUp(self):
+  @mock.patch.object(renderer.wrapper.core,
+                     '_estimate_max_renderable_geoms',
+                     return_value=1000)
+  @mock.patch.object(renderer.wrapper.core.mjlib, 'mjv_makeScene')
+  def setUp(self, mock_make_scene, _):
     super(RaycastsTest, self).setUp()
     self.model = mock.MagicMock()
     self.data = mock.MagicMock()
     self.options = mock.MagicMock()
 
     self.viewport = mock.MagicMock()
-    with mock.patch.object(renderer.wrapper.core.mjlib,
-                           'mjv_makeScene') as mock_make_scene:
-      self.camera = renderer.SceneCamera(self.model, self.data, self.options)
+    self.camera = renderer.SceneCamera(self.model, self.data, self.options)
     mock_make_scene.assert_called_once()
     self.initialize_camera(True)
 
diff --git a/setup.py b/setup.py
@@ -177,7 +177,7 @@ def is_excluded(s):
 
 setup(
     name='dm_control',
-    version='0.0.319404054',
+    version='0.0.319497192',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
