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

Reuse sparse Dof-Dof matrix to make compressed sparse M in mj_addM.

Browse Source 
PiperOrigin-RevId: 556823353
Change-Id: I650ac83767d60c69859aae6e38aa331f43814378
This commit is contained in:
Kyle Bayes
2023-08-14 09:48:35 -07:00
committed by Copybara-Service
parent 7e1c65b10e
commit f34cb22bd3
4 changed files with 115 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_solver.c
+7 -7
View File
@@ -1364,12 +1364,12 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// sparse
if (mj_isSparse(m)) {
// create sparse inertia matrix M (uncompressed)
mjtNum* M = mj_stackAlloc(d, nv*nv);
int* M_rownnz = (int*) mj_stackAlloc(d, nv);
int* M_rowadr = (int*) mj_stackAlloc(d, nv);
int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
// create sparse inertia matrix M
int nnz = m->nD; // use sparse dof-dof matrix
int* M_rownnz = (int*) mj_stackAlloc(d, nv); // actual nnz count
int* M_colind = (int*) mj_stackAlloc(d, nnz);
mjtNum* M = mj_stackAlloc(d, nnz);
mj_makeMSparse(m, d, M, M_rownnz, NULL, M_colind);
// compute H = J'*D*J
@@ -1384,7 +1384,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// compute H = M + J'*D*J
mj_addMSparse(m, d, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind,
M, M_rownnz, M_rowadr, M_colind);
M, M_rownnz, NULL, M_colind);
// factorize H, uncompressed layout
int rank = mju_cholFactorSparse(ctx->H, nv, mjMINVAL,
src/engine/engine_support.c
+56 -50
View File
@@ -950,14 +950,15 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
if (rownnz && rowadr && colind) {
int nv = m->nv;
mjMARKSTACK;
// create sparse inertia matrix M (uncompressed)
mjtNum* M = mj_stackAlloc(d, nv*nv);
int* M_rownnz = (int*) mj_stackAlloc(d, nv);
int* M_rowadr = (int*) mj_stackAlloc(d, nv);
int* M_colind = (int*) mj_stackAlloc(d, nv*nv);
mj_createMSparse(m, d, M, M_rownnz, M_rowadr, M_colind);
// create sparse inertia matrix M
int nnz = m->nD; // use sparse dof-dof matrix
int* M_rownnz = (int*) mj_stackAlloc(d, nv); // actual nnz count
int* M_colind = (int*) mj_stackAlloc(d, nnz);
mjtNum* M = mj_stackAlloc(d, nnz);
mj_makeMSparse(m, d, M, M_rownnz, NULL, M_colind);
mj_addMSparse(m, d, dst, rownnz, rowadr, colind, M,
M_rownnz, M_rowadr, M_colind);
M_rownnz, NULL, M_colind);
mjFREESTACK;
}
@@ -969,65 +970,64 @@ void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
// create inertia matrix M (uncompressed)
void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind) {
int adr, adr1, nv = m->nv;
// make inertia matrix M
void mj_makeMSparse(const mjModel* m, mjData* d, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind) {
int nv = m->nv;
// currently the sparse dof-dof matrix D row addresses are used, since D has
// the same predetermined sparsity structure as M, however with simple bodies
// M has less non-zeros and can be precounted for further memory reduction
if (M_rowadr == NULL) {
M_rowadr = d->D_rowadr;
}
// build M into sparse format, lower-triangular
for (int i=0; i < nv; i++) {
M_rowadr[i] = i*nv;
adr = m->dof_Madr[i];
// build M into sparse format, lower triangle
for (int i = 0; i < nv; i++) {
int Madr = m->dof_Madr[i];
// simple, fill diagonal only
if (m->dof_simplenum[i]) {
M[i*nv] = d->qM[adr];
M_colind[i*nv] = i;
M_rownnz[i] = 1;
M[M_rowadr[i]] = d->qM[Madr];
M_colind[M_rowadr[i]] = i;
continue;
}
// backward pass over dofs: construct M_row(i) in reverse order
int j = i;
adr1 = 0;
while (j >= 0) {
// assign
M[i*nv+adr1] = d->qM[adr];
M_colind[i*nv+adr1] = j;
// count columns
adr1++;
// advance
adr++;
j = m->dof_parentid[j];
int col = M_rowadr[i]; // current column in row i
for (int j = i; j >= 0; j = m->dof_parentid[j]) {
M[col] = d->qM[Madr++];
M_colind[col++] = j;
}
// assign row descriptors
M_rownnz[i] = adr1;
// track nnz of lower triangle for row i
int nnz = M_rownnz[i] = col - M_rowadr[i];
// reverse order
for (int k=0; k < adr1/2; k++) {
mjtNum tmp = M[i*nv+k];
M[i*nv+k] = M[i*nv+adr1-1-k];
M[i*nv+adr1-1-k] = tmp;
int end = nnz >> 1;
for (int j = 0; j < end; j++) {
int a1 = M_rowadr[i] + j; // address 1
int a2 = (M_rowadr[i] + nnz - 1) - j; // address 2
int tmpi = M_colind[i*nv+k];
M_colind[i*nv+k] = M_colind[i*nv+adr1-1-k];
M_colind[i*nv+adr1-1-k] = tmpi;
// swap M data on row i
mjtNum val = M[a1];
M[a1] = M[a2];
M[a2] = val;
// swap M column indices on row i
int ind = M_colind[a1];
M_colind[a1] = M_colind[a2];
M_colind[a2] = ind;
}
}
// make symmetric
for (int i=1; i < nv; i++) {
if (m->dof_simplenum[i]) {
continue;
}
for (int k=nv*i; k < nv*i+M_rownnz[i]-1; k++) {
// add to row given by column index
adr1 = nv*M_colind[k] + M_rownnz[M_colind[k]]++;
M[adr1] = M[k];
M_colind[adr1] = i;
// fill upper triangle
for (int i = 1; i < nv; i++) {
int end = M_rowadr[i] + M_rownnz[i] - 1;
for (int j = M_rowadr[i]; j < end; j++) {
int a = M_rowadr[M_colind[j]] + M_rownnz[M_colind[j]]++;
M[a] = M[j];
M_colind[a] = i;
}
}
}
@@ -1039,6 +1039,12 @@ void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind) {
int nv = m->nv;
// currently the sparse dof-dof matrix D row addresses are used, since D has
// the same predetermined sparsity structure as M, however with simple bodies
// M has less non-zeros and can be precounted for further memory reduction
if (M_rowadr == NULL) {
M_rowadr = d->D_rowadr;
}
mjMARKSTACK;
int* buf_ind = (int*) mj_stackAlloc(d, nv);
src/engine/engine_support.h
+3 -3
View File
@@ -116,9 +116,9 @@ MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
int* rownnz, int* rowadr, int* colind);
// create inertia matrix M (uncompressed)
MJAPI void mj_createMSparse(const mjModel* m, mjData* d, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind);
// make inertia matrix M
MJAPI void mj_makeMSparse(const mjModel* m, mjData* d, mjtNum* M,
int* M_rownnz, int* M_rowadr, int* M_colind);
// add inertia matrix to sparse destination matrix
MJAPI void mj_addMSparse(const mjModel* m, mjData* d, mjtNum* dst,
test/engine/engine_support_test.cc
+49 -1
View File
@@ -16,10 +16,10 @@
#include <random>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mujoco.h>
#include "test/fixture.h"
@@ -34,6 +34,7 @@ using ::testing::DoubleNear;
using ::testing::ContainsRegex;
using ::testing::MatchesRegex;
using ::testing::Pointwise;
using ::testing::ElementsAreArray;
using JacobianTest = MujocoTest;
static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
@@ -405,5 +406,52 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
mj_deleteModel(model);
}
using AddMTest = MujocoTest;
TEST_F(AddMTest, DenseSameAsSparse) {
mjModel* m = LoadModelFromPath("humanoid100/humanoid100.xml");
mjData* d = mj_makeData(m);
int nv = m->nv;
// force use of sparse matrices
m->opt.jacobian = mjJAC_SPARSE;
// warm-up rollout to get a typical state
while (d->time < 2) {
mj_step(m, d);
}
// dense zero matrix
std::vector<mjtNum> dst_sparse = std::vector(nv * nv, 0.0);
// sparse zero matrix
std::vector<mjtNum> dst_dense = std::vector(nv * nv, 0.0);
std::vector<int> rownnz = std::vector(nv, nv);
std::vector<int> rowadr = std::vector(nv, 0);
std::vector<int> colind = std::vector(nv * nv, 0);
// set sparse structure
for (int i = 0; i < nv; i++) {
rowadr[i] = i * nv;
for (int j = 0; j < nv; j++) {
colind[rowadr[i] + j] = j;
}
}
// sparse addM
mj_addM(m, d, dst_sparse.data(), rownnz.data(),
rowadr.data(), colind.data());
// dense addM
mj_addM(m, d, dst_dense.data(), NULL, NULL, NULL);
// dense comparison, should be same matrix
EXPECT_THAT(dst_dense, ElementsAreArray(dst_sparse));
// clean up
mj_deleteData(d);
mj_deleteModel(m);
}
} // namespace
} // namespace mujoco