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

Convert qLD to CSR format.

Browse Source 
PiperOrigin-RevId: 723955038
Change-Id: I30c3dc7f59739e89ae5fff8841432bc74717ec1b
This commit is contained in:
Yuval Tassa
2025-02-06 08:58:45 -08:00
committed by Copybara-Service
parent cb1696eb34
commit c27d3758c2
22 changed files with 151 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/engine/engine_core_smooth_test.cc
+44 -43
View File
@@ -449,9 +449,8 @@ TEST_F(CoreSmoothTest, SolveMIsland) {
// expect corresponding values to match
for (int j=0; j < dofnum; j++) {
EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-14));
EXPECT_THAT(res_i[j], DoubleNear(res[dofind[j]], 1e-12));
}
mju_free(res_i);
}
@@ -475,21 +474,21 @@ TEST_F(CoreSmoothTest, FactorI) {
// dense L matrix
int nv = model->nv;
vector<mjtNum> Ldense(nv*nv);
mj_fullM(model, Ldense.data(), data->qLD);
// clear upper triangle, set diagonal to 1
vector<mjtNum> Ldense(nv*nv, 0);
mju_sparse2dense(Ldense.data(), data->qLD, nv, nv,
data->C_rownnz, data->C_rowadr, data->C_colind);
for (int i=0; i < nv; i++) {
for (int j=i; j < nv; j++) {
Ldense[i*nv+j] = i == j ? 1 : 0;
}
// set diagonal to 1
Ldense[i*nv+i] = 1;
}
// dense D matrix
vector<mjtNum> Ddense(nv*nv);
mj_fullM(model, Ddense.data(), data->qLD);
// clear everything but the diagonal
mju_sparse2dense(Ddense.data(), data->qLD, nv, nv,
data->C_rownnz, data->C_rowadr, data->C_colind);
for (int i=0; i < nv; i++) {
for (int j=0; j < nv; j++) {
// zero everything except the diagonal
if (i != j) Ddense[i*nv+j] = 0;
}
}
@@ -521,20 +520,21 @@ TEST_F(CoreSmoothTest, SolveLDs) {
mj_forward(m, d);
int nv = m->nv;
int nM = m->nM;
int nC = m->nC;
// copy LD into LDs: CSR format
vector<mjtNum> LDs(nC);
// copy M into LD: Legacy format
vector<mjtNum> LDlegacy(nM, 0);
for (int i=0; i < nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]];
LDlegacy[d->mapM2C[i]] = d->qLD[i];
}
// compare LD and LDs densified matrices
vector<mjtNum> LDdense(nv*nv);
mju_sparse2dense(LDdense.data(), LDs.data(), nv, nv,
mju_sparse2dense(LDdense.data(), d->qLD, nv, nv,
d->C_rownnz, d->C_rowadr, d->C_colind);
vector<mjtNum> LDdense2(nv*nv);
mj_fullM(m, LDdense2.data(), d->qLD);
mj_fullM(m, LDdense2.data(), LDlegacy.data());
// expect lower triangles to match exactly
for (int i=0; i < nv; i++) {
@@ -543,14 +543,14 @@ TEST_F(CoreSmoothTest, SolveLDs) {
}
}
// compare LD and LDs vector solve
// compare legacy and CSR LD vector solve
vector<mjtNum> vec(nv);
vector<mjtNum> vec2(nv);
for (int i=0; i < nv; i++) vec[i] = vec2[i] = 20 + 30*i;
for (int i=0; i < nv; i+=2) vec[i] = vec2[i] = 0;
mj_solveLD(m, vec.data(), 1, d->qLD, d->qLDiagInv);
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, 1,
mj_solveLD(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv);
mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, 1,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect vectors to match up to floating point precision
@@ -572,12 +572,13 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
mj_forward(m, d);
int nv = m->nv;
int nM = m->nM;
int nC = m->nC;
// copy LD into LDs: CSR format
vector<mjtNum> LDs(nC);
// copy LD into LDlegacy: Legacy format
vector<mjtNum> LDlegacy(nM, 0);
for (int i=0; i < nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]];
LDlegacy[d->mapM2C[i]] = d->qLD[i];
}
// compare n LD and LDs vector solve
@@ -587,8 +588,8 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) {
for (int i=0; i < nv*n; i++) vec[i] = vec2[i] = 2 + 3*i;
for (int i=0; i < nv*n; i+=3) vec[i] = vec2[i] = 0;
mj_solveLD(m, vec.data(), n, d->qLD, d->qLDiagInv);
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
mj_solveLD(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv);
mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect vectors to match up to floating point precision
@@ -609,19 +610,12 @@ TEST_F(CoreSmoothTest, SolveM2) {
mjData* d = mj_makeData(m);
mj_forward(m, d);
int nv = m->nv;
int nC = m->nC;
// copy LD into LDs: CSR format
vector<mjtNum> LDs(nC);
for (int i=0; i < nC; i++) {
LDs[i] = d->qLD[d->mapM2C[i]];
}
// inverse square root of D from inertia LDL decomposition
int nv = m->nv;
vector<mjtNum> sqrtInvD(nv);
for (int i=0; i < nv; i++) {
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
int diag = d->C_rowadr[i] + d->C_rownnz[i] - 1;
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]);
}
// compare full solve and half solve
@@ -633,7 +627,7 @@ TEST_F(CoreSmoothTest, SolveM2) {
vector<mjtNum> res(nv*n);
mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n);
mj_solveLDs(vec2.data(), LDs.data(), d->qLDiagInv, nv, n,
mj_solveLDs(vec2.data(), d->qLD, d->qLDiagInv, nv, n,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v)
@@ -655,25 +649,32 @@ TEST_F(CoreSmoothTest, FactorIs) {
mjData* d = mj_makeData(m);
mj_forward(m, d);
int nC = m->nC, nv = m->nv;
int nC = m->nC, nM = m->nM, nv = m->nv;
// copy qM into LDs, qLD into qLDexpected: CSR format
vector<mjtNum> qLDsExpected(nC);
vector<mjtNum> qLDs(nC);
// copy qM into into qLDlegacy and factorize
vector<mjtNum> qLDlegacy(nM);
mj_factorI(m, d, d->qM, qLDlegacy.data(), d->qLDiagInv);
// copy qLDlegacy into qLDexpected: CSR format
vector<mjtNum> qLDexpected(nC);
for (int i=0; i < nC; i++) {
int index = d->mapM2C[i];
qLDs[i] = d->qM[index]; // mj_factorIs is in-place
qLDsExpected[i] = d->qLD[index];
qLDexpected[i] = qLDlegacy[d->mapM2C[i]]; // mj_factorIs is in-place
}
// copy qM into qLD: CSR format
vector<mjtNum> qLD(nC);
for (int i=0; i < nC; i++) {
qLD[i] = d->qM[d->mapM2C[i]]; // mj_factorIs is in-place
}
vector<mjtNum> qLDiagInvExpected(d->qLDiagInv, d->qLDiagInv + nv);
vector<mjtNum> qLDiagInv(nv, 0);
mj_factorIs(qLDs.data(), qLDiagInv.data(), nv,
mj_factorIs(qLD.data(), qLDiagInv.data(), nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// expect outputs to match to floating point precision
EXPECT_THAT(qLDs, Pointwise(DoubleNear(1e-12), qLDsExpected));
EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected));
EXPECT_THAT(qLDiagInv, Pointwise(DoubleNear(1e-12), qLDiagInvExpected));
/* uncomment for debugging
test/engine/engine_derivative_test.cc
+4 -2
View File
@@ -436,7 +436,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
Ac[i*nv + i] = -m->jnt_stiffness[i];
Ac[nv*nv + i*nv + i] = -m->dof_damping[i];
}
mj_solveLD(m, Ac, 2*nv, d->qH, d->qHDiagInv);
mj_solveLDs(Ac, d->qH, d->qHDiagInv, nv, 2*nv,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
// A = [dt*Ac; Ac]
mju_transpose(A, Ac, 2*nv, nv);
@@ -463,7 +464,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
mjtNum *BcT = mj_stackAllocNum(d, nv*nu);
mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
d->moment_rowadr, d->moment_colind);
mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv);
mj_solveLDs(Bc, d->qH, d->qHDiagInv, nv, nu,
d->C_rownnz, d->C_rowadr, m->dof_simplenum, d->C_colind);
mju_transpose(BcT, Bc, nu, nv);
mju_scl(B, BcT, dt*dt, nu*nv);
mju_scl(B+nu*nv, BcT, dt, nu*nv);