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

Allocate qM CSR matrix structure on mjData (currently empty).

Browse Source 
PiperOrigin-RevId: 733247419
Change-Id: I814e574e3161875e89433a6e03f2c9e200d8b1d6
This commit is contained in:
Yuval Tassa
2025-03-04 02:36:54 -08:00
committed by Copybara-Service
parent 8f37ae1bf2
commit c72c83aa3b
7 changed files with 60 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+4
View File
@@ -339,6 +339,10 @@ def make_data(
'B_rownnz': (m.nbody, jp.int32),
'B_rowadr': (m.nbody, jp.int32),
'B_colind': (m.nB, jp.int32),
'M_rownnz': (m.nv, jp.int32),
'M_rowadr': (m.nv, jp.int32),
'M_colind': (m.nM, jp.int32),
'mapM2M': (m.nM, jp.int32),
'C_rownnz': (m.nv, jp.int32),
'C_rowadr': (m.nv, jp.int32),
'C_colind': (m.nC, jp.int32),
mjx/mujoco/mjx/_src/types.py
+8
View File
@@ -1316,6 +1316,10 @@ class Data(PyTreeNode):
B_rownnz: body-dof: non-zeros in each row (nbody,)
B_rowadr: body-dof: address of each row in B_colind (nbody,)
B_colind: body-dof: column indices of non-zeros (nB,)
M_rownnz: inertia: non-zeros in each row (nv,)
M_rowadr: inertia: address of each row in M_colind (nv,)
M_colind: inertia: column indices of non-zeros (nM,)
mapM2M: index mapping from M (legacy) to M (CSR) (nM,)
C_rownnz: reduced dof-dof: non-zeros in each row (nv,)
C_rowadr: reduced dof-dof: address of each row in C_colind (nv,)
C_colind: reduced dof-dof: column indices of non-zeros (nC,)
@@ -1445,6 +1449,10 @@ class Data(PyTreeNode):
B_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
B_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
B_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
M_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
M_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
M_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
mapM2M: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
C_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
C_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
C_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name