get_data is slow with mjx #2718

jstmn · 2025-07-04T02:16:31Z

jstmn
Jul 4, 2025

Intro

Hi!

I am a graduate student at USC, I use MuJoCo for my research on robotic manipulation.

My setup

ubuntu 20.04
cuda 12.6
python 3.10
mujoco                                3.3.2         h89c7902_2                17 KiB     conda  https://conda.anaconda.org/conda-forge/
mujoco-mjx                            3.3.2         pyhd8ed1ab_0              5.3 MiB    conda  https://conda.anaconda.org/conda-forge/

My question

I've noticed that mjx.get_data() is really slow. Here's some timing data:

with n=100 envs, 100 sim steps, no get_data:
# time:          0.12 s
# fps:           849.45
# ms per step:   1.18 ms

With n=100 envs, 100 sim steps, with get_data:
# time:          9.23 s
# fps:           10.84
# ms per step:   92.26 ms

I'll need to retrieve the sim's state after every few steps, so this will be prohibitively expensive. Is there some way to speed this call up?

Is there a jitted version of mjx.get_data?

Minimal model and/or code that explain my question

This code is from the mjx tutorial notebook https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb#scrollTo=rrdrcKRVK6w9

from time import time
import mediapy as media
import jax
from jax import numpy as jp
import mujoco
from mujoco import mjx

xml = """
<mujoco>
  <worldbody>
    <light name="top" pos="0 0 1"/>
    <body name="box_and_sphere" euler="0 0 -30">
      <joint name="swing" type="hinge" axis="1 -1 0" pos="-.2 -.2 -.2"/>
      <geom name="red_box" type="box" size=".2 .2 .2" rgba="1 0 0 1"/>
      <geom name="green_sphere" pos=".2 .2 .2" size=".1" rgba="0 1 0 1"/>
    </body>
  </worldbody>
</mujoco>
"""

scene_option = mujoco.MjvOption()
scene_option.flags[mujoco.mjtVisFlag.mjVIS_JOINT] = True

duration = 0.1  # (seconds)
framerate = 60  # (Hz)

# Make model, data, and renderer
mj_model = mujoco.MjModel.from_xml_string(xml)
mj_data = mujoco.MjData(mj_model)
renderer = mujoco.Renderer(mj_model)
mjx_model = mjx.put_model(mj_model)
mjx_data = mjx.put_data(mj_model, mj_data)

batch_size = 100
rng = jax.random.PRNGKey(0)
rng = jax.random.split(rng, batch_size)
batch = jax.vmap(lambda rng: mjx_data.replace(qpos=jax.random.uniform(rng, (1,))))(rng)
jit_step = jax.jit(jax.vmap(mjx.step, in_axes=(None, 0)))


frames = []
mujoco.mj_resetData(mj_model, mj_data)
mjx_data = mjx.put_data(mj_model, mj_data)
n_steps = 0


print("Taking jit warm up step")
jit_step(mjx_model, batch)
mjx.get_data(mj_model, batch)

print("Starting simulation")
t0 = time()
while n_steps < 100:
    batch = jit_step(mjx_model, batch)
    batched_mj_data = mjx.get_data(mj_model, batch)
    n_steps += 1
    print(f"n_steps: {n_steps}")
runtime = time() - t0

print(f"n_steps:       {n_steps}")
print(f"batch_size:    {batch_size}")
print(f"time:          {runtime:.2f} s")
print(f"fps:           {n_steps / runtime:.2f}")
print(f"ms per step:   {(runtime* 1000) / n_steps :.2f} ms")

Confirmations

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

jstmn · 2025-08-15T19:23:45Z

jstmn
Aug 15, 2025
Author

erikfrey · 2025-08-15T19:32:43Z

erikfrey
Aug 15, 2025
Maintainer

Hi @jstmn - can you say a bit more about why you want to transfer mjx.Data back to host every few steps? This is less a problem with MJX and more of an architectural issue - transfer between CPU and GPU is slow unless you have specialized compute hardware like Apple's unified memory.

You have two options:

Only fetch the arrays that you actually need back onto the host, rather than the entire mjx.Data. This might help things a bit.
Rethink your approach such that whatever computation you are doing host-side can be done on the GPU.

0 replies

jstmn · 2025-08-27T23:53:36Z

jstmn
Aug 27, 2025
Author

Hey, sorry for the delay.

can you say a bit more about why you want to transfer mjx.Data back to host every few steps?

Yep - I'll be using mjx in a sampling based MPC application. Ballpark numbers, I need to have 100 parallel environments each taking ~16 simulation steps. This will be happening as fast as possible to control a robot. Each of the resulting 100*16 simulation states will be evaluated by a reward function r: s -> R^1 to determine the sequence of actions that led to the best rollout. Small note: I may instead just evaluate the final state of each rollout, meaning they'll only be 100 evaluations. But this should still run entirely on the gpu.

A more representative task will a Panda pick-cube task like this one: https://github.com/google-deepmind/mujoco_menagerie/blob/main/franka_emika_panda/mjx_single_cube.xml

Only fetch the arrays that you actually need back onto the host, rather than the entire mjx.Data. This might help things a bit.

Getting only the robots joint angles and the pose of the objects would work, but option 2 sounds more appropriate

Rethink your approach such that whatever computation you are doing host-side can be done on the GPU.

This would work for me! I'm new to mjx / mujoco though. Would you mind showing me how for example I can get a batch of joint angles of the robot?

Thanks

0 replies

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

get_data is slow with mjx #2718

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Replies: 3 comments

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{title}}

Uh oh!

Select a reply

Uh oh!

get_data is slow with mjx #2718

Uh oh!

Uh oh!

jstmn Jul 4, 2025

Intro

My setup

My question

Minimal model and/or code that explain my question

Confirmations

Replies: 3 comments

Uh oh!

jstmn Aug 15, 2025 Author

Uh oh!

erikfrey Aug 15, 2025 Maintainer

Uh oh!

jstmn Aug 27, 2025 Author

jstmn
Jul 4, 2025

jstmn
Aug 15, 2025
Author

erikfrey
Aug 15, 2025
Maintainer

jstmn
Aug 27, 2025
Author