Shape mismatch in MJX's put_data function

[rustam-e]

Intro

Hi!

I am a masters student at Vrije Universiteit Amsterdam, I use MuJoCo for my research on Modular Robotics.

My setup

Currently, Mujoco - trying to convert existing Evolutionary Computing for modular robotics framework Revolve2 (https://github.com/ci-group/revolve2) from Mujoco to MJX 3.2.3

What's happening? What did you expect?

I am following the provided tutorial on https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb#scrollTo=Jtz7j1PDOnw5 and convert currently working code running the mujoco.mj_step to run using MJX.

I get an error related efc_J:

Steps for reproduction

1. checkout https://github.com/rustam-e/revolve2/tree/develop
2. dependency installation - https://ci-group.github.io/revolve2/installation/index.html
3. run the experiment with python examples/4_example_experiment_setups/4c_robot_bodybrain_ea-mjx/main.py

the only differences from a functional examples/4_example_experiment_setups/4c_robot_bodybrain_ea/main.py can be found in rustam-e/revolve2@30bfc7f

Minimal model for reproduction

the issue is consistent with all of the experiements we ran

Code required for reproduction

running any of the examples by changing
from:
from revolve2.simulators.mujoco_simulator import LocalSimulator

to:
from revolve2.simulators.mjx_simulator import LocalSimulator

would throw consistently the same error like:
ValueError: could not broadcast input array from shape (0,10) into shape (4,10)

the only diference across examples being the second dimension

Confirmations

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

[rustam-e]

The issue seems to be due to the fact the the Mujoco did not instantiate efc_J property of the model, but running the mj_step once before the put_data call seems to have resolved the issue

[yuvaltassa]

It's good that you found the solution, but put_data should either handle this situation (being called before any arena allocations) gracefully or, if that's not possible, give the user an informative error.

That, I think, is a bug. WDYT @erikfrey ?