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