From dbe18f574006adbc410f1c0551a53a24e6156c6e Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Wed, 17 Jul 2024 10:03:53 -0700
Subject: [PATCH] Add additional MuJoCo fields to MJX structs.

PiperOrigin-RevId: 653270676
Change-Id: I5f7a62d3bd542abc0d618e63d2426663e7d6406f
---
 doc/changelog.rst             |   8 +
 mjx/mujoco/mjx/_src/io.py     | 190 ++++++++-----
 mjx/mujoco/mjx/_src/smooth.py |   3 +
 mjx/mujoco/mjx/_src/types.py  | 510 +++++++++++++++++++++++++++++-----
 4 files changed, 569 insertions(+), 142 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index 8b2080c2..6173007f 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+MJX
+^^^
+
+1. Added more fields to ``mjx.Model`` and ``mjx.Data`` for further compatibility with the corresponding MuJoCo structs.
+
 Version 3.2.0 (Jul 15, 2024)
 ----------------------------
 
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index efa1da86..1cdb015f 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -15,7 +15,7 @@
 """Functions to initialize, load, or save data."""
 
 import copy
-from typing import List, Tuple, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import jax
 from jax import numpy as jp
@@ -110,8 +110,6 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
   mj_field_names = {f.name for f in types.Model.fields()} - mjx_only
   fields = {f: getattr(m, f) for f in mj_field_names}
   fields['geom_rbound_hfield'] = fields['geom_rbound']
-  fields['geom_rgba'] = fields['geom_rgba'].reshape((-1, 4))
-  fields['mat_rgba'] = fields['mat_rgba'].reshape((-1, 4))
   fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
   fields['opt'] = _make_option(m.opt)
   fields['stat'] = _make_statistic(m.stat)
@@ -137,33 +135,19 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   ne, nf, nl, nc = constraint.counts(efc_type)
   ncon, nefc = dim.size, ne + nf + nl + nc
 
-  zero_0 = jp.zeros(0, dtype=float)
-  zero_nv = jp.zeros(m.nv, dtype=float)
-  zero_nv_6 = jp.zeros((m.nv, 6), dtype=float)
-  zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=float)
-  zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=float)
-  zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=float)
-  zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=float)
-  zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=float)
-  zero_nefc = jp.zeros(nefc, dtype=float)
-  zero_na = jp.zeros(m.na, dtype=float)
-  zero_nu = jp.zeros(m.nu, dtype=float)
-  zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=float)
-  zero_nm = jp.zeros(m.nM, dtype=float)
-
   contact = types.Contact(
-      dist=jp.zeros(ncon),
-      pos=jp.zeros((ncon, 3)),
-      frame=jp.zeros((ncon, 3, 3)),
-      includemargin=jp.zeros(ncon),
-      friction=jp.zeros((ncon, 5)),
-      solref=jp.zeros((ncon, mujoco.mjNREF)),
-      solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
-      solimp=jp.zeros((ncon, mujoco.mjNIMP)),
+      dist=jp.zeros((ncon,), dtype=float),
+      pos=jp.zeros((ncon, 3), dtype=float),
+      frame=jp.zeros((ncon, 3, 3), dtype=float),
+      includemargin=jp.zeros((ncon,), dtype=float),
+      friction=jp.zeros((ncon, 5), dtype=float),
+      solref=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
+      solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float),
+      solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float),
       dim=dim,
-      geom1=jp.zeros(ncon, dtype=int) - 1,
-      geom2=jp.zeros(ncon, dtype=int) - 1,
-      geom=jp.zeros((ncon, 2), dtype=int) - 1,
+      geom1=jp.full((ncon,), -1, dtype=int),
+      geom2=jp.full((ncon,), -1, dtype=int),
+      geom=jp.full((ncon, 2), -1, dtype=int),
       efc_address=efc_address,
   )
 
@@ -173,59 +157,118 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       nl=nl,
       nefc=nefc,
       ncon=ncon,
-      solver_niter=jp.array(0, dtype=int),
-      time=jp.array(0.0, dtype=float),
+      solver_niter=jp.zeros((), dtype=int),
+      time=jp.zeros((), dtype=float),
       qpos=jp.array(m.qpos0),
-      qvel=zero_nv,
-      act=zero_na,
-      qacc_warmstart=zero_nv,
-      ctrl=zero_nu,
-      qfrc_applied=zero_nv,
-      xfrc_applied=zero_nbody_6,
-      eq_active=jp.zeros(m.neq, dtype=jp.uint8),
-      qacc=zero_nv,
-      act_dot=zero_na,
-      xpos=zero_nbody_3,
+      qvel=jp.zeros((m.nv,), dtype=float),
+      act=jp.zeros((m.na,), dtype=float),
+      qacc_warmstart=jp.zeros((m.nv,), dtype=float),
+      ctrl=jp.zeros((m.nu,), dtype=float),
+      qfrc_applied=jp.zeros((m.nv,), dtype=float),
+      xfrc_applied=jp.zeros((m.nbody, 6), dtype=float),
+      eq_active=jp.zeros((m.neq,), dtype=jp.uint8),
+      mocap_pos=jp.zeros((m.nmocap, 3), dtype=float),
+      mocap_quat=jp.zeros((m.nmocap, 4), dtype=float),
+      qacc=jp.zeros((m.nv,), dtype=float),
+      act_dot=jp.zeros((m.na,), dtype=float),
+      userdata=jp.zeros((m.nuserdata,), dtype=float),
+      sensordata=jp.zeros((m.nsensordata,), dtype=float),
+      xpos=jp.zeros((m.nbody, 3), dtype=float),
       xquat=jp.zeros((m.nbody, 4), dtype=float),
-      xmat=zero_nbody_3_3,
-      xipos=zero_nbody_3,
-      ximat=zero_nbody_3_3,
-      xanchor=zero_njnt_3,
-      xaxis=zero_njnt_3,
+      xmat=jp.zeros((m.nbody, 3, 3), dtype=float),
+      xipos=jp.zeros((m.nbody, 3), dtype=float),
+      ximat=jp.zeros((m.nbody, 3, 3), dtype=float),
+      xanchor=jp.zeros((m.njnt, 3), dtype=float),
+      xaxis=jp.zeros((m.njnt, 3), dtype=float),
       geom_xpos=jp.zeros((m.ngeom, 3), dtype=float),
       geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=float),
       site_xpos=jp.zeros((m.nsite, 3), dtype=float),
       site_xmat=jp.zeros((m.nsite, 3, 3), dtype=float),
       cam_xpos=jp.zeros((m.ncam, 3), dtype=float),
       cam_xmat=jp.zeros((m.ncam, 3, 3), dtype=float),
