Free the Physics object when Environment.close is called.

Make physics.free() idempotent since sometimes it's called by __del__ functions.

PiperOrigin-RevId: 717526328
Change-Id: I1df3c0141cf0d5440420ee48f805a8a71580bb13

Author: nimrod-gileadi

dm_control/composer/environment.py

@@ -249,8 +249,9 @@ def _recompile_physics_and_update_observables(self):
 
   def _recompile_physics(self):
     """Creates a new Physics using the latest MJCF model from the task."""
-    if getattr(self, '_physics', None):
-      self._physics.free()
+    physics = getattr(self, '_physics', None)
+    if physics:
+      physics.free()
     self._physics = mjcf.Physics.from_mjcf_model(
         self._task.root_entity.mjcf_model)
     self._physics.legacy_step = self._legacy_step
@@ -267,9 +268,10 @@ def physics(self):
     """Returns a `weakref.ProxyType` pointing to the current `mjcf.Physics`.
 
     Note that the underlying `mjcf.Physics` will be destroyed whenever the MJCF
-    model is recompiled. It is therefore unsafe for external objects to hold a
-    reference to `environment.physics`. Attempting to access attributes of a
-    dead `Physics` instance will result in a `ReferenceError`.
+    model is recompiled or environment.close() is called. It is therefore unsafe
+    for external objects to hold a reference to `environment.physics`.
+    Attempting to access attributes of a dead `Physics` instance will result in
+    a `ReferenceError`.
     """
     return self._physics_proxy
 
@@ -462,6 +464,11 @@ def _substep(self, action):
     self._physics.step()
     self._hooks.after_substep(self._physics_proxy, self._random_state)
 
+  def close(self):
+    super().close()
+    self._physics.free()
+    self._physics = None
+
   def action_spec(self):
     """Returns the action specification for this environment."""
     return self._task.action_spec(self._physics_proxy)

dm_control/mujoco/engine.py

@@ -417,7 +417,8 @@ def free(self):
     with self._contexts_lock:
       if self._contexts:
         self._free_rendering_contexts()
-    del self._data
+    if hasattr(self, '_data'):
+      del self._data
 
   @classmethod
   def from_model(cls, model):

dm_control/mujoco/engine_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ============================================================================
 
-"""Tests for `engine`."""
-
 import copy
 import pickle
 
@@ -202,6 +200,10 @@ def testRenderFlags(self):
         np.all(disabled == enabled),
         msg='Images are identical regardless of whether wireframe is enabled.')
 
+  def testFreeIdempotent(self):
+    self._physics.free()
+    self._physics.free()
+
   @parameterized.parameters(((0.5, 0.5), (1, 3)),  # pole
                             ((0.5, 0.1), (0, 0)),  # ground
                             ((0.9, 0.9), (None, None)),  # sky