Add userdata support to MJX.

PiperOrigin-RevId: 622948369
Change-Id: I5eca0afcb8cd50f53eacc57a78bfdf68ccde12e7

doc/changelog.rst

@@ -36,10 +36,11 @@ MJX
 5. Fixed bug for sphere/capsule-convex deep penetration.
 6. Fixed bug where ``mjx.Data`` produced by ``mjx.put_data`` had different treedef than ``mjx.make_data``.
 7. Throw an error for margin/gap for convex mesh collisions, since they are not supported.
+8. Added support for userdata.
 
 Simulate
 ^^^^^^^^
-8. Fixed bug in order of enable flag strings. Before this change, using the simulate UI to toggle the
+9. Fixed bug in order of enable flag strings. Before this change, using the simulate UI to toggle the
    :ref:`invdiscrete<option-flag-invdiscrete>` or the (now removed) ``sensornoise`` flags would actually toggle the
    other flag.
 
@@ -50,7 +51,7 @@ Python bindings
    :align: right
    :width: 240px
 
-9. Added the ``mujoco.minimize`` Python module for nonlinear least-squares, designed for System Identification (sysID).
+10. Added the ``mujoco.minimize`` Python module for nonlinear least-squares, designed for System Identification (sysID).
    The sysID tutorial is work in progress, but a pedagogical colab notebook with examples, including Inverse
    Kinematics, is available here: |ls_colab|
    |br| The video on the right shows example clips from the tutorial.

mjx/mujoco/mjx/_src/io.py

@@ -219,6 +219,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       qfrc_constraint=zero_nv,
       qfrc_inverse=zero_nv,
       efc_force=zero_nefc,
+      userdata=jp.zeros(m.nuserdata, dtype=float),
   )
 
   return d

mjx/mujoco/mjx/_src/types.py

@@ -310,6 +310,7 @@ class Model(PyTreeNode):
     nexclude: number of excluded geom pairs
     neq: number of equality constraints
     nnumeric: number of numeric custom fields
+    nuserdata: size of userdata array
     nM: number of non-zeros in sparse inertia matrix
     opt: physics options
     stat: model statistics
@@ -454,6 +455,7 @@ class Model(PyTreeNode):
   nexclude: int
   neq: int
   nnumeric: int
+  nuserdata: int
   nM: int  # pylint:disable=invalid-name
   opt: Option
   stat: Statistic
@@ -685,6 +687,7 @@ class Data(PyTreeNode):
     qfrc_inverse: net external force; should equal:               (nv,)
       qfrc_applied + J'*xfrc_applied + qfrc_actuator
     efc_force: constraint force in constraint space               (nefc,)
+    userdata: user data, not touched by engine                    (nuserdata,)
   """
   # solver statistics:
   solver_niter: jax.Array
@@ -703,6 +706,8 @@ class Data(PyTreeNode):
   # dynamics:
   qacc: jax.Array
   act_dot: jax.Array
+  # user data:
+  userdata: jax.Array
   # position dependent:
   xpos: jax.Array
   xquat: jax.Array