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

Store index of diagonal element for C and D in mjData.

Browse Source 
PiperOrigin-RevId: 704287869
Change-Id: I15d1faa0b5c3ff54721f692ec6329726456ec318
This commit is contained in:
Yuval Tassa
2024-12-09 08:00:29 -08:00
committed by Copybara-Service
parent b8344c191f
commit 3a14ada308
10 changed files with 67 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+2
View File
@@ -332,10 +332,12 @@ def make_data(
'B_colind': (m.nB, jp.int32),
'C_rownnz': (m.nv, jp.int32),
'C_rowadr': (m.nv, jp.int32),
'C_diag': (m.nv, jp.int32),
'C_colind': (m.nC, jp.int32),
'mapM2C': (m.nC, jp.int32),
'D_rownnz': (m.nv, jp.int32),
'D_rowadr': (m.nv, jp.int32),
'D_diag': (m.nv, jp.int32),
'D_colind': (m.nD, jp.int32),
'mapM2D': (m.nD, jp.int32),
'mapD2M': (m.nM, jp.int32),
mjx/mujoco/mjx/_src/types.py
+4
View File
@@ -1273,10 +1273,12 @@ class Data(PyTreeNode):
B_colind: body-dof: column indices of non-zeros (nB,)
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_diag: reduced dof-dof: index of diagonal element (nv,)
C_colind: reduced dof-dof: column indices of non-zeros (nC,)
mapM2C: index mapping from M to C (nC,)
D_rownnz: dof-dof: non-zeros in each row (nv,)
D_rowadr: dof-dof: address of each row in D_colind (nv,)
D_diag: dof-dof: index of diagonal element (nv,)
D_colind: dof-dof: column indices of non-zeros (nD,)
mapM2D: index mapping from M to D (nD,)
mapD2M: index mapping from D to M (nM,)
@@ -1402,10 +1404,12 @@ class Data(PyTreeNode):
B_colind: 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_diag: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
C_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
mapM2C: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_diag: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
D_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
mapM2D: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
mapD2M: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name