chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add additional MuJoCo fields to MJX structs.

Browse Source 
PiperOrigin-RevId: 653270676
Change-Id: I5f7a62d3bd542abc0d618e63d2426663e7d6406f
This commit is contained in:
Saran Tunyasuvunakool
2024-07-17 10:03:53 -07:00
committed by Copybara-Service
parent dd25973481
commit dbe18f5740
4 changed files with 569 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+8
View File
@@ -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)
----------------------------
mjx/mujoco/mjx/_src/io.py
+120 -70
View File
@@ -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)
mjx/mujoco/mjx/_src/smooth.py
+3
View File
@@ -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
mjx/mujoco/mjx/_src/types.py
+438 -72
View File
@@ -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