Fix #2212.

PiperOrigin-RevId: 693294401 Change-Id: Ifef2b268f61b90f557580f61316786c50dc9698d
google-deepmind · Nov 5, 2024 · f9569cd · f9569cd
1 parent b7bbb2b
commit f9569cd
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 3 deletions.
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+
+Upcoming version (not yet released)
+-----------------------------------
+
+Bug fixes
+^^^^^^^^^
+1. Fixed :github:issue:`2212`, type error in ```mjx.get_data``.
+
 Version 3.2.5 (Nov 4, 2024)
 ---------------------------
 

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
@@ -432,9 +432,9 @@ def get_data_into(
 
       # MuJoCo actuator_moment is sparse, MJX uses a dense representation.
       if field.name == 'actuator_moment' and m.nu:
-        moment_rownnz = np.zeros(m.nu, dtype=int)
-        moment_rowadr = np.zeros(m.nu, dtype=int)
-        moment_colind = np.zeros(m.nu * m.nv, dtype=int)
+        moment_rownnz = np.zeros(m.nu, dtype=np.int32)
+        moment_rowadr = np.zeros(m.nu, dtype=np.int32)
+        moment_colind = np.zeros(m.nu * m.nv, dtype=np.int32)
         actuator_moment = np.zeros(m.nu * m.nv)
         mujoco.mju_dense2sparse(
             actuator_moment,

diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
@@ -422,6 +422,26 @@ def test_get_data(self):
     np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
     np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
 
+  def test_get_data_runs(self):
+    xml = """
+      <mujoco>
+        <compiler autolimits="true"/>
+        <worldbody>
+          <body name="box">
+            <joint name="slide1" type="slide" axis="1 0 0" />
+            <geom type="box" size=".05 .05 .05" mass="1"/>
+          </body>
+        </worldbody>
+        <actuator>
+          <motor joint="slide1"/>
+        </actuator>
+      </mujoco>
+    """
+    m = mujoco.MjModel.from_xml_string(xml)
+    d = mujoco.MjData(m)
+    dx = mjx.put_data(m, d)
+    mjx.get_data(m, dx)
+
   def test_get_data_batched(self):
     """Test that get_data makes correct List[MjData] for batched Data."""