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

@@ -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