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

Refactor islands to be memory contiguous.

Browse Source 
PiperOrigin-RevId: 755803476
Change-Id: I41972b07e0d5ef5d0117c94f565b93367b87458b
This commit is contained in:
Yuval Tassa
2025-05-07 05:05:34 -07:00
committed by Copybara-Service
parent 449de73430
commit ecb769fc3a
30 changed files with 1742 additions and 1116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+41 -151
View File
@@ -378,50 +378,6 @@ void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
// multiply Jacobian by vector, for one 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);
return;
}
// sizes
int vecnnz = d->island_dofnum[island];
int resnnz = d->island_efcnum[island];
// indices
int* vecind = d->island_dofind + d->island_dofadr[island];
int* resind = d->island_efcind + d->island_efcadr[island];
// sparse Jacobian
if (mj_isSparse(m)) {
for (int i=0; i < resnnz; i++) {
int row = resind[i];
int Jnnz = d->efc_J_rownnz[row];
int Jrowadr = d->efc_J_rowadr[row];
int* Jind = d->efc_J_colind + Jrowadr;
mjtNum* J = d->efc_J + Jrowadr;
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse2(J, vec, Jnnz, Jind, vecnnz, vecind, flg_vecunc);
}
}
// dense Jacobian
else {
int nv = m->nv;
for (int i=0; i < resnnz; i++) {
int row = resind[i];
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse(vec, d->efc_J + nv*row, vecnnz, vecind, flg_vecunc);
}
}
}
// multiply JacobianT by vector
void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
// exit if no constraints
@@ -443,50 +399,6 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
// multiply Jacobian transpose by vector, for one 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);
return;
}
// sizes
int vecnnz = d->island_efcnum[island];
int resnnz = d->island_dofnum[island];
// indices
int* vecind = d->island_efcind + d->island_efcadr[island];
int* resind = d->island_dofind + d->island_dofadr[island];
// sparse Jacobian
if (mj_isSparse(m)) {
for (int i=0; i < resnnz; i++) {
int row = resind[i];
int JTnnz = d->efc_JT_rownnz[row];
int JTrowadr = d->efc_JT_rowadr[row];
int* JTind = d->efc_JT_colind + JTrowadr;
mjtNum* JT = d->efc_JT + JTrowadr;
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse2(JT, vec, JTnnz, JTind, vecnnz, vecind, flg_vecunc);
}
}
// dense Jacobian
else {
int nefc = d->nefc;
for (int i=0; i < resnnz; i++) {
int row = resind[i];
int j = flg_resunc ? row : i;
res[j] = mju_dotSparse(vec, d->efc_JT + nefc*row, vecnnz, vecind, flg_vecunc);
}
}
}
//--------------------- instantiate constraints by type --------------------------------------------
// equality constraints
@@ -2102,10 +2014,6 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
// supernodes of JT
mju_superSparse(m->nv, d->efc_JT_rowsuper,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind);
} else {
if (mjENABLED(mjENBL_ISLAND)) {
mju_transpose(d->efc_JT, d->efc_J, d->nefc, m->nv);
}
}
// compute diagApprox
@@ -2377,25 +2285,17 @@ void mj_referenceConstraint(const mjModel* m, mjData* d) {
//---------------------------- update constraint state ---------------------------------------------
// compute efc_state, efc_force, qfrc_constraint, optionally restricted to one island
// island < 0: update all d->nefc constraints
// island >= 0: update only d->island_efcnum[island] constraints
// jar = Jac*qacc-aref is restricted to the island, in the above sense
// compute efc_state, efc_force
// optional: cost(qacc) = shat(jar); cone Hessians
void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
mjtNum cost[1], int flg_coneHessian, int island) {
int ne = d->ne, nf = d->nf;
const mjtNum *D = d->efc_D, *R = d->efc_R, *floss = d->efc_frictionloss;
mjtNum* force = d->efc_force;
void mj_constraintUpdate_impl(int ne, int nf, int nefc,
const mjtNum* D, const mjtNum* R, const mjtNum* floss,
const mjtNum* jar, const int* type, const int* id,
mjContact* contact, int* state, mjtNum* force, mjtNum cost[1],
int flg_coneHessian) {
mjtNum s = 0;
int nefc = island < 0 ? d->nefc : d->island_efcnum[island];
int* efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island];
// no constraints: clear qfrc_constraint and cost, return
// no constraints: clear cost, return
if (!nefc) {
// can only occur for island == -1
mju_zero(d->qfrc_constraint, m->nv);
if (cost) {
*cost = 0;
}
@@ -2403,55 +2303,49 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
}
// compute unconstrained efc_force
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
force[i] = -D[i]*jar[c];
for (int i=0; i < nefc; i++) {
force[i] = -D[i]*jar[i];
}
// update constraints
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
for (int i=0; i < nefc; i++) {
// ==== equality
if (i < ne) {
if (cost) {
s += 0.5*D[i]*jar[c]*jar[c];
s += 0.5*D[i]*jar[i]*jar[i];
}
d->efc_state[i] = mjCNSTRSTATE_QUADRATIC;
state[i] = mjCNSTRSTATE_QUADRATIC;
continue;
}
// ==== friction
if (i < ne + nf) {
// linear negative
if (jar[c] <= -R[i]*floss[i]) {
if (jar[i] <= -R[i]*floss[i]) {
if (cost) {
s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[c];
s += -0.5*R[i]*floss[i]*floss[i] - floss[i]*jar[i];
}
force[i] = floss[i];
d->efc_state[i] = mjCNSTRSTATE_LINEARNEG;
state[i] = mjCNSTRSTATE_LINEARNEG;
}
// linear positive
else if (jar[c] >= R[i]*floss[i]) {
else if (jar[i] >= R[i]*floss[i]) {
if (cost) {
s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[c];
s += -0.5*R[i]*floss[i]*floss[i] + floss[i]*jar[i];
}
force[i] = -floss[i];
d->efc_state[i] = mjCNSTRSTATE_LINEARPOS;
state[i] = mjCNSTRSTATE_LINEARPOS;
}
// quadratic
else {
if (cost) {
s += 0.5*D[i]*jar[c]*jar[c];
s += 0.5*D[i]*jar[i]*jar[i];
}
d->efc_state[i] = mjCNSTRSTATE_QUADRATIC;
state[i] = mjCNSTRSTATE_QUADRATIC;
}
continue;
}
@@ -2459,36 +2353,35 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
// ==== contact
// non-negative constraint
if (d->efc_type[i] != mjCNSTR_CONTACT_ELLIPTIC) {
if (type[i] != mjCNSTR_CONTACT_ELLIPTIC) {
// constraint is satisfied: no cost
if (jar[c] >= 0) {
if (jar[i] >= 0) {
force[i] = 0;
d->efc_state[i] = mjCNSTRSTATE_SATISFIED;
state[i] = mjCNSTRSTATE_SATISFIED;
}
// quadratic
else {
if (cost) {
s += 0.5*D[i]*jar[c]*jar[c];
s += 0.5*D[i]*jar[i]*jar[i];
}
d->efc_state[i] = mjCNSTRSTATE_QUADRATIC;
state[i] = mjCNSTRSTATE_QUADRATIC;
}
}
// contact with elliptic cone
else {
// get contact
mjContact* con = d->contact + d->efc_id[i];
mjContact* con = contact + id[i];
mjtNum mu = con->mu, *friction = con->friction;
int dim = con->dim;
// map to regular dual cone space
mjtNum U[6];
U[0] = jar[c]*mu;
U[0] = jar[i]*mu;
for (int j=1; j < dim; j++) {
U[j] = jar[c+j]*friction[j-1];
U[j] = jar[i+j]*friction[j-1];
}
// decompose into normal and tangent
@@ -2498,19 +2391,17 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
// top zone
if (N >= mu*T || (T <= 0 && N >= 0)) {
mju_zero(force+i, dim);
d->efc_state[i] = mjCNSTRSTATE_SATISFIED;
state[i] = mjCNSTRSTATE_SATISFIED;
}
// bottom zone
else if (mu*N+T <= 0 || (T <= 0 && N < 0)) {
if (cost) {
for (int j=0; j < dim; j++) {
s += 0.5*D[i+j]*jar[c+j]*jar[c+j];
s += 0.5*D[i+j]*jar[i+j]*jar[i+j];
}
}
d->efc_state[i] = mjCNSTRSTATE_QUADRATIC;
state[i] = mjCNSTRSTATE_QUADRATIC;
}
// middle zone
@@ -2530,12 +2421,12 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
}
// set state
d->efc_state[i] = mjCNSTRSTATE_CONE;
state[i] = mjCNSTRSTATE_CONE;
// cone Hessian
if (flg_coneHessian) {
// get Hessian pointer
mjtNum* H = d->contact[d->efc_id[i]].H;
mjtNum* H = contact[id[i]].H;
// set first row: (1, -mu/T * U)
mjtNum scl = -mu/T;
@@ -2546,10 +2437,11 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
// set upper block: mu*N/T^3 * U*U'
scl = mu*N/(T*T*T);
for (int k=1; k < dim; k++)
for (int k=1; k < dim; k++) {
for (int j=k; j < dim; j++) {
H[k*dim+j] = scl*U[j]*U[k];
}
}
// add to diagonal: (mu^2 - mu*N/T) * I
scl = mu*mu - mu*N/T;
@@ -2576,19 +2468,14 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
// replicate state in all cone dimensions
for (int j=1; j < dim; j++) {
d->efc_state[i+j] = d->efc_state[i];
state[i+j] = state[i];
}
// advance to end of contact
c += (dim-1);
i += (dim-1);
}
}
// compute qfrc_constraint
int flg_vecunc = 1;
int flg_resunc = 1;
mj_mulJacTVec_island(m, d, d->qfrc_constraint, d->efc_force, island, flg_vecunc, flg_resunc);
// assign cost
if (cost) {
*cost = s;
@@ -2601,5 +2488,8 @@ void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians
void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
mjtNum cost[1], int flg_coneHessian) {
mj_constraintUpdate_island(m, d, jar, cost, flg_coneHessian, -1);
mj_constraintUpdate_impl(d->ne, d->nf, d->nefc, d->efc_D, d->efc_R, d->efc_frictionloss,
jar, d->efc_type, d->efc_id, d->contact, d->efc_state, d->efc_force,
cost, flg_coneHessian);
mj_mulJacTVec(m, d, d->qfrc_constraint, d->efc_force);
}