Shape mismatch when calling "mjx.get_data" #2905

BissonnetteVincent · 2025-10-17T21:18:30Z

BissonnetteVincent
Oct 17, 2025

Intro

Hi, I'm currently working on a project where I try to use a physics engine to validate some physical configurations. I though I could use MJX to try different actions on a physical env and validate those in parallel

My setup

I'm using Mujoco 3.3.6
I'm also using JAX_cuda_13
I'm using WSL which is Ubuntu 22.04.05 LTS

My question

When I try to use "mjx.get_data" to retrieve the data from my device I'm getting an error regarding a shape mismatch for qLD. Im not sure what could cause this issue. As you can see the qPos are printed when iterating over the batch but as soon as I try to retrieve the data the error occurs

Something interesting is that if I remove body named "box_10" the error no longer occurs.

Model:

minimal XML

<?xml version="1.0" encoding="UTF-8"?>
<mujoco model="json_model">
	<compiler angle="degree"/>
	<option gravity="0 0 -9.81"/>
	<worldbody>
		<geom name="ground" type="plane" pos="0 0 0" size="5 5 0.1" rgba="0.8 0.8 0.8 1"/>
    <body name="pallet" pos="0 0 86">
      <geom type="box" size="605 508 86" mass="20" rgba="0.55 0.27 0.07 1" friction="0.8 0.5 0.5" solref="0.02 2" solimp="0.95 0.99 0.01"/>
      <joint type="free"/>
    </body>
    <body name="box_0" pos="519 350 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_1" pos="519 0 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_2" pos="519 -350 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_3" pos="311 350 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_4" pos="311 0 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_5" pos="311 -350 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
		<body name="box_9" pos="-103 350 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_10" pos="-103 0 257">
      <geom type="box" size="170 103 85" mass="16.3" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint type="free"/>
    </body>
    <body name="box_22" pos="-155 266 384">
      <geom type="box" size="254 155 42" mass="26" friction="0.8 0.1 0.1" solref="0.02 2" solimp="0.95 0.99 0.01" rgba="0.7 0.3 0.3 1" quat="0.7071 0 0 0.7071"/>
      <joint name="box_22_joint" type="free"/>
    </body>
	</worldbody>
	<visual>
		<map znear="0.01" zfar="10"/>
	</visual>
</mujoco>

Code:

import jax
import mujoco
import mujoco.mjx as mjx

mj_model = mujoco.MjModel.from_xml_string(XML)
mj_data = mujoco.MjData(mj_model)
mjx_model = mjx.put_model(mj_model)
mjx_data = mjx.put_data(mj_model, mj_data)

joint_id = mujoco.mj_name2id(mj_model, mujoco.mjtObj.mjOBJ_JOINT, "box_22_joint")
joint_addr = mj_model.jnt_qposadr[joint_id]

qPos = jax.numpy.arange(1.0, 10.0, 1)
batch = jax.vmap(lambda newPos: mjx_data.replace(qpos=mjx_data.qpos.at[joint_addr].set(newPos)))(qPos)

jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))
batch = jit_step(mjx_model, batch)

for idx, qpos in enumerate(batch.qpos):
     print(f"======== qPos {idx} ========")
     print(qpos[joint_addr])

mj_data = mjx.get_data(mj_model, batch)

Confirmations

I searched the latest documentation thoroughly before posting.
I searched previous Issues and Discussions, I am certain this has not been raised before.

yuvaltassa · 2025-11-17T13:20:31Z

yuvaltassa
Nov 17, 2025
Maintainer

I this still an issue with 3.3.7?

1 reply

BissonnetteVincent Nov 17, 2025
Author

I'll need to test and get back to you

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Shape mismatch when calling "mjx.get_data" #2905

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Confirmations

Replies: 1 comment 1 reply

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

Shape mismatch when calling "mjx.get_data" #2905

Uh oh!

Uh oh!

BissonnetteVincent Oct 17, 2025

Intro

My setup

My question

Confirmations

Replies: 1 comment · 1 reply

Uh oh!

yuvaltassa Nov 17, 2025 Maintainer

Uh oh!

BissonnetteVincent Nov 17, 2025 Author

BissonnetteVincent
Oct 17, 2025

Replies: 1 comment 1 reply

yuvaltassa
Nov 17, 2025
Maintainer

BissonnetteVincent Nov 17, 2025
Author