Skip to content

Commit

Permalink
Switch mj_solveM_island to use CSR representation
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 717684720
Change-Id: I7d74a5d4aef4aa5b0c0acddf88c5b8591d63e88d
  • Loading branch information
@yuvaltassa @copybara-github
yuvaltassa authored and copybara-github committed Jan 21, 2025
1 parent a5ab7a9 commit 2624d52
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 73 deletions.
62 changes: 32 additions & 30 deletions src/engine/engine_core_smooth.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1701,7 +1701,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) {

// in-place sparse backsubstitution for one island: x = inv(L'*D*L)*x
// L is in lower triangle of qLD; D is on diagonal of qLD
void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int island) {
void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* restrict x, int island) {
// if no islands, call mj_solveLD
const mjtNum* qLD = d->qLD;
const mjtNum* qLDiagInv = d->qLDiagInv;
Expand All @@ -1710,10 +1710,19 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
return;
}

// local constants: general
const int* Madr = m->dof_Madr;
const int* parentid = m->dof_parentid;
const int* simplenum = m->dof_simplenum;
// local copies of key variables
const int* rownnz = d->C_rownnz;
const int* rowadr = d->C_rowadr;
const int* colind = d->C_colind;
const int* diagnum = m->dof_simplenum;

// temporary: make local CSR version of qLD
int nC = m->nC;
mj_markStack(d);
mjtNum* qLDs = mjSTACKALLOC(d, nC, mjtNum);
for (int i=0; i < nC; i++) {
qLDs[i] = d->qLD[d->mapM2C[i]];
}

// local constants: island specific
int ndof = d->island_dofnum[island];
Expand All @@ -1723,18 +1732,12 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
// x <- inv(L') * x; skip simple, exploit sparsity of input vector
for (int k=ndof-1; k >= 0; k--) {
int i = dofind[k];
if (!simplenum[i] && x[k]) {
// init
int Madr_ij = Madr[i]+1;
int j = parentid[i];

// traverse ancestors backwards
// read directly from x[l] since j cannot be a parent of itself
while (j >= 0) {
x[islandind[j]] -= qLD[Madr_ij++]*x[k]; // x(j) -= L(i,j) * x(i)

// advance to parent
j = parentid[j];
mjtNum x_k;
if (!diagnum[i] && (x_k = x[k])) {
int start = rowadr[i];
int end = start + rownnz[i] - 1;
for (int adr=end-1; adr >= start; adr--) {
x[islandind[colind[adr]]] -= qLDs[adr] * x_k;
}
}
}
Expand All @@ -1747,21 +1750,20 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
// x <- inv(L) * x; skip simple
for (int k=0; k < ndof; k++) {
int i = dofind[k];
if (!simplenum[i]) {
// init
int Madr_ij = Madr[i]+1;
int j = parentid[i];

// traverse ancestors backwards
// write directly in x[i] since i cannot be a parent of itself
while (j >= 0) {
x[k] -= qLD[Madr_ij++]*x[islandind[j]]; // x(i) -= L(i,j) * x(j)

// advance to parent
j = parentid[j];
}

// skip diagonal rows
if (diagnum[i]) {
continue;
}

int start = rowadr[i];
int end = start + rownnz[i] - 1;
for (int adr=end-1; adr >= start; adr--) {
x[k] -= x[islandind[colind[adr]]] * qLDs[adr];
}
}

mj_freeStack(d);
}


Expand Down
3 changes: 2 additions & 1 deletion src/engine/engine_core_smooth.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,8 +71,9 @@ MJAPI void mj_solveLDs(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, i
// sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d
MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);

// TODO(tassa): Restore mjData const-ness.
// sparse backsubstitution for one island: x = inv(L'*D*L)*x, use factorization in d
MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island);
MJAPI void mj_solveM_island(const mjModel* m, mjData* d, mjtNum* x, int island);

// half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y
MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
Expand Down
4 changes: 2 additions & 2 deletions src/engine/engine_solver.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -895,9 +895,9 @@ static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
}



// TODO(tassa): Restore mjData const-ness.
// update grad, Mgrad
static void CGupdateGradient(const mjModel* m, const mjData* d, mjCGContext* ctx) {
static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = ctx->nv;
const int* dofind = ctx->dofind;

Expand Down
85 changes: 45 additions & 40 deletions test/engine/engine_core_smooth_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -353,58 +353,63 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) {

static const char* const kIlslandEfcPath =
"engine/testdata/island/island_efc.xml";
static const char* const kModelPath =
"testdata/model.xml";

TEST_F(CoreSmoothTest, SolveMIsland) {
const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
mjData* data = mj_makeData(model);
int nv = model->nv;
for (auto model_path : {kModelPath, kIlslandEfcPath}) {
const std::string xml_path = GetTestDataFilePath(model_path);
mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
mjData* data = mj_makeData(model);
int nv = model->nv;

// allocate vec, fill with arbitrary values, copy to sol
mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
for (int i=0; i < nv; i++) {
vec[i] = 0.2 + 0.3*i;
}
mju_copy(res, vec, nv);

// allocate vec, fill with arbitrary values, copy to sol
mjtNum* vec = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
mjtNum* res = (mjtNum*) mju_malloc(sizeof(mjtNum) * nv);
for (int i=0; i < nv; i++) {
vec[i] = 0.2 + 0.3*i;
}
mju_copy(res, vec, nv);
if (model->nkey > 0) mj_resetDataKeyframe(model, data, 0);

// simulate for 0.2 seconds
mj_resetData(model, data);
while (data->time < 0.2) {
mj_step(model, data);
}
mj_forward(model, data);
for (int i=0; i < 6; i++) {
mj_step(model, data);
}

// divide by mass matrix: sol = M^-1 * vec
mj_solveM(model, data, res, res, 1);
mj_forward(model, data);

// iterate over islands
for (int i=0; i < data->nisland; i++) {
// allocate dof vectors for island
int dofnum = data->island_dofnum[i];
mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);
// divide by mass matrix: sol = M^-1 * vec
mj_solveM(model, data, res, res, 1);

// copy values into sol_i
int* dofind = data->island_dofind + data->island_dofadr[i];
for (int j=0; j < dofnum; j++) {
res_i[j] = vec[dofind[j]];
}
// iterate over islands
for (int i=0; i < data->nisland; i++) {
// allocate dof vectors for island
int dofnum = data->island_dofnum[i];
mjtNum* res_i = (mjtNum*)mju_malloc(sizeof(mjtNum) * dofnum);

// copy values into sol_i
int* dofind = data->island_dofind + data->island_dofadr[i];
for (int j=0; j < dofnum; j++) {
res_i[j] = vec[dofind[j]];
}

// divide by mass matrix, for this island
mj_solveM_island(model, data, res_i, i);
// divide by mass matrix, for this island
mj_solveM_island(model, data, res_i, i);

// expect corresponding values to match
for (int j=0; j < dofnum; j++) {
EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12));
// expect corresponding values to match
for (int j=0; j < dofnum; j++) {
EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14));
}

mju_free(res_i);
}

mju_free(res_i);
mju_free(res);
mju_free(vec);
mj_deleteData(data);
mj_deleteModel(model);
}

mju_free(res);
mju_free(vec);
mj_deleteData(data);
mj_deleteModel(model);
}

static const char* const kInertiaPath = "engine/testdata/inertia.xml";
Expand Down

0 comments on commit 2624d52

Please sign in to comment.