Real-Time Requirements & Broken MjxState

[Nickick-ICRS]

Hi,

First off, great paper, though it took me 6 hours to get this setup due to various issues with cuda and jax (please lock the version of jax and brax in the requirements!)

1. Why did I have to change line 60 and 61 in dial_plan.py?
   From

    data = _reformat_contact(sys, data.data)
    return MjxState(q=q, qd=qd, x=x, xd=xd, **data.__dict__)

To

    contact = _reformat_contact(sys, data.contact)
    return MjxState(q=q, qd=qd, x=x, xd=xd, contact=contact, **data.__dict__)

2. In your paper, you claim a 50hz running frequency for the go2 trotting simulation. What GPU/other machine stats are you using for this? With a standard laptop GPU (3060, 4gb VRAM) my machine cannot even converge in sim, with average delays of 100ms. I guess that's 10hz?

Cheers

[tc2882]

So what is the version of jax and brax should I use? I got the following error:

🚀Creating environment
/home/linux/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/brax/io/mjcf.py:480: UserWarning: Brax System, piplines and environments are not actively being maintained. Please see MJX for a well maintained JAX-based physics engine: https://github.com/google-deepmind/mujoco/tree/main/mjx. For a host of environments that use MJX, see: https://github.com/google-deepmind/mujoco_playground.
  warnings.warn(
[0.5904899  0.6560999  0.7289999  0.80999994 0.9        1.        ]
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/linux/anaconda3/envs/dial-mpc/bin/dial-mpc-plan", line 33, in <module>
    sys.exit(load_entry_point('dial-mpc', 'console_scripts', 'dial-mpc-plan')())
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 288, in main
    mbd_publisher.main_loop()
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 165, in main_loop
    state = self.init_mjx_state(
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 143, in init_mjx_state
    pipeline_state = self.pipeline_init_jit(self.env.sys, q, qd)
  File "/home/linux/dial-mpc/dial_mpc/deploy/dial_plan.py", line 60, in pipeline_init
    data = _reformat_contact(sys, data)
  File "/home/linux/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/brax/mjx/pipeline.py", line 32, in _reformat_contact
    elasticity = jp.zeros(contact.pos.shape[0])
AttributeError: 'Data' object has no attribute 'pos'. Did you mean: 'qpos'?
/home/linux/anaconda3/envs/dial-mpc/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 6 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

Thank you very much!

[Singularhdp]

I have the same problem with you, could you please tell me how to solve it?

[tc2882]

@Singularhdp This setting works for me.

    install_requires=[
        "numpy<2.0.0",
        "matplotlib",
        "tqdm",
        "tyro",
        "jax[cuda12]==0.4.33",
        "jax-cosmo",
        "mujoco<=3.2.4",
        "brax<=0.11",
        "art",
        "emoji",
        "scienceplots",
    ]

[Singularhdp]

@tc2882 Thank you so much for helping me!