Python bindings: Store a reference to MjModel in MjDataWrapper.

Before this change, MjDataWrapper contained a custom data structure called MjDataMetadata, which contained the sizes of various arrays needed to create MjData. This was used for serialization and deserialization, as well as copying MjData instances.

This worked fine for native MuJoCo models, but as soon as plugins were used, there was information in the model that was needed and not available in MjDataMetadata. Since plugins are so general, keeping MjDataMetadata was untenable.

PiperOrigin-RevId: 578205355
Change-Id: I11b9ee797d1da5aaf1fecbbe170b3b0f1f1a3e6c

python/mujoco/bindings_test.py

@@ -74,10 +74,21 @@ TEST_XML_SENSOR = r"""
 """
 
 TEST_XML_PLUGIN = r"""
-<mujoco model="test">
+<mujoco>
+  <option gravity="0 0 0"/>
   <extension>
     <plugin plugin="mujoco.elasticity.cable"/>
   </extension>
+  <worldbody>
+    <composite type="cable" curve="s" count="41 1 1" size="1" offset="0 0 1" initial="none">
+      <plugin plugin="mujoco.elasticity.cable">
+        <config key="twist" value="1e6"/>
+        <config key="bend" value="1e9"/>
+      </plugin>
+      <joint kind="main" damping="2"/>
+      <geom type="capsule" size=".005" density="1"/>
+    </composite>
+  </worldbody>
 </mujoco>
 """
 
@@ -1048,10 +1059,24 @@ Euler integrator, semi-implicit in velocity.
     while data_instances:
       d = data_instances.pop()
       self.assertEqual(sys.getrefcount(d), 2)
+      del d
     while model_instances:
       m = model_instances.pop()
       self.assertEqual(sys.getrefcount(m), 2)
 
+  # This test is disabled on PyPy as it uses sys.getrefcount
+  # However PyPy is not officially supported by MuJoCo
+  @absltest.skipIf(sys.implementation.name == 'pypy',
+                   reason='requires sys.getrefcount')
+  def test_mjdata_holds_ref_to_model(self):
+    data = mujoco.MjData(mujoco.MjModel.from_xml_string('<mujoco/>'))
+    model = data.model
+    # references: one in `data.model, one in `model`, one in the temporary ref
+    # passed to getrefcount.
+    self.assertEqual(sys.getrefcount(data.model), 3)
+    del data
+    self.assertEqual(sys.getrefcount(model), 2)
+
   def test_can_initialize_mjv_structs(self):
     self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
     self.assertIsInstance(mujoco.MjvCamera(), mujoco.MjvCamera)
@@ -1287,6 +1312,34 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(data.geom(3).xpos[2], 4)
     self.assertEqual(data.geom(4).xpos[2], 5)
 
+  def test_load_plugin(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
+    data = mujoco.MjData(model)
+    mujoco.mj_forward(model, data)
+
+  def test_copy_mjdata_with_plugin(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
+    data1 = mujoco.MjData(model)
+    self.assertIs(data1.model, model)
+    mujoco.mj_step(model, data1)
+    data2 = copy.copy(data1)
+    mujoco.mj_step(model, data1)
+    mujoco.mj_step(model, data2)
+    np.testing.assert_array_equal(data1.qpos, data2.qpos)
+    self.assertIs(data1.model, data2.model)
+
+  def test_deepcopy_mjdata_with_plugin(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
+    data1 = mujoco.MjData(model)
+    self.assertIs(data1.model, model)
+    mujoco.mj_step(model, data1)
+    data2 = copy.deepcopy(data1)
+    mujoco.mj_step(model, data1)
+    mujoco.mj_step(model, data2)
+    np.testing.assert_array_equal(data1.qpos, data2.qpos)
+    self.assertIsNot(data1.model, data2.model)
+    self.assertNotEqual(data1.model._address, data2.model._address)
+
   def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
     for name in attr_to_compare:
       actual_value = getattr(actual_obj, name)
@@ -1300,9 +1353,6 @@ Euler integrator, semi-implicit in velocity.
         self.fail("Attribute '{}' differs from expected value: {}".format(
             name, str(e)))
 
-  def test_load_plugin(self):
-    mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
-
 
 if __name__ == '__main__':
   absltest.main()