why is MJX pipeline_step so slow

[gautica]

Intro

Hi!

I am a researcher in legged robot locomotion.

My setup

MJX 3.2.3

My question

I am very interested in applying RL methods with MJX, but I tried the Unitree H1 robot and found that it takes for a single environment about 3-5 milliseconds to run one step pipeline_step function (on a single GPU NVIDIA RTX A4000) and 0.5 seconds for 500 environments, which is much slower than the reported performance in the documentation. In my H1 model, the contact is simplified to be only at the feet. Is it normal that it is so slow? Thank you for any help.

Minimal model and/or code that explain my question

If you encountered the issue in a complex model, please simplify it as much as possible (while still reproducing the issue).

Model:

minimal XML

<mujoco>
  <worldbody>
    <light pos="0 0 1"/>
    <geom type="sphere" size="1" rgba="1 0 0 1"/>
  </worldbody>
</mujoco>

Code:

import mujoco
import mediapy as media
model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)

with mujoco.Renderer(model) as renderer:
  mujoco.mj_forward(model, data)
  renderer.update_scene(data)

  media.show_image(renderer.render())

Confirmations

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