-      subtree_com=zero_nbody_3,
-      cdof=zero_nv_6,
-      cinert=zero_nbody_10,
-      actuator_length=zero_nu,
+      light_xpos=jp.zeros((m.nlight, 3), dtype=float),
+      light_xdir=jp.zeros((m.nlight, 3), dtype=float),
+      subtree_com=jp.zeros((m.nbody, 3), dtype=float),
+      cdof=jp.zeros((m.nv, 6), dtype=float),
+      cinert=jp.zeros((m.nbody, 10), dtype=float),
+      flexvert_xpos=jp.zeros((m.nflexvert, 3), dtype=float),
+      flexelem_aabb=jp.zeros((m.nflexelem, 6), dtype=float),
+      flexedge_J_rownnz=jp.zeros((m.nflexedge,), dtype=int),
+      flexedge_J_rowadr=jp.zeros((m.nflexedge,), dtype=int),
+      flexedge_J_colind=jp.zeros((m.nflexedge, m.nv), dtype=int),
+      flexedge_J=jp.zeros((m.nflexedge, m.nv), dtype=float),
+      flexedge_length=jp.zeros((m.nflexedge,), dtype=float),
+      ten_wrapadr=jp.zeros((m.ntendon,), dtype=int),
+      ten_wrapnum=jp.zeros((m.ntendon,), dtype=int),
+      ten_J_rownnz=jp.zeros((m.ntendon,), dtype=int),
+      ten_J_rowadr=jp.zeros((m.ntendon,), dtype=int),
+      ten_J_colind=jp.zeros((m.ntendon, m.nv), dtype=int),
+      ten_J=jp.zeros((m.ntendon, m.nv), dtype=float),
+      ten_length=jp.zeros((m.ntendon,), dtype=float),
+      wrap_obj=jp.zeros((m.nwrap, 2), dtype=int),
+      wrap_xpos=jp.zeros((m.nwrap, 6), dtype=float),
+      actuator_length=jp.zeros((m.nu,), dtype=float),
       actuator_moment=jp.zeros((m.nu, m.nv), dtype=float),
-      crb=zero_nbody_10,
-      qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
-      qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
-      qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
+      crb=jp.zeros((m.nbody, 10), dtype=float),
+      qM=(
+          jp.zeros((m.nM,), dtype=float)
+          if support.is_sparse(m)
+          else jp.zeros((m.nv, m.nv), dtype=float)
+      ),
+      qLD=(
+          jp.zeros((m.nM,), dtype=float)
+          if support.is_sparse(m)
+          else jp.zeros((m.nv, m.nv), dtype=float)
+      ),
+      qLDiagInv=(
+          jp.zeros((m.nv,), dtype=float) if support.is_sparse(m)
+          else jp.zeros((0,), dtype=float)
+      ),
+      qLDiagSqrtInv=jp.zeros((m.nv,), dtype=float),
+      bvh_aabb_dyn=jp.zeros((m.nbvhdynamic, 6), dtype=float),
+      bvh_active=jp.zeros((m.nbvh,), dtype=jp.uint8),
+      flexedge_velocity=jp.zeros((m.nflexedge,), dtype=float),
+      ten_velocity=jp.zeros((m.ntendon,), dtype=float),
+      actuator_velocity=jp.zeros((m.nu,), dtype=float),
+      cvel=jp.zeros((m.nbody, 6), dtype=float),
+      cdof_dot=jp.zeros((m.nv, 6), dtype=float),
+      qfrc_bias=jp.zeros((m.nv,), dtype=float),
+      qfrc_spring=jp.zeros((m.nv,), dtype=float),
+      qfrc_damper=jp.zeros((m.nv,), dtype=float),
+      qfrc_gravcomp=jp.zeros((m.nv,), dtype=float),
+      qfrc_fluid=jp.zeros((m.nv,), dtype=float),
+      qfrc_passive=jp.zeros((m.nv,), dtype=float),
+      subtree_linvel=jp.zeros((m.nbody, 3), dtype=float),
+      subtree_angmom=jp.zeros((m.nbody, 3), dtype=float),
+      qH=jp.zeros((m.nM,), dtype=float),
+      qHDiagInv=jp.zeros((m.nv,), dtype=float),
+      D_rownnz=jp.zeros((m.nv,), dtype=int),
+      D_rowadr=jp.zeros((m.nv,), dtype=int),
+      D_colind=jp.zeros((m.nD,), dtype=int),
+      B_rownnz=jp.zeros((m.nbody,), dtype=int),
+      B_rowadr=jp.zeros((m.nbody,), dtype=int),
+      B_colind=jp.zeros((m.nB,), dtype=int),
+      qDeriv=jp.zeros((m.nD,), dtype=float),
+      qLU=jp.zeros((m.nD,), dtype=float),
+      actuator_force=jp.zeros((m.nu,), dtype=float),
+      qfrc_actuator=jp.zeros((m.nv,), dtype=float),
+      qfrc_smooth=jp.zeros((m.nv,), dtype=float),
+      qacc_smooth=jp.zeros((m.nv,), dtype=float),
+      qfrc_constraint=jp.zeros((m.nv,), dtype=float),
+      qfrc_inverse=jp.zeros((m.nv,), dtype=float),
+      cacc=jp.zeros((m.nbody, 6), dtype=float),
+      cfrc_int=jp.zeros((m.nbody, 6), dtype=float),
+      cfrc_ext=jp.zeros((m.nbody, 6), dtype=float),
       contact=contact,
       efc_type=efc_type,
       efc_J=jp.zeros((nefc, m.nv), dtype=float),
-      efc_frictionloss=zero_nefc,
-      efc_D=zero_nefc,
-      actuator_velocity=zero_nu,
-      cvel=zero_nbody_6,
-      cdof_dot=zero_nv_6,
-      qfrc_bias=zero_nv,
-      qfrc_gravcomp=zero_nv,
-      qfrc_passive=zero_nv,
-      efc_aref=zero_nefc,
-      qfrc_actuator=zero_nv,
-      qfrc_smooth=zero_nv,
-      qacc_smooth=zero_nv,
-      qfrc_constraint=zero_nv,
-      qfrc_inverse=zero_nv,
-      efc_force=zero_nefc,
-      userdata=jp.zeros(m.nuserdata, dtype=float),
+      efc_frictionloss=jp.zeros((nefc,), dtype=float),
+      efc_D=jp.zeros((nefc,), dtype=float),
+      efc_aref=jp.zeros((nefc,), dtype=float),
+      efc_force=jp.zeros((nefc,), dtype=float),
+      _qM_sparse=jp.zeros((m.nM), dtype=float),
+      _qLD_sparse=jp.zeros((m.nM), dtype=float),
+      _qLDiagInv_sparse=jp.zeros((m.nv,), dtype=float),
   )
 
   return d
@@ -296,6 +339,9 @@ def get_data_into(
     result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
 
     for field in types.Data.fields():
+      if field.name.startswith('_') and field.name.endswith('_sparse'):
+        continue
+
       if field.name == 'contact':
         _get_contact(result_i.contact, d_i.contact)
         # efc_address must be updated because rows were deleted above:
@@ -379,7 +425,8 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
     if d_val > val:
       raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
 
-  fields = {f.name: getattr(d, f.name) for f in types.Data.fields()}
+  fields = {f.name: getattr(d, f.name) for f in types.Data.fields()
+            if not f.name.endswith('_sparse')}
 
   # MJX prefers square matrices for these fields:
   for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'):
@@ -427,6 +474,9 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
     fields[fname] = value
 
   # convert qM and qLD if jacobian is dense
+  fields['_qM_sparse'] = fields['qM']
+  fields['_qLD_sparse'] = fields['qLD']
+  fields['_qLDiagInv_sparse'] = fields['qLDiagInv']
   if not support.is_sparse(m):
     fields['qM'] = np.zeros((m.nv, m.nv))
     mujoco.mj_fullM(m, fields['qM'], d.qM)
diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
index a3fe0019..7f105199 100644
--- a/mjx/mujoco/mjx/_src/smooth.py
+++ b/mjx/mujoco/mjx/_src/smooth.py
@@ -282,6 +282,8 @@ def crb(m: Model, d: Data) -> Data:
   crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
   qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
   d = d.replace(qM=qm)
+  if support.is_sparse(m):
+    d = d.replace(_qM_sparse=qm)
 
   return d
 
@@ -340,6 +342,7 @@ def factor_m(m: Model, d: Data) -> Data:
   qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
 
   d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag)
+  d = d.replace(_qLD_sparse=d.qLD, _qLDiagInv_sparse=d.qLDiagInv)
 
   return d
 
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index c85701d4..f20f1cbb 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -334,25 +334,49 @@ class Model(PyTreeNode):
     nu: number of actuators/controls = dim(ctrl)
     na: number of activation states = dim(act)
     nbody: number of bodies
+    nbvh: number of total bounding volumes in all bodies
+    nbvhstatic: number of static bounding volumes (aabb stored in mjModel)
+    nbvhdynamic: number of dynamic bounding volumes (aabb stored in mjData)
     njnt: number of joints
     ngeom: number of geoms
     nsite: number of sites
     ncam: number of cameras
+    nlight: number of lights
+    nflex: number of flexes
+    nflexvert: number of vertices in all flexes
+    nflexedge: number of edges in all flexes
+    nflexelem: number of elements in all flexes
+    nflexelemdata: number of element vertex ids in all flexes
+    nflexshelldata: number of shell fragment vertex ids in all flexes
+    nflexevpair: number of element-vertex pairs in all flexes
+    nflextexcoord: number of vertices with texture coordinates
     nmesh: number of meshes
     nmeshvert: number of vertices in all meshes
+    nmeshnormal: number of normals in all meshes
+    nmeshtexcoord: number of texcoords in all meshes
     nmeshface: number of triangular faces in all meshes
+    nmeshgraph: number of ints in mesh auxiliary data
     nhfield: number of heightfields
+    nhfielddata: number of data points in all heightfields
     nmat: number of materials
     npair: number of predefined geom pairs
     nexclude: number of excluded geom pairs
     neq: number of equality constraints
-    ngravcomp: number of bodies with nonzero gravcomp
+    ntendon: number of tendons
+    nwrap: number of wrap objects in all tendon paths
+    nsensor: number of sensors
     nnumeric: number of numeric custom fields
     ntuple: number of tuple custom fields
-    nsensor: number of sensors
     nkey: number of keyframes
-    nuserdata: size of userdata array
+    nmocap: number of mocap bodies
     nM: number of non-zeros in sparse inertia matrix
+    nD: number of non-zeros in sparse dof-dof matrix
+    nB: number of non-zeros in sparse body-dof matrix
+    ntree: number of kinematic trees under world body
+    ngravcomp: number of bodies with nonzero gravcomp
+    nuserdata: size of userdata array
+    nsensordata: number of mjtNums in sensor data vector
+    narena: number of bytes in the mjData arena (inclusive of stack)
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose                        (nq,)
@@ -364,8 +388,10 @@ class Model(PyTreeNode):
     body_jntadr: start addr of joints; -1: no joints          (nbody,)
     body_dofnum: number of motion degrees of freedom          (nbody,)
     body_dofadr: start addr of dofs; -1: no dofs              (nbody,)
+    body_treeid: id of body's kinematic tree; -1: static      (nbody,)
     body_geomnum: number of geoms                             (nbody,)
     body_geomadr: start addr of geoms; -1: no geoms           (nbody,)
+    body_simple: 1: diag M; 2: diag M, sliders only           (nbody,)
     body_pos: position offset rel. to parent body             (nbody, 3)
     body_quat: orientation offset rel. to parent body         (nbody, 4)
     body_ipos: local position of center of mass               (nbody, 3)
@@ -374,6 +400,14 @@ class Model(PyTreeNode):
     body_subtreemass: mass of subtree starting at this body   (nbody,)
     body_inertia: diagonal inertia in ipos/iquat frame        (nbody, 3)
     body_gravcomp: antigravity force, units of body weight    (nbody,)
+    body_margin: MAX over all geom margins                    (nbody,)
+    body_contype: OR over all geom contypes                   (nbody,)
+    body_conaffinity: OR over all geom conaffinities          (nbody,)
+    body_bvhadr: address of bvh root                          (nbody,)
+    body_bvhnum: number of bounding volumes                   (nbody,)
+    bvh_child: left and right children in tree                (nbvh, 2)
+    bvh_nodeid: geom or elem id of node; -1: non-leaf         (nbvh,)
+    bvh_aabb: local bounding box (center, size)               (nbvhstatic, 6)
     body_invweight0: mean inv inert in qpos0 (trn, rot)       (nbody, 2)
     jnt_type: type of joint (mjtJoint)                        (njnt,)
     jnt_qposadr: start addr in 'qpos' for joint's data        (njnt,)
@@ -394,7 +428,9 @@ class Model(PyTreeNode):
     dof_bodyid: id of dof's body                              (nv,)
     dof_jntid: id of dof's joint                              (nv,)
     dof_parentid: id of dof's parent; -1: none                (nv,)
+    dof_treeid: id of dof's kinematic tree                    (nv,)
     dof_Madr: dof address in M-diagonal                       (nv,)
+    dof_simplenum: number of consecutive simple dofs          (nv,)
     dof_solref: constraint solver reference:frictionloss      (nv, mjNREF)
     dof_solimp: constraint solver impedance:frictionloss      (nv, mjNIMP)
     dof_frictionloss: dof friction loss                       (nv,)
@@ -415,6 +451,7 @@ class Model(PyTreeNode):
     geom_solref: constraint solver reference: contact         (ngeom, mjNREF)
     geom_solimp: constraint solver impedance: contact         (ngeom, mjNIMP)
     geom_size: geom-specific size parameters                  (ngeom, 3)
+    geom_aabb: bounding box, (center, size)                   (ngeom, 6)
     geom_rbound: radius of bounding sphere                    (ngeom,)
     geom_rbound_hfield: static rbound for hfield grid bounds  (ngeom,)
     geom_pos: local position offset rel. to body              (ngeom, 3)
@@ -432,14 +469,71 @@ class Model(PyTreeNode):
     cam_pos:  position rel. to body frame                     (ncam, 3)
     cam_quat:  orientation rel. to body frame                 (ncam, 4)
     cam_poscom0:  global position rel. to sub-com in qpos0    (ncam, 3)
-    cam_pos0:  global position rel. to body in qpos0          (ncam, 3)
-    cam_mat0:  global orientation in qpos0                    (ncam, 9)
+    cam_pos0: global position rel. to body in qpos0           (ncam, 3)
+    cam_mat0: global orientation in qpos0                     (ncam, 3, 3)
+    cam_fovy: y field-of-view                                 (ncam,)
+    cam_resolution: resolution: pixels                        (ncam, 2)
+    cam_sensorsize: sensor size: length                       (ncam, 2)
+    cam_intrinsic: [focal length; principal point]            (ncam, 4)
+    light_mode: light tracking mode (mjtCamLight)             (nlight,)
+    light_bodyid: id of light's body                          (nlight,)
+    light_targetbodyid: id of targeted body; -1: none         (nlight,)
+    light_pos: position rel. to body frame                    (nlight, 3)
+    light_dir: direction rel. to body frame                   (nlight, 3)
+    light_poscom0: global position rel. to sub-com in qpos0   (nlight, 3)
+    light_pos0: global position rel. to body in qpos0         (nlight, 3)
+    light_dir0: global direction in qpos0                     (nlight, 3)
+    flex_contype: flex contact type                           (nflex,)
+    flex_conaffinity: flex contact affinity                   (nflex,)
+    flex_condim: contact dimensionality (1, 3, 4, 6)          (nflex,)
+    flex_priority: flex contact priority                      (nflex,)
+    flex_solmix: mix coef for solref/imp in contact pair      (nflex,)
+    flex_solref: constraint solver reference: contact         (nflex, mjNREF)
+    flex_solimp: constraint solver impedance: contact         (nflex, mjNIMP)
+    flex_friction: friction for (slide, spin, roll)           (nflex,)
+    flex_margin: detect contact if dist<margin                (nflex,)
+    flex_gap: include in solver if dist<margin-gap            (nflex,)
+    flex_internal: internal flex collision enabled            (nflex,)
+    flex_selfcollide: self collision mode (mjtFlexSelf)       (nflex,)
+    flex_activelayers: number of active element layers, 3D only  (nflex,)
+    flex_dim: 1: lines, 2: triangles, 3: tetrahedra           (nflex,)
+    flex_vertadr: first vertex address                        (nflex,)
+    flex_vertnum: number of vertices                          (nflex,)
+    flex_edgeadr: first edge address                          (nflex,)
+    flex_edgenum: number of edges                             (nflex,)
+    flex_elemadr: first element address                       (nflex,)
+    flex_elemnum: number of elements                          (nflex,)
+    flex_elemdataadr: first element vertex id address         (nflex,)
+    flex_evpairadr: first evpair address                      (nflex,)
+    flex_evpairnum: number of evpairs                         (nflex,)
+    flex_vertbodyid: vertex body ids                          (nflex,)
+    flex_edge: edge vertex ids (2 per edge)                   (nflexedge, 2)
+    flex_elem: element vertex ids (dim+1 per elem)            (nflexelemdata,)
+    flex_elemlayer: element distance from surface, 3D only    (nflexelem,)
+    flex_evpair: (element, vertex) collision pairs            (nflexevpair, 2)
+    flex_vert: vertex positions in local body frames          (nflexvert, 3)
+    flexedge_length0: edge lengths in qpos0                   (nflexedge,)
+    flexedge_invweight0: edge inv. weight in qpos0            (nflexedge,)
+    flex_radius: radius around primitive element              (nflex,)
+    flex_edgestiffness: edge stiffness                        (nflex,)
+    flex_edgedamping: edge damping                            (nflex,)
+    flex_edgeequality: is edge equality constraint defined    (nflex,)
+    flex_rigid: are all verices in the same body              (nflex,)
+    flexedge_rigid: are both edge vertices in same body       (nflexedge,)
+    flex_centered: are all vertex coordinates (0,0,0)         (nflex,)
+    flex_bvhadr: address of bvh root; -1: no bvh              (nflex,)
+    flex_bvhnum: number of bounding volumes                   (nflex,)
     mesh_vertadr: first vertex address                        (nmesh,)
+    mesh_vertnum: number of vertices                          (nmesh,)
     mesh_faceadr: first face address                          (nmesh,)
+    mesh_bvhadr: address of bvh root                          (nmesh,)
+    mesh_bvhnum: number of bvh                                (nmesh,)
     mesh_graphadr: graph data address; -1: no graph           (nmesh,)
     mesh_vert: vertex positions for all meshes                (nmeshvert, 3)
     mesh_face: vertex face data                               (nmeshface, 3)
     mesh_graph: convex graph data                             (nmeshgraph,)
+    mesh_pos: translation applied to asset vertices           (nmesh, 3)
+    mesh_quat: rotation applied to asset vertices             (nmesh, 4)
     mesh_convex: pre-compiled convex mesh info for MJX        (nmesh,)
     hfield_size: (x, y, z_top, z_bottom)                      (nhfield,)
     hfield_nrow: number of rows in grid                       (nhfield,)
@@ -450,6 +544,7 @@ class Model(PyTreeNode):
     pair_dim: contact dimensionality                          (npair,)
     pair_geom1: id of geom1                                   (npair,)
     pair_geom2: id of geom2                                   (npair,)
+    pair_signature: body1 << 16 + body2                       (npair,)
     pair_solref: solver reference: contact normal             (npair, mjNREF)
     pair_solreffriction: solver reference: contact friction   (npair, mjNREF)
     pair_solimp: solver impedance: contact                    (npair, mjNIMP)
@@ -464,6 +559,24 @@ class Model(PyTreeNode):
     eq_solref: constraint solver reference                    (neq, mjNREF)
     eq_solimp: constraint solver impedance                    (neq, mjNIMP)
     eq_data: numeric data for constraint                      (neq, mjNEQDATA)
+    tendon_adr: address of first object in tendon's path      (ntendon,)
+    tendon_num: number of objects in tendon's path            (ntendon,)
+    tendon_limited: does tendon have length limits            (ntendon,)
+    tendon_solref_lim: constraint solver reference: limit     (ntendon, mjNREF)
+    tendon_solimp_lim: constraint solver impedance: limit     (ntendon, mjNIMP)
+    tendon_solref_fri: constraint solver reference: friction  (ntendon, mjNREF)
+    tendon_solimp_fri: constraint solver impedance: friction  (ntendon, mjNIMP)
+    tendon_range: tendon length limits                        (ntendon, 2)
+    tendon_margin: min distance for limit detection           (ntendon,)
+    tendon_stiffness: stiffness coefficient                   (ntendon,)
+    tendon_damping: damping coefficient                       (ntendon,)
+    tendon_frictionloss: loss due to friction                 (ntendon,)
+    tendon_lengthspring: spring resting length range          (ntendon, 2)
+    tendon_length0: tendon length in qpos0                    (ntendon,)
+    tendon_invweight0: inv. weight in qpos0                   (ntendon,)
+    wrap_type: wrap object type (mjtWrap)                     (nwrap,)
+    wrap_objid: object id: geom, site, joint                  (nwrap,)
+    wrap_prm: divisor, joint coef, or site id                 (nwrap,)
     actuator_trntype: transmission type (mjtTrn)              (nu,)
     actuator_dyntype: dynamics type (mjtDyn)                  (nu,)
     actuator_gaintype: gain type (mjtGain)                    (nu,)
@@ -471,16 +584,31 @@ class Model(PyTreeNode):
     actuator_trnid: transmission id: joint, tendon, site      (nu, 2)
     actuator_actadr: first activation address; -1: stateless  (nu,)
     actuator_actnum: number of activation variables           (nu,)
+    actuator_group: group for visibility                      (nu,)
     actuator_ctrllimited: is control limited                  (nu,)
     actuator_forcelimited: is force limited                   (nu,)
     actuator_actlimited: is activation limited                (nu,)
     actuator_dynprm: dynamics parameters                      (nu, mjNDYN)
     actuator_gainprm: gain parameters                         (nu, mjNGAIN)
     actuator_biasprm: bias parameters                         (nu, mjNBIAS)
+    actuator_actearly: step activation before force           (nu,)
     actuator_ctrlrange: range of controls                     (nu, 2)
     actuator_forcerange: range of forces                      (nu, 2)
     actuator_actrange: range of activations                   (nu, 2)
     actuator_gear: scale length and transmitted force         (nu, 6)
+    actuator_cranklength: crank length for slider-crank       (nu,)
+    actuator_acc0: acceleration from unit force in qpos0      (nu,)
+    actuator_lengthrange: feasible actuator length range      (nu, 2)
+    sensor_type: sensor type (mjtSensor)                      (nsensor,)
+    sensor_datatype: numeric data type (mjtDataType)          (nsensor,)
+    sensor_needstage: required compute stage (mjtStage)       (nsensor,)
+    sensor_objtype: type of sensorized object (mjtObj)        (nsensor,)
+    sensor_objid: id of sensorized object                     (nsensor,)
+    sensor_reftype: type of reference frame (mjtObj)          (nsensor,)
+    sensor_refid: id of reference frame; -1: global frame     (nsensor,)
+    sensor_dim: number of scalar outputs                      (nsensor,)
+    sensor_adr: address in sensor array                       (nsensor,)
+    sensor_cutoff: cutoff for real and positive; 0: ignore    (nsensor,)
     numeric_adr: address of field in numeric_data             (nnumeric,)
     numeric_data: array of all numeric fields                 (nnumericdata,)
     tuple_adr: address of text in text_data                   (ntuple,)
@@ -508,38 +636,66 @@ class Model(PyTreeNode):
   nu: int
   na: int
   nbody: int
+  nbvh: int
+  nbvhstatic: int
+  nbvhdynamic: int
   njnt: int
   ngeom: int
   nsite: int
   ncam: int
+  nlight: int
+  nflex: int
+  nflexvert: int
+  nflexedge: int
+  nflexelem: int
+  nflexelemdata: int
+  nflexshelldata: int
+  nflexevpair: int
+  nflextexcoord: int
   nmesh: int
   nmeshvert: int
+  nmeshnormal: int
+  nmeshtexcoord: int
   nmeshface: int
+  nmeshgraph: int
   nhfield: int
+  nhfielddata: int
   nmat: int
   npair: int
   nexclude: int
   neq: int
-  ngravcomp: int
-  nnumeric: int
-  nuserdata: int
-  ntuple: int
+  ntendon: int
+  nwrap: int
   nsensor: int
+  nnumeric: int
+  ntuple: int
   nkey: int
+  nmocap: int
   nM: int  # pylint:disable=invalid-name
+  nD: int  # pylint:disable=invalid-name
+  nB: int  # pylint:disable=invalid-name
+  ntree: int
+  ngravcomp: int
+  nuserdata: int
+  nsensordata: int
+  narena: int
   opt: Option
   stat: Statistic
   qpos0: jax.Array
   qpos_spring: jax.Array
   body_parentid: np.ndarray
+  body_mocapid: np.ndarray
   body_rootid: np.ndarray
   body_weldid: np.ndarray
   body_jntnum: np.ndarray
   body_jntadr: np.ndarray
+  body_sameframe: np.ndarray
   body_dofnum: np.ndarray
   body_dofadr: np.ndarray
+  body_treeid: np.ndarray
   body_geomnum: np.ndarray
   body_geomadr: np.ndarray
+  body_simple: np.ndarray
   body_pos: jax.Array
   body_quat: jax.Array
   body_ipos: jax.Array
@@ -548,6 +704,14 @@ class Model(PyTreeNode):
   body_subtreemass: jax.Array
   body_inertia: jax.Array
   body_gravcomp: jax.Array
+  body_margin: np.ndarray
+  body_contype: np.ndarray
+  body_conaffinity: np.ndarray
+  body_bvhadr: np.ndarray
+  body_bvhnum: np.ndarray
+  bvh_child: np.ndarray
+  bvh_nodeid: np.ndarray
+  bvh_aabb: np.ndarray
   body_invweight0: jax.Array
   jnt_type: np.ndarray
   jnt_qposadr: np.ndarray
@@ -567,7 +731,9 @@ class Model(PyTreeNode):
   dof_bodyid: np.ndarray
   dof_jntid: np.ndarray
   dof_parentid: np.ndarray
+  dof_treeid: np.ndarray
   dof_Madr: np.ndarray  # pylint:disable=invalid-name
+  dof_simplenum: np.ndarray
   dof_solref: jax.Array
   dof_solimp: jax.Array
   dof_frictionloss: jax.Array
@@ -580,6 +746,7 @@ class Model(PyTreeNode):
   geom_conaffinity: np.ndarray
   geom_condim: np.ndarray
   geom_bodyid: np.ndarray
+  geom_sameframe: np.ndarray
   geom_dataid: np.ndarray
   geom_group: np.ndarray
   geom_matid: np.ndarray
@@ -588,6 +755,7 @@ class Model(PyTreeNode):
   geom_solref: jax.Array
   geom_solimp: jax.Array
   geom_size: jax.Array
+  geom_aabb: np.ndarray
   geom_rbound: jax.Array
   geom_rbound_hfield: np.ndarray
   geom_pos: jax.Array
@@ -595,8 +763,12 @@ class Model(PyTreeNode):
   geom_friction: jax.Array
   geom_margin: jax.Array
   geom_gap: jax.Array
+  geom_fluid: np.ndarray
   geom_rgba: np.ndarray
+  site_type: np.ndarray
   site_bodyid: np.ndarray
+  site_sameframe: np.ndarray
+  site_size: np.ndarray
   site_pos: jax.Array
   site_quat: jax.Array
   cam_mode: np.ndarray
@@ -607,12 +779,69 @@ class Model(PyTreeNode):
   cam_poscom0: jax.Array
   cam_pos0: jax.Array
   cam_mat0: jax.Array
+  cam_fovy: np.ndarray
+  cam_resolution: np.ndarray
+  cam_sensorsize: np.ndarray
+  cam_intrinsic: np.ndarray
+  light_mode: np.ndarray
+  light_bodyid: np.ndarray
+  light_targetbodyid: np.ndarray
+  light_pos: np.ndarray
+  light_dir: np.ndarray
+  light_poscom0: np.ndarray
+  light_pos0: np.ndarray
+  light_dir0: np.ndarray
+  flex_contype: np.ndarray
+  flex_conaffinity: np.ndarray
+  flex_condim: np.ndarray
+  flex_priority: np.ndarray
+  flex_solmix: np.ndarray
+  flex_solref: np.ndarray
+  flex_solimp: np.ndarray
+  flex_friction: np.ndarray
+  flex_margin: np.ndarray
+  flex_gap: np.ndarray
+  flex_internal: np.ndarray
+  flex_selfcollide: np.ndarray
+  flex_activelayers: np.ndarray
+  flex_dim: np.ndarray
+  flex_vertadr: np.ndarray
+  flex_vertnum: np.ndarray
+  flex_edgeadr: np.ndarray
+  flex_edgenum: np.ndarray
+  flex_elemadr: np.ndarray
+  flex_elemnum: np.ndarray
+  flex_elemdataadr: np.ndarray
+  flex_evpairadr: np.ndarray
+  flex_evpairnum: np.ndarray
+  flex_vertbodyid: np.ndarray
+  flex_edge: np.ndarray
+  flex_elem: np.ndarray
+  flex_elemlayer: np.ndarray
+  flex_evpair: np.ndarray
+  flex_vert: np.ndarray
+  flexedge_length0: np.ndarray
+  flexedge_invweight0: np.ndarray
+  flex_radius: np.ndarray
+  flex_edgestiffness: np.ndarray
+  flex_edgedamping: np.ndarray
+  flex_edgeequality: np.ndarray
+  flex_rigid: np.ndarray
+  flexedge_rigid: np.ndarray
+  flex_centered: np.ndarray
+  flex_bvhadr: np.ndarray
+  flex_bvhnum: np.ndarray
   mesh_vertadr: np.ndarray
+  mesh_vertnum: np.ndarray
   mesh_faceadr: np.ndarray
+  mesh_bvhadr: np.ndarray
+  mesh_bvhnum: np.ndarray
   mesh_graphadr: np.ndarray
   mesh_vert: np.ndarray
   mesh_face: np.ndarray
   mesh_graph: np.ndarray
+  mesh_pos: np.ndarray
+  mesh_quat: np.ndarray
   mesh_convex: Tuple[ConvexMesh, ...]
   hfield_size: np.ndarray
   hfield_nrow: np.ndarray
@@ -623,6 +852,7 @@ class Model(PyTreeNode):
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
   pair_geom2: np.ndarray
+  pair_signature: np.ndarray
   pair_solref: jax.Array
   pair_solreffriction: jax.Array
   pair_solimp: jax.Array
@@ -637,6 +867,24 @@ class Model(PyTreeNode):
   eq_solref: jax.Array
   eq_solimp: jax.Array
   eq_data: jax.Array
+  tendon_adr: np.ndarray
+  tendon_num: np.ndarray
+  tendon_limited: np.ndarray
+  tendon_solref_lim: np.ndarray
+  tendon_solimp_lim: np.ndarray
+  tendon_solref_fri: np.ndarray
+  tendon_solimp_fri: np.ndarray
+  tendon_range: np.ndarray
+  tendon_margin: np.ndarray
+  tendon_stiffness: np.ndarray
+  tendon_damping: np.ndarray
+  tendon_frictionloss: np.ndarray
+  tendon_lengthspring: np.ndarray
+  tendon_length0: np.ndarray
+  tendon_invweight0: np.ndarray
+  wrap_type: np.ndarray
+  wrap_objid: np.ndarray
+  wrap_prm: np.ndarray
   actuator_trntype: np.ndarray
   actuator_dyntype: np.ndarray
   actuator_gaintype: np.ndarray
@@ -644,16 +892,31 @@ class Model(PyTreeNode):
   actuator_trnid: np.ndarray
   actuator_actadr: np.ndarray
   actuator_actnum: np.ndarray
+  actuator_group: np.ndarray
   actuator_ctrllimited: np.ndarray
   actuator_forcelimited: np.ndarray
   actuator_actlimited: np.ndarray
   actuator_dynprm: jax.Array
   actuator_gainprm: jax.Array
   actuator_biasprm: jax.Array
+  actuator_actearly: np.ndarray
   actuator_ctrlrange: jax.Array
   actuator_forcerange: jax.Array
   actuator_actrange: jax.Array
   actuator_gear: jax.Array
+  actuator_cranklength: np.ndarray
+  actuator_acc0: np.ndarray
+  actuator_lengthrange: np.ndarray
+  sensor_type: np.ndarray
+  sensor_datatype: np.ndarray
+  sensor_needstage: np.ndarray
+  sensor_objtype: np.ndarray
+  sensor_objid: np.ndarray
+  sensor_reftype: np.ndarray
+  sensor_refid: np.ndarray
+  sensor_dim: np.ndarray
+  sensor_adr: np.ndarray
+  sensor_cutoff: np.ndarray
   numeric_adr: np.ndarray
   numeric_data: np.ndarray
   tuple_adr: np.ndarray
@@ -724,61 +987,109 @@ class Data(PyTreeNode):
     ncon: number of contacts
     solver_niter: number of solver iterations
     time: simulation time
-    qpos: position                                                (nq,)
-    qvel: velocity                                                (nv,)
-    act: actuator activation                                      (na,)
-    qacc_warmstart: acceleration used for warmstart               (nv,)
-    ctrl: control                                                 (nu,)
-    qfrc_applied: applied generalized force                       (nv,)
-    xfrc_applied: applied Cartesian force/torque                  (nbody, 6)
-    eq_active: enable/disable constraints                         (neq,)
-    qacc: acceleration                                            (nv,)
-    act_dot: time-derivative of actuator activation               (na,)
-    xpos:  Cartesian position of body frame                       (nbody, 3)
-    xquat: Cartesian orientation of body frame                    (nbody, 4)
-    xmat:  Cartesian orientation of body frame                    (nbody, 3, 3)
-    xipos: Cartesian position of body com                         (nbody, 3)
-    ximat: Cartesian orientation of body inertia                  (nbody, 3, 3)
-    xanchor: Cartesian position of joint anchor                   (njnt, 3)
-    xaxis: Cartesian joint axis                                   (njnt, 3)
-    geom_xpos: Cartesian geom position                            (ngeom, 3)
-    geom_xmat: Cartesian geom orientation                         (ngeom, 3, 3)
-    site_xpos: Cartesian site position                            (nsite, 3)
-    site_xmat: Cartesian site orientation                         (nsite, 9)
-    cam_xpos: Cartesian camera position                           (ncam, 3)
-    cam_xmat: Cartesian camera orientation                        (ncam, 9)
-    subtree_com: center of mass of each subtree                   (nbody, 3)
-    cdof: com-based motion axis of each dof                       (nv, 6)
-    cinert: com-based body inertia and mass                       (nbody, 10)
-    actuator_length: actuator lengths                             (nu,)
-    actuator_moment: actuator moments                             (nu, nv)
-    crb: com-based composite inertia and mass                     (nbody, 10)  \
-    qM: total inertia                                  if sparse: (nM,)
-                                                       if dense:  (nv, nv)
-    qLD: L'*D*L (or Cholesky) factorization of M.      if sparse: (nM,)
-                                                       if dense:  (nv, nv)
-    qLDiagInv: 1/diag(D)                               if sparse: (nv,)
-                                                       if dense:  (0,)
-    contact: all detected contacts                                (ncon,)
-    efc_type: constraint type                                     (nefc,)
-    efc_J: constraint Jacobian                                    (nefc, nv)
-    efc_frictionloss: frictionloss (friction)                     (nefc,)
-    efc_D: constraint mass                                        (nefc,)
-    actuator_velocity: actuator velocities                        (nu,)
-    cvel: com-based velocity [3D rot; 3D tran]                    (nbody, 6)
-    cdof_dot: time-derivative of cdof                             (nv, 6)
-    qfrc_bias: C(qpos,qvel)                                       (nv,)
-    qfrc_gravcomp: passive gravity compensation force             (nv,)
-    qfrc_passive: passive force                                   (nv,)
-    efc_aref: reference pseudo-acceleration                       (nefc,)
-    qfrc_actuator: actuator force                                 (nv,)
-    qfrc_smooth: net unconstrained force                          (nv,)
-    qacc_smooth: unconstrained acceleration                       (nv,)
-    qfrc_constraint: constraint force                             (nv,)
-    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,)
+    qpos: position                                              (nq,)
+    qvel: velocity                                              (nv,)
+    act: actuator activation                                    (na,)
+    qacc_warmstart: acceleration used for warmstart             (nv,)
+    ctrl: control                                               (nu,)
+    qfrc_applied: applied generalized force                     (nv,)
+    xfrc_applied: applied Cartesian force/torque                (nbody, 6)
+    eq_active: enable/disable constraints                       (neq,)
+    mocap_pos: positions of mocap bodies                        (nmocap x 3)
+    mocap_quat: orientations of mocap bodies                    (nmocap x 4)
+    qacc: acceleration                                          (nv,)
+    act_dot: time-derivative of actuator activation             (na,)
+    userdata: user data, not touched by engine                  (nuserdata,)
+    sensordata: sensor data array                               (nsensordata,)
+    xpos:  Cartesian position of body frame                     (nbody, 3)
+    xquat: Cartesian orientation of body frame                  (nbody, 4)
+    xmat:  Cartesian orientation of body frame                  (nbody, 3, 3)
+    xipos: Cartesian position of body com                       (nbody, 3)
+    ximat: Cartesian orientation of body inertia                (nbody, 3, 3)
+    xanchor: Cartesian position of joint anchor                 (njnt, 3)
+    xaxis: Cartesian joint axis                                 (njnt, 3)
+    geom_xpos: Cartesian geom position                          (ngeom, 3)
+    geom_xmat: Cartesian geom orientation                       (ngeom, 3, 3)
+    site_xpos: Cartesian site position                          (nsite, 3)
+    site_xmat: Cartesian site orientation                       (nsite, 9)
+    cam_xpos: Cartesian camera position                         (ncam, 3)
+    cam_xmat: Cartesian camera orientation                      (ncam, 3, 3)
+    light_xpos: Cartesian light position                        (nlight, 3)
+    light_xdir: Cartesian light direction                       (nlight, 3)
+    subtree_com: center of mass of each subtree                 (nbody, 3)
+    cdof: com-based motion axis of each dof                     (nv, 6)
+    cinert: com-based body inertia and mass                     (nbody, 10)
+    flexvert_xpos: Cartesian flex vertex positions              (nflexvert, 3)
+    flexelem_aabb: flex element bounding boxes (center, size)   (nflexelem, 6)
+    flexedge_J_rownnz: number of non-zeros in Jacobian row      (nflexedge,)
+    flexedge_J_rowadr: row start address in colind array        (nflexedge,)
+    flexedge_J_colind: column indices in sparse Jacobian        (nflexedge, nv)
+    flexedge_J: flex edge Jacobian                              (nflexedge, nv)
+    flexedge_length: flex edge lengths                          (nflexedge,)
+    ten_wrapadr: start address of tendon's path                 (ntendon,)
+    ten_wrapnum: number of wrap points in path                  (ntendon,)
+    ten_J_rownnz: number of non-zeros in Jacobian row           (ntendon,)
+    ten_J_rowadr: row start address in colind array             (ntendon,)
+    ten_J_colind: column indices in sparse Jacobian             (ntendon, nv)
+    ten_J: tendon Jacobian                                      (ntendon, nv)
+    ten_length: tendon lengths                                  (ntendon,)
+    wrap_obj: geom id; -1: site; -2: pulley                     (nwrap*2,)
+    wrap_xpos: Cartesian 3D points in all path                  (nwrap*2, 3)
+    actuator_length: actuator lengths                           (nu,)
+    actuator_moment: actuator moments                           (nu, nv)
+    crb: com-based composite inertia and mass                   (nbody, 10)
+    qM: total inertia                                if sparse: (nM,)
+                                                     if dense:  (nv, nv)
+    qLD: L'*D*L (or Cholesky) factorization of M.    if sparse: (nM,)
+                                                     if dense:  (nv, nv)
+    qLDiagInv: 1/diag(D)                             if sparse: (nv,)
+                                                     if dense:  (0,)
+    qLDiagSqrtInv: 1/sqrt(diag(D))                              (nv,)
+    bvh_aabb_dyn: global bounding box (center, size)            (nbvhdynamic, 6)
+    bvh_active: volume has been added to collisions             (nbvh,)
+    flexedge_velocity: flex edge velocities                     (nflexedge,)
+    ten_velocity: tendon velocities                             (ntendon,)
+    actuator_velocity: actuator velocities                      (nu,)
+    cvel: com-based velocity [3D rot; 3D tran]                  (nbody, 6)
+    cdof_dot: time-derivative of cdof                           (nv, 6)
+    qfrc_bias: C(qpos,qvel)                                     (nv,)
+    qfrc_spring: passive spring force                           (nv,)
+    qfrc_damper: passive damper force                           (nv,)
+    qfrc_gravcomp: passive gravity compensation force           (nv,)
+    qfrc_fluid: passive fluid force                             (nv,)
+    qfrc_passive: total passive force                           (nv,)
+    subtree_linvel: linear velocity of subtree com              (nbody, 3)
+    subtree_angmom: angular momentum about subtree com          (nbody, 3)
+    qH: L'*D*L factorization of modified M                      (nM,)
+    qHDiagInv: 1/diag(D) of modified M                          (nv,)
+    D_rownnz: non-zeros in each row                             (nv,)
+    D_rowadr: address of each row in D_colind                   (nv,)
+    D_colind: column indices of non-zeros                       (nD,)
+    B_rownnz: non-zeros in each row                             (nbody,)
+    B_rowadr: address of each row in B_colind                   (nbody,)
+    B_colind: column indices of non-zeros                       (nB,)
+    qDeriv: d (passive + actuator - bias) / d qvel              (nD,)
+    qLU: sparse LU of (qM - dt*qDeriv)                          (nD,)
+    actuator_force: actuator force in actuation space           (nu,)
+    qfrc_actuator: actuator force                               (nv,)
+    qfrc_smooth: net unconstrained force                        (nv,)
+    qacc_smooth: unconstrained acceleration                     (nv,)
+    qfrc_constraint: constraint force                           (nv,)
+    qfrc_inverse: net external force; should equal:             (nv,)
+                  qfrc_applied + J'*xfrc_applied + qfrc_actuator
+    cacc: com-based acceleration                                (nbody, 6)
+    cfrc_int: com-based interaction force with parent           (nbody, 6)
+    cfrc_ext: com-based external force on body                  (nbody, 6)
+    contact: all detected contacts                              (ncon,)
+    efc_type: constraint type                                   (nefc,)
+    efc_J: constraint Jacobian                                  (nefc, nv)
+    efc_frictionloss: frictionloss (friction)                   (nefc,)
+    efc_D: constraint mass                                      (nefc,)
+    efc_aref: reference pseudo-acceleration                     (nefc,)
+    efc_force: constraint force in constraint space             (nefc,)
+    _qM_sparse: qM in sparse representation                     (nM,)
+    _qLD_sparse: qLD in sparse representation                   (nM,)
+    _qLDiagInv_sparse: qLDiagInv in sparse representation       (nv,)
   """
   # constant sizes:
   ne: int
@@ -800,11 +1111,15 @@ class Data(PyTreeNode):
   qfrc_applied: jax.Array
   xfrc_applied: jax.Array
   eq_active: jax.Array
+  # mocap data:
+  mocap_pos: jax.Array
+  mocap_quat: jax.Array
   # dynamics:
   qacc: jax.Array
   act_dot: jax.Array
   # user data:
   userdata: jax.Array
+  sensordata: jax.Array
   # position dependent:
   xpos: jax.Array
   xquat: jax.Array
@@ -819,32 +1134,83 @@ class Data(PyTreeNode):
   site_xmat: jax.Array
   cam_xpos: jax.Array
   cam_xmat: jax.Array
+  light_xpos: jax.Array
+  light_xdir: jax.Array
   subtree_com: jax.Array
   cdof: jax.Array
   cinert: jax.Array
-  crb: jax.Array
+  flexvert_xpos: jax.Array
+  flexelem_aabb: jax.Array
+  flexedge_J_rownnz: jax.Array  # pylint:disable=invalid-name
+  flexedge_J_rowadr: jax.Array  # pylint:disable=invalid-name
+  flexedge_J_colind: jax.Array  # pylint:disable=invalid-name
+  flexedge_J: jax.Array  # pylint:disable=invalid-name
+  flexedge_length: jax.Array
+  ten_wrapadr: jax.Array
+  ten_wrapnum: jax.Array
+  ten_J_rownnz: jax.Array  # pylint:disable=invalid-name
+  ten_J_rowadr: jax.Array  # pylint:disable=invalid-name
+  ten_J_colind: jax.Array  # pylint:disable=invalid-name
+  ten_J: jax.Array  # pylint:disable=invalid-name
+  ten_length: jax.Array
+  wrap_obj: jax.Array
+  wrap_xpos: jax.Array
   actuator_length: jax.Array
   actuator_moment: jax.Array
+  crb: jax.Array
   qM: jax.Array  # pylint:disable=invalid-name
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
-  contact: Contact
-  efc_type: np.ndarray
-  efc_J: jax.Array  # pylint:disable=invalid-name
-  efc_frictionloss: jax.Array
-  efc_D: jax.Array  # pylint:disable=invalid-name
+  qLDiagSqrtInv: jax.Array
+  bvh_aabb_dyn: jax.Array
+  bvh_active: jax.Array
   # position, velocity dependent:
+  flexedge_velocity: jax.Array
+  ten_velocity: jax.Array
   actuator_velocity: jax.Array
   cvel: jax.Array
   cdof_dot: jax.Array
   qfrc_bias: jax.Array
-  qfrc_passive: jax.Array
+  qfrc_spring: jax.Array
+  qfrc_damper: jax.Array
   qfrc_gravcomp: jax.Array
-  efc_aref: jax.Array
+  qfrc_fluid: jax.Array
+  qfrc_passive: jax.Array
+  subtree_linvel: jax.Array
+  subtree_angmom: jax.Array
+  qH: jax.Array  # pylint:disable=invalid-name
+  qHDiagInv: jax.Array  # pylint:disable=invalid-name
+  D_rownnz: jax.Array  # pylint:disable=invalid-name
+  D_rowadr: jax.Array  # pylint:disable=invalid-name
+  D_colind: jax.Array  # pylint:disable=invalid-name
+  B_rownnz: jax.Array  # pylint:disable=invalid-name
+  B_rowadr: jax.Array  # pylint:disable=invalid-name
+  B_colind: jax.Array  # pylint:disable=invalid-name
+  qDeriv: jax.Array  # pylint:disable=invalid-name
+  qLU: jax.Array  # pylint:disable=invalid-name
   # position, velocity, control & acceleration dependent:
   qfrc_actuator: jax.Array
+  actuator_force: jax.Array
   qfrc_smooth: jax.Array
   qacc_smooth: jax.Array
   qfrc_constraint: jax.Array
   qfrc_inverse: jax.Array
+  cacc: jax.Array
+  cfrc_int: jax.Array
+  cfrc_ext: jax.Array
+  # dynamically sized
+  contact: Contact
+  # dynamically sized - position dependent:
+  efc_type: jax.Array
+  efc_J: jax.Array  # pylint:disable=invalid-name
+  efc_frictionloss: jax.Array
+  efc_D: jax.Array  # pylint:disable=invalid-name
+  # dynamically sized - position & velocity dependent:
+  efc_aref: jax.Array
+  # dynamically sized - position, velocity, control & acceleration dependent:
   efc_force: jax.Array
+  # sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo
+  # when in dense mode
+  _qM_sparse: jax.Array
+  _qLD_sparse: jax.Array
+  _qLDiagInv_sparse: jax.Array