From b637d7eb761708872afe9945dee231437bad67d3 Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Mon, 9 Feb 2026 10:28:05 -0800
Subject: [PATCH] Call kinematics in case of state mismatch in `get_data_into`
 with CPP backend.

PiperOrigin-RevId: 867667583
Change-Id: Ieec5ff86f6768c11459795068abf17f69bd26f3c
---
 mjx/mujoco/mjx/_src/io.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index ba8c91c4..67038a63 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -1660,18 +1660,24 @@ def _get_data_into_cpp(
   # mjx.Model which we don't have access to in this function.
   fields_to_check = ['qpos', 'qvel', 'act', 'mocap_pos', 'mocap_quat']
   for i in range(batch_size):
-    d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
+    d_i: types.Data = jax.tree_util.tree_map(
+        lambda x, i=i: x[i], d) if batched else d
     src_data = mj_data_list[i]
 
+    needs_syncing = False
     for field in fields_to_check:
       mj_value = getattr(src_data, field)
       mjx_value = np.asarray(getattr(d_i, field))
       if not np.allclose(mj_value, mjx_value):
-        raise ValueError(
-            f'State mismatch at batch index {i}, field {field}: underlying '
-            'MjData does not match mjx.Data. The mjx.Data may have been '
-            'modified without updating the underlying MjData.'
-        )
+        needs_syncing = True
+        break
+    if needs_syncing:
+      src_data.qpos[:] = d_i.qpos
+      src_data.qvel[:] = d_i.qvel
+      src_data.act[:] = d_i.act
+      src_data.mocap_pos[:] = d_i.mocap_pos
+      src_data.mocap_quat[:] = d_i.mocap_quat
+      mujoco.mj_kinematics(m, src_data)
 
   for i in range(batch_size):
     result_i = result[i] if batched else result