chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 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 -----------------------------------------------------