MJX Support for Linearized Dynamics Estimation

[jc-bao]

Intro

Hi!

I am a graduate student at CMU, I use MuJoCo for my research on optimal control.

My setup

I am using MJX.

My question

I noticed that using JAX autograd with MJX dynamics is relatively slow. Besides, to make sure the Jacobian is calculated without JAX error, I had to set the iteration number to 1, which avoiding the use of while_loop for MuJoCo inner solver.

I am wondering if MJX provides a similar function to mjd_transitionFD for fast linearized dynamic estimation. If not, is it possible to implement such functionality with mj_jac?

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.

[yuvaltassa]

Why would you want to use MJX for this?

[jc-bao]

I tend to run multiple iLQR instance or use randomized gradient estimators.

[yuvaltassa]

Generally speaking autodiff is not what you want for iLQR, since that's a serial operation. What you want it to calculate the Jacobians in parallel and the do the backpass (which has to be serial), so I'm surprised you'd want MJX.

I take it you're familiar with the MJPC codebase?

[jc-bao]

Thanks for the response. I used MJPC a lot and what I want is to get A, B matrix from MJX. But when I try to run jax.jac(mjx.step), it would through out an error since there is while loop in MJX. Also, the jacobian calculation is super slow (mainly the jit time). I am wondering what is the recommended why to linearize dynamics in MJX. Or MJX is not designed to do so.