From a64f1fd33526944a9a32391b98f7820d4e2bcb4b Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Tue, 17 Jun 2025 13:20:14 -0700
Subject: [PATCH] Compute M in mjx._get_data_into.

PiperOrigin-RevId: 772594295
Change-Id: I7da992e9e5590c6060fbe1e2c1e40ccbd3c9b2b9
---
 mjx/mujoco/mjx/_src/io.py      |  4 ++++
 mjx/mujoco/mjx/_src/io_test.py | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index d5743082..82551c86 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -1210,6 +1210,10 @@ def _get_data_into(
       else:
         setattr(result_i, field.name, value)
 
+    # TODO(taylorhowell): remove mapping once qM is deprecated
+    # map inertia (sparse) to reduced inertia (compressed sparse) representation
+    result_i.M[:] = result_i.qM[result_i.mapM2M]
+
     # recalculate qLD and qLDiagInv as MJX and MuJoCo have different
     # representations of the Cholesky decomposition.
     mujoco.mj_factorM(m, result_i)
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 4bac5e3c..c8670597 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/mjx/mujoco/mjx/_src/io_test.py
@@ -28,6 +28,7 @@ from mujoco.mjx._src import test_util
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import ConeType
 from mujoco.mjx._src.types import Impl
+from mujoco.mjx._src.types import JacobianType
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -656,6 +657,25 @@ class DataIOTest(parameterized.TestCase):
     with self.assertRaises(NotImplementedError):
       mjx.put_model(m, impl='jax')
 
+  @parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
+  def test_qm_mapm2m(self, jacobian):
+    """Test that qM is mapped to M."""
+    m = test_util.load_test_file('humanoid/humanoid.xml')
+    m.opt.jacobian = jacobian
+    d = mujoco.MjData(m)
+    mx = mjx.put_model(m, impl='jax')
+    dx = mjx.make_data(m, impl='jax')
+    dx = mjx.forward(mx, dx)
+
+    mjx.get_data_into(d, m, dx)
+
+    res_mj = np.zeros((1, m.nv))
+    mujoco.mj_solveM(m, d, res_mj, np.ones((1, m.nv)))
+
+    res = mjx._src.smooth.solve_m(mx, dx, jp.ones(m.nv))
+
+    np.testing.assert_allclose(res_mj[0], res, rtol=1e-3, atol=1e-3)
+
 
 class FullCompatTest(parameterized.TestCase):
   """Tests for the _full_compat flag."""