Evaluating a jacobian of a mass matrix (manipulator dynamics) with Jax and MJX

[danielpmorton]

Hi, I am a graduate student interested in taking derivatives of the manipulator dynamics with Jax and MJX. I'm new to Mujoco but have worked extensively with other simulators, and have used Jax in a limited capacity.

Ideally, I'd like to use automatic differentiation with the robot dynamics (for instance, compiling a function for the derivative of the mass matrix, with respect to the joint angles).

I've done a preliminary test (below) and I can access the qM attribute of the mjx_step Data instance, but I am unsure where the function to compute this is located, so that I can pass this into jax.jacrev or jax.jacfwd. Currently the qM I get is just a Jax Array

So to summarize, my main question is:
Where can I find a MJX function for the robot mass matrix, and how can I pass this into Jax to evaluate a derivative?

import mujoco
from mujoco import mjx

# The mujoco menagerie model required commenting out the "implicitfast" integrator
mj_model = mujoco.MjModel.from_xml_path(
    "mujoco_menagerie/universal_robots_ur5e/ur5e.xml"
)
mj_data = mujoco.MjData(mj_model)

mjx_model = mjx.put_model(mj_model)
mjx_data = mjx.make_data(mjx_model)
qM = mjx.step(mjx_model, mjx_data).qM
print("qM from MJX step return: \n", qM)

Returns: A Jax Array with numeric values

[yuvaltassa]

here it is. I can't help you with the second part, since JAX is not my strength, I just searched the MJX source for "qM", it wasn't hard to find.