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

Merge pull request #1446 from jaro-sevcik:mjx-solve-m-parallel

Browse Source 
PiperOrigin-RevId: 610984251
Change-Id: Ifd69480643df236fe2d6b9f77d1b029d2650acff
This commit is contained in:
Copybara-Service
2024-02-27 22:21:11 -08:00
parent 2e43f18bee 7688e03eea
commit a5ea87b0e6
1 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/smooth.py
+12 -8
View File
@@ -339,6 +339,10 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
if not support.is_sparse(m):
return jax.scipy.linalg.cho_solve((d.qLD, False), x)
depth = []
for i in range(m.nv):
depth.append(depth[m.dof_parentid[i]] + 1 if m.dof_parentid[i] != -1 else 0)
updates_i, updates_j = {}, {}
for i in range(m.nv):
madr_ij, j = m.dof_Madr[i], i
@@ -346,21 +350,21 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
madr_ij, j = madr_ij + 1, m.dof_parentid[j]
if j == -1:
break
updates_i.setdefault(i, []).append((madr_ij, j))
updates_j.setdefault(j, []).append((madr_ij, i))
updates_i.setdefault(depth[i], []).append((i, madr_ij, j))
updates_j.setdefault(depth[j], []).append((j, madr_ij, i))
# x <- inv(L') * x
for j, vals in sorted(updates_j.items(), reverse=True):
madr_ij, i = jp.array(vals).T
x = x.at[j].add(-jp.sum(d.qLD[madr_ij] * x[i]))
for _, vals in sorted(updates_j.items(), reverse=True):
j, madr_ij, i = np.array(vals).T
x = x.at[j].add(-d.qLD[madr_ij] * x[i])
# x <- inv(D) * x
x = x * d.qLDiagInv
# x <- inv(L) * x
for i, vals in sorted(updates_i.items()):
madr_ij, j = jp.array(vals).T
x = x.at[i].add(-jp.sum(d.qLD[madr_ij] * x[j]))
for _, vals in sorted(updates_i.items()):
i, madr_ij, j = np.array(vals).T
x = x.at[i].add(-d.qLD[madr_ij] * x[j])
return x