From 02c62c118ca8577fffcb9dc8658a07841684c11c Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Thu, 14 Mar 2024 15:50:16 -0700
Subject: [PATCH] Fix bug where treedef for mjx.make_data doesn't match
 mjx.put_data.

PiperOrigin-RevId: 615928973
Change-Id: I08af97ac0300b9d51c1f4b46d48dde98d24197fa
---
 doc/changelog.rst              |  1 +
 mjx/mujoco/mjx/_src/io.py      |  5 ++++-
 mjx/mujoco/mjx/_src/io_test.py | 15 +++++++++++++++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index bb66f206..fdfafeef 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -10,6 +10,7 @@ MJX
 
 1. Improved performance of SAT for convex collisions.
 2. Fixed bug for sphere/capsule-convex deep penetration.
+3. Fixed bug where ``mjx.Data`` produced by ``mjx.put_data`` had different treedef than ``mjx.make_data``.
 
 Version 3.1.3 (March 5th, 2024)
 -----------------------------------
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 4b0c2f4f..ddcd8907 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -172,7 +172,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       ctrl=zero_nu,
       qfrc_applied=zero_nv,
       xfrc_applied=zero_nbody_6,
-      eq_active=jp.zeros(m.neq, dtype=int),
+      eq_active=jp.zeros(m.neq, dtype=jp.uint8),
       qacc=zero_nv,
       act_dot=zero_na,
       xpos=zero_nbody_3,
@@ -374,6 +374,9 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
   for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
     fields[fname] = fields[fname].reshape((-1, 3, 3))
 
+  # MJX does not support islanding, so only transfer the first solver_niter
+  fields['solver_niter'] = fields['solver_niter'][0]
+
   # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
   # efc_J is also optionally column-sparse (typically for large nv).  MJX is
   # neither: it contains zeros for inactive constraints, and efc_J is always
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 9ceca6b0..e3f414df 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/mjx/mujoco/mjx/_src/io_test.py
@@ -434,5 +434,20 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d_2.contact.frame.shape, (1, 9))
     np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
 
+  def test_make_matches_put(self):
+    """Test that make_data produces a pytree that matches put_data."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+
+    step_fn = lambda d: d.replace(time=d.time + 1)
+    step_fn_jit = jax.jit(step_fn).lower(dx).compile()
+
+    # placing an MjData onto device should yield the same treedef mjx.Data as
+    # calling make_data.  they should be interchangeable for jax functions:
+    step_fn_jit(mjx.make_data(m))
+
 if __name__ == '__main__':
   absltest.main()