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

Allow mj_mulJac(T)Vec_island to read and write from uncompressed vectors (length nv or nefc).

Browse Source 
PiperOrigin-RevId: 562547310
Change-Id: Id517c1842e711639f4113d23cc364945839a3759
This commit is contained in:
Yuval Tassa
2023-09-04 06:50:45 -07:00
committed by Copybara-Service
parent 0e36f00857
commit 5cef4708f2
3 changed files with 95 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+16 -8
View File
@@ -386,8 +386,9 @@ void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
// multiply Jacobian by vector, for one island
void mj_mulJacVec_island(const mjModel* m, const mjData* d,
mjtNum* res, const mjtNum* vec, int island) {
// flg_resunc and flg_vecunc denote whether res/vec are uncompressed
void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
int island, int flg_resunc, int flg_vecunc) {
// no island, call regular function
if (island < 0) {
mj_mulJacVec(m, d, res, vec);
@@ -410,7 +411,8 @@ void mj_mulJacVec_island(const mjModel* m, const mjData* d,
int Jrowadr = d->efc_J_rowadr[row];
int* Jind = d->efc_J_colind + Jrowadr;
mjtNum* J = d->efc_J + Jrowadr;
res[i] = mju_dotSparse2(vec, J, vecnnz, vecind, Jnnz, Jind, /*flg_unc2=*/0);
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse2(J, vec, Jnnz, Jind, vecnnz, vecind, flg_vecunc);
}
}
@@ -418,7 +420,9 @@ void mj_mulJacVec_island(const mjModel* m, const mjData* d,
else {
int nv = m->nv;
for (int i=0; i < resnnz; i++) {
res[i] = mju_dotSparse(vec, d->efc_J + nv*resind[i], vecnnz, vecind, /*flg_unc1=*/0);
int row = resind[i];
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse(vec, d->efc_J + nv*row, vecnnz, vecind, flg_vecunc);
}
}
}
@@ -447,8 +451,9 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
// multiply Jacobian transpose by vector, for one island
void mj_mulJacTVec_island(const mjModel* m, const mjData* d,
mjtNum* res, const mjtNum* vec, int island) {
// flg_resunc and flg_vecunc denote whether res/vec are uncompressed
void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
int island, int flg_resunc, int flg_vecunc) {
// no island, call regular function
if (island < 0) {
mj_mulJacTVec(m, d, res, vec);
@@ -471,7 +476,8 @@ void mj_mulJacTVec_island(const mjModel* m, const mjData* d,
int JTrowadr = d->efc_JT_rowadr[row];
int* JTind = d->efc_JT_colind + JTrowadr;
mjtNum* JT = d->efc_JT + JTrowadr;
res[i] = mju_dotSparse2(vec, JT, vecnnz, vecind, JTnnz, JTind, /*flg_unc2=*/0);
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse2(JT, vec, JTnnz, JTind, vecnnz, vecind, flg_vecunc);
}
}
@@ -479,7 +485,9 @@ void mj_mulJacTVec_island(const mjModel* m, const mjData* d,
else {
int nefc = d->nefc;
for (int i=0; i < resnnz; i++) {
res[i] = mju_dotSparse(vec, d->efc_JT + nefc*resind[i], vecnnz, vecind, /*flg_unc1=*/0);
int row = resind[i];
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse(vec, d->efc_JT + nefc*row, vecnnz, vecind, flg_vecunc);
}
}
}
src/engine/engine_core_constraint.h
+4 -4
View File
@@ -39,15 +39,15 @@ MJAPI int mj_isDual(const mjModel* m);
MJAPI void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
// multiply Jacobian by vector, for one island
MJAPI void mj_mulJacVec_island(const mjModel* m, const mjData* d,
mjtNum* res, const mjtNum* vec, int island);
MJAPI void mj_mulJacVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
int island, int flg_resunc, int flg_vecunc);
// multiply JacobianT by vector
MJAPI void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
// multiply JacobianT by vector, for one island
MJAPI void mj_mulJacTVec_island(const mjModel* m, const mjData* d,
mjtNum* res, const mjtNum* vec, int island);
MJAPI void mj_mulJacTVec_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
int island, int flg_resunc, int flg_vecunc);
//-------------------------- utility functions -----------------------------------------------------
test/engine/engine_core_constraint_test.cc
+75 -8
View File
@@ -15,6 +15,7 @@
// Tests for engine/engine_core_constraint.c.
#include <cstddef>
#include <cstring>
#include <string>
#include <vector>
@@ -225,6 +226,7 @@ TEST_F(CoreConstraintTest, MulJacVecIsland) {
// multiply by Jacobian: vec_nefc = J * vec_nv
mjtNum* vec_nefc = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc);
mj_mulJacVec(model, data, vec_nefc, vec_nv);
mjtNum* vec_nefc_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*data->nefc);
// iterate over islands
for (int i=0; i < data->nisland; i++) {
@@ -234,25 +236,57 @@ TEST_F(CoreConstraintTest, MulJacVecIsland) {
int efcnum = data->island_efcnum[i];
mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum);
// copy values into vec_nvi
// get indices
int* dofind = data->island_dofind + data->island_dofadr[i];
int* efcind = data->island_efcind + data->island_efcadr[i];
// copy values into vec_nvi
for (int j=0; j < dofnum; j++) {
vec_nvi[j] = vec_nv[dofind[j]];
}
// multiply by Jacobian, for this island
mj_mulJacVec_island(model, data, vec_nefci, vec_nvi, i);
// ===== both compressed
int flg_resunc = 0;
int flg_vecunc = 0;
mju_zero(vec_nefci, efcnum); // clear output
mj_mulJacVec_island(model, data, vec_nefci, vec_nvi,
i, flg_resunc, flg_vecunc);
// expect corresponding values to match
int* efcind = data->island_efcind + data->island_efcadr[i];
for (int j=0; j < efcnum; j++) {
EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12));
}
// ===== input uncompressed: read from vec_nv
flg_resunc = 0;
flg_vecunc = 1;
mju_zero(vec_nefci, efcnum); // clear output
mj_mulJacVec_island(model, data, vec_nefci, vec_nv,
i, flg_resunc, flg_vecunc);
// expect corresponding values to match
for (int j=0; j < efcnum; j++) {
EXPECT_THAT(vec_nefci[j], DoubleNear(vec_nefc[efcind[j]], 1e-12));
}
// ===== output uncompressed: write to vec_nefc_tmp
flg_resunc = 1;
flg_vecunc = 0;
mju_zero(vec_nefc_tmp, data->nefc); // clear output
mj_mulJacVec_island(model, data, vec_nefc_tmp, vec_nvi,
i, flg_resunc, flg_vecunc);
// expect corresponding values to match
for (int j=0; j < efcnum; j++) {
EXPECT_THAT(vec_nefc_tmp[efcind[j]],
DoubleNear(vec_nefc[efcind[j]], 1e-12));
}
mju_free(vec_nvi);
mju_free(vec_nefci);
}
mju_free(vec_nefc_tmp);
mju_free(vec_nefc);
}
@@ -268,6 +302,7 @@ TEST_F(CoreConstraintTest, MulJacTVecIsland) {
// allocate vec_nv
mjtNum* vec_nv = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv);
mjtNum* vec_nv_tmp = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nv);
// iterate through dense and sparse
for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) {
@@ -297,27 +332,59 @@ TEST_F(CoreConstraintTest, MulJacTVecIsland) {
int efcnum = data->island_efcnum[i];
mjtNum* vec_nefci = (mjtNum*)mju_malloc(sizeof(mjtNum) * efcnum);
// copy values into vec_nefci
// get indices
int* efcind = data->island_efcind + data->island_efcadr[i];
int* dofind = data->island_dofind + data->island_dofadr[i];
// copy values into vec_nefci
for (int j=0; j < efcnum; j++) {
vec_nefci[j] = vec_nefc[efcind[j]];
}
// multiply by Jacobian, for this island
mj_mulJacTVec_island(model, data, vec_nvi, vec_nefci, i);
// ==== both compressed
int flg_resunc = 0;
int flg_vecunc = 0;
mju_zero(vec_nvi, dofnum); // clear output
mj_mulJacTVec_island(model, data, vec_nvi, vec_nefci,
i, flg_resunc, flg_vecunc);
// expect corresponding values to match
int* dofind = data->island_dofind + data->island_dofadr[i];
for (int j=0; j < dofnum; j++) {
EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12));
}
// ===== input uncompressed: read from vec_nefc
flg_resunc = 0;
flg_vecunc = 1;
mju_zero(vec_nvi, dofnum); // clear output
mj_mulJacTVec_island(model, data, vec_nvi, vec_nefc,
i, flg_resunc, flg_vecunc);
// expect corresponding values to match
for (int j=0; j < dofnum; j++) {
EXPECT_THAT(vec_nvi[j], DoubleNear(vec_nv[dofind[j]], 1e-12));
}
// ===== output uncompressed: write to vec_nv_tmp
flg_resunc = 1;
flg_vecunc = 0;
mju_zero(vec_nv_tmp, model->nv); // clear output
mj_mulJacTVec_island(model, data, vec_nv_tmp, vec_nefci,
i, flg_resunc, flg_vecunc);
// expect corresponding values to match
for (int j=0; j < dofnum; j++) {
EXPECT_THAT(vec_nv_tmp[dofind[j]],
DoubleNear(vec_nv[dofind[j]], 1e-12));
}
mju_free(vec_nvi);
mju_free(vec_nefci);
}
mju_free(vec_nefc);
}
mju_free(vec_nv_tmp);
mju_free(vec_nv);
mj_deleteData(data);
mj_deleteModel(model);