Dynamically increase offscreen framebuffer size.

kevinzakka · Jan 18, 2023 · 3b4be04 · 3b4be04
1 parent 0804655
commit 3b4be04
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/dm_env_wrappers/_src/mujoco/dm_control_video.py b/dm_env_wrappers/_src/mujoco/dm_control_video.py
@@ -41,6 +41,14 @@ def __init__(
         self._height = height
         self._width = width
 
+        # Ensure the offscreen framebuffer is large enough to accommodate the requested
+        # resolution.
+        new_offwidth = max(self.physics.model.vis.global_.offwidth, width)
+        new_offheight = max(self.physics.model.vis.global_.offheight, height)
+        mjcf_model = self._task.root_entity.mjcf_model
+        mjcf_model.visual.__getattr__("global").offheight = new_offheight
+        mjcf_model.visual.__getattr__("global").offwidth = new_offwidth
+
     # Helper methods.
 
     def _render_frame(self, observation) -> np.ndarray: