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);
}
src/engine/engine_core_constraint.h
+10 -10
View File
@@ -24,6 +24,7 @@
extern "C" {
#endif
//-------------------------- Jacobian-related ------------------------------------------------------
// determine type of friction cone
@@ -38,16 +39,9 @@ MJAPI int mj_isDual(const mjModel* m);
// multiply Jacobian by vector
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, 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, int flg_resunc, int flg_vecunc);
//-------------------------- utility functions -----------------------------------------------------
@@ -90,6 +84,7 @@ void mj_diagApprox(const mjModel* m, mjData* d);
// compute efc_R, efc_D, efc_KDIP, adjust diagApprox
void mj_makeImpedance(const mjModel* m, mjData* d);
//---------------------------- top-level API for constraint construction ---------------------------
// main driver: call all functions above
@@ -101,14 +96,19 @@ MJAPI void mj_projectConstraint(const mjModel* m, mjData* d);
// compute efc_vel, efc_aref
MJAPI void mj_referenceConstraint(const mjModel* m, mjData* d);
// compute efc_state, efc_force
// optional: cost(qacc) = shat(jar); cone Hessians
MJAPI 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);
// compute efc_state, efc_force, qfrc_constraint
// optional: cost(qacc) = shat(jar) where jar = Jac*qacc-aref; cone Hessians
MJAPI void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
mjtNum cost[1], int flg_coneHessian);
// compute efc_state, efc_force, qfrc_constraint for one island
MJAPI void mj_constraintUpdate_island(const mjModel* m, mjData* d, const mjtNum* jar,
mjtNum cost[1], int flg_coneHessian, int island);
#ifdef __cplusplus
}
src/engine/engine_core_smooth.c
+5 -64
View File
@@ -1803,7 +1803,7 @@ void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n,
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n,
void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind) {
// x <- L^-T x
for (int i=nv-1; i > 0; i--) {
@@ -1819,7 +1819,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv,
int start = rowadr[i];
int end = start + rownnz[i] - 1;
for (int adr=start; adr < end; adr++) {
x[colind[adr]] -= qLDs[adr] * x_i;
x[colind[adr]] -= qLD[adr] * x_i;
}
}
}
@@ -1832,7 +1832,7 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv,
mjtNum x_i;
if ((x_i = x[i+offset])) {
for (int adr=start; adr < end; adr++) {
x[offset + colind[adr]] -= qLDs[adr] * x_i;
x[offset + colind[adr]] -= qLD[adr] * x_i;
}
}
}
@@ -1870,13 +1870,13 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLDs, const mjtNum* qLDiagInv,
// one vector
if (n == 1) {
x[i] -= mju_dotSparse(qLDs+adr, x, d, colind+adr, /*flg_unc1=*/0);
x[i] -= mju_dotSparse(qLD+adr, x, d, colind+adr, /*flg_unc1=*/0);
}
// multiple vectors
else {
for (int offset=0; offset < n*nv; offset+=nv) {
x[i+offset] -= mju_dotSparse(qLDs+adr, x+offset, d, colind+adr, /*flg_unc1=*/0);
x[i+offset] -= mju_dotSparse(qLD+adr, x+offset, d, colind+adr, /*flg_unc1=*/0);
}
}
}
@@ -1896,65 +1896,6 @@ 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) {
// if no islands, call mj_solveLD
const mjtNum* qLD = d->qLD;
const mjtNum* qLDiagInv = d->qLDiagInv;
if (island < 0) {
mj_solveLD(x, qLD, qLDiagInv, m->nv, 1,
d->M_rownnz, d->M_rowadr, m->dof_simplenum, d->M_colind);
return;
}
// local copies of key variables
const int* rownnz = d->M_rownnz;
const int* rowadr = d->M_rowadr;
const int* colind = d->M_colind;
const int* diagnum = m->dof_simplenum;
// local constants: island specific
int ndof = d->island_dofnum[island];
const int* dofind = d->island_dofind + d->island_dofadr[island];
const int* islandind = d->dof_islandind;
// x <- inv(L') * x; skip simple, exploit sparsity of input vector
for (int k=ndof-1; k >= 0; k--) {
int i = dofind[k];
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]]] -= qLD[adr] * x_k;
}
}
}
// x <- inv(D) * x
for (int k=ndof-1; k >= 0; k--) {
x[k] *= qLDiagInv[dofind[k]]; // x(i) /= L(i,i)
}
// x <- inv(L) * x; skip simple
for (int k=0; k < ndof; k++) {
int i = dofind[k];
// 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]]] * qLD[adr];
}
}
}
// half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
src/engine/engine_core_smooth.h
+1 -4
View File
@@ -71,15 +71,12 @@ MJAPI void mj_solveLD_legacy(const mjModel* m, mjtNum* x, int n,
// in-place sparse backsubstitution: x = inv(L'*D*L)*x
// handle n vectors at once
MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLDs, const mjtNum* qLDiagInv, int nv, int n,
MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n,
const int* rownnz, const int* rowadr, const int* diagnum, const int* colind);
// 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);
// 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);
// 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,
const mjtNum* sqrtInvD, int n);
src/engine/engine_forward.c
+24 -9
View File
@@ -631,10 +631,10 @@ static void warmstart(const mjModel* m, mjData* d) {
// have island structure: unconstrained qacc = qacc_smooth
if (d->nisland > 0) {
for (int i=0; i < nv; i++) {
if (d->dof_island[i] < 0) {
d->qacc[i] = d->qacc_smooth[i];
}
// loop over unconstrained dofs in map_idof2dof[nidof, nv)
for (int i=d->nidof; i < nv; i++) {
int dof = d->map_idof2dof[i];
d->qacc[dof] = d->qacc_smooth[dof];
}
}
@@ -723,22 +723,37 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
// check if islands are supported
int islands_supported = mjENABLED(mjENBL_ISLAND) &&
d->nisland > 0 &&
nisland > 0 &&
m->opt.solver == mjSOL_CG &&
m->opt.noslip_iterations == 0;
// run solver over constraint islands
if (islands_supported) {
// no threadpool, loop over islands
int nidof = d->nidof;
// copy CG inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
mju_gather(d->ifrc_smooth, d->qfrc_smooth, d->map_idof2dof, nidof);
mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->iacc_smooth, d->qacc_smooth, d->map_idof2dof, nidof);
mju_gather(d->iacc, d->qacc, d->map_idof2dof, nidof);
mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc);
mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc);
// solve per island
if (!d->threadpool) {
// no threadpool, loop over islands
for (int island=0; island < nisland; island++) {
mj_solCG_island(m, d, island, m->opt.iterations);
}
}
else {
// solve using threads
} else {
// have threadpool, solve using threads
mj_solCG_island_multithreaded(m, d);
}
// copy back solver outputs (scatter dofs since ni <= nv)
mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof);
mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
}
// run solver over all constraints
src/engine/engine_io.c
+1
View File
@@ -1917,6 +1917,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
d->nJ = 0;
d->nA = 0;
d->nisland = 0;
d->nidof = 0;
// clear global properties
d->time = 0;
src/engine/engine_island.c
+216 -92
View File
@@ -16,6 +16,7 @@
#include <stdio.h>
#include <stddef.h>
#include <string.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
@@ -26,12 +27,65 @@
#include "engine/engine_support.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_sparse.h"
#ifdef MEMORY_SANITIZER
#include <sanitizer/msan_interface.h>
#endif
//-------------------------- local utilities -------------------------------------------------------
// clear island-related arena pointers in mjData
static void clearIsland(mjData* d, size_t parena) {
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS_ISLAND
#undef X
d->nefc = 0;
d->nisland = 0;
d->nidof = 0;
d->parena = parena;
// poison remaining memory
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
(char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
#endif
}
// allocate island arrays on arena, return 1 on success, 0 on failure
static int arenaAllocIsland(const mjModel* m, mjData* d) {
#undef MJ_M
#define MJ_M(n) m->n
#undef MJ_D
#define MJ_D(n) d->n
size_t parena_old = d->parena;
#define X(type, name, nr, nc) \
d->name = mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \
if (!d->name) { \
mj_warning(d, mjWARN_CNSTRFULL, d->narena); \
clearIsland(d, parena_old); \
return 0; \
}
MJDATA_ARENA_POINTERS_ISLAND
#undef X
#undef MJ_M
#define MJ_M(n) n
#undef MJ_D
#define MJ_D(n) n
return 1;
}
//-------------------------- flood-fill and graph construction ------------------------------------
// find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
// arguments:
@@ -87,54 +141,6 @@ int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, cons
// clear island-related arena pointers in mjData
static void clearIsland(mjData* d, size_t parena) {
#define X(type, name, nr, nc) d->name = NULL;
MJDATA_ARENA_POINTERS_ISLAND
#undef X
d->nefc = 0;
d->nisland = 0;
d->parena = parena;
// poison remaining memory
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
(char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
#endif
}
// allocate island arrays on arena, return 1 on success, 0 on failure
static int arenaAllocIsland(const mjModel* m, mjData* d) {
#undef MJ_M
#define MJ_M(n) m->n
#undef MJ_D
#define MJ_D(n) d->n
size_t parena_old = d->parena;
#define X(type, name, nr, nc) \
d->name = mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), _Alignof(type)); \
if (!d->name) { \
mj_warning(d, mjWARN_CNSTRFULL, d->narena); \
clearIsland(d, parena_old); \
return 0; \
}
MJDATA_ARENA_POINTERS_ISLAND
#undef X
#undef MJ_M
#define MJ_M(n) n
#undef MJ_D
#define MJ_D(n) n
return 1;
}
// return upper bound on number of tree-tree edges
static int countMaxEdge(const mjModel* m, const mjData* d) {
int nedge_max = 0;
@@ -411,14 +417,17 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
//-------------------------- main entry-point -----------------------------------------------------
// discover islands:
// nisland, island_dofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext
// nisland, island_idofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext
void mj_island(const mjModel* m, mjData* d) {
int nv = m->nv, nefc = d->nefc, ntree=m->ntree;
// no constraints: quick return
if (!nefc || m->nflex) { // TODO: add flex support to island discovery
d->nisland = 0;
d->nidof = 0;
return;
}
@@ -454,86 +463,201 @@ void mj_island(const mjModel* m, mjData* d) {
int* stack = mjSTACKALLOC(d, nedge, int);
d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack);
// no islands found: quick return
if (!d->nisland) {
d->nidof = 0;
mj_freeStack(d);
return;
}
// count ni: total number of dofs in islands
int nidof = 0;
for (int i=0; i < nv; i++) {
nidof += (tree_island[m->dof_treeid[i]] >= 0);
}
d->nidof = nidof;
// allocate island arrays on arena
if (!arenaAllocIsland(m, d)) {
mj_freeStack(d);
return;
}
int nisland = d->nisland; // local copy
// local copy
int nisland = d->nisland;
// compute dof_island, island_dofnum
int num_dof_unc = 0; // number of unconstrained dofs
mju_zeroInt(d->island_dofnum, nisland);
// ------------------------------------- degrees of freedom --------------------------------------
// compute dof_island, island_nv
mju_zeroInt(d->island_nv, nisland);
for (int i=0; i < nv; i++) {
// dof_island
int island = tree_island[m->dof_treeid[i]];
// assign dofs to islands
int island = tree_island[m->dof_treeid[i]]; // -1 if unconstrained
d->dof_island[i] = island;
// island_dofnum
// increment island_nv
if (island >= 0) {
d->island_dofnum[island]++;
} else {
num_dof_unc++;
d->island_nv[island]++;
}
}
// compute island_dofadr
if (nisland) d->island_dofadr[0] = 0;
// compute island_idofadr (cumsum of island_nv)
d->island_idofadr[0] = 0;
for (int i=1; i < nisland; i++) {
d->island_dofadr[i] = d->island_dofadr[i-1] + d->island_dofnum[i-1];
d->island_idofadr[i] = d->island_idofadr[i-1] + d->island_nv[i-1];
}
// reset island_dofnum
mju_zeroInt(d->island_dofnum, nisland);
// compute dof_islandind, island_dofind
int num_dof_island = 0;
for (int i=0; i < nv; i++) {
int island = d->dof_island[i];
// compute dof <-> idof maps
int* island_nv2 = mjSTACKALLOC(d, nisland + 1, int); // last element counts unconstrained dofs
mju_zeroInt(island_nv2, nisland + 1);
for (int dof=0; dof < nv; dof++) {
int island = d->dof_island[dof];
int idof;
if (island >= 0) {
d->island_dofind[d->island_dofadr[island] + d->island_dofnum[island]] = i;
d->dof_islandind[i] = d->island_dofnum[island]++;
num_dof_island++;
// constrained dof
idof = d->island_idofadr[island] + island_nv2[island]++;
} else {
d->dof_islandind[i] = -1;
// unconstrained dof
idof = nidof + island_nv2[nisland]++;
}
d->map_dof2idof[dof] = idof;
d->map_idof2dof[idof] = dof; // only the first ni elements of map_idof2dof are in some island
}
// sanity check, SHOULD NOT OCCUR
if (num_dof_island + num_dof_unc != nv) {
mjERROR("not all islands assigned to dofs");
// SHOULD NOT OCCUR
if (!mju_compare(island_nv2, d->island_nv, nisland)) mjERROR("island_nv miscount");
if (nidof + island_nv2[nisland] != nv) mjERROR("miscount of unconstrained dofs");
// compute island_dofadr (used for visualization)
for (int i=0; i < nisland; i++) {
d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]];
}
// finalize dof_islandind: set remaining indices to -1
for (int i=num_dof_island; i < nv; i++) {
d->island_dofind[i] = -1;
// local CSR copy of qM
mjtNum* qM = mjSTACKALLOC(d, m->nM, mjtNum);
mju_gather(qM, d->qM, d->mapM2M, m->nM);
// inertia: block-diagonalize both iLD <- qLD and iM <- qM
mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind,
d->qLD, d->M_rownnz, d->M_rowadr, d->M_colind,
nidof, nisland,
d->map_idof2dof, d->map_dof2idof,
d->island_idofadr, d->island_idofadr,
d->iM, qM);
mju_gather(d->iLDiagInv, d->qLDiagInv, d->map_idof2dof, nidof);
// compute iM_diagnum (dof_simplenum per island)
int count = 0;
int dof_next = d->map_idof2dof[nidof-1];
for (int i=nidof-1; i >= 0; i--) {
// check if island boundary was crossed
int dof = d->map_idof2dof[i];
int island_boundary = (d->dof_island[dof] != d->dof_island[dof_next]);
dof_next = dof;
// accumulate and set simple dof (diagonal row) counter
if (m->dof_simplenum[dof] && !island_boundary) {
count++; // increment counter
} else {
count = 0; // reset
}
d->iM_diagnum[i] = count;
}
// compute efc_island, island_efcnum
mju_zeroInt(d->island_efcnum, nisland);
// ------------------------------------- constraints ---------------------------------------------
// compute efc_island, island_{ne,nf,nefc}
mju_zeroInt(d->island_ne, nisland);
mju_zeroInt(d->island_nf, nisland);
mju_zeroInt(d->island_nefc, nisland);
for (int i=0; i < nefc; i++) {
int tree[2];
treeFirst(m, d, tree, i);
int island = tree_island[tree[0]];
d->efc_island[i] = island;
d->island_efcnum[island]++;
d->island_nefc[island]++;
switch (d->efc_type[i]) {
case mjCNSTR_EQUALITY:
d->island_ne[island]++;
break;
case mjCNSTR_FRICTION_DOF:
case mjCNSTR_FRICTION_TENDON:
d->island_nf[island]++;
break;
default:
break;
}
}
// compute island_efcadr
if (nisland) d->island_efcadr[0] = 0;
// compute island_iefcadr (cumsum of island_nefc)
d->island_iefcadr[0] = 0;
for (int i=1; i < nisland; i++) {
d->island_efcadr[i] = d->island_efcadr[i-1] + d->island_efcnum[i-1];
d->island_iefcadr[i] = d->island_iefcadr[i-1] + d->island_nefc[i-1];
}
// reset island_efcnum
mju_zeroInt(d->island_efcnum, nisland);
// compute efc_islandind
for (int i=0; i < nefc; i++) {
int island = d->efc_island[i];
d->island_efcind[d->island_efcadr[island] + (d->island_efcnum[island]++)] = i;
// compute efc <-> iefc maps
int* island_nefc2 = island_nv2; // reuse island_nv2
mju_zeroInt(island_nefc2, nisland);
for (int c=0; c < nefc; c++) {
int island = d->efc_island[c];
int ic = d->island_iefcadr[island] + island_nefc2[island]++;
d->map_efc2iefc[c] = ic;
d->map_iefc2efc[ic] = c;
}
// SHOULD NOT OCCUR
if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount");
// dense: block-diagonalize Jacobian
if (!mj_isSparse(m)) {
mju_blockDiag(d->iefc_J, d->efc_J,
nv, nidof, nisland,
d->map_iefc2efc, d->map_idof2dof,
d->island_nefc, d->island_nv,
d->island_iefcadr, d->island_idofadr);
}
// sparse
else {
// block-diagonalize Jacobian
mju_blockDiagSparse(d->iefc_J, d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind,
d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
nefc, nisland,
d->map_iefc2efc, d->map_dof2idof,
d->island_iefcadr, d->island_idofadr, NULL, NULL);
// recompute rowsuper per island
for (int island=0; island < nisland; island++) {
int adr = d->island_iefcadr[island];
mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr,
d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind);
}
// block-diagonalize Jacobian-transpose
mju_blockDiagSparse(d->iefc_JT, d->iefc_JT_rownnz, d->iefc_JT_rowadr, d->iefc_JT_colind,
d->efc_JT, d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind,
nidof, nisland,
d->map_idof2dof, d->map_efc2iefc,
d->island_idofadr, d->island_iefcadr, NULL, NULL);
// recompute rowsuper per island
for (int island=0; island < nisland; island++) {
int adr = d->island_idofadr[island];
mju_superSparse(d->island_nv[island], d->iefc_JT_rowsuper + adr,
d->iefc_JT_rownnz + adr, d->iefc_JT_rowadr + adr, d->iefc_JT_colind);
}
}
// copy position-dependent efc vectors required by solver
mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc);
mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc);
mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, nefc);
mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, nefc);
mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, nefc);
mj_freeStack(d);
}
src/engine/engine_print.c
+29 -17
View File
@@ -1392,27 +1392,30 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
}
fprintf(fp, "\n\n");
fprintf(fp, NAME_FORMAT, "ISLAND_DOFNUM");
fprintf(fp, NAME_FORMAT, "ISLAND_NV");
for (int i = 0; i < d->nisland; i++) {
fprintf(fp, " %d", d->island_dofnum[i]);
fprintf(fp, " %d", d->island_nv[i]);
}
fprintf(fp, "\n\n");
fprintf(fp, NAME_FORMAT, "ISLAND_DOFADR");
fprintf(fp, NAME_FORMAT, "ISLAND_IDOFADR");
for (int i = 0; i < d->nisland; i++) {
fprintf(fp, " %d", d->island_dofadr[i]);
fprintf(fp, " %d", d->island_idofadr[i]);
}
fprintf(fp, "\n\n");
fprintf(fp, NAME_FORMAT, "ISLAND_DOFIND");
fprintf(fp, NAME_FORMAT, "MAP_IDOF2DOF");
for (int i = 0; i < m->nv; i++) {
fprintf(fp, " %d", d->island_dofind[i]);
}
fprintf(fp, "\n\n");
int dof = d->map_idof2dof[i];
if (i > 0) {
int dofprev = d->map_idof2dof[i-1];
fprintf(fp, NAME_FORMAT, "DOF_ISLANDIND");
for (int i = 0; i < m->nv; i++) {
fprintf(fp, " %d", d->dof_islandind[i]);
// print '|' at island boundaries
if (d->dof_island[dof] != d->dof_island[dofprev]) {
fprintf(fp, " |");
}
}
fprintf(fp, " %d", dof);
}
fprintf(fp, "\n\n");
@@ -1422,21 +1425,30 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
}
fprintf(fp, "\n\n");
fprintf(fp, NAME_FORMAT, "ISLAND_EFCNUM");
fprintf(fp, NAME_FORMAT, "ISLAND_NEFC");
for (int i = 0; i < d->nisland; i++) {
fprintf(fp, " %d", d->island_efcnum[i]);
fprintf(fp, " %d", d->island_nefc[i]);
}
fprintf(fp, "\n\n");
fprintf(fp, NAME_FORMAT, "ISLAND_EFCADR");
fprintf(fp, NAME_FORMAT, "ISLAND_IEFCADR");
for (int i = 0; i < d->nisland; i++) {
fprintf(fp, " %d", d->island_efcadr[i]);
fprintf(fp, " %d", d->island_iefcadr[i]);
}
fprintf(fp, "\n\n");
fprintf(fp, NAME_FORMAT, "ISLAND_EFCIND");
fprintf(fp, NAME_FORMAT, "MAP_IEFC2EFC");
for (int i = 0; i < d->nefc; i++) {
fprintf(fp, " %d", d->island_efcind[i]);
int efc = d->map_iefc2efc[i];
if (i > 0) {
int efcprev = d->map_iefc2efc[i-1];
// print '|' at island boundaries
if (d->efc_island[efc] != d->efc_island[efcprev]) {
fprintf(fp, " |");
}
}
fprintf(fp, " %d", efc);
}
fprintf(fp, "\n\n");
}
src/engine/engine_solver.c
+279 -136
View File
@@ -766,13 +766,55 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
// CG context
struct _mjCGContext {
int flg_Newton; // 1: Newton, 0: CG
// island-related
int island; // current island index, -1 if monolithic
// sizes
int nv; // number of dofs
int nefc; // number of constraints
int* dofind; // dof indices of this island, NULL if monolithic
int* efcind; // constraint indices of this island, NULL if monolithic
int ne; // number of equalities
int nf; // number of friction constraints
int nefc; // number of all constraints
// contact array
mjContact* contact;
// dof arrays
const mjtNum* qfrc_smooth;
const mjtNum* qacc_smooth;
mjtNum* qfrc_constraint;
mjtNum* qacc;
// inertia
const int* M_rownnz;
const int* M_rowadr;
const int* M_diagnum;
const int* M_colind;
const int* dof_Madr;
const int* dof_parentid;
const mjtNum* qM;
const mjtNum* qLD;
const mjtNum* qLDiagInv;
// efc arrays
const mjtNum* efc_D;
const mjtNum* efc_R;
const mjtNum* efc_frictionloss;
const mjtNum* efc_aref;
const int* efc_id;
const int* efc_type;
mjtNum* efc_force;
int* efc_state;
// Jacobians
const int* J_rownnz;
const int* J_rowadr;
const int* J_rowsuper;
const int* J_colind;
const int* JT_rownnz;
const int* JT_rowadr;
const int* JT_rowsuper;
const int* JT_colind;
const mjtNum* J;
const mjtNum* JT;
// common arrays (CGallocate)
mjtNum* Jaref; // Jac*qacc - aref (nefc x 1)
@@ -793,7 +835,7 @@ struct _mjCGContext {
int* L_rownnz; // Hessian factor row nonzeros (nv x 1)
int* L_rowadr; // Hessian factor row addresses (nv x 1)
// Newton arrays, computed-size (HessianMake)
// Newton arrays, computed-size (MakeHessian)
int nH; // number of nonzeros in Hessian H
int* H_colind; // Hessian column indices (nH x 1)
mjtNum* H; // Hessian (nH x 1)
@@ -818,23 +860,130 @@ struct _mjCGContext {
typedef struct _mjCGContext mjCGContext;
// set sizes and pointers to mjData arrays in mjCGContext
static void CGpointers(const mjModel* m, const mjData* d, mjCGContext* ctx, int island) {
int is_sparse = mj_isSparse(m);
ctx->contact = d->contact;
ctx->island = island;
// set sizes and pointers (monolithic)
if (island < 0) {
// sizes
ctx->nv = m->nv;
ctx->ne = d->ne;
ctx->nf = d->nf;
ctx->nefc = d->nefc;
// dof arrays
ctx->qfrc_smooth = d->qfrc_smooth;
ctx->qfrc_constraint = d->qfrc_constraint;
ctx->qacc_smooth = d->qacc_smooth;
ctx->qacc = d->qacc;
// inertia
ctx->M_rownnz = d->M_rownnz;
ctx->M_rowadr = d->M_rowadr;
ctx->M_diagnum = m->dof_simplenum;
ctx->M_colind = d->M_colind;
ctx->dof_Madr = m->dof_Madr;
ctx->dof_parentid = m->dof_parentid;
ctx->qM = d->qM;
ctx->qLD = d->qLD;
ctx->qLDiagInv = d->qLDiagInv;
// efc arrays
ctx->efc_D = d->efc_D;
ctx->efc_R = d->efc_R;
ctx->efc_frictionloss = d->efc_frictionloss;
ctx->efc_aref = d->efc_aref;
ctx->efc_id = d->efc_id;
ctx->efc_type = d->efc_type;
ctx->efc_force = d->efc_force;
ctx->efc_state = d->efc_state;
// Jacobians
ctx->J = d->efc_J;
if (is_sparse) {
ctx->J_rownnz = d->efc_J_rownnz;
ctx->J_rowadr = d->efc_J_rowadr;
ctx->J_rowsuper = d->efc_J_rowsuper;
ctx->J_colind = d->efc_J_colind;
ctx->JT_rownnz = d->efc_JT_rownnz;
ctx->JT_rowadr = d->efc_JT_rowadr;
ctx->JT_rowsuper = d->efc_JT_rowsuper;
ctx->JT_colind = d->efc_JT_colind;
ctx->JT = d->efc_JT;
}
}
// set sizes and pointers (per-island)
else {
// sizes
ctx->nv = d->island_nv[island];
ctx->ne = d->island_ne[island];
ctx->nf = d->island_nf[island];
ctx->nefc = d->island_nefc[island];
// dof arrays
int idofadr = d->island_idofadr[island];
ctx->qfrc_smooth = d->ifrc_smooth + idofadr;
ctx->qfrc_constraint = d->ifrc_constraint + idofadr;
ctx->qacc_smooth = d->iacc_smooth + idofadr;
ctx->qacc = d->iacc + idofadr;
// inertia
ctx->M_rownnz = d->iM_rownnz + idofadr;
ctx->M_rowadr = d->iM_rowadr + idofadr;
ctx->M_diagnum = d->iM_diagnum + idofadr;
ctx->M_colind = d->iM_colind;
ctx->qM = d->iM;
ctx->qLD = d->iLD;
ctx->qLDiagInv = d->iLDiagInv + idofadr;
// efc arrays
int iefcadr = d->island_iefcadr[island];
ctx->efc_D = d->iefc_D + iefcadr;
ctx->efc_R = d->iefc_R + iefcadr;
ctx->efc_frictionloss = d->iefc_frictionloss + iefcadr;
ctx->efc_aref = d->iefc_aref + iefcadr;
ctx->efc_id = d->iefc_id + iefcadr;
ctx->efc_type = d->iefc_type + iefcadr;
ctx->efc_force = d->iefc_force + iefcadr;
ctx->efc_state = d->iefc_state + iefcadr;
// Jacobians
if (!is_sparse) {
ctx->J = d->iefc_J + d->nidof * iefcadr;
} else {
ctx->J_rownnz = d->iefc_J_rownnz + iefcadr;
ctx->J_rowadr = d->iefc_J_rowadr + iefcadr;
ctx->J_rowsuper = d->iefc_J_rowsuper + iefcadr;
ctx->J_colind = d->iefc_J_colind;
ctx->JT_rownnz = d->iefc_JT_rownnz + idofadr;
ctx->JT_rowadr = d->iefc_JT_rowadr + idofadr;
ctx->JT_rowsuper = d->iefc_JT_rowsuper + idofadr;
ctx->JT_colind = d->iefc_JT_colind;
ctx->J = d->iefc_J;
ctx->JT = d->iefc_JT;
}
}
}
// allocate fixed-size arrays in mjCGContext
// mj_{mark/free}Stack in calling function!
static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
int island, int flg_Newton) {
static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx, int island, int flg_Newton) {
// clear everything
memset(ctx, 0, sizeof(mjCGContext));
// get sizes
int nv = island < 0 ? m->nv : d->island_dofnum[island];
int nefc = island < 0 ? d->nefc : d->island_efcnum[island];
// set sizes and pointers
CGpointers(m, d, ctx, island);
// island-related
ctx->island = island;
ctx->nv = nv;
ctx->nefc = nefc;
ctx->dofind = island < 0 ? NULL : d->island_dofind + d->island_dofadr[island];
ctx->efcind = island < 0 ? NULL : d->island_efcind + d->island_efcadr[island];
// local sizes
int nv = ctx->nv;
int nefc = ctx->nefc;
// common arrays
ctx->Jaref = mjSTACKALLOC(d, nefc, mjtNum);
@@ -849,7 +998,7 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
// Newton only, known-size arrays
ctx->flg_Newton = flg_Newton;
if (flg_Newton) {
ctx->D = mjSTACKALLOC(d, nefc, mjtNum);
ctx->D = mjSTACKALLOC(d, nefc, mjtNum);
// sparse Newton only
if (mj_isSparse(m)) {
@@ -866,28 +1015,35 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
// update efc_force, qfrc_constraint, cost-related
static void CGupdateConstraint(const mjModel* m, mjData* d, mjCGContext* ctx) {
static void CGupdateConstraint(mjCGContext* ctx) {
int nefc = ctx->nefc, nv = ctx->nv;
const int* dofind = ctx->dofind;
const int* efcind = ctx->efcind;
// update constraints
mj_constraintUpdate_island(m, d, ctx->Jaref, &(ctx->cost), ctx->flg_Newton, ctx->island);
mj_constraintUpdate_impl(ctx->ne, ctx->nf, ctx->nefc, ctx->efc_D, ctx->efc_R,
ctx->efc_frictionloss, ctx->Jaref, ctx->efc_type, ctx->efc_id,
ctx->contact, ctx->efc_state, ctx->efc_force,
&(ctx->cost), ctx->flg_Newton);
// compute qfrc_constraint (dense or sparse)
if (!ctx->JT) {
mju_mulMatTVec(ctx->qfrc_constraint, ctx->J, ctx->efc_force, nefc, nv);
} else {
mju_mulMatVecSparse(ctx->qfrc_constraint, ctx->JT, ctx->efc_force, nv,
ctx->JT_rownnz, ctx->JT_rowadr, ctx->JT_colind, ctx->JT_rowsuper);
}
// count active and cone
ctx->nactive = 0;
ctx->ncone = 0;
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
ctx->nactive += (d->efc_state[i] != mjCNSTRSTATE_SATISFIED);
ctx->ncone += (d->efc_state[i] == mjCNSTRSTATE_CONE);
for (int i=0; i < nefc; i++) {
ctx->nactive += (ctx->efc_state[i] != mjCNSTRSTATE_SATISFIED);
ctx->ncone += (ctx->efc_state[i] == mjCNSTRSTATE_CONE);
}
// add Gauss cost, set in quadratic[0]
mjtNum Gauss = 0;
for (int c=0; c < nv; c++) {
int i = dofind ? dofind[c] : c;
Gauss += 0.5 * (ctx->Ma[c] - d->qfrc_smooth[i]) * (d->qacc[i] - d->qacc_smooth[i]);
for (int i=0; i < nv; i++) {
Gauss += 0.5 * (ctx->Ma[i] - ctx->qfrc_smooth[i]) * (ctx->qacc[i] - ctx->qacc_smooth[i]);
}
ctx->quadGauss[0] = Gauss;
@@ -895,22 +1051,20 @@ 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, mjData* d, mjCGContext* ctx) {
static void CGupdateGradient(mjCGContext* ctx) {
int nv = ctx->nv;
const int* dofind = ctx->dofind;
// grad = M*qacc - qfrc_smooth - qfrc_constraint
for (int c=0; c < nv; c++) {
int i = dofind ? dofind[c] : c;
ctx->grad[c] = ctx->Ma[c] - d->qfrc_smooth[i] - d->qfrc_constraint[i];
for (int i=0; i < nv; i++) {
ctx->grad[i] = ctx->Ma[i] - ctx->qfrc_smooth[i] - ctx->qfrc_constraint[i];
}
// Newton: Mgrad = H \ grad
// TODO: b/295296178 - add island support to Newton solver
if (ctx->flg_Newton) {
if (mj_isSparse(m)) {
if (ctx->L_rowadr) {
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Lcone : ctx->L),
ctx->grad, nv, ctx->L_rownnz, ctx->L_rowadr, ctx->L_colind);
} else {
@@ -921,44 +1075,32 @@ static void CGupdateGradient(const mjModel* m, mjData* d, mjCGContext* ctx) {
// CG: Mgrad = M \ grad
else {
mju_copy(ctx->Mgrad, ctx->grad, nv);
mj_solveM_island(m, d, ctx->Mgrad, ctx->island);
mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
}
}
// prepare quadratic polynomials and contact cone quantities
static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc, island = ctx->island;
const int* dofind = ctx->dofind;
const int* efcind = ctx->efcind;
static void CGprepare(mjCGContext* ctx) {
int nv = ctx->nv, nefc = ctx->nefc;
const mjtNum* v = ctx->search;
// Gauss: alpha^2*0.5*v'*M*v + alpha*v'*(Ma-qfrc_smooth) + 0.5*(a-qacc_smooth)'*(Ma-qfrc_smooth)
// quadGauss[0] already computed in CGupdateConstraint
mjtNum v_dot_smooth;
if (island < 0) {
v_dot_smooth = mju_dot(d->qfrc_smooth, v, nv);
} else {
v_dot_smooth = 0;
for (int c=0; c < nv; c++) {
v_dot_smooth += d->qfrc_smooth[dofind[c]] * v[c];
}
}
ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - v_dot_smooth;
ctx->quadGauss[1] = mju_dot(v, ctx->Ma, nv) - mju_dot(ctx->qfrc_smooth, v, nv);
ctx->quadGauss[2] = 0.5*mju_dot(v, ctx->Mv, nv);
// process constraints
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
for (int i=0; i < nefc; i++) {
// pointers to numeric data
const mjtNum* Jv = ctx->Jv + c;
const mjtNum* Jaref = ctx->Jaref + c;
const mjtNum* D = d->efc_D + i;
const mjtNum* Jv = ctx->Jv + i;
const mjtNum* Jaref = ctx->Jaref + i;
const mjtNum* D = ctx->efc_D + i;
// pointer to this quadratic
mjtNum* quad = ctx->quad + 3*c;
mjtNum* quad = ctx->quad + 3*i;
// init with scalar quadratic
mjtNum DJ0 = D[0]*Jaref[0];
@@ -967,12 +1109,12 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
quad[2] = Jv[0]*D[0]*Jv[0];
// elliptic cone: extra processing
if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) {
if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) {
// extract contact info
mjContact* con = d->contact + d->efc_id[i];
const mjContact* con = ctx->contact + ctx->efc_id[i];
int dim = con->dim;
mjtNum U[6], V[6], UU = 0, UV = 0, VV = 0, mu = con->mu;
mjtNum* friction = con->friction;
const mjtNum* friction = con->friction;
// complete vector quadratic (for bottom zone)
for (int j=1; j < dim; j++) {
@@ -1006,7 +1148,7 @@ static void CGprepare(const mjModel* m, const mjData* d, mjCGContext* ctx) {
quad[8] = D[0] / ((mu*mu) * (1 + (mu*mu)));
// advance to next constraint
c += (dim-1);
i += (dim-1);
}
// apply scaling
@@ -1028,9 +1170,8 @@ typedef struct _mjCGPnt mjCGPnt;
// evaluate linesearch cost, return first and second derivatives
static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt* p) {
int ne = d->ne, nf = d->nf, nefc = ctx->nefc;
const int* efcind = ctx->efcind;
static void CGeval(mjCGContext* ctx, mjCGPnt* p) {
int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc;
// clear result
mjtNum cost = 0, alpha = p->alpha;
@@ -1041,26 +1182,24 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
mju_copy3(quadTotal, ctx->quadGauss);
// process 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) {
mju_addTo3(quadTotal, ctx->quad+3*c);
mju_addTo3(quadTotal, ctx->quad+3*i);
continue;
}
// friction
if (i < ne + nf) {
// search point, friction loss, bound (Rf)
mjtNum start = ctx->Jaref[c], dir = ctx->Jv[c];
mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i];
mjtNum x = start + alpha*dir;
mjtNum f = d->efc_frictionloss[i];
mjtNum Rf = d->efc_R[i]*f;
mjtNum f = ctx->efc_frictionloss[i];
mjtNum Rf = ctx->efc_R[i]*f;
// -bound < x < bound : quadratic
if (-Rf < x && x < Rf) {
mju_addTo3(quadTotal, ctx->quad+3*c);
mju_addTo3(quadTotal, ctx->quad+3*i);
}
// x < -bound : linear negative
@@ -1078,10 +1217,10 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
}
// limit and contact
if (d->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone
if (ctx->efc_type[i] == mjCNSTR_CONTACT_ELLIPTIC) { // elliptic cone
// extract contact info
mjContact* con = d->contact + d->efc_id[i];
mjtNum* quad = ctx->quad + 3*c;
const mjContact* con = ctx->contact + ctx->efc_id[i];
mjtNum* quad = ctx->quad + 3*i;
int dim = con->dim;
mjtNum mu = con->mu;
@@ -1137,14 +1276,14 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
}
// advance to next constraint
c += (dim-1);
i += (dim-1);
} else { // inequality
// search point
mjtNum x = ctx->Jaref[c] + alpha*ctx->Jv[c];
mjtNum x = ctx->Jaref[i] + alpha*ctx->Jv[i];
// active
if (x < 0) {
mju_addTo3(quadTotal, ctx->quad+3*c);
mju_addTo3(quadTotal, ctx->quad+3*i);
}
}
}
@@ -1170,7 +1309,7 @@ static void CGeval(const mjModel* m, const mjData* d, mjCGContext* ctx, mjCGPnt*
// update bracket point given 3 candidate points
static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
static int updateBracket(mjCGContext* ctx,
mjCGPnt* p, const mjCGPnt candidates[3], mjCGPnt* pnext) {
int flag = 0;
for (int i=0; i < 3; i++) {
@@ -1192,7 +1331,7 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
// compute next point if updated
if (flag) {
pnext->alpha = p->alpha - p->deriv[0]/p->deriv[1];
CGeval(m, d, ctx, pnext);
CGeval(ctx, pnext);
}
return flag;
@@ -1201,8 +1340,8 @@ static int updateBracket(const mjModel* m, const mjData* d, mjCGContext* ctx,
// line search
static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
int nv = ctx->nv;
static mjtNum CGsearch(mjCGContext* ctx, mjtNum tolerance, mjtNum ls_iterations) {
int nv = ctx->nv, nefc = ctx->nefc;
mjCGPnt p0, p1, p2, pmid, p1next, p2next;
// clear results
@@ -1218,23 +1357,36 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
}
// compute scaled gradtol and slope scaling
mjtNum gtol = m->opt.tolerance * m->opt.ls_tolerance * snorm / ctx->scale;
mjtNum gtol = tolerance * snorm / ctx->scale;
mjtNum slopescl = ctx->scale / snorm;
// compute Mv, Jv
mj_mulM_island(m, d, ctx->Mv, ctx->search, ctx->island, /*flg_vecunc=*/0);
mj_mulJacVec_island(m, d, ctx->Jv, ctx->search, ctx->island, /*flg_resunc=*/0, /*flg_vecunc=*/0);
// compute Mv = M * v (island or monolithic)
if (ctx->island >= 0) {
mju_mulSymVecSparse(ctx->Mv, ctx->qM, ctx->search, nv,
ctx->M_rownnz, ctx->M_rowadr, ctx->M_diagnum, ctx->M_colind);
} else {
mj_mulM_impl(ctx->Mv, ctx->search, nv, ctx->qM,
ctx->dof_Madr, ctx->dof_parentid, ctx->M_diagnum);
}
// compute Jv = J * search (dense or sparse)
if (!ctx->J_rowadr) {
mju_mulMatVec(ctx->Jv, ctx->J, ctx->search, nefc, nv);
} else {
mju_mulMatVecSparse(ctx->Jv, ctx->J, ctx->search, nefc,
ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind, ctx->J_rowsuper);
}
// prepare quadratics and cones
CGprepare(m, d, ctx);
CGprepare(ctx);
// init at alpha = 0, save
p0.alpha = 0;
CGeval(m, d, ctx, &p0);
CGeval(ctx, &p0);
// always attempt one Newton step
p1.alpha = p0.alpha - p0.deriv[0]/p0.deriv[1];
CGeval(m, d, ctx, &p1);
CGeval(ctx, &p1);
if (p0.cost < p1.cost) {
p1 = p0;
}
@@ -1289,14 +1441,14 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
// one-sided search
int p2update = 0;
while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < m->opt.ls_iterations) {
while (p1.deriv[0]*dir <= -gtol && ctx->LSiter < ls_iterations) {
// save current
p2 = p1;
p2update = 1;
// move to Newton point w.r.t current
p1.alpha -= p1.deriv[0]/p1.deriv[1];
CGeval(m, d, ctx, &p1);
CGeval(ctx, &p1);
// check for convergence
if (mju_abs(p1.deriv[0]) < gtol) {
@@ -1306,7 +1458,7 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
}
// check for failure to bracket
if (ctx->LSiter >= m->opt.ls_iterations) {
if (ctx->LSiter >= ls_iterations) {
ctx->LSresult = 3; // could not bracket
ctx->LSslope = mju_abs(p1.deriv[0])*slopescl;
return p1.alpha;
@@ -1322,13 +1474,13 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
// compute next-points for bracket
p2next = p1;
p1next.alpha = p1.alpha - p1.deriv[0]/p1.deriv[1];
CGeval(m, d, ctx, &p1next);
CGeval(ctx, &p1next);
// bracketed search
while (ctx->LSiter < m->opt.ls_iterations) {
while (ctx->LSiter < ls_iterations) {
// evaluate at midpoint
pmid.alpha = 0.5*(p1.alpha + p2.alpha);
CGeval(m, d, ctx, &pmid);
CGeval(ctx, &pmid);
// make list of candidates
mjCGPnt candidates[3] = {p1next, p2next, pmid};
@@ -1349,8 +1501,8 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
}
// update brackets
int b1 = updateBracket(m, d, ctx, &p1, candidates, &p1next);
int b2 = updateBracket(m, d, ctx, &p2, candidates, &p2next);
int b1 = updateBracket(ctx, &p1, candidates, &p1next);
int b2 = updateBracket(ctx, &p2, candidates, &p2next);
// no update possible: numerical accuracy reached, use midpoint
if (!b1 && !b2) {
@@ -1730,8 +1882,6 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// local copies
int nv = ctx.nv;
int nefc = ctx.nefc;
const int* dofind = ctx.dofind;
const int* efcind = ctx.efcind;
// allocate local storage
if (!flg_Newton) {
@@ -1741,27 +1891,32 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
}
int* oldstate = mjSTACKALLOC(d, nefc, int);
// initialize matrix-vector products
int flg_vecunc = 1; // d->qacc is uncompressed
mj_mulM_island(m, d, ctx.Ma, d->qacc, island, flg_vecunc);
int flg_resunc = 0; // ctx.Jaref is compressed
mj_mulJacVec_island(m, d, ctx.Jaref, d->qacc, island, flg_resunc, flg_vecunc);
if (island < 0) {
mju_subFrom(ctx.Jaref, d->efc_aref, nefc);
// compute Ma = M * qacc (island or monolithic)
if (island >= 0) {
mju_mulSymVecSparse(ctx.Ma, ctx.qM, ctx.qacc, nv,
ctx.M_rownnz, ctx.M_rowadr, ctx.M_diagnum, ctx.M_colind);
} else {
for (int c=0; c < nefc; c++) {
ctx.Jaref[c] -= d->efc_aref[efcind[c]];
}
mj_mulM_impl(ctx.Ma, ctx.qacc, nv, ctx.qM,
ctx.dof_Madr, ctx.dof_parentid, ctx.M_diagnum);
}
// compute Jaref = J * qacc - aref (dense or sparse)
if (!ctx.J_rownnz) {
mju_mulMatVec(ctx.Jaref, ctx.J, ctx.qacc, nefc, nv);
} else {
mju_mulMatVecSparse(ctx.Jaref, ctx.J, ctx.qacc, nefc,
ctx.J_rownnz, ctx.J_rowadr, ctx.J_colind, ctx.J_rowsuper);
}
mju_subFrom(ctx.Jaref, ctx.efc_aref, nefc);
// first update
CGupdateConstraint(m, d, &ctx);
CGupdateConstraint(&ctx);
if (flg_Newton) {
// compute and factorize Hessian
MakeHessian(m, d, &ctx);
FactorizeHessian(m, d, &ctx, /*flg_recompute=*/0);
}
CGupdateGradient(m, d, &ctx);
CGupdateGradient(&ctx);
// start both with preconditioned gradient
mju_scl(ctx.search, ctx.Mgrad, -1, nv);
@@ -1772,8 +1927,9 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv));
} else {
mjtNum island_inertia = 0;
for (int c=0; c < nv; c++) {
island_inertia += d->qM[m->dof_Madr[dofind[c]]];
for (int i=0; i < nv; i++) {
int* map2dof = d->map_idof2dof + d->island_idofadr[island];
island_inertia += d->qM[m->dof_Madr[map2dof[i]]];
}
scale = 1 / island_inertia;
}
@@ -1782,7 +1938,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// main loop
while (iter < maxiter) {
// perform linesearch
alpha = CGsearch(m, d, &ctx);
alpha = CGsearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations);
// no improvement: done
if (alpha == 0) {
@@ -1790,13 +1946,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
}
// move to new solution
if (island < 0) {
mju_addToScl(d->qacc, ctx.search, alpha, nv);
} else {
for (int c=0; c < nv; c++) {
d->qacc[dofind[c]] += alpha * ctx.search[c];
}
}
mju_addToScl(ctx.qacc, ctx.search, alpha, nv);
mju_addToScl(ctx.Ma, ctx.Mv, alpha, nv);
mju_addToScl(ctx.Jaref, ctx.Jv, alpha, nefc);
@@ -1805,27 +1955,20 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
mju_copy(gradold, ctx.grad, nv);
mju_copy(Mgradold, ctx.Mgrad, nv);
}
if (island < 0) {
mju_copyInt(oldstate, d->efc_state, nefc);
} else {
for (int c=0; c < nefc; c++) {
oldstate[c] = d->efc_state[efcind[c]];
}
}
mju_copyInt(oldstate, ctx.efc_state, nefc);
mjtNum oldcost = ctx.cost;
// update
CGupdateConstraint(m, d, &ctx);
CGupdateConstraint(&ctx);
if (flg_Newton) {
HessianIncremental(m, d, &ctx, oldstate);
}
CGupdateGradient(m, d, &ctx);
CGupdateGradient(&ctx);
// count state changes
int nchange = 0;
for (int c=0; c < nefc; c++) {
int i = efcind ? efcind[c] : c;
nchange += (d->efc_state[i] != oldstate[c]);
for (int i=0; i < nefc; i++) {
nchange += (ctx.efc_state[i] != oldstate[i]);
}
// scale improvement, gradient, save stats
@@ -1857,8 +2000,8 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
}
// update
for (int c=0; c < nv; c++) {
ctx.search[c] = -ctx.Mgrad[c] + beta*ctx.search[c];
for (int i=0; i < nv; i++) {
ctx.search[i] = -ctx.Mgrad[i] + beta*ctx.search[i];
}
}
}
src/engine/engine_support.c
+6 -66
View File
@@ -972,14 +972,9 @@ void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M) {
// multiply vector by inertia matrix
void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
int nv = m->nv;
const mjtNum* M = d->qM;
const int* Madr = m->dof_Madr;
const int* parentid = m->dof_parentid;
const int* simplenum = m->dof_simplenum;
// multiply vector by inertia matrix (implementation)
void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M,
const int* Madr, const int* parentid, const int* simplenum) {
mju_zero(res, nv);
for (int i=0; i < nv; i++) {
@@ -1031,64 +1026,9 @@ void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
// multiply vector by inertia matrix for one dof island
void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
int island, int flg_vecunc) {
// if no island, call regular function
if (island < 0) {
mj_mulM(m, d, res, vec);
return;
}
// local constants: general
const mjtNum* M = d->qM;
const int* Madr = m->dof_Madr;
const int* parentid = m->dof_parentid;
const int* simplenum = m->dof_simplenum;
// local constants: island specific
int ndof = d->island_dofnum[island];
const int* dofind = d->island_dofind + d->island_dofadr[island];
const int* islandind = d->dof_islandind;
mju_zero(res, ndof);
for (int k=0; k < ndof; k++) {
// address in full dof vector
int i = dofind[k];
// address in M
int adr = Madr[i];
// diagonal
if (flg_vecunc) {
res[k] = M[adr]*vec[i];
} else {
res[k] = M[adr]*vec[k];
}
// simple dof: continue
if (simplenum[i]) {
continue;
}
// off-diagonal
int j = parentid[i];
while (j >= 0) {
adr++;
int l = islandind[j];
if (flg_vecunc) {
res[k] += M[adr]*vec[j];
res[l] += M[adr]*vec[i];
} else {
res[k] += M[adr]*vec[l];
res[l] += M[adr]*vec[k];
}
// advance to parent
j = parentid[j];
}
}
// multiply vector by inertia matrix
void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
mj_mulM_impl(res, vec, m->nv, d->qM, m->dof_Madr, m->dof_parentid, m->dof_simplenum);
}
src/engine/engine_support.h
+4 -4
View File
@@ -120,13 +120,13 @@ MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
// convert sparse inertia matrix M into full matrix
MJAPI void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
// multiply vector by inertia matrix (implementation)
MJAPI void mj_mulM_impl(mjtNum* res, const mjtNum* vec, int nv, const mjtNum* M,
const int* Madr, const int* parentid, const int* simplenum);
// multiply vector by inertia matrix
MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
// multiply vector by inertia matrix for one dof island
MJAPI void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec,
int island, int flg_vecunc);
// multiply vector by (inertia matrix)^(1/2)
MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
src/engine/engine_util_misc.c
+18
View File
@@ -1408,6 +1408,24 @@ void mju_scatter(mjtNum* restrict res, const mjtNum* restrict vec, const int* re
// gather integers
void mju_gatherInt(int* restrict res, const int* restrict vec, const int* restrict ind, int n) {
for (int i=0; i < n; i++) {
res[i] = vec[ind[i]];
}
}
// scatter integers
void mju_scatterInt(int* restrict res, const int* restrict vec, const int* restrict ind, int n) {
for (int i=0; i < n; i++) {
res[ind[i]] = vec[i];
}
}
// insertion sort, increasing order
void mju_insertionSort(mjtNum* list, int n) {
for (int i=1; i < n; i++) {
src/engine/engine_util_misc.h
+8 -2
View File
@@ -156,12 +156,18 @@ MJAPI void mju_d2n(mjtNum* res, const double* vec, int n);
// convert from mjtNum to double
MJAPI void mju_n2d(double* res, const mjtNum* vec, int n);
// gather
// gather mjtNums
MJAPI void mju_gather(mjtNum* res, const mjtNum* vec, const int* ind, int n);
// scatter
// scatter mjtNums
MJAPI void mju_scatter(mjtNum* res, const mjtNum* vec, const int* ind, int n);
// gather integers
MJAPI void mju_gatherInt(int* res, const int* vec, const int* ind, int n);
// scatter integers
MJAPI void mju_scatterInt(int* res, const int* vec, const int* ind, int n);
// insertion sort, increasing order
MJAPI void mju_insertionSort(mjtNum* list, int n);
src/engine/engine_vis_state.c
-6
View File
@@ -98,7 +98,6 @@ void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnsta
// buffer space required for islands
scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->island_dofadr) * m->ntree);
scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->island_dofind) * m->nv);
scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->dof_island) * m->nv);
scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->efc_island) * maxgeom * condimmax);
scnstate->nbuffer += roundUpToCacheLine(sizeof(*d->tendon_efcadr) * m->ntendon);
@@ -136,9 +135,6 @@ void mjv_makeSceneState(const mjModel* m, const mjData* d, mjvSceneState* scnsta
scnstate->data.island_dofadr = (int*)ptr;
ptr += roundUpToCacheLine(sizeof(*scnstate->data.island_dofadr) * scnstate->model.ntree);
scnstate->data.island_dofind = (int*)ptr;
ptr += roundUpToCacheLine(sizeof(*scnstate->data.island_dofind) * scnstate->model.nv);
scnstate->data.dof_island = (int*)ptr;
ptr += roundUpToCacheLine(sizeof(*scnstate->data.dof_island) * scnstate->model.nv);
@@ -224,7 +220,6 @@ void mjv_assignFromSceneState(const mjvSceneState* scnstate, mjModel* m, mjData*
d->contact = scnstate->data.contact;
d->efc_force = scnstate->data.efc_force;
d->island_dofadr = scnstate->data.island_dofadr;
d->island_dofind = scnstate->data.island_dofind;
d->dof_island = scnstate->data.dof_island;
d->efc_island = scnstate->data.efc_island;
d->tendon_efcadr = scnstate->data.tendon_efcadr;
@@ -385,7 +380,6 @@ void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
scnstate->data.nisland = d->nisland;
if (d->nisland) {
memcpy(scnstate->data.island_dofadr, d->island_dofadr, sizeof(*d->island_dofadr) * d->nisland);
memcpy(scnstate->data.island_dofind, d->island_dofind, sizeof(*d->island_dofind) * m->nv);
memcpy(scnstate->data.dof_island, d->dof_island, sizeof(*d->dof_island) * m->nv);
memcpy(scnstate->data.tendon_efcadr, d->tendon_efcadr, sizeof(*d->tendon_efcadr) * m->ntendon);
}
src/engine/engine_vis_visualize.c
+6 -6
View File
@@ -91,9 +91,9 @@ static void makeLabel(const mjModel* m, mjtObj type, int id, char* label) {
// assign pseudo-random rgba to constraint island using Halton sequence
static void islandColor(float rgba[4], int islanddofadr) {
rgba[0] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 2);
rgba[1] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 3);
rgba[2] = 0.1f + 0.8f*mju_Halton(islanddofadr + 1, 5);
rgba[0] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 2);
rgba[1] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 3);
rgba[2] = 0.1f + 0.9f*mju_Halton(islanddofadr + 1, 5);
rgba[3] = 1;
}
@@ -152,7 +152,7 @@ static void addContactGeom(const mjModel* m, mjData* d, const mjtByte* flags,
// override standard colors if visualizing islands
if (vopt->flags[mjVIS_ISLAND] && d->nisland && efc_adr >= 0) {
// set color using island's first dof
islandColor(thisgeom->rgba, d->island_dofind[d->island_dofadr[d->efc_island[efc_adr]]]);
islandColor(thisgeom->rgba, d->island_dofadr[d->efc_island[efc_adr]]);
}
// otherwise regular colors (different for included and excluded contacts)
@@ -1344,7 +1344,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
int island = d->dof_island[m->body_dofadr[weld_id]];
if (island > -1) {
// color using island's first dof
islandColor(rgba_island, d->island_dofind[d->island_dofadr[island]]);
islandColor(rgba_island, d->island_dofadr[island]);
}
}
}
@@ -1835,7 +1835,7 @@ void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
if (d->tendon_efcadr[i] != -1) {
// set color using island's first dof
int island = d->efc_island[d->tendon_efcadr[i]];
islandColor(rgba_island, d->island_dofind[d->island_dofadr[island]]);
islandColor(rgba_island, d->island_dofadr[island]);
}
}
setMaterial(m, thisgeom, tendon_matid, rgba, vopt->flags);