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

Fix sparse matmul logic in addJTBJSparse.

Browse Source 
Also rename arguments and loop variables in `addJTBJSparse` for additional clarity.

PiperOrigin-RevId: 534908360
Change-Id: I143ef5aef7ff90091369d408740f7f44d9e19378
This commit is contained in:
Saran Tunyasuvunakool
2023-05-24 11:01:58 -07:00
committed by Copybara-Service
parent d82465d710
commit 55c7f45dc5
2 changed files with 23 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+3
View File
@@ -17,6 +17,9 @@ Bug fixes
invalid. In such case, now MuJoCo only accepts a non-watertight geometry if :ref:`shellinertia<body-geom-shellinertia>`
is equal to ``true``.
- Fixed the sparse Jacobian multiplication logic that is used to compute derivatives for tendon damping and fluid force,
which affects the behaviour of the :ref:`implicit and implicitfast integrators<geIntegration>`.
Plugins
^^^^^^^
src/engine/engine_derivative.c
+20 -18
View File
@@ -708,39 +708,41 @@ static void addJTBJ(const mjModel* m, mjData* d, const mjtNum* J, const mjtNum*
// add J'*B*J to qDeriv, sparse version
static void addJTBJSparse(const mjModel* m, mjData* d, const mjtNum* J,
const mjtNum* B, int n, int offset,
const int* rownnz, const int* rowadr, const int* colind) {
static void addJTBJSparse(
const mjModel* m, mjData* d, const mjtNum* J,
const mjtNum* B, int n, int offset,
const int* J_rownnz, const int* J_rowadr, const int* J_colind) {
int nv = m->nv;
// allocate row
mjMARKSTACK;
mjtNum* row = mj_stackAlloc(d, nv);
mju_zero(row, nv);
// process non-zero elements of B
for (int i=0; i<n; i++) {
for (int j=0; j<n; j++) {
// compute qDeriv(k,p) += sum_{i,j} ( J(i,k)*B(i,j)*J(j,p) )
for (int i = 0; i < n; i++) {
for (int j = 0; j < n; j++) {
if (!B[i*n+j]) {
continue;
}
// process non-zero elements of J(i,k)
for (int k=0; k<rownnz[offset+i]; k++) {
int ik = rowadr[offset+i] + k;
mju_zero(row, nv);
// loop over non-zero elements of J(i,:)
for (int k = 0; k < J_rownnz[offset+i]; k++) {
int ik = J_rowadr[offset+i] + k;
mjtNum scl = J[ik]*B[i*n+j];
// process non-zero elements of J(j,p)
for (int p=0; p<rownnz[offset+j]; p++) {
int jp = rowadr[offset+j] + p;
// row[p] = J(i,k)*B(i,j)*J(j,p)
row[p] = scl * J[jp];
// loop over non-zero elements of J(j,:)
// (pJ is the sparse column index into J)
for (int pJ = 0; pJ < J_rownnz[offset+j]; pJ++) {
int adr = J_rowadr[offset+j] + pJ;
row[J_colind[adr]] = scl * J[adr];
}
// add row to qDeriv(k,:)
for (int s=0; s<d->D_rownnz[k]; s++) {
int adr = d->D_rowadr[k] + s;
// (pD is the sparse column index into qDeriv)
for (int pD = 0; pD < d->D_rownnz[k]; pD++) {
int adr = d->D_rowadr[k] + pD;
d->qDeriv[adr] += row[d->D_colind[adr]];
}
}