diff --git a/STYLEGUIDE.md b/STYLEGUIDE.md index 85d598e1..2c80956e 100644 --- a/STYLEGUIDE.md +++ b/STYLEGUIDE.md @@ -82,70 +82,70 @@ are required. #### Braces -- MuJoCo uses -[attached K&R braces](https://en.wikipedia.org/wiki/Indentation_style#Variant:_mandatory_braces), -including for one-line blocks: +- MuJoCo uses + [attached K&R braces](https://en.wikipedia.org/wiki/Indentation_style#Variant:_mandatory_braces), + including for one-line blocks: - ```C - // transpose matrix - void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc) { - for (int i=0; i < nr; i++) { - for (int j=0; j < nc; j++) { - res[j*nr+i] = mat[i*nc+j]; + ```c + // transpose matrix + void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc) { + for (int i=0; i < nr; i++) { + for (int j=0; j < nc; j++) { + res[j*nr+i] = mat[i*nc+j]; + } } } - } - ``` + ``` -- Brace-less single line statements are allowed outside of `engine/` code, for -similar, repeated blocks, that do not contain flow control statements (`return`, -`continue`, etc.). For an example of this exception, inspect the [`mjCModel` -destructor](https://github.com/google-deepmind/mujoco/search?q=repo%3Adeepmind%2Fmujoco+filename%3Auser_model.cc). +- Brace-less single line statements are allowed with the exception of `return` + and `break` statements. For an example of this exception, inspect the + [`mjCModel` destructor](https://github.com/google-deepmind/mujoco/search?q=repo%3Adeepmind%2Fmujoco+filename%3Auser_model.cc). -- Unattached braces are allowed in `if/else` blocks, when inserting a comment -before the `else`: +- Unattached braces are allowed in `if/else` blocks, when inserting an + explanatory comment above the `else`: - ```C - // rotate vector by quaternion - void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]) { - // null quat: copy vec - if (quat[0] == 1 && quat[1] == 0 && quat[2] == 0 && quat[3] == 0) { - mju_copy3(res, vec); + ```c + // rotate vector by quaternion + void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]) { + // null quat: copy vec + if (quat[0] == 1 && quat[1] == 0 && quat[2] == 0 && quat[3] == 0) { + mju_copy3(res, vec); + } + + // regular processing + else { + mjtNum mat[9]; + mju_quat2Mat(mat, quat); + mju_mulMatVec3(res, mat, vec); + } } - - // regular processing - else { - mjtNum mat[9]; - mju_quat2Mat(mat, quat); - mju_mulMatVec3(res, mat, vec); - } - } - ``` + ``` #### Spacing -- MuJoCo encourages judicious use of spacing around operators to promote -readability. For example below, note the lack of spaces around the -multiplication operator, and the aligning spaces in the second and fourth -assignments: +- MuJoCo encourages judicious use of spacing around operators to promote + readability. For example below, note the lack of spaces around the + multiplication operator, and the aligning spaces in the second and fourth + assignments: - ```C - // time-derivative of quaternion, given 3D rotational velocity - void mju_derivQuat(mjtNum res[4], const mjtNum quat[4], const mjtNum vel[3]) { - res[0] = 0.5*(-vel[0]*quat[1] - vel[1]*quat[2] - vel[2]*quat[3]); - res[1] = 0.5*( vel[0]*quat[0] + vel[1]*quat[3] - vel[2]*quat[2]); - res[2] = 0.5*(-vel[0]*quat[3] + vel[1]*quat[0] + vel[2]*quat[1]); - res[3] = 0.5*( vel[0]*quat[2] - vel[1]*quat[1] + vel[2]*quat[0]); - } - ``` + ```c + // time-derivative of quaternion, given 3D rotational velocity + void mju_derivQuat(mjtNum res[4], const mjtNum quat[4], const mjtNum vel[3]) { + res[0] = 0.5*(-vel[0]*quat[1] - vel[1]*quat[2] - vel[2]*quat[3]); + res[1] = 0.5*( vel[0]*quat[0] + vel[1]*quat[3] - vel[2]*quat[2]); + res[2] = 0.5*(-vel[0]*quat[3] + vel[1]*quat[0] + vel[2]*quat[1]); + res[3] = 0.5*( vel[0]*quat[2] - vel[1]*quat[1] + vel[2]*quat[0]); + } + ``` -- Spaces are required around comparison operators. +- Spaces are required around comparison operators. -- Spaces are not allowed around operators in array subscripts `[]` or in - variable initialisation in `for` loops. For example, inspect the - `mju_transpose` implementation above. +- Spaces are not allowed around operators in array subscripts `[]` or in + variable initialisation in `for` loops. For example, inspect the + `mju_transpose` implementation above. -- Two blank lines are required between function implementations in source files. +- Two blank lines are required between function implementations in source + files. #### Variable declarations diff --git a/simulate/simulate.cc b/simulate/simulate.cc index a266c57a..cd31a36e 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -408,10 +408,12 @@ void UpdateProfiler(mj::Simulate* sim, const mjModel* m, const mjData* d) { } sqrt_nnz = mju_sqrt(sqrt_nnz); - // get sizes: nv, nbody, nefc, sqrt(nnz), ncont, iter + // get sizes: nv, nbody, nefc, sqrt(nnz), ncon, iter + int nv = mjENABLED(mjENBL_SLEEP) ? d->nv_awake : m->nv; + int nbody = mjENABLED(mjENBL_SLEEP) ? d->nbody_awake : m->nbody; float sdata[6] = { - static_cast(m->nv), - static_cast(m->nbody), + static_cast(nv), + static_cast(nbody), static_cast(d->nefc), static_cast(sqrt_nnz), static_cast(d->ncon), @@ -708,6 +710,7 @@ void MakePhysicsSection(mj::Simulate* sim) { {mjITEM_EDITNUM, "Noslip Tol", 2, &(opt->noslip_tolerance), "1 0 1"}, {mjITEM_EDITINT, "CCD Iter", 2, &(opt->ccd_iterations), "1 0 1000"}, {mjITEM_EDITNUM, "CCD Tol", 2, &(opt->ccd_tolerance), "1 0 1"}, + {mjITEM_EDITNUM, "Sleep Tol", 2, &(opt->sleep_tolerance), "1 0 1"}, {mjITEM_EDITINT, "SDF Iter", 2, &(opt->sdf_iterations), "1 1 20"}, {mjITEM_EDITINT, "SDF Init", 2, &(opt->sdf_initpoints), "1 1 100"}, {mjITEM_SEPARATOR, "Physical Parameters", mjPRESERVE}, diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c index f474bce5..e02e8190 100644 --- a/src/engine/engine_collision_driver.c +++ b/src/engine/engine_collision_driver.c @@ -156,14 +156,25 @@ static int mj_filterSphere(const mjModel* m, mjData* d, int g1, int g2, mjtNum m } -// filter body pair: 1- discard, 0- proceed -static int filterBodyPair(int weldbody1, int weldparent1, int weldbody2, - int weldparent2, int dsbl_filterparent) { +// filter body pair; 1: discard, 0: proceed +static int filterBodyPair(int weldbody1, int weldparent1, int asleep1, + int weldbody2, int weldparent2, int asleep2, + int dsbl_filterparent) { // same weldbody check if (weldbody1 == weldbody2) { return 1; } + // both asleep check + if (asleep1 && asleep2) { + return 1; + } + + // asleep and static check + if ((asleep1 && !weldbody2) || (asleep2 && !weldbody1)) { + return 1; + } + // weldparent check if ((!dsbl_filterparent && weldbody1 != 0 && weldbody2 != 0) && (weldbody1 == weldparent2 || weldbody2 == weldparent1)) { @@ -1126,6 +1137,7 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) { int npair = 0, nbody = m->nbody, ngeom = m->ngeom; int nvert = m->nflexvert, nflex = m->nflex, nbodyflex = m->nbody + m->nflex; int dsbl_filterparent = mjDISABLED(mjDSBL_FILTERPARENT); + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < nbody; mjtNum cov[9], cen[3], eigval[3], frame[9], quat[4]; // init with pairs involving always-colliding bodies @@ -1138,7 +1150,7 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) { // b1 is world body with geoms, or world-welded body with plane if ((b1 == 0 && m->body_geomnum[b1] > 0) || (m->body_weldid[b1] == 0 && hasPlane(m, b1))) { - // add b1:body pairs that are not welded together + // add b1:b2 pairs that are not welded together for (int b2=0; b2 < nbody; b2++) { // cannot collide if (!canCollide(m, b2)) { @@ -1148,7 +1160,8 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) { // welded together int weld2 = m->body_weldid[b2]; int parent_weld2 = m->body_weldid[m->body_parentid[weld2]]; - if (filterBodyPair(0, 0, weld2, parent_weld2, dsbl_filterparent)) { + int asleep2 = sleep_filter ? d->body_awake[b2] == mjS_ASLEEP : 0; + if (filterBodyPair(0, 0, 1, weld2, parent_weld2, asleep2, dsbl_filterparent)) { continue; } @@ -1230,14 +1243,17 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) { int bf1 = bfid[sappair[i] >> 16]; int bf2 = bfid[sappair[i] & 0xFFFF]; - // body pair: prune based on weld filter + // body pair: prune based on sleep filter and weld filter if (bf1 < nbody && bf2 < nbody) { + int asleep1 = sleep_filter ? d->body_awake[bf1] == mjS_ASLEEP : 0; + int asleep2 = sleep_filter ? d->body_awake[bf2] == mjS_ASLEEP : 0; int weld1 = m->body_weldid[bf1]; int weld2 = m->body_weldid[bf2]; int parent_weld1 = m->body_weldid[m->body_parentid[weld1]]; int parent_weld2 = m->body_weldid[m->body_parentid[weld2]]; - if (filterBodyPair(weld1, parent_weld1, weld2, parent_weld2, + if (filterBodyPair(weld1, parent_weld1, asleep1, + weld2, parent_weld2, asleep2, dsbl_filterparent)) { continue; } @@ -1421,6 +1437,15 @@ void mj_collideGeoms(const mjModel* m, mjData* d, int g1, int g2) { if (ipair >= 0) { g1 = m->pair_geom1[ipair]; g2 = m->pair_geom2[ipair]; + + // sleep filtering for explicit pairs + if (mjENABLED(mjENBL_SLEEP)) { + int b1 = m->geom_bodyid[g1]; + int b2 = m->geom_bodyid[g2]; + if (d->body_awake[b1] != mjS_AWAKE && d->body_awake[b2] != mjS_AWAKE) { + return; + } + } } // order geoms by type diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 3915745f..f8df1155 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -26,6 +26,7 @@ #include "engine/engine_core_util.h" #include "engine/engine_core_smooth.h" #include "engine/engine_memory.h" +#include "engine/engine_sleep.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" @@ -377,6 +378,9 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { return; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + mj_markStack(d); // allocate space @@ -392,10 +396,16 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // find active equality constraints for (int i=0; i < m->neq; i++) { + // skip inactive if (!d->eq_active[i]) { continue; } + // skip sleeping + if (sleep_filter && mj_sleepState(m, d, mjOBJ_EQUALITY, i) == mjS_ASLEEP) { + continue; + } + // get constraint data data = m->eq_data + mjNEQDATA*i; id[0] = m->eq_obj1id[i]; @@ -649,6 +659,9 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) { return; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + mj_markStack(d); // allocate Jacobian @@ -656,21 +669,30 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) { // find frictional dofs for (int i=0; i < nv; i++) { - if (m->dof_frictionloss[i] > 0) { - // prepare Jacobian: sparse or dense - if (issparse) { - jac[0] = 1; - } else { - mju_zero(jac, nv); - jac[i] = 1; - } - - // add constraint - mj_addConstraint(m, d, jac, 0, 0, m->dof_frictionloss[i], - 1, mjCNSTR_FRICTION_DOF, i, - issparse ? 1 : 0, - issparse ? &i : NULL); + // no friction loss: skip + if (!m->dof_frictionloss[i]) { + continue; } + + // sleeping tree: skip + if (sleep_filter && mj_sleepState(m, d, mjOBJ_DOF, i) == mjS_ASLEEP) { + continue; + } + + // prepare Jacobian: sparse or dense + if (issparse) { + jac[0] = 1; + } else { + mju_zero(jac, nv); + jac[i] = 1; + } + + // add constraint + mj_addConstraint(m, d, jac, 0, 0, m->dof_frictionloss[i], + 1, mjCNSTR_FRICTION_DOF, i, + issparse ? 1 : 0, + issparse ? &i : NULL); + } // find frictional tendons @@ -696,7 +718,7 @@ void mj_instantiateFriction(const mjModel* m, mjData* d) { // joint and tendon limits void mj_instantiateLimit(const mjModel* m, mjData* d) { - int side, nv = m->nv, issparse = mj_isSparse(m); + int nv = m->nv, issparse = mj_isSparse(m); mjtNum margin, value, dist, angleAxis[3]; mjtNum *jac; @@ -705,6 +727,9 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { return; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + mj_markStack(d); // allocate Jacobian @@ -712,82 +737,90 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { // find joint limits for (int i=0; i < m->njnt; i++) { - if (m->jnt_limited[i]) { - // get margin - margin = m->jnt_margin[i]; + // no limit: skip + if (!m->jnt_limited[i]) { + continue; + } - // HINGE or SLIDE joint - if (m->jnt_type[i] == mjJNT_SLIDE || m->jnt_type[i] == mjJNT_HINGE) { - // get joint value - value = d->qpos[m->jnt_qposadr[i]]; + // sleeping tree: skip + if (sleep_filter && mj_sleepState(m, d, mjOBJ_JOINT, i) == mjS_ASLEEP) { + continue; + } - // process lower and upper limits - for (side=-1; side <= 1; side+=2) { - // compute distance (negative: penetration) - dist = side * (m->jnt_range[2*i+(side+1)/2] - value); + // get margin + margin = m->jnt_margin[i]; - // detect joint limit - if (dist < margin) { - // prepare Jacobian: sparse or dense - if (issparse) { - jac[0] = -(mjtNum)side; - } else { - mju_zero(jac, nv); - jac[m->jnt_dofadr[i]] = -(mjtNum)side; - } + // HINGE or SLIDE joint + if (m->jnt_type[i] == mjJNT_SLIDE || m->jnt_type[i] == mjJNT_HINGE) { + // get joint value + value = d->qpos[m->jnt_qposadr[i]]; - // add constraint - mj_addConstraint(m, d, jac, &dist, &margin, 0, - 1, mjCNSTR_LIMIT_JOINT, i, - issparse ? 1 : 0, - issparse ? m->jnt_dofadr+i : NULL); - } - } - } - - // BALL joint - else if (m->jnt_type[i] == mjJNT_BALL) { - // convert joint quaternion to axis-angle - int adr = m->jnt_qposadr[i]; - mjtNum quat[4] = {d->qpos[adr], d->qpos[adr+1], d->qpos[adr+2], d->qpos[adr+3]}; - mju_normalize4(quat); - mju_quat2Vel(angleAxis, quat, 1); - - // get rotation angle, normalize - value = mju_normalize3(angleAxis); - - // compute distance, using max of range (negative: penetration) - dist = mju_max(m->jnt_range[2*i], m->jnt_range[2*i+1]) - value; + // process lower and upper limits + for (int side=-1; side <= 1; side+=2) { + // compute distance (negative: penetration) + dist = side * (m->jnt_range[2*i+(side+1)/2] - value); // detect joint limit if (dist < margin) { - // sparse + // prepare Jacobian: sparse or dense if (issparse) { - // prepare dof index array - int chain[3] = { - m->jnt_dofadr[i], - m->jnt_dofadr[i] + 1, - m->jnt_dofadr[i] + 2 - }; - - // prepare Jacobian - mju_scl3(jac, angleAxis, -1); - - // add constraint - mj_addConstraint(m, d, jac, &dist, &margin, 0, - 1, mjCNSTR_LIMIT_JOINT, i, 3, chain); - } - - // dense - else { - // prepare Jacobian + jac[0] = -(mjtNum)side; + } else { mju_zero(jac, nv); - mju_scl3(jac + m->jnt_dofadr[i], angleAxis, -1); - - // add constraint - mj_addConstraint(m, d, jac, &dist, &margin, 0, - 1, mjCNSTR_LIMIT_JOINT, i, 0, 0); + jac[m->jnt_dofadr[i]] = -(mjtNum)side; } + + // add constraint + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_JOINT, i, + issparse ? 1 : 0, + issparse ? m->jnt_dofadr+i : NULL); + } + } + } + + // BALL joint + else if (m->jnt_type[i] == mjJNT_BALL) { + // convert joint quaternion to axis-angle + int adr = m->jnt_qposadr[i]; + mjtNum quat[4] = {d->qpos[adr], d->qpos[adr+1], d->qpos[adr+2], d->qpos[adr+3]}; + mju_normalize4(quat); + mju_quat2Vel(angleAxis, quat, 1); + + // get rotation angle, normalize + value = mju_normalize3(angleAxis); + + // compute distance, using max of range (negative: penetration) + dist = mju_max(m->jnt_range[2*i], m->jnt_range[2*i+1]) - value; + + // detect joint limit + if (dist < margin) { + // sparse + if (issparse) { + // prepare dof index array + int chain[3] = { + m->jnt_dofadr[i] + 0, + m->jnt_dofadr[i] + 1, + m->jnt_dofadr[i] + 2 + }; + + // prepare Jacobian + mju_scl3(jac, angleAxis, -1); + + // add constraint + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_JOINT, i, 3, chain); + } + + // dense + else { + // prepare Jacobian + mju_zero(jac, nv); + mju_scl3(jac + m->jnt_dofadr[i], angleAxis, -1); + + // add constraint + mj_addConstraint(m, d, jac, &dist, &margin, 0, + 1, mjCNSTR_LIMIT_JOINT, i, 0, 0); } } } @@ -801,7 +834,7 @@ void mj_instantiateLimit(const mjModel* m, mjData* d) { margin = m->tendon_margin[i]; // process lower and upper limits - for (side=-1; side <= 1; side+=2) { + for (int side=-1; side <= 1; side+=2) { // compute distance (negative: penetration) dist = side * (m->tendon_range[2*i+(side+1)/2] - value); @@ -913,6 +946,7 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di } } + // frictionless and frictional contacts void mj_instantiateContact(const mjModel* m, mjData* d) { int ispyramid = mj_isPyramidal(m), issparse = mj_isSparse(m), ncon = d->ncon; @@ -1560,6 +1594,9 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { return 0; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + mj_markStack(d); if (nnz) { @@ -1569,110 +1606,118 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) { // find active equality constraints for (int i=0; i < neq; i++) { - if (d->eq_active[i]) { - id[0] = m->eq_obj1id[i]; - id[1] = m->eq_obj2id[i]; - size = 0; - NV = 0; - NV2 = 0; + // skip inactive + if (!d->eq_active[i]) { + continue; + } - // process according to type - switch ((mjtEq) m->eq_type[i]) { - case mjEQ_CONNECT: - size = 3; - if (!nnz) { - break; - } + // skip sleeping + if (sleep_filter && mj_sleepState(m, d, mjOBJ_EQUALITY, i) == mjS_ASLEEP) { + continue; + } - // get body ids if using site semantics - if (m->eq_objtype[i] == mjOBJ_SITE) { - id[0] = m->site_bodyid[id[0]]; - id[1] = m->site_bodyid[id[1]]; - } + id[0] = m->eq_obj1id[i]; + id[1] = m->eq_obj2id[i]; + size = 0; + NV = 0; + NV2 = 0; - NV = mj_jacDifPairCount(m, chain, id[1], id[0], issparse); + // process according to type + switch ((mjtEq) m->eq_type[i]) { + case mjEQ_CONNECT: + size = 3; + if (!nnz) { break; - - case mjEQ_WELD: - size = 6; - if (!nnz) { - break; - } - - // get body ids if using site semantics - if (m->eq_objtype[i] == mjOBJ_SITE) { - id[0] = m->site_bodyid[id[0]]; - id[1] = m->site_bodyid[id[1]]; - } - - NV = mj_jacDifPairCount(m, chain, id[1], id[0], issparse); - break; - - case mjEQ_JOINT: - case mjEQ_TENDON: - size = 1; - if (!nnz) { - break; - } - - for (int j=0; j < 1+(id[1] >= 0); j++) { - if (m->eq_type[i] == mjEQ_JOINT) { - if (!j) { - NV = 1; - chain[0] = m->jnt_dofadr[id[j]]; - } else { - NV2 = 1; - chain2[0] = m->jnt_dofadr[id[j]]; - } - } else { - if (!j) { - NV = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); - } else { - NV2 = d->ten_J_rownnz[id[j]]; - mju_copyInt(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2); - } - } - } - - if (id[1] >= 0) { - NV = mju_combineSparseCount(NV, NV2, chain, chain2); - } - break; - - case mjEQ_FLEX: - flex_edgeadr = m->flex_edgeadr[id[0]]; - flex_edgenum = m->flex_edgenum[id[0]]; - - // init with all edges, subract rigid later - size = flex_edgenum; - - // process edges of this flex - for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { - // rigid: reduce size and skip - if (m->flexedge_rigid[e]) { - size--; - continue; - } - - // accumulate NV if needed - if (nnz) { - int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]]; - int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]]; - NV += mj_jacDifPairCount(m, chain, b1, b2, issparse); - } - } - break; - - default: - // might occur in case of the now-removed distance equality constraint - mjERROR("unknown constraint type type %d", m->eq_type[i]); // SHOULD NOT OCCUR } - // accumulate counts; flex NV already accumulated - ne += mj_addConstraintCount(m, size, NV); - nnze += (m->eq_type[i] == mjEQ_FLEX) ? NV : size*NV; + // get body ids if using site semantics + if (m->eq_objtype[i] == mjOBJ_SITE) { + id[0] = m->site_bodyid[id[0]]; + id[1] = m->site_bodyid[id[1]]; + } + + NV = mj_jacDifPairCount(m, chain, id[1], id[0], issparse); + break; + + case mjEQ_WELD: + size = 6; + if (!nnz) { + break; + } + + // get body ids if using site semantics + if (m->eq_objtype[i] == mjOBJ_SITE) { + id[0] = m->site_bodyid[id[0]]; + id[1] = m->site_bodyid[id[1]]; + } + + NV = mj_jacDifPairCount(m, chain, id[1], id[0], issparse); + break; + + case mjEQ_JOINT: + case mjEQ_TENDON: + size = 1; + if (!nnz) { + break; + } + + for (int j=0; j < 1+(id[1] >= 0); j++) { + if (m->eq_type[i] == mjEQ_JOINT) { + if (!j) { + NV = 1; + chain[0] = m->jnt_dofadr[id[j]]; + } else { + NV2 = 1; + chain2[0] = m->jnt_dofadr[id[j]]; + } + } else { + if (!j) { + NV = d->ten_J_rownnz[id[j]]; + mju_copyInt(chain, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV); + } else { + NV2 = d->ten_J_rownnz[id[j]]; + mju_copyInt(chain2, d->ten_J_colind+d->ten_J_rowadr[id[j]], NV2); + } + } + } + + if (id[1] >= 0) { + NV = mju_combineSparseCount(NV, NV2, chain, chain2); + } + break; + + case mjEQ_FLEX: + flex_edgeadr = m->flex_edgeadr[id[0]]; + flex_edgenum = m->flex_edgenum[id[0]]; + + // init with all edges, subract rigid later + size = flex_edgenum; + + // process edges of this flex + for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { + // rigid: reduce size and skip + if (m->flexedge_rigid[e]) { + size--; + continue; + } + + // accumulate NV if needed + if (nnz) { + int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]]; + int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]]; + NV += mj_jacDifPairCount(m, chain, b1, b2, issparse); + } + } + break; + + default: + // might occur in case of the now-removed distance equality constraint + mjERROR("unknown constraint type type %d", m->eq_type[i]); // SHOULD NOT OCCUR } + + // accumulate counts; flex NV already accumulated + ne += mj_addConstraintCount(m, size, NV); + nnze += (m->eq_type[i] == mjEQ_FLEX) ? NV : size*NV; } if (nnz) { @@ -1693,11 +1738,22 @@ static int mj_nf(const mjModel* m, const mjData* d, int *nnz) { return 0; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + for (int i=0; i < nv; i++) { - if (m->dof_frictionloss[i] > 0) { - nf += mj_addConstraintCount(m, 1, 1); - if (nnz) *nnz += 1; + // no friction loss: skip + if (!m->dof_frictionloss[i]) { + continue; } + + // sleeping tree: skip + if (sleep_filter && !d->tree_awake[m->dof_treeid[i]]) { + continue; + } + + nf += mj_addConstraintCount(m, 1, 1); + if (nnz) *nnz += 1; } for (int i=0; i < ntendon; i++) { @@ -1723,15 +1779,22 @@ static int mj_nl(const mjModel* m, const mjData* d, int *nnz) { return 0; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; for (int i=0; i < m->njnt; i++) { if (!m->jnt_limited[i]) { continue; } + // sleeping tree: skip + if (sleep_filter && !d->tree_awake[m->dof_treeid[m->jnt_dofadr[i]]]) { + continue; + } + margin = m->jnt_margin[i]; - // slider and hinge joint limits can be bilateral, check both sides + // SLIDE and HINGE joint limits can be bilateral, check both sides if (m->jnt_type[i] == mjJNT_SLIDE || m->jnt_type[i] == mjJNT_HINGE) { value = d->qpos[m->jnt_qposadr[i]]; for (side=-1; side <= 1; side+=2) { @@ -1742,6 +1805,8 @@ static int mj_nl(const mjModel* m, const mjData* d, int *nnz) { } } } + + // BALL joint limits are always unilateral else if (m->jnt_type[i] == mjJNT_BALL) { mjtNum angleAxis[3]; int adr = m->jnt_qposadr[i]; @@ -1757,19 +1822,12 @@ static int mj_nl(const mjModel* m, const mjData* d, int *nnz) { } } + // tendon limits for (int i=0; i < ntendon; i++) { - if (m->tendon_limited[i]) { - value = d->ten_length[i]; - margin = m->tendon_margin[i]; - - // tendon limits can be bilateral, check both sides - for (side=-1; side <= 1; side+=2) { - dist = side * (m->tendon_range[2*i+(side+1)/2] - value); - if (dist < margin) { - nl += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); - if (nnz) *nnz += d->ten_J_rownnz[i]; - } - } + int count = tendonLimit(m, d->ten_length, i); + for (int j = 0; j < count; j++) { + nl += mj_addConstraintCount(m, 1, d->ten_J_rownnz[i]); + if (nnz) *nnz += d->ten_J_rownnz[i]; } } @@ -1786,6 +1844,9 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) { return 0; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + mj_markStack(d); int *chain = mjSTACKALLOC(d, m->nv, int); @@ -1804,6 +1865,21 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) { continue; } + // check for contact with sleeping tree; SHOULD NOT OCCUR + if (sleep_filter) { + int g1 = con->geom[0]; + int g2 = con->geom[1]; + if (g1 >= 0 && g2 >= 0) { + int b1 = m->body_weldid[m->geom_bodyid[g1]]; + int b2 = m->body_weldid[m->geom_bodyid[g2]]; + int asleep1 = d->body_awake[b1] == mjS_ASLEEP; + int asleep2 = d->body_awake[b2] == mjS_ASLEEP; + if (asleep1 || asleep2) { + mjERROR("contact %d involves sleeping geom %d", i, asleep1 ? g1 : g2); + } + } + } + // compute NV only if nnz requested int NV = 0; if (nnz) { @@ -1908,8 +1984,7 @@ void mj_makeConstraint(const mjModel* m, mjData* d) { d->tendon_efcadr[i] = -1; } - // reset nefc for the instantiation functions, - // and instantiate all elements of Jacobian + // reset nefc for the instantiation functions, instantiate all elements of Jacobian d->nefc = 0; mj_instantiateEquality(m, d); mj_instantiateFriction(m, d); diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index ac4129b6..6ec5af52 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -25,17 +25,19 @@ #include "engine/engine_crossplatform.h" #include "engine/engine_macro.h" #include "engine/engine_memory.h" +#include "engine/engine_sleep.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" #include "engine/engine_util_sparse.h" #include "engine/engine_util_spatial.h" + //--------------------------- position ------------------------------------------------------------- -// forward kinematics -void mj_kinematics(const mjModel* m, mjData* d) { - int nbody = m->nbody, nsite = m->nsite, ngeom = m->ngeom; +// forward kinematics part 1: bodies +void mj_kinematics1(const mjModel* m, mjData* d) { + int nbody = m->nbody; // set world position and orientation mju_zero3(d->xpos); @@ -46,8 +48,15 @@ void mj_kinematics(const mjModel* m, mjData* d) { d->xmat[0] = d->xmat[4] = d->xmat[8] = 1; d->ximat[0] = d->ximat[4] = d->ximat[8] = 1; + int sleep_filter = mjENABLED(mjENBL_SLEEP); + // compute global cartesian positions and orientations of all bodies for (int i=1; i < nbody; i++) { + // skip static bodies + if (sleep_filter) { + if (d->body_awake[i] == mjS_STATIC) continue; + } + mjtNum xpos[3], xquat[4]; int jntadr = m->body_jntadr[i]; int jntnum = m->body_jntnum[i]; @@ -138,7 +147,7 @@ void mj_kinematics(const mjModel* m, mjData* d) { break; default: - mjERROR("unknown joint type %d", jtype); // SHOULD NOT OCCUR + mjERROR("unknown joint type %d", jtype); // SHOULD NOT OCCUR } // assign xanchor and xaxis @@ -147,53 +156,124 @@ void mj_kinematics(const mjModel* m, mjData* d) { } } - // assign xquat and xpos, construct xmat + // normalize quaternion mju_normalize4(xquat); + + // sleeping body, check for mismatch + if (sleep_filter && jntnum && d->body_awake[i] == mjS_ASLEEP) { + // compare new and existing xpos and xquat + const mjtNum* pos = d->xpos+3*i; + const mjtNum* xq = d->xquat+4*i; + int match = xpos[0] == pos[0] && xpos[1] == pos[1] && xpos[2] == pos[2] && + xquat[0] == xq[0] && xquat[1] == xq[1] && xquat[2] == xq[2] && xquat[3] == xq[3]; + + // match: continue to next body + if (match) { + continue; + } + + // mismatch: mark the tree for waking later (in mj_wake) + else { + d->tree_awake[m->body_treeid[i]] = 1; + } + } + + // assign xquat and xpos, construct xmat mju_copy4(d->xquat+4*i, xquat); mju_copy3(d->xpos+3*i, xpos); mju_quat2Mat(d->xmat+9*i, xquat); } +} + + +// forward kinematics part 2: body inertias, geoms and sites +void mj_kinematics2(const mjModel* m, mjData* d) { + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; // compute/copy Cartesian positions and orientations of body inertial frames - for (int i=1; i < nbody; i++) { + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + mj_local2Global(d, d->xipos+3*i, d->ximat+9*i, m->body_ipos+3*i, m->body_iquat+4*i, i, m->body_sameframe[i]); } // compute/copy Cartesian positions and orientations of geoms - for (int i=0; i < ngeom; i++) { - mj_local2Global(d, d->geom_xpos+3*i, d->geom_xmat+9*i, - m->geom_pos+3*i, m->geom_quat+4*i, - m->geom_bodyid[i], m->geom_sameframe[i]); + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + + // skip geom in sleeping or static body + if (sleep_filter && d->body_awake[i] != mjS_AWAKE) continue; + + int start = m->body_geomadr[i]; + int end = start + m->body_geomnum[i]; + for (int g=start; g < end; g++) { + mj_local2Global(d, d->geom_xpos+3*g, d->geom_xmat+9*g, + m->geom_pos+3*g, m->geom_quat+4*g, + m->geom_bodyid[g], m->geom_sameframe[g]); + } } // compute/copy Cartesian positions and orientations of sites + int nsite = m->nsite; for (int i=0; i < nsite; i++) { + int bodyid = m->site_bodyid[i]; + + // skip site in sleeping or static body + if (sleep_filter && d->body_awake[bodyid] != mjS_AWAKE) continue; + mj_local2Global(d, d->site_xpos+3*i, d->site_xmat+9*i, m->site_pos+3*i, m->site_quat+4*i, - m->site_bodyid[i], m->site_sameframe[i]); + bodyid, m->site_sameframe[i]); } } +// forward kinematics +void mj_kinematics(const mjModel* m, mjData* d) { + mj_kinematics1(m, d); + if (mj_wake(m, d)) { + mj_updateSleep(m, d); + } + mj_kinematics2(m, d); +} + + // map inertias and motion dofs to global frame centered at subtree-CoM void mj_comPos(const mjModel* m, mjData* d) { - int nbody = m->nbody, njnt = m->njnt; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + int nparent = sleep_filter ? d->nparent_awake : m->nbody; // subtree_com: initialize with body moment - for (int i=0; i < nbody; i++) { + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + mju_scl3(d->subtree_com+3*i, d->xipos+3*i, m->body_mass[i]); } // subtree_com: accumulate to parent in backward pass - for (int i=nbody-1; i > 0; i--) { - int j = m->body_parentid[i]; - mju_addTo3(d->subtree_com+3*j, d->subtree_com+3*i); + for (int b=nparent-1; b >= 0; b--) { + int i = sleep_filter ? d->parent_awake_ind[b] : b; + if (!i) continue; + + // accumulate moment to parent, rescale if sleeping + int parent = m->body_parentid[i]; + if (sleep_filter && d->body_awake[i] == mjS_ASLEEP) { + mjtNum child_moment[3]; + mju_scl3(child_moment, d->subtree_com+3*i, m->body_subtreemass[i]); + mju_addTo3(d->subtree_com+3*parent, child_moment); + } else { + mju_addTo3(d->subtree_com+3*parent, d->subtree_com+3*i); + } } // subtree_com: normalize - for (int i=0; i < nbody; i++) { + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + if (m->body_subtreemass[i] < mjMINVAL) { mju_copy3(d->subtree_com+3*i, d->xipos+3*i); } else { @@ -205,55 +285,64 @@ void mj_comPos(const mjModel* m, mjData* d) { mju_zero(d->cinert, 10); // map inertias to frame centered at subtree_com - for (int i=1; i < nbody; i++) { + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + mjtNum offset[3]; mju_sub3(offset, d->xipos+3*i, d->subtree_com+3*m->body_rootid[i]); - mju_inertCom(d->cinert+10*i, m->body_inertia+3*i, d->ximat+9*i, - offset, m->body_mass[i]); + mju_inertCom(d->cinert+10*i, m->body_inertia+3*i, d->ximat+9*i, offset, m->body_mass[i]); } // map motion dofs to global frame centered at subtree_com - for (int j=0; j < njnt; j++) { - // get dof address, body index - int da = 6*m->jnt_dofadr[j]; - int bi = m->jnt_bodyid[j]; + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; - // compute com-anchor vector - mjtNum offset[3], axis[3]; - mju_sub3(offset, d->subtree_com+3*m->body_rootid[bi], d->xanchor+3*j); + int jntnum = m->body_jntnum[i]; + if (!jntnum) continue; - // create motion dof - int skip = 0; - switch ((mjtJoint) m->jnt_type[j]) { - case mjJNT_FREE: - // translation components: x, y, z in global frame - mju_zero(d->cdof+da, 18); - for (int i=0; i < 3; i++) { - d->cdof[da+3+7*i] = 1; + int start = m->body_jntadr[i]; + int end = start + jntnum; + for (int j=start; j < end; j++) { + // get cdof address + int da = 6*m->jnt_dofadr[j]; + + // compute com-anchor vector + mjtNum offset[3], axis[3]; + mju_sub3(offset, d->subtree_com+3*m->body_rootid[i], d->xanchor+3*j); + + // create motion dof + int skip = 0; + switch ((mjtJoint) m->jnt_type[j]) { + case mjJNT_FREE: + // translation components: x, y, z in global frame + mju_zero(d->cdof+da, 18); + d->cdof[da+3+7*0] = 1; + d->cdof[da+3+7*1] = 1; + d->cdof[da+3+7*2] = 1; + + // rotation components: same as ball + skip = 18; + mjFALLTHROUGH; + + case mjJNT_BALL: + for (int k=0; k < 3; k++) { + // I_3 rotation in child frame (assume no subsequent rotations) + axis[0] = d->xmat[9*i + k + 0]; + axis[1] = d->xmat[9*i + k + 3]; + axis[2] = d->xmat[9*i + k + 6]; + + mju_dofCom(d->cdof+da+skip+6*k, axis, offset); + } + break; + + case mjJNT_SLIDE: + mju_dofCom(d->cdof+da, d->xaxis+3*j, 0); + break; + + case mjJNT_HINGE: + mju_dofCom(d->cdof+da, d->xaxis+3*j, offset); + break; } - - // rotation components: same as ball - skip = 18; - mjFALLTHROUGH; - - case mjJNT_BALL: - for (int i=0; i < 3; i++) { - // I_3 rotation in child frame (assume no subsequent rotations) - axis[0] = d->xmat[9*bi+i+0]; - axis[1] = d->xmat[9*bi+i+3]; - axis[2] = d->xmat[9*bi+i+6]; - - mju_dofCom(d->cdof+da+skip+6*i, axis, offset); - } - break; - - case mjJNT_SLIDE: - mju_dofCom(d->cdof+da, d->xaxis+3*j, 0); - break; - - case mjJNT_HINGE: - mju_dofCom(d->cdof+da, d->xaxis+3*j, offset); - break; } } } @@ -261,18 +350,26 @@ void mj_comPos(const mjModel* m, mjData* d) { // compute camera and light positions and orientations void mj_camlight(const mjModel* m, mjData* d) { - mjtNum pos[3], matT[9]; + int ncam = m->ncam, nlight = m->nlight; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; // compute Cartesian positions and orientations of cameras - for (int i=0; i < m->ncam; i++) { - // default processing for fixed mode - mj_local2Global(d, d->cam_xpos+3*i, d->cam_xmat+9*i, - m->cam_pos+3*i, m->cam_quat+4*i, m->cam_bodyid[i], 0); - + for (int i=0; i < ncam; i++) { // get camera body id and target body id int id = m->cam_bodyid[i]; int id1 = m->cam_targetbodyid[i]; + // skip camera if both body and target body are asleep or static + if (sleep_filter && d->body_awake[id] != mjS_AWAKE) { + if (id1 < 0 || d->body_awake[id1] != mjS_AWAKE) { + continue; + } + } + + // default processing for fixed mode + mj_local2Global(d, d->cam_xpos+3*i, d->cam_xmat+9*i, + m->cam_pos+3*i, m->cam_quat+4*i, id, 0); + // adjust for mode switch ((mjtCamLight) m->cam_mode[i]) { case mjCAMLIGHT_FIXED: @@ -297,6 +394,7 @@ void mj_camlight(const mjModel* m, mjData* d) { case mjCAMLIGHT_TARGETBODYCOM: // only if target body is specified if (id1 >= 0) { + mjtNum pos[3]; // get position to look at if (m->cam_mode[i] == mjCAMLIGHT_TARGETBODY) { mju_copy3(pos, d->xpos+3*id1); @@ -305,6 +403,7 @@ void mj_camlight(const mjModel* m, mjData* d) { } // zaxis = -desired camera direction, in global frame + mjtNum matT[9]; mju_sub3(matT+6, d->cam_xpos+3*i, pos); mju_normalize3(matT+6); @@ -326,15 +425,22 @@ void mj_camlight(const mjModel* m, mjData* d) { } // compute Cartesian positions and directions of lights - for (int i=0; i < m->nlight; i++) { - // default processing for fixed mode - mj_local2Global(d, d->light_xpos+3*i, 0, m->light_pos+3*i, 0, m->light_bodyid[i], 0); - mju_rotVecQuat(d->light_xdir+3*i, m->light_dir+3*i, d->xquat+4*m->light_bodyid[i]); - + for (int i=0; i < nlight; i++) { // get light body id and target body id int id = m->light_bodyid[i]; int id1 = m->light_targetbodyid[i]; + // skip light if both body and target body are asleep or static + if (sleep_filter && d->body_awake[id] != mjS_AWAKE) { + if (id1 < 0 || d->body_awake[id1] != mjS_AWAKE) { + continue; + } + } + + // default processing for fixed mode + mj_local2Global(d, d->light_xpos+3*i, 0, m->light_pos+3*i, 0, id, 0); + mju_rotVecQuat(d->light_xdir+3*i, m->light_dir+3*i, d->xquat+4*id); + // adjust for mode switch ((mjtCamLight) m->light_mode[i]) { case mjCAMLIGHT_FIXED: @@ -360,14 +466,15 @@ void mj_camlight(const mjModel* m, mjData* d) { // only if target body is specified if (id1 >= 0) { // get position to look at + mjtNum lookat[3]; if (m->light_mode[i] == mjCAMLIGHT_TARGETBODY) { - mju_copy3(pos, d->xpos+3*id1); + mju_copy3(lookat, d->xpos+3*id1); } else { - mju_copy3(pos, d->subtree_com+3*id1); + mju_copy3(lookat, d->subtree_com+3*id1); } // set dir - mju_sub3(d->light_xdir+3*i, pos, d->light_xpos+3*i); + mju_sub3(d->light_xdir+3*i, lookat, d->light_xpos+3*i); } } @@ -665,9 +772,17 @@ void mj_tendon(const mjModel* m, mjData* d) { mju_zero(J, nten*nv); } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + // loop over tendons int wrapcount = 0; for (int i=0; i < nten; i++) { + // skip sleeping tendon + if (sleep_filter && mj_sleepState(m, d, mjOBJ_TENDON, i) == mjS_ASLEEP) { + continue; + } + // initialize tendon path int adr = m->tendon_adr[i]; d->ten_wrapadr[i] = wrapcount; @@ -994,11 +1109,19 @@ void mj_transmission(const mjModel* m, mjData* d) { // define stack variables required for site transmission, don't allocate mjtNum *jacref = NULL, *moment_tmp = NULL; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < nv; + // compute lengths and moments for (int i=0; i < nu; i++) { rowadr[i] = i == 0 ? 0 : rowadr[i-1] + rownnz[i-1]; int nnz, adr = rowadr[i]; + // skip sleeping actuator + if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) { + rownnz[i] = 0; + continue; + } + // extract info int id = m->actuator_trnid[2*i]; mjtNum* gear = m->actuator_gear+6*i; @@ -1457,9 +1580,16 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { const int* M_rowadr = m->M_rowadr; const int* M_colind = m->M_colind; - for (int k=0; k < ntendon; k++) { - mjtNum armature = m->tendon_armature[k]; + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < nv; + for (int k=0; k < ntendon; k++) { + // skip sleeping tendon + if (sleep_filter && mj_sleepState(m, d, mjOBJ_TENDON, k) == mjS_ASLEEP) { + continue; + } + + mjtNum armature = m->tendon_armature[k]; if (!armature) { continue; } @@ -1512,39 +1642,57 @@ void mj_tendonArmature(const mjModel* m, mjData* d) { // composite rigid body inertia algorithm void mj_crb(const mjModel* m, mjData* d) { - int nv = m->nv, nbody = m->nbody; - // outputs mjtNum* crb = d->crb; mjtNum* M = d->M; // inputs - const mjtNum* cinert = d->cinert; - const mjtNum* cdof = d->cdof; - const mjtNum* dof_M0 = m->dof_M0; - const mjtNum* dof_armature = m->dof_armature; - const int* rownnz = m->M_rownnz; - const int* rowadr = m->M_rowadr; - const int* body_parentid = m->body_parentid; - const int* dof_parentid = m->dof_parentid; - const int* dof_simplenum = m->dof_simplenum; - const int* dof_bodyid = m->dof_bodyid; + const mjtNum* cinert = d->cinert; + const mjtNum* cdof = d->cdof; + const mjtNum* dof_M0 = m->dof_M0; + const mjtNum* dof_armature = m->dof_armature; + const int* body_awake_ind = d->body_awake_ind; + const int* parent_awake_ind = d->parent_awake_ind; + const int* dof_awake_ind = d->dof_awake_ind; + const int* rownnz = m->M_rownnz; + const int* rowadr = m->M_rowadr; + const int* body_parentid = m->body_parentid; + const int* dof_parentid = m->dof_parentid; + const int* dof_simplenum = m->dof_simplenum; + const int* dof_bodyid = m->dof_bodyid; + + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + int nparent = sleep_filter ? d->nparent_awake : m->nbody; + int nv = sleep_filter ? d->nv_awake : m->nv; // crb = cinert - mju_copy(crb, cinert, 10*nbody); + if (!sleep_filter) { + mju_copy(crb, cinert, 10*nbody); + } else { + mju_copyRows(crb, cinert, body_awake_ind, nbody, 10); + } // backward pass over bodies, accumulate composite inertias - for (int i=nbody - 1; i > 0; i--) { - if (body_parentid[i]) { + for (int b = nparent - 1; b >= 0; b--) { + int i = sleep_filter ? parent_awake_ind[b] : b; + if (body_parentid[i] > 0) { mju_addTo(crb + 10*body_parentid[i], crb + 10*i, 10); } } // clear M - mju_zero(M, m->nC); + if (!sleep_filter) { + mju_zero(M, m->nC); + } else { + mju_zeroSparse(M, rownnz, rowadr, dof_awake_ind, nv); + } // dense forward pass over dofs - for (int i=0; i < nv; i++) { + for (int v=0; v < nv; v++) { + int i = sleep_filter ? dof_awake_ind[v] : v; + // simple dof: fixed diagonal inertia int adr = rowadr[i]; if (dof_simplenum[i]) { @@ -1573,7 +1721,7 @@ void mj_makeM(const mjModel* m, mjData* d) { TM_START; mj_crb(m, d); mj_tendonArmature(m, d); - mju_scatter(d->qM, d->M, m->mapM2M, m->nC); + mju_scatter(d->qM, d->M, m->mapM2M, m->nC); // TODO(tassa): scatter only awake dofs TM_END(mjTIMER_POS_INERTIA); } @@ -1644,17 +1792,41 @@ void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd void mj_factorM(const mjModel* m, mjData* d) { TM_START; - mju_copy(d->qLD, d->M, m->nC); - mj_factorI(d->qLD, d->qLDiagInv, m->nv, m->M_rownnz, m->M_rowadr, m->M_colind); + + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + const int* index; + int nv; + + // no sleep filtering: copy everything + if (!sleep_filter) { + index = NULL; + nv = m->nv; + mju_copy(d->qLD, d->M, m->nC); + } + + // sleep filtering: copy only awake dofs + else { + index = d->dof_awake_ind; + nv = d->nv_awake; + mju_copySparse(d->qLD, d->M, m->M_rownnz, m->M_rowadr, d->dof_awake_ind, d->nv_awake); + } + + // factorize + mj_factorI(d->qLD, d->qLDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, index); + TM_ADD(mjTIMER_POS_INERTIA); } -// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd +// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd (with dof skipping) void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv, - const int* rownnz, const int* rowadr, const int* colind) { + const int* rownnz, const int* rowadr, const int* colind, + const int* index) { // backward loop over rows - for (int k=nv-1; k >= 0; k--) { + for (int j=nv-1; j >= 0; j--) { + int k = index ? index[j] : j; + // get row k's address, diagonal index, inverse diagonal value int start = rowadr[k]; int diag = rownnz[k] - 1; @@ -1787,11 +1959,13 @@ void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n, } -// in-place sparse backsubstitution: x = inv(L'*D*L)*x +// in-place sparse backsubstitution: x = inv(L'*D*L)*x (with dof skipping) void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n, - const int* rownnz, const int* rowadr, const int* colind) { + const int* rownnz, const int* rowadr, const int* colind, const int* index) { // x <- L^-T x - for (int i=nv-1; i > 0; i--) { + for (int k = nv - 1; k >= 0; k--) { + int i = index ? index[k] : k; + // skip diagonal rows if (rownnz[i] == 1) { continue; @@ -1825,7 +1999,9 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, } // x <- D^-1 x - for (int i=0; i < nv; i++) { + for (int k = 0; k < nv; k++) { + int i = index ? index[k] : k; + mjtNum invD_i = qLDiagInv[i]; // one vector @@ -1842,7 +2018,9 @@ void mj_solveLD(mjtNum* restrict x, const mjtNum* qLD, const mjtNum* qLDiagInv, } // x <- L^-1 x - for (int i=1; i < nv; i++) { + for (int k = 0; k < nv; k++) { + int i = index ? index[k] : k; + // skip diagonal rows if (rownnz[i] == 1) { continue; @@ -1874,8 +2052,7 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { if (x != y) { mju_copy(x, y, n*m->nv); } - mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n, - m->M_rownnz, m->M_rowadr, m->M_colind); + mj_solveLD(x, d->qLD, d->qLDiagInv, m->nv, n, m->M_rownnz, m->M_rowadr, m->M_colind, NULL); } @@ -1930,15 +2107,15 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, // compute cvel, cdof_dot void mj_comVel(const mjModel* m, mjData* d) { - int nbody = m->nbody; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; // set world vel to 0 mju_zero(d->cvel, 6); // forward pass over bodies - for (int i=1; i < nbody; i++) { - // get body's first dof address - int bda = m->body_dofadr[i]; + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; // cvel = cvel_parent mjtNum cvel[6]; @@ -1946,6 +2123,7 @@ void mj_comVel(const mjModel* m, mjData* d) { // cvel = cvel_parent + cdof * qvel, cdofdot = cvel x cdof int dofnum = m->body_dofnum[i]; + int bda = m->body_dofadr[i]; mjtNum cdofdot[36]; for (int j=0; j < dofnum; j++) { mjtNum tmp[6]; @@ -1966,9 +2144,9 @@ void mj_comVel(const mjModel* m, mjData* d) { case mjJNT_BALL: // compute all 3 cdofdots using parent velocity - for (int k=0; k < 3; k++) { - mju_crossMotion(cdofdot+6*(j+k), cvel, d->cdof+6*(bda+j+k)); - } + mju_crossMotion(cdofdot+6*(j+0), cvel, d->cdof+6*(bda+j+0)); + mju_crossMotion(cdofdot+6*(j+1), cvel, d->cdof+6*(bda+j+1)); + mju_crossMotion(cdofdot+6*(j+2), cvel, d->cdof+6*(bda+j+2)); // update velocity mju_mulDofVec(tmp, d->cdof+6*(bda+j), d->qvel+bda+j, 3); @@ -1999,13 +2177,16 @@ void mj_comVel(const mjModel* m, mjData* d) { // subtree linear velocity and angular momentum void mj_subtreeVel(const mjModel* m, mjData* d) { - int nbody = m->nbody; - mjtNum dx[3], dv[3], dp[3], dL[3]; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + mj_markStack(d); mjtNum* body_vel = mjSTACKALLOC(d, 6*m->nbody, mjtNum); // bodywise quantities - for (int i=0; i < nbody; i++) { + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + // compute and save body velocity mj_objectVelocity(m, d, mjOBJ_BODY, i, body_vel+6*i, 0); @@ -2013,6 +2194,7 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { mju_scl3(d->subtree_linvel+3*i, body_vel+6*i+3, m->body_mass[i]); // body angular momentum + mjtNum dv[3]; mju_mulMatTVec3(dv, d->ximat+9*i, body_vel+6*i); dv[0] *= m->body_inertia[3*i]; dv[1] *= m->body_inertia[3*i+1]; @@ -2020,8 +2202,10 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { mju_mulMatVec3(d->subtree_angmom+3*i, d->ximat+9*i, dv); } - // subtree linvel - for (int i=nbody-1; i >= 0; i--) { + // subtree linear velocity + for (int b=nbody-1; b >= 0; b--) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + // non-world: add linear momentum to parent if (i) { mju_addTo3(d->subtree_linvel+3*m->body_parentid[i], d->subtree_linvel+3*i); @@ -2029,14 +2213,17 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { // convert linear momentum to linear velocity mju_scl3(d->subtree_linvel+3*i, d->subtree_linvel+3*i, - 1/mjMAX(mjMINVAL, m->body_subtreemass[i])); + 1/mju_max(mjMINVAL, m->body_subtreemass[i])); } - // subtree angmom - for (int i=nbody-1; i > 0; i--) { + // subtree angular momentum + for (int b=nbody-1; b > 0; b--) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + int parent = m->body_parentid[i]; // momentum wrt body i + mjtNum dx[3], dv[3], dp[3], dL[3]; mju_sub3(dx, d->xipos+3*i, d->subtree_com+3*i); mju_sub3(dv, body_vel+6*i+3, d->subtree_linvel+3*i); mju_scl3(dp, dv, m->body_mass[i]); @@ -2066,8 +2253,11 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { // RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { - int nbody = m->nbody, nv = m->nv; - mjtNum tmp[6], tmp1[6]; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + int nparent = sleep_filter ? d->nparent_awake : m->nbody; + int nv = sleep_filter ? d->nv_awake : m->nv; + mj_markStack(d); mjtNum* loc_cacc = mjSTACKALLOC(d, m->nbody*6, mjtNum); mjtNum* loc_cfrc_body = mjSTACKALLOC(d, m->nbody*6, mjtNum); @@ -2079,11 +2269,14 @@ void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { } // forward pass over bodies: accumulate cacc, set cfrc_body - for (int i=1; i < nbody; i++) { + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + // get body's first dof address int bda = m->body_dofadr[i]; // cacc = cacc_parent + cdofdot * qvel + mjtNum tmp[6]; mju_mulDofVec(tmp, d->cdof_dot+6*bda, d->qvel+bda, m->body_dofnum[i]); mju_add(loc_cacc+6*i, loc_cacc+6*m->body_parentid[i], tmp, 6); @@ -2096,22 +2289,27 @@ void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { // cfrc_body = cinert * cacc + cvel x (cinert * cvel) mju_mulInertVec(loc_cfrc_body+6*i, d->cinert+10*i, loc_cacc+6*i); mju_mulInertVec(tmp, d->cinert+10*i, d->cvel+6*i); + mjtNum tmp1[6]; mju_crossForce(tmp1, d->cvel+6*i, tmp); mju_addTo(loc_cfrc_body+6*i, tmp1, 6); } - // clear world cfrc_body, for style + // clear world cfrc_body mju_zero(loc_cfrc_body, 6); // backward pass over bodies: accumulate cfrc_body from children - for (int i=nbody-1; i > 0; i--) { - if (m->body_parentid[i]) { - mju_addTo(loc_cfrc_body+6*m->body_parentid[i], loc_cfrc_body+6*i, 6); + for (int b=nparent-1; b > 0; b--) { + int i = sleep_filter ? d->parent_awake_ind[b] : b; + int j = m->body_parentid[i]; + + if (j) { + mju_addTo(loc_cfrc_body+6*j, loc_cfrc_body+6*i, 6); } } // result = cdof * cfrc_body - for (int i=0; i < nv; i++) { + for (int v=0; v < nv; v++) { + int i = sleep_filter ? d->dof_awake_ind[v] : v; result[i] = mju_dot(d->cdof+6*i, loc_cfrc_body+6*m->dof_bodyid[i], 6); } @@ -2308,12 +2506,18 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { // add bias force due to tendon armature void mj_tendonBias(const mjModel* m, mjData* d, mjtNum* qfrc) { + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; int ntendon = m->ntendon, nv = m->nv, issparse = mj_isSparse(m); mjtNum* ten_Jdot = NULL; mj_markStack(d); // add bias term due to tendon armature for (int i=0; i < ntendon; i++) { + // skip sleeping tendon + if (sleep_filter && mj_sleepState(m, d, mjOBJ_TENDON, i) == mjS_ASLEEP) { + continue; + } + mjtNum armature = m->tendon_armature[i]; // no armature: skip diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index 48b0a94c..6e18fe1a 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -24,6 +24,12 @@ extern "C" { #endif //-------------------------- position -------------------------------------------------------------- +// forward kinematics part 1: bodies +void mj_kinematics1(const mjModel* m, mjData* d); + +// forward kinematics part 2: body inertias, geoms and sites +void mj_kinematics2(const mjModel* m, mjData* d); + // forward kinematics MJAPI void mj_kinematics(const mjModel* m, mjData* d); @@ -61,9 +67,9 @@ MJAPI void mj_makeM(const mjModel* m, mjData* d); MJAPI void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv); -// sparse L'*D*L factorizaton of inertia-like matrix +// sparse L'*D*L factorizaton of inertia-like matrix (only dofs in index, if given) MJAPI void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv, - const int* rownnz, const int* rowadr, const int* colind); + const int* rownnz, const int* rowadr, const int* colind, const int* index); // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd MJAPI void mj_factorM(const mjModel* m, mjData* d); @@ -72,10 +78,10 @@ MJAPI void mj_factorM(const mjModel* m, mjData* d); MJAPI void mj_solveLD_legacy(const mjModel* m, mjtNum* x, int n, const mjtNum* qLD, const mjtNum* qLDiagInv); -// in-place sparse backsubstitution: x = inv(L'*D*L)*x +// in-place sparse backsubstitution (only dofs in index, if given): x = inv(L'*D*L)*x // handle n vectors at once MJAPI void mj_solveLD(mjtNum* x, const mjtNum* qLD, const mjtNum* qLDiagInv, int nv, int n, - const int* rownnz, const int* rowadr, const int* colind); + const int* rownnz, const int* rowadr, const int* colind, const int* index); // 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); diff --git a/src/engine/engine_core_util.c b/src/engine/engine_core_util.c index abb27188..2eaa4823 100644 --- a/src/engine/engine_core_util.c +++ b/src/engine/engine_core_util.c @@ -726,6 +726,8 @@ void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) { } +//-------------------------- spatial frame utilities ----------------------------------------------- + // compute object 6D velocity in object-centered frame, world/local orientation void mj_objectVelocity(const mjModel* m, const mjData* d, int objtype, int objid, mjtNum res[6], int flg_local) { @@ -894,6 +896,8 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], } +//-------------------------- miscellaneous utilities ----------------------------------------------- + // extract 6D force:torque for one contact, in contact frame void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) { mjContact* con; @@ -915,6 +919,26 @@ void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6] } +// count the number of length limit violations for tendon i (0, 1 or 2) +int tendonLimit(const mjModel* m, const mjtNum* ten_length, int i) { + if (!m->tendon_limited[i]) { + return 0; + } + + int nl = 0; + mjtNum value = ten_length[i]; + mjtNum margin = m->tendon_margin[i]; + + // tendon limits can be bilateral, check both sides + for (int side = -1; side <= 1; side += 2) { + mjtNum dist = side * (m->tendon_range[2 * i + (side + 1) / 2] - value); + if (dist < margin) nl++; + } + + return nl; +} + + // count warnings, print only the first time void mj_warning(mjData* d, int warning, int info) { // check type diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h index 91a735d6..3b4d7465 100644 --- a/src/engine/engine_core_util.h +++ b/src/engine/engine_core_util.h @@ -114,17 +114,20 @@ MJAPI void mj_objectVelocity(const mjModel* m, const mjData* d, MJAPI void mj_objectAcceleration(const mjModel* m, const mjData* d, int objtype, int objid, mjtNum res[6], int flg_local); - -//-------------------------- miscellaneous --------------------------------------------------------- - // map from body local to global Cartesian coordinates MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], const mjtNum pos[3], const mjtNum quat[4], int body, mjtByte sameframe); + +//-------------------------- miscellaneous --------------------------------------------------------- + // extract 6D force:torque for one contact, in contact frame MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]); +// count the number of length limit violations for tendon i (0, 1 or 2) +int tendonLimit(const mjModel* m, const mjtNum* ten_length, int i); + // high-level warning function: count warnings in mjData, print only the first time MJAPI void mj_warning(mjData* d, int warning, int info); diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 242bf497..7feecb28 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -21,6 +21,7 @@ #include "engine/engine_crossplatform.h" #include "engine/engine_memory.h" #include "engine/engine_passive.h" +#include "engine/engine_sleep.h" #include "engine/engine_support.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" @@ -521,12 +522,16 @@ static void addToParent(const mjModel* m, mjData* d, mjtNum* mat, int n) { // derivative of cvel, cdof_dot w.r.t qvel static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* Dcdofdot) { - int nv = m->nv, nbody = m->nbody; + int nv = m->nv, nM = m->nM; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; int* Badr = m->B_rowadr, * Dadr = m->D_rowadr; mjtNum mat[36], matT[36]; // 6x6 matrices // forward pass over bodies: accumulate Dcvel, set Dcdofdot - for (int i = 1; i < nbody; i++) { + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + // Dcvel = Dcvel_parent copyFromParent(m, d, Dcvel, i); @@ -534,7 +539,7 @@ static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* D int doflast = m->body_dofadr[i] + m->body_dofnum[i]; for (int j = m->body_dofadr[i]; j < doflast; j++) { // number of dof ancestors of dof j - int Jadr = (j < nv - 1 ? m->dof_Madr[j + 1] : m->nM) - (m->dof_Madr[j] + 1); + int Jadr = (j < nv - 1 ? m->dof_Madr[j + 1] : nM) - (m->dof_Madr[j] + 1); // Dcvel += D(cdof * qvel), Dcdofdot = D(cvel x cdof) switch ((mjtJoint) m->jnt_type[m->dof_jntid[j]]) { @@ -589,7 +594,13 @@ static void mjd_comVel_vel(const mjModel* m, mjData* d, mjtNum* Dcvel, mjtNum* D // subtract d qfrc_bias / d qvel from qDeriv static void mjd_rne_vel(const mjModel* m, mjData* d) { - int nv = m->nv, nbody = m->nbody; + int nM = m->nM; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + int nparent = sleep_filter ? d->nparent_awake : m->nbody; + int mnv = m->nv; + int nv = sleep_filter ? d->nv_awake : mnv; + const int* Badr = m->B_rowadr; const int* Dadr = m->D_rowadr; const int* Bnnz = m->B_rownnz; @@ -601,19 +612,37 @@ static void mjd_rne_vel(const mjModel* m, mjData* d) { mjtNum* Dcvel = mjSTACKALLOC(d, 6*m->nB, mjtNum); mjtNum* Dcacc = mjSTACKALLOC(d, 6*m->nB, mjtNum); mjtNum* Dcfrcbody = mjSTACKALLOC(d, 6*m->nB, mjtNum); - mjtNum* row = mjSTACKALLOC(d, nv, mjtNum); + mjtNum* row = mjSTACKALLOC(d, m->nv, mjtNum); // clear - mju_zero(Dcdofdot, 6*m->nD); - mju_zero(Dcvel, 6*m->nB); - mju_zero(Dcacc, 6*m->nB); - mju_zero(Dcfrcbody, 6*m->nB); + if (!sleep_filter) { + mju_zero(Dcdofdot, 6*m->nD); + mju_zero(Dcvel, 6*m->nB); + mju_zero(Dcacc, 6*m->nB); + mju_zero(Dcfrcbody, 6*m->nB); + } else { + for (int i = 0; i < nv; i++) { + int dof = d->dof_awake_ind[i]; + mju_zero(Dcdofdot + 6*m->D_rowadr[dof], 6*m->D_rownnz[dof]); + } + + for (int i = 0; i < nbody; i++) { + int body = d->body_awake_ind[i]; + int adr = 6*m->B_rowadr[body]; + int nnz = 6*m->B_rownnz[body]; + mju_zero(Dcvel + adr, nnz); + mju_zero(Dcacc + adr, nnz); + mju_zero(Dcfrcbody + adr, nnz); + } + } // compute Dcvel and Dcdofdot mjd_comVel_vel(m, d, Dcvel, Dcdofdot); // forward pass over bodies: accumulate Dcacc, set Dcfrcbody - for (int i=1; i < nbody; i++) { + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + // Dcacc = Dcacc_parent copyFromParent(m, d, Dcacc, i); @@ -621,7 +650,7 @@ static void mjd_rne_vel(const mjModel* m, mjData* d) { int doflast = m->body_dofadr[i] + m->body_dofnum[i]; for (int j=m->body_dofadr[i]; j < doflast; j++) { // number of dof ancestors of dof j - int Jadr = (j < nv - 1 ? m->dof_Madr[j + 1] : m->nM) - (m->dof_Madr[j] + 1); + int Jadr = (j < mnv - 1 ? m->dof_Madr[j + 1] : nM) - (m->dof_Madr[j] + 1); // Dcacc += cdofdot * (D qvel) mju_addTo(Dcacc + 6*(Badr[i] + Jadr), d->cdof_dot + 6*j, 6); @@ -655,12 +684,15 @@ static void mjd_rne_vel(const mjModel* m, mjData* d) { mju_zero(Dcfrcbody, 6*Bnnz[0]); // backward pass over bodies: accumulate Dcfrcbody - for (int i=m->nbody-1; i > 0; i--) { + for (int b=nparent-1; b > 0; b--) { + int i = sleep_filter ? d->parent_awake_ind[b] : b; addToParent(m, d, Dcfrcbody, i); } // process all dofs, update qDeriv - for (int j=0; j < nv; j++) { + for (int v=0; v < nv; v++) { + int j = sleep_filter ? d->dof_awake_ind[v] : v; + // get body index int i = m->dof_bodyid[j]; @@ -797,6 +829,7 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang // add (d qfrc_actuator / d qvel) to qDeriv void mjd_actuator_vel(const mjModel* m, mjData* d) { int nu = m->nu; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; // disabled: nothing to add if (mjDISABLED(mjDSBL_ACTUATION)) { @@ -810,6 +843,11 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) { continue; } + // skip if sleeping + if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) { + continue; + } + mjtNum bias_vel = 0, gain_vel = 0; // affine bias @@ -1401,10 +1439,14 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { return; } + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + // fluid drag model, either body-level (inertia box) or geom-level (ellipsoid) if (m->opt.viscosity > 0 || m->opt.density > 0) { - int nbody = m->nbody; - for (int i=1; i < nbody; i++) { + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + if (m->body_mass[i] < mjMINVAL) { continue; } @@ -1430,7 +1472,9 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { // dof damping int nv = m->nv; - for (int i=0; i < nv; i++) { + int nv_awake = sleep_filter ? d->nv_awake : nv; + for (int j = 0; j < nv_awake; j++) { + int i = sleep_filter ? d->dof_awake_ind[j] : j; d->qDeriv[m->D_rowadr[i] + m->D_diag[i]] -= m->dof_damping[i]; } @@ -1464,6 +1508,15 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { // tendon damping int ntendon = m->ntendon; for (int i=0; i < ntendon; i++) { + // skip tendon in one or two sleeping trees + if (sleep_filter) { + int treenum = m->tendon_treenum[i]; + int id1 = m->tendon_treeid[2*i]; + if (treenum == 1 && !d->tree_awake[id1]) continue; + int id2 = m->tendon_treeid[2*i+1]; + if (treenum == 2 && !d->tree_awake[id1] && !d->tree_awake[id2]) continue; + } + mjtNum B = -m->tendon_damping[i]; if (!B) { @@ -1485,8 +1538,14 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { // analytical derivative of smooth forces w.r.t velocities: // d->qDeriv = d (qfrc_actuator + qfrc_passive - [qfrc_bias]) / d qvel void mjd_smooth_vel(const mjModel* m, mjData* d, int flg_bias) { + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + // clear qDeriv - mju_zero(d->qDeriv, m->nD); + if (!sleep_filter) { + mju_zero(d->qDeriv, m->nD); + } else { + mju_zeroSparse(d->qDeriv, m->D_rownnz, m->D_rowadr, d->dof_awake_ind, d->nv_awake); + } // qDeriv += d qfrc_actuator / d qvel mjd_actuator_vel(m, d); diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index a68521e6..93570699 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -29,12 +29,12 @@ #include "engine/engine_derivative.h" #include "engine/engine_inverse.h" #include "engine/engine_island.h" -#include "engine/engine_io.h" #include "engine/engine_macro.h" #include "engine/engine_memory.h" #include "engine/engine_passive.h" #include "engine/engine_plugin.h" #include "engine/engine_sensor.h" +#include "engine/engine_sleep.h" #include "engine/engine_solver.h" #include "engine/engine_support.h" #include "engine/engine_util_blas.h" @@ -51,8 +51,10 @@ // check positions, reset if bad void mj_checkPos(const mjModel* m, mjData* d) { - for (int i=0; i < m->nq; i++) { - if (mju_isBad(d->qpos[i])) { + int nq = m->nq; + const mjtNum* qpos = d->qpos; + for (int i=0; i < nq; i++) { + if (mju_isBad(qpos[i])) { mj_warning(d, mjWARN_BADQPOS, i); if (!mjDISABLED(mjDSBL_AUTORESET)) { mj_resetData(m, d); @@ -67,7 +69,12 @@ void mj_checkPos(const mjModel* m, mjData* d) { // check velocities, reset if bad void mj_checkVel(const mjModel* m, mjData* d) { - for (int i=0; i < m->nv; i++) { + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nv = sleep_filter ? d->nv_awake : m->nv; + + for (int j=0; j < nv; j++) { + int i = sleep_filter ? d->dof_awake_ind[j] : j; + if (mju_isBad(d->qvel[i])) { mj_warning(d, mjWARN_BADQVEL, i); if (!mjDISABLED(mjDSBL_AUTORESET)) { @@ -83,7 +90,12 @@ void mj_checkVel(const mjModel* m, mjData* d) { // check accelerations, reset if bad void mj_checkAcc(const mjModel* m, mjData* d) { - for (int i=0; i < m->nv; i++) { + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nv = sleep_filter ? d->nv_awake : m->nv; + + for (int j=0; j < nv; j++) { + int i = sleep_filter ? d->dof_awake_ind[j] : j; + if (mju_isBad(d->qacc[i])) { mj_warning(d, mjWARN_BADQACC, i); if (!mjDISABLED(mjDSBL_AUTORESET)) { @@ -135,6 +147,10 @@ void mj_fwdPosition(const mjModel* m, mjData* d) { mj_camlight(m, d); mj_flex(m, d); mj_tendon(m, d); + if (mj_wakeTendon(m, d)) { + mj_updateSleep(m, d); + } + TM_END(mjTIMER_POS_KINEMATICS); // no threadpool: inertia and collision on main thread @@ -168,6 +184,15 @@ void mj_fwdPosition(const mjModel* m, mjData* d) { mju_taskJoin(&tasks[1]); } + if (mj_wakeCollision(m, d)) { + mj_updateSleep(m, d); + mj_collision(m, d); + } + + if (mj_wakeEquality(m, d)) { + mj_updateSleep(m, d); + } + TM_RESTART; mj_makeConstraint(m, d); mj_island(m, d); @@ -277,6 +302,8 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { // clear actuator_force mju_zero(force, nu); + int sleep_filter = mjENABLED(mjENBL_SLEEP); + // disabled or no actuation: return if (nu == 0 || mjDISABLED(mjDSBL_ACTUATION)) { mju_zero(d->qfrc_actuator, nv); @@ -305,6 +332,10 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { // act_dot for stateful actuators for (int i=0; i < nu; i++) { + if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) { + continue; + } + int act_first = m->actuator_actadr[i]; if (act_first < 0) { continue; @@ -371,6 +402,11 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { // force = gain .* [ctrl/act] + bias for (int i=0; i < nu; i++) { + // skip if sleeping + if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) { + continue; + } + // skip if disabled if (mj_actuatorDisabled(m, i)) { continue; @@ -548,16 +584,38 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { // add up all non-constraint forces, compute qacc_smooth void mj_fwdAcceleration(const mjModel* m, mjData* d) { - int nv = m->nv; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nv; + const int* index; - // qfrc_smooth = sum of all non-constraint forces - mju_sub(d->qfrc_smooth, d->qfrc_passive, d->qfrc_bias, nv); // qfrc_bias is negative - mju_addTo(d->qfrc_smooth, d->qfrc_applied, nv); - mju_addTo(d->qfrc_smooth, d->qfrc_actuator, nv); + // qfrc_smooth = qfrc_passive - qfrc_bias + qfrc_applied + qfrc_actuator + if (!sleep_filter) { + nv = m->nv; + index = NULL; + mju_sub(d->qfrc_smooth, d->qfrc_passive, d->qfrc_bias, nv); + mju_addTo(d->qfrc_smooth, d->qfrc_applied, nv); + mju_addTo(d->qfrc_smooth, d->qfrc_actuator, nv); + } else { + nv = d->nv_awake; + index = d->dof_awake_ind; + mju_subInd(d->qfrc_smooth, d->qfrc_passive, d->qfrc_bias, index, nv); + mju_addToInd(d->qfrc_smooth, d->qfrc_applied, index, nv); + mju_addToInd(d->qfrc_smooth, d->qfrc_actuator, index, nv); + } + + // qfrc_smooth += project(xfrc_applied) mj_xfrcAccumulate(m, d, d->qfrc_smooth); + // copy for in-place solve: qacc_smooth = qfrc_smooth + if (!sleep_filter) { + mju_copy(d->qacc_smooth, d->qfrc_smooth, nv); + } else { + mju_copyInd(d->qacc_smooth, d->qfrc_smooth, index, nv); + } + // qacc_smooth = M \ qfrc_smooth - mj_solveM(m, d, d->qacc_smooth, d->qfrc_smooth, 1); + mj_solveLD(d->qacc_smooth, d->qLD, d->qLDiagInv, nv, 1, + m->M_rownnz, m->M_rowadr, m->M_colind, index); } @@ -746,7 +804,6 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) { solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON); } - // 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); @@ -800,11 +857,27 @@ static void mj_advance(const mjModel* m, mjData* d, } } + // put islands to sleep according to velocity tolerance + if (mj_sleep(m, d)) { + // if any trees put to sleep (qvel set to 0), recompute all velocity-dependent quantities + mj_forwardSkip(m, d, mjSTAGE_POS, 0); + + // update sleep indices + mj_updateSleep(m, d); + } + // advance velocities - mju_addToScl(d->qvel, qacc, m->opt.timestep, m->nv); + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + if (sleep_filter) { + mju_addToSclInd(d->qvel, qacc, d->dof_awake_ind, m->opt.timestep, d->nv_awake); + } else { + mju_addToScl(d->qvel, qacc, m->opt.timestep, m->nv); + } // advance positions with qvel if given, d->qvel otherwise (semi-implicit) - mj_integratePos(m, d->qpos, qvel ? qvel : d->qvel, m->opt.timestep); + const int* index = sleep_filter ? d->body_awake_ind : NULL; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + mj_integratePosInd(m, d->qpos, qvel ? qvel : d->qvel, m->opt.timestep, index, nbody); // advance time d->time += m->opt.timestep; @@ -831,15 +904,20 @@ static void mj_advance(const mjModel* m, mjData* d, // Euler integrator, semi-implicit in velocity, possibly skipping factorisation void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { TM_START; - int nv = m->nv, nC = m->nC; mj_markStack(d); - mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); - mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum); + mjtNum* qfrc = mjSTACKALLOC(d, m->nv, mjtNum); + mjtNum* qacc = mjSTACKALLOC(d, m->nv, mjtNum); + + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nv = sleep_filter ? d->nv_awake : m->nv; + const int* dof_awake_ind = sleep_filter ? d->dof_awake_ind : NULL; // check for dof damping if disable flag is not set int dof_damping = 0; if (!mjDISABLED(mjDSBL_EULERDAMP) && !mjDISABLED(mjDSBL_DAMPER)) { - for (int i=0; i < nv; i++) { + for (int v=0; v < nv; v++) { + int i = sleep_filter ? dof_awake_ind[v] : v; if (m->dof_damping[i] > 0) { dof_damping = 1; break; @@ -849,27 +927,43 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { // no damping or disabled: explicit velocity integration if (!dof_damping) { - mju_copy(qacc, d->qacc, nv); + if (sleep_filter) { + mju_copyInd(qacc, d->qacc, dof_awake_ind, nv); + } else { + mju_copy(qacc, d->qacc, nv); + } } // damping: integrate implicitly else { if (!skipfactor) { - // qH = M + h*diag(B) - mju_copy(d->qH, d->M, nC); - for (int i=0; i < nv; i++) { + // qH = M + if (sleep_filter) { + mju_copySparse(d->qH, d->M, m->M_rownnz, m->M_rowadr, dof_awake_ind, d->nv_awake); + } else { + mju_copy(d->qH, d->M, m->nC); + } + + // qH += h*diag(B) + for (int v=0; v < nv; v++) { + int i = sleep_filter ? dof_awake_ind[v] : v; d->qH[m->M_rowadr[i] + m->M_rownnz[i] - 1] += m->opt.timestep * m->dof_damping[i]; } // factorize in-place - mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind); + mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind); } // solve - mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); - mju_copy(qacc, qfrc, m->nv); + if (sleep_filter) { + mju_addInd(qfrc, d->qfrc_smooth, d->qfrc_constraint, dof_awake_ind, nv); + mju_copyInd(qacc, qfrc, dof_awake_ind, nv); + } else { + mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); + mju_copy(qacc, qfrc, nv); + } mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind); } // advance state and time @@ -995,14 +1089,23 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) { // fully implicit in velocity, possibly skipping factorization void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { TM_START; - int nv = m->nv, nD = m->nD, nC = m->nC; + int nD = m->nD, nC = m->nC; mj_markStack(d); - mjtNum* qfrc = mjSTACKALLOC(d, nv, mjtNum); - mjtNum* qacc = mjSTACKALLOC(d, nv, mjtNum); + mjtNum* qfrc = mjSTACKALLOC(d, m->nv, mjtNum); + mjtNum* qacc = mjSTACKALLOC(d, m->nv, mjtNum); + + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nv = sleep_filter ? d->nv_awake : m->nv; + const int* dof_awake_ind = sleep_filter ? d->dof_awake_ind : NULL; // set qfrc = qfrc_smooth + qfrc_constraint - mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); + if (sleep_filter) { + mju_addInd(qfrc, d->qfrc_smooth, d->qfrc_constraint, dof_awake_ind, nv); + } else { + mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); + } // IMPLICIT if (m->opt.integrator == mjINT_IMPLICIT) { @@ -1018,11 +1121,12 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { // factorize qLU int* scratch = mjSTACKALLOC(d, nv, int); - mju_factorLUSparse(d->qLU, nv, scratch, m->D_rownnz, m->D_rowadr, m->D_colind); + mju_factorLUSparse(d->qLU, nv, scratch, m->D_rownnz, m->D_rowadr, m->D_colind, dof_awake_ind); } // solve for qacc: (M - dt*qDeriv) * qacc = qfrc - mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind); + mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind, + dof_awake_ind); } // IMPLICITFAST @@ -1038,13 +1142,17 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { mju_addScl(d->qH, d->M, d->qH, -m->opt.timestep, nC); // factorize in-place - mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind); + mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind); } // solve for qacc: (M - dt*qDeriv) * qacc = qfrc - mju_copy(qacc, qfrc, nv); + if (sleep_filter) { + mju_copyInd(qacc, qfrc, dof_awake_ind, nv); + } else { + mju_copy(qacc, qfrc, nv); + } mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind); } else { mjERROR("integrator must be implicit or implicitfast"); diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 53bc690b..a2ae5544 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -26,6 +26,8 @@ #include #include // IWYU pragma: keep #include +#include "engine/engine_core_smooth.h" +#include "engine/engine_forward.h" #include "engine/engine_init.h" #include "engine/engine_macro.h" #include "engine/engine_memory.h" @@ -1087,6 +1089,11 @@ void mj_makeRawData(mjData** dest, const mjModel* m) { // clear nplugin (overwritten by _initPlugin) d->nplugin = 0; + // set awake array sizes to default (all awake) + d->ntree_awake = m->ntree; + d->nbody_awake = d->nparent_awake = m->nbody; + d->nv_awake = m->nv; + // copy pointer if allocated here if (allocate) { *dest = d; @@ -1369,6 +1376,68 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { d->tree_asleep[i] = kAwake; } + // sleep enabled: handle static bodies and trees marked as mjSLEEP_INIT + if (mjENABLED(mjENBL_SLEEP)) { + // count trees initialized as asleep + int num_asleep_init = 0; + for (int i=0; i < m->ntree; i++) { + num_asleep_init += (m->tree_sleep_policy[i] == mjSLEEP_INIT); + } + + // update sleep arrays, treat static bodies as awake + mj_updateSleepInit(m, d, /*flg_staticawake*/ 1); + + // partial mj_fwdPosition, functions that update STATIC values + if (!num_asleep_init) { + mj_kinematics(m, d); + mj_comPos(m, d); + mj_camlight(m, d); + mj_tendon(m, d); + } + + // if any trees initialized as sleeping call entire mj_forward, put them to sleep + else { + mj_forward(m, d); + + // mark asleep-init trees as ready to sleep + for (int i=0; i < m->ntree; i++) { + int init = m->tree_sleep_policy[i] == mjSLEEP_INIT; + d->tree_asleep[i] = init ? -1 : kAwake; + } + + int nslept = mj_sleep(m, d); + + // raise error if any failed to sleep + if (nslept != num_asleep_init) { + // find root body of the first tree that could not be slept + int root = -1; + for (int i=0; i < m->ntree; i++) { + if (m->tree_sleep_policy[i] == mjSLEEP_INIT && d->tree_asleep[i] < 0) { + root = m->tree_bodyadr[i]; + break; + } + } + + // free all memory held by d just before aborting + mj_deleteData(d); + + // raise error and abort + const char* hasname = mj_id2name(m, mjOBJ_BODY, root); + const char* name = hasname ? hasname : ""; + mjERROR("%d trees were marked as sleep='init' but only %d could be slept.\n" + "Body '%s' (id=%d) is the root of the first tree that could not be slept.", + num_asleep_init, nslept, name, root); + } + + // clear arrays to avoid MSAN errors upon mid-step wake + mju_zero(d->qacc_smooth, m->nv); + mju_zero(d->qfrc_smooth, m->nv); + + // clear arena + mj_clearEfc(d); + } + } + // update sleep arrays and counters mj_updateSleep(m, d); diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 95aa16af..52387fbf 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -25,6 +25,7 @@ #include "engine/engine_crossplatform.h" #include "engine/engine_memory.h" #include "engine/engine_plugin.h" +#include "engine/engine_sleep.h" #include "engine/engine_support.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" @@ -58,63 +59,72 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3], // spring and damper forces static void mj_springdamper(const mjModel* m, mjData* d) { - int nv = m->nv, njnt = m->njnt, ntendon = m->ntendon; + int nv = m->nv, ntendon = m->ntendon; int has_spring = !mjDISABLED(mjDSBL_SPRING); int has_damping = !mjDISABLED(mjDSBL_DAMPER); int issparse = mj_isSparse(m); + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; // joint-level springs if (has_spring) { - for (int i=0; i < njnt; i++) { - mjtNum stiffness = m->jnt_stiffness[i]; + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; + int jnt_start = m->body_jntadr[i]; + int jnt_end = jnt_start + m->body_jntnum[i]; + for (int j=jnt_start; j < jnt_end; j++) { + mjtNum stiffness = m->jnt_stiffness[j]; - // disabled : nothing to do - if (stiffness == 0) { - continue; - } - - int padr = m->jnt_qposadr[i]; - int dadr = m->jnt_dofadr[i]; - - switch ((mjtJoint) m->jnt_type[i]) { - case mjJNT_FREE: - // apply force - d->qfrc_spring[dadr+0] = -stiffness*(d->qpos[padr+0] - m->qpos_spring[padr+0]); - d->qfrc_spring[dadr+1] = -stiffness*(d->qpos[padr+1] - m->qpos_spring[padr+1]); - d->qfrc_spring[dadr+2] = -stiffness*(d->qpos[padr+2] - m->qpos_spring[padr+2]); - - // continue with rotations - dadr += 3; - padr += 3; - mjFALLTHROUGH; - - case mjJNT_BALL: - { - // convert quaternion difference into angular "velocity" - mjtNum dif[3], quat[4]; - mju_copy4(quat, d->qpos+padr); - mju_normalize4(quat); - mju_subQuat(dif, quat, m->qpos_spring + padr); - - // apply torque - d->qfrc_spring[dadr+0] = -stiffness*dif[0]; - d->qfrc_spring[dadr+1] = -stiffness*dif[1]; - d->qfrc_spring[dadr+2] = -stiffness*dif[2]; + // disabled : nothing to do + if (stiffness == 0) { + continue; } - break; - case mjJNT_SLIDE: - case mjJNT_HINGE: - // apply force or torque - d->qfrc_spring[dadr] = -stiffness*(d->qpos[padr] - m->qpos_spring[padr]); - break; + int padr = m->jnt_qposadr[j]; + int dadr = m->jnt_dofadr[j]; + + switch ((mjtJoint) m->jnt_type[j]) { + case mjJNT_FREE: + // apply force + d->qfrc_spring[dadr+0] = -stiffness*(d->qpos[padr+0] - m->qpos_spring[padr+0]); + d->qfrc_spring[dadr+1] = -stiffness*(d->qpos[padr+1] - m->qpos_spring[padr+1]); + d->qfrc_spring[dadr+2] = -stiffness*(d->qpos[padr+2] - m->qpos_spring[padr+2]); + + // continue with rotations + dadr += 3; + padr += 3; + mjFALLTHROUGH; + + case mjJNT_BALL: + { + // convert quaternion difference into angular "velocity" + mjtNum dif[3], quat[4]; + mju_copy4(quat, d->qpos+padr); + mju_normalize4(quat); + mju_subQuat(dif, quat, m->qpos_spring + padr); + + // apply torque + d->qfrc_spring[dadr+0] = -stiffness*dif[0]; + d->qfrc_spring[dadr+1] = -stiffness*dif[1]; + d->qfrc_spring[dadr+2] = -stiffness*dif[2]; + } + break; + + case mjJNT_SLIDE: + case mjJNT_HINGE: + // apply force or torque + d->qfrc_spring[dadr] = -stiffness*(d->qpos[padr] - m->qpos_spring[padr]); + break; + } } } } // dof-level dampers if (has_damping) { - for (int i=0; i < m->nv; i++) { + int nv_awake = sleep_filter ? d->nv_awake : nv; + for (int j = 0; j < nv_awake; j++) { + int i = sleep_filter ? d->dof_awake_ind[j] : j; mjtNum damping = m->dof_damping[i]; if (damping != 0) { d->qfrc_damper[i] = -damping*d->qvel[i]; @@ -419,6 +429,11 @@ static void mj_springdamper(const mjModel* m, mjData* d) { // tendon-level spring-dampers for (int i=0; i < ntendon; i++) { + // skip sleeping or static tendon + if (sleep_filter && mj_sleepState(m, d, mjOBJ_TENDON, i) != mjS_AWAKE) { + continue; + } + mjtNum stiffness = m->tendon_stiffness[i] * has_spring; mjtNum damping = m->tendon_damping[i] * has_damping; @@ -466,11 +481,14 @@ static int mj_gravcomp(const mjModel* m, mjData* d) { return 0; } - int nbody = m->nbody, has_gravcomp = 0; + int has_gravcomp = 0; mjtNum force[3], torque[3]={0}; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; // apply per-body gravity compensation - for (int i=1; i < nbody; i++) { + for (int b=1; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; if (m->body_gravcomp[i]) { has_gravcomp = 1; mju_scl3(force, m->opt.gravity, -(m->body_mass[i]*m->body_gravcomp[i])); @@ -484,32 +502,37 @@ static int mj_gravcomp(const mjModel* m, mjData* d) { // fluid forces static int mj_fluid(const mjModel* m, mjData* d) { - int has_fluid = m->opt.viscosity > 0 || m->opt.density > 0; + // no fluid forces: early return + if (!m->opt.viscosity && !m->opt.density) { + return 0; + } - if (has_fluid) { - int nbody = m->nbody; - for (int i=1; i < nbody; i++) { - if (m->body_mass[i] < mjMINVAL) { - continue; - } + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + int nbody = sleep_filter ? d->nbody_awake : m->nbody; - // if any child geom uses the ellipsoid model, inertia-box model is disabled for parent body - int use_ellipsoid_model = 0; - int geomnum = m->body_geomnum[i]; - for (int j=0; j < geomnum && use_ellipsoid_model == 0; j++) { - const int geomid = m->body_geomadr[i] + j; - use_ellipsoid_model += (m->geom_fluid[mjNFLUID*geomid] > 0); - } + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; - if (use_ellipsoid_model) { - mj_ellipsoidFluidModel(m, d, i); - } else { - mj_inertiaBoxFluidModel(m, d, i); - } + if (m->body_mass[i] < mjMINVAL) { + continue; + } + + // if any child geom uses the ellipsoid model, inertia-box model is disabled for parent body + int use_ellipsoid_model = 0; + int geomnum = m->body_geomnum[i]; + for (int j=0; j < geomnum && use_ellipsoid_model == 0; j++) { + const int geomid = m->body_geomadr[i] + j; + use_ellipsoid_model += (m->geom_fluid[mjNFLUID*geomid] > 0); + } + + if (use_ellipsoid_model) { + mj_ellipsoidFluidModel(m, d, i); + } else { + mj_inertiaBoxFluidModel(m, d, i); } } - return has_fluid; + return 1; } @@ -597,14 +620,24 @@ int mj_contactPassive(const mjModel* m, mjData* d) { // all passive forces void mj_passive(const mjModel* m, mjData* d) { - int nv = m->nv; + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv; + int nv = sleep_filter ? d->nv_awake : m->nv; + const int* dof_awake_ind = sleep_filter ? d->dof_awake_ind : NULL; - // clear all passive force vectors - mju_zero(d->qfrc_spring, nv); - mju_zero(d->qfrc_damper, nv); - mju_zero(d->qfrc_gravcomp, nv); - mju_zero(d->qfrc_fluid, nv); - mju_zero(d->qfrc_passive, nv); + // clear passive force vectors for awake dofs + if (sleep_filter) { + mju_zeroInd(d->qfrc_spring, nv, dof_awake_ind); + mju_zeroInd(d->qfrc_damper, nv, dof_awake_ind); + mju_zeroInd(d->qfrc_gravcomp, nv, dof_awake_ind); + mju_zeroInd(d->qfrc_fluid, nv, dof_awake_ind); + mju_zeroInd(d->qfrc_passive, nv, dof_awake_ind); + } else { + mju_zero(d->qfrc_spring, nv); + mju_zero(d->qfrc_damper, nv); + mju_zero(d->qfrc_gravcomp, nv); + mju_zero(d->qfrc_fluid, nv); + mju_zero(d->qfrc_passive, nv); + } // both spring and damping disabled: skip all passive forces if (mjDISABLED(mjDSBL_SPRING) && mjDISABLED(mjDSBL_DAMPER)) { @@ -624,39 +657,49 @@ void mj_passive(const mjModel* m, mjData* d) { mj_contactPassive(m, d); // add passive forces into qfrc_passive - mju_add(d->qfrc_passive, d->qfrc_spring, d->qfrc_damper, nv); - if (has_fluid) { - mju_addTo(d->qfrc_passive, d->qfrc_fluid, nv); + if (sleep_filter) { + mju_addInd(d->qfrc_passive, d->qfrc_spring, d->qfrc_damper, dof_awake_ind, nv); + } else { + mju_add(d->qfrc_passive, d->qfrc_spring, d->qfrc_damper, nv); } + + if (has_fluid) { + if (sleep_filter) { + mju_addToInd(d->qfrc_passive, d->qfrc_fluid, dof_awake_ind, nv); + } else { + mju_addTo(d->qfrc_passive, d->qfrc_fluid, nv); + } + } + if (has_gravcomp) { - int njnt = m->njnt; - for (int i=0; i < njnt; i++) { - // skip if gravcomp added via actuators - if (m->jnt_actgravcomp[i]) { - continue; - } + int nbody = sleep_filter ? d->nbody_awake : m->nbody; + for (int b=0; b < nbody; b++) { + int i = sleep_filter ? d->body_awake_ind[b] : b; - // get number of dofs for this joint - int dofnum; - switch (m->jnt_type[i]) { - case mjJNT_HINGE: - case mjJNT_SLIDE: - dofnum = 1; - break; + // skip if no joints + int jntnum = m->body_jntnum[i]; + if (!jntnum) continue; - case mjJNT_BALL: - dofnum = 3; - break; + // skip if no gravity compensation + if (!m->body_gravcomp[i]) continue; - case mjJNT_FREE: - dofnum = 6; - break; - } + int start = m->body_jntadr[i]; + int end = start + jntnum; + for (int j=start; j < end; j++) { + // skip if gravity compensation added via actuators + if (m->jnt_actgravcomp[j]) { + continue; + } - // add gravcomp force - int dofadr = m->jnt_dofadr[i]; - for (int j=0; j < dofnum; j++) { - d->qfrc_passive[dofadr+j] += d->qfrc_gravcomp[dofadr+j]; + // get number of dofs for this joint + const int jnt_dofnum[4] = {6, 3, 1, 1}; + int dofnum = jnt_dofnum[m->jnt_type[j]]; + + // add gravity compensation force + int dofadr = m->jnt_dofadr[j]; + for (int k=0; k < dofnum; k++) { + d->qfrc_passive[dofadr+k] += d->qfrc_gravcomp[dofadr+k]; + } } } } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 67dd9d5f..aeae37e0 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -80,7 +80,7 @@ static void printArr(FILE* fp, const char* name, const float* data, int n, const // print 2D array of mjtNum into file static void printArray2d(const char* str, int nr, int nc, const mjtNum* data, FILE* fp, - const char* float_format) { + const char* float_format) { if (!data) { return; } diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index c218a1dc..9a76d56a 100644 --- a/src/engine/engine_sensor.c +++ b/src/engine/engine_sensor.c @@ -28,6 +28,7 @@ #include "engine/engine_memory.h" #include "engine/engine_plugin.h" #include "engine/engine_ray.h" +#include "engine/engine_sleep.h" #include "engine/engine_sort.h" #include "engine/engine_support.h" #include "engine/engine_util_blas.h" @@ -394,10 +395,18 @@ void mj_sensorPos(const mjModel* m, mjData* d) { return; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + // process sensors matching stage for (int i=0; i < nsensor; i++) { mjtSensor type = (mjtSensor) m->sensor_type[i]; + // skip sleeping sensor + if (sleep_filter && mj_sleepState(m, d, mjOBJ_SENSOR, i) == mjS_ASLEEP) { + continue; + } + // skip sensor plugins -- these are handled after builtin sensor types if (type == mjSENS_PLUGIN) { continue; @@ -698,6 +707,9 @@ void mj_sensorVel(const mjModel* m, mjData* d) { return; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + // process sensors matching stage int subtreeVel = 0; for (int i=0; i < m->nsensor; i++) { @@ -706,6 +718,11 @@ void mj_sensorVel(const mjModel* m, mjData* d) { continue; } + // skip sleeping sensor + if (sleep_filter && mj_sleepState(m, d, mjOBJ_SENSOR, i) == mjS_ASLEEP) { + continue; + } + if (m->sensor_needstage[i] == mjSTAGE_VEL) { // get sensor info mjtSensor type = m->sensor_type[i]; @@ -883,9 +900,17 @@ void mj_sensorAcc(const mjModel* m, mjData* d) { return; } + // sleep filtering + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + // process sensors matching stage int rnePost = 0; for (int i=0; i < m->nsensor; i++) { + // skip sleeping sensor + if (sleep_filter && mj_sleepState(m, d, mjOBJ_SENSOR, i) == mjS_ASLEEP) { + continue; + } + // skip sensor plugins -- these are handled after builtin sensor types if (m->sensor_type[i] == mjSENS_PLUGIN) { continue; @@ -1412,35 +1437,47 @@ void mj_energyPos(const mjModel* m, mjData* d) { } } + int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody; + // add joint-level springs if (!mjDISABLED(mjDSBL_SPRING)) { - for (int i=0; i < m->njnt; i++) { - stiffness = m->jnt_stiffness[i]; - padr = m->jnt_qposadr[i]; + int nbody = m->nbody; + for (int b=1; b < nbody; b++) { + if (sleep_filter && d->body_awake[b] != mjS_AWAKE) continue; - switch ((mjtJoint) m->jnt_type[i]) { - case mjJNT_FREE: - mju_sub3(dif, d->qpos+padr, m->qpos_spring+padr); - d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif); + int jnt_start = m->body_jntadr[b]; + int jnt_end = jnt_start + m->body_jntnum[b]; + for (int j=jnt_start; j < jnt_end; j++) { + stiffness = m->jnt_stiffness[j]; + if (stiffness == 0) { + continue; + } + padr = m->jnt_qposadr[j]; - // continue with rotations - padr += 3; - mjFALLTHROUGH; + switch ((mjtJoint) m->jnt_type[j]) { + case mjJNT_FREE: + mju_sub3(dif, d->qpos+padr, m->qpos_spring+padr); + d->energy[0] += 0.5 * stiffness * mju_dot3(dif, dif); - case mjJNT_BALL: - // convert quaternion difference into angular "velocity" - mju_copy4(quat, d->qpos+padr); - mju_normalize4(quat); - mju_subQuat(dif, d->qpos + padr, m->qpos_spring + padr); - d->energy[0] += 0.5*stiffness*mju_dot3(dif, dif); - break; + // continue with rotations + padr += 3; + mjFALLTHROUGH; - case mjJNT_SLIDE: - case mjJNT_HINGE: - d->energy[0] += 0.5*stiffness* - (d->qpos[padr] - m->qpos_spring[padr])* - (d->qpos[padr] - m->qpos_spring[padr]); - break; + case mjJNT_BALL: + // convert quaternion difference into angular "velocity" + mju_copy4(quat, d->qpos+padr); + mju_normalize4(quat); + mju_subQuat(dif, d->qpos + padr, m->qpos_spring + padr); + d->energy[0] += 0.5 * stiffness * mju_dot3(dif, dif); + break; + + case mjJNT_SLIDE: + case mjJNT_HINGE: + d->energy[0] += 0.5 * stiffness * + (d->qpos[padr] - m->qpos_spring[padr]) * + (d->qpos[padr] - m->qpos_spring[padr]); + break; + } } } } @@ -1448,6 +1485,11 @@ void mj_energyPos(const mjModel* m, mjData* d) { // add tendon-level springs if (!mjDISABLED(mjDSBL_SPRING)) { for (int i=0; i < m->ntendon; i++) { + // skip sleeping or static tendon + if (sleep_filter && mj_sleepState(m, d, mjOBJ_TENDON, i) != mjS_AWAKE) { + continue; + } + stiffness = m->tendon_stiffness[i]; mjtNum length = d->ten_length[i]; mjtNum displacement = 0; diff --git a/src/engine/engine_sleep.c b/src/engine/engine_sleep.c index 99231859..0c8d55a6 100644 --- a/src/engine/engine_sleep.c +++ b/src/engine/engine_sleep.c @@ -19,7 +19,13 @@ #include #include +#include "engine/engine_core_util.h" +#include "engine/engine_util_blas.h" +#include "engine/engine_util_errmem.h" +#include "engine/engine_util_misc.h" +// uncomment to print sleep/wake events +// #define MJ_DEBUG_SLEEP //-------------------------------- update ---------------------------------------------------------- @@ -97,3 +103,676 @@ void mj_updateSleepInit(const mjModel* m, mjData* d, int flg_staticawake) { void mj_updateSleep(const mjModel* m, mjData* d) { mj_updateSleepInit(m, d, /*flg_staticawake*/0); } + + +//-------------------------------- utilities ------------------------------------------------------- + +// return 1 if the weighted infinity norm of vec is smaller than tol, 0 otherwise +static int isSmaller(const mjtNum* vec, const mjtNum* weight, int n, mjtNum tol) { + mjtNum max = 0; + for (int i=0; i < n; i++) { + max = mju_max(max, weight[i] * mju_abs(vec[i])); + if (max >= tol) { + return 0; + } + } + return 1; +} + + +// return 1 if tree i can sleep, 0 otherwise +static int treeCanSleep(const mjModel* m, const mjData* d, int i, mjtNum tol) { + // check sleep policy + if (m->tree_sleep_policy[i] == mjSLEEP_NEVER || + m->tree_sleep_policy[i] == mjSLEEP_AUTO_NEVER) { + return 0; + } + + // check xfrc_applied + int adr = m->tree_bodyadr[i]; + int num = m->tree_bodynum[i]; + if (!mju_isZeroByte((const unsigned char*)(d->xfrc_applied+6*adr), 6*num*sizeof(mjtNum))) { + return 0; + } + + // check qfrc_applied + adr = m->tree_dofadr[i]; + num = m->tree_dofnum[i]; + if (!mju_isZeroByte((const unsigned char*)(d->qfrc_applied+adr), num*sizeof(mjtNum))) { + return 0; + } + + // check qvel + if (tol) { + return isSmaller(d->qvel+adr, m->dof_length+adr, num, tol); + } else { + return mju_isZeroByte((const unsigned char*)(d->qvel+adr), num*sizeof(mjtNum)); + } +} + + +// return the first tree in the sleep cycle that starts at i, -1 if error +int mj_sleepCycle(const int* tree_asleep, int ntree, int i) { + if (i < 0 || i >= ntree) { + return -1; // index i out of bounds + } + + int smallest = i; + int current = i; + int count = 0; + + do { + if (count > ntree) { + return -1; // cycle detection failed (too many steps) + } + + int next = tree_asleep[current]; + + if (next < 0 || next >= ntree) { + return -1; // next index out of bounds + } + + if (next < smallest) { + smallest = next; + } + + current = next; + count++; + } while (current != i); + + return smallest; +} + + +//-------------------------------- wake ------------------------------------------------------------ + +// wake tree i and its associated cycle, return number of woke trees +int mj_wakeTree(int* tree_asleep, int ntree, int i, int wakeval) { + int nwoke = 0; + + // i is invalid; SHOULD NOT OCCUR + if (i < 0 || i >= ntree) { + mjERROR("invalid tree %d", i); + return nwoke; + } + + // tree i already awake: set to wakeval if larger than current value + int asleep_val = tree_asleep[i]; + if (asleep_val < 0) { + tree_asleep[i] = mjMIN(wakeval, asleep_val); + return nwoke; + } + + // tree i asleep: wake up tree and its island cycle + else { + int current = i; + do { + // get the index of the next tree in the cycle + int next = tree_asleep[current]; + + // next is invalid; SHOULD NOT OCCUR + if (next < 0 || next >= ntree) { + mjERROR("invalid sleep state index %d when waking tree %d", next, i); + return 0; + } + + // wake the current tree, increment count, advance to next + tree_asleep[current] = wakeval; + nwoke++; + current = next; + } while (current != i && nwoke < ntree); + + // did not come back to tree i, not a cycle; SHOULD NOT OCCUR + if (current != i) { + mjERROR("tree %d is not in a cycle", i); + return 0; + } + } + + return nwoke; +} + + +static int kAwake = -(1+mjMINAWAKE); // tree_asleep value for fully awake tree + +// wake sleeping trees due to changes by user, return number of woke trees +int mj_wake(const mjModel* m, mjData* d) { + int ntree = m->ntree, nwoke = 0; + + // sleep disabled + if (!mjENABLED(mjENBL_SLEEP)) { + // sleep disabled but some trees still asleep: wake all + if (d->ntree_awake < ntree) { + for (int i=0; i < ntree; i++) d->tree_asleep[i] = kAwake; + } + return ntree - d->ntree_awake; + } + + // sweep over trees, wake if required + for (int i=0; i < ntree; i++) { + int asleep = d->tree_asleep[i] >= 0; + + // awake: nothing to do + if (!asleep) { + continue; + } + + // if qpos mismatch or cannot sleep: wake up + if (d->tree_awake[i] || !treeCanSleep(m, d, i, 0)) { + int woke = mj_wakeTree(d->tree_asleep, ntree, i, kAwake); + if (woke) { + nwoke += woke; + + #ifdef MJ_DEBUG_SLEEP + printf("woke tree %d due to perturbation at t=%g\n", i, d->time); + #endif + } + } + } + + return nwoke; +} + + +// wake sleeping trees that touch awake trees, return number of woke trees +int mj_wakeCollision(const mjModel* m, mjData* d) { + int ntree = m->ntree, ncon = d->ncon, nwoke = 0; + + if (!mjENABLED(mjENBL_SLEEP)) { + return nwoke; + } + + // sweep over contacts, wake trees if required + for (int i=0; i < ncon; i++) { + const mjContact* con = d->contact + i; + + // only geom-geom contacts are handled + if (con->geom[0] < 0 || con->geom[1] < 0) { + continue; + } + + int b1 = m->geom_bodyid[con->geom[0]]; + int b2 = m->geom_bodyid[con->geom[1]]; + int tree1 = m->body_treeid[b1]; + int tree2 = m->body_treeid[b2]; + + // contact with static body, nothing to do + if (tree1 < 0 || tree2 < 0) { + continue; + } + + int awake1 = d->tree_awake[tree1]; + int awake2 = d->tree_awake[tree2]; + + // both trees awake, nothing to do + if (awake1 && awake2) { + continue; + } + + // both trees asleep; SHOULD NOT OCCUR + if (!awake1 && !awake2) { + mjERROR("contact between sleeping bodies %d and %d", b1, b2); + } + + // wake sleeping tree + int sleeping_tree = awake1 ? tree2 : tree1; + int wakeval = awake1 ? d->tree_asleep[tree1] : d->tree_asleep[tree2]; + nwoke += mj_wakeTree(d->tree_asleep, ntree, sleeping_tree, wakeval); + + #ifdef MJ_DEBUG_SLEEP + printf("woke tree %d due to contact at t=%g\n", sleeping_tree, d->time); + #endif + } + + return nwoke; +} + + +// wake sleeping trees with a constrained tendon to a waking tree, return number of woke trees +int mj_wakeTendon(const mjModel* m, mjData* d) { + int ntendon = m->ntendon, nwoke = 0; + + if (!mjENABLED(mjENBL_SLEEP)) { + return nwoke; + } + + // sweep over tendons, wake trees if required + for (int i=0; i < ntendon; i++) { + if (m->tendon_treenum[i] != 2 || !tendonLimit(m, d->ten_length, i)) { + continue; + } + + int tree1 = m->tendon_treeid[2*i]; + int tree2 = m->tendon_treeid[2*i + 1]; + int awake1 = d->tree_awake[tree1]; + int awake2 = d->tree_awake[tree2]; + if (awake1 != awake2) { + int sleeping_tree = awake1 ? tree2 : tree1; + int wakeval = awake1 ? d->tree_asleep[tree1] : d->tree_asleep[tree2]; + nwoke += mj_wakeTree(d->tree_asleep, m->ntree, sleeping_tree, wakeval); + + #ifdef MJ_DEBUG_SLEEP + printf("woke tree %d due to tendon constraint at t=%g\n", sleeping_tree, d->time); + #endif + } + } + + return nwoke; +} + + +// wake sleeping trees with an equality to a waking tree, return number of woke trees +int mj_wakeEquality(const mjModel* m, mjData* d) { + int neq = m->neq, nwoke = 0; + + if (!mjENABLED(mjENBL_SLEEP)) { + return nwoke; + } + + // sweep over equalities, wake trees if required + for (int i=0; i < neq; i++) { + // skip inactive + if (!d->eq_active[i]) continue; + + mjtEq eqtype = m->eq_type[i]; + int id1 = m->eq_obj1id[i]; + int id2 = m->eq_obj2id[i]; + int tree1, tree2; + + switch (eqtype) { + case mjEQ_CONNECT: + case mjEQ_WELD: + if (m->eq_objtype[i] == mjOBJ_BODY) { + tree1 = m->body_treeid[id1]; + tree2 = m->body_treeid[id2]; + } else { + tree1 = m->body_treeid[m->site_bodyid[id1]]; + tree2 = m->body_treeid[m->site_bodyid[id2]]; + } + break; + case mjEQ_JOINT: + tree1 = id1 >= 0 ? m->body_treeid[m->jnt_bodyid[id1]] : -1; + tree2 = id2 >= 0 ? m->body_treeid[m->jnt_bodyid[id2]] : -1; + break; + case mjEQ_TENDON: + mjERROR("tendon equality does not yet support sleeping"); + continue; + case mjEQ_FLEX: + mjERROR("flex equality does not yet support sleeping"); + continue; + default: + continue; + } + + // get sleep state + mjtSleepState s1 = tree1 >= 0 ? d->tree_awake[tree1] : mjS_STATIC; + mjtSleepState s2 = tree2 >= 0 ? d->tree_awake[tree2] : mjS_STATIC; + + // neither is asleep, nothing to do + if (s1 != mjS_ASLEEP && s2 != mjS_ASLEEP) { + continue; + } + + // one is static, nothing to do + if (s1 == mjS_STATIC || s2 == mjS_STATIC) { + continue; + } + + // equality within the same tree, nothing to do + if (tree1 == tree2) { + continue; + } + + // both are asleep, wake if in different islands + if (s1 == mjS_ASLEEP && s2 == mjS_ASLEEP) { + int cycle1 = mj_sleepCycle(d->tree_asleep, m->ntree, tree1); + int cycle2 = mj_sleepCycle(d->tree_asleep, m->ntree, tree2); + if (cycle1 != cycle2) { + int nwoke1 = mj_wakeTree(d->tree_asleep, m->ntree, tree1, kAwake); + int nwoke2 = mj_wakeTree(d->tree_asleep, m->ntree, tree2, kAwake); + + #ifdef MJ_DEBUG_SLEEP + printf("woke trees %d, %d due to equality %d at t=%g\n", tree1, tree2, i, d->time); + #endif + + nwoke += nwoke1 + nwoke2; + } + continue; + } + + // one is asleep and one is awake, wake the sleeping tree + int sleeping_tree = s1 == mjS_ASLEEP ? tree1 : tree2; + nwoke += mj_wakeTree(d->tree_asleep, m->ntree, sleeping_tree, kAwake); + + #ifdef MJ_DEBUG_SLEEP + printf("woke tree %d due to equality %d at t=%g\n", sleeping_tree, i, d->time); + #endif + } + + return nwoke; +} + + +//-------------------------------- sleep ----------------------------------------------------------- + +// put n trees to sleep (create cycle), set their velocity and acceleration to zero +static inline void sleepTrees(const mjModel* m, mjData* d, const int* tree, int n) { + for (int i=0; i < n; i++) { + // create cycle + int current = tree[i]; + int next = (i == n - 1) ? tree[0] : tree[i + 1]; + if (d->tree_asleep[current] == -1) { + d->tree_asleep[current] = next; + } + + // SHOULD NOT OCCUR + else if (d->tree_asleep[current] >= 0) { + mjERROR("trying to sleep tree %d which is already asleep", i); + } else { + mjERROR("trying to sleep tree %d which is not ready to sleep", i); + } + + // set tree velocity and acceleration to zero + int adr = m->tree_dofadr[current]; + int num = m->tree_dofnum[current]; + mju_zero(d->qvel+adr, num); + mju_zero(d->qacc+adr, num); + } + + #ifdef MJ_DEBUG_SLEEP + if (n == 1) { + printf("tree %d put to sleep at t=%g\n", tree[0], d->time); + } else if (n > 1) { + printf("trees "); + for (int i = 0; i < n; i++) { + printf("%d%s", tree[i], (i == n - 1) ? "" : ", "); + } + printf(" put to sleep at t=%g\n", d->time); + } + #endif +} + + +// put trees to sleep according to tolerance, return number of slept trees +int mj_sleep(const mjModel* m, mjData* d) { + int ntree = m->ntree, nisland = d->nisland, nslept = 0; + + // sleep disabled: nothing to do + if (!mjENABLED(mjENBL_SLEEP)) { + return nslept; + } + + // have constraints but no island structure: can't sleep + if (d->nefc && !nisland) { + return nslept; + } + + // sweep over awake trees, increment tree_asleep if under tolerance + for (int i=0; i < ntree; i++) { + // skip sleeping tree + if (d->tree_asleep[i] >= 0) { + continue; + } + + // increment tree_asleep if tree can sleep, otherwise wake up + if (treeCanSleep(m, d, i, m->opt.sleep_tolerance)) { + d->tree_asleep[i] += (d->tree_asleep[i] < -1); + } else { + d->tree_asleep[i] = -(1+mjMINAWAKE); + } + } + + // sweep over islands, put to sleep if all trees are under tolerance + for (int i=0; i < nisland; i++) { + // check if all trees in the island can sleep + int can_sleep = 1; + int start = d->island_itreeadr[i]; + int end = start + d->island_ntree[i]; + for (int j=start; j < end; j++) { + int tree_asleep = d->tree_asleep[d->map_itree2tree[j]]; + if (tree_asleep < -1) { + can_sleep = 0; + break; + } + + // sleeping tree in an island; SHOULD NOT OCCUR + else if (tree_asleep >= 0) { + mjERROR("found sleeping tree %d in island %d", d->map_itree2tree[j], i); + } + } + + // put island to sleep + if (can_sleep) { + const int* tree = d->map_itree2tree + start; + int n = d->island_ntree[i]; + sleepTrees(m, d, tree, n); + nslept += n; + } + } + + // sleep unconstrained trees (with or without island structure) + int start = nisland ? d->island_itreeadr[nisland-1] + d->island_ntree[nisland-1] : 0; + for (int j=start; j < ntree; j++) { + int i = nisland ? d->map_itree2tree[j] : j; + if (d->tree_asleep[i] == -1) { + sleepTrees(m, d, &i, 1); + nslept++; + } + } + + return nslept; +} + + +//-------------------------------- sleep state ----------------------------------------------------- + +// return sleep state of tendon i +static mjtSleepState mj_tendonSleepState(const mjModel* m, const mjData* d, int i) { + int treenum = m->tendon_treenum[i]; + + // no trees: tendon is static + if (treenum == 0) { + return mjS_STATIC; + } + + // single tree: awake if tree is awake, asleep otherwise + int id1 = m->tendon_treeid[2*i]; + if (treenum == 1) { + return d->tree_awake[id1] ? mjS_AWAKE : mjS_ASLEEP; + } + + // two trees: asleep only if both are asleep + int id2 = m->tendon_treeid[2*i+1]; + if (treenum == 2) { + return (d->tree_awake[id1] || d->tree_awake[id2]) ? mjS_AWAKE : mjS_ASLEEP; + } + + return mjS_AWAKE; +} + + +// return sleep state of actuator i +static mjtSleepState mj_actuatorSleepState(const mjModel* m, const mjData* d, int i) { + mjtSleepState s1, s2; + int trnid = m->actuator_trnid[i*2]; + + switch ((mjtTrn)m->actuator_trntype[i]) { + case mjTRN_JOINT: + case mjTRN_JOINTINPARENT: + return mj_sleepState(m, d, mjOBJ_JOINT, trnid); + + case mjTRN_SLIDERCRANK: + s1 = mj_sleepState(m, d, mjOBJ_SITE, trnid); + s2 = mj_sleepState(m, d, mjOBJ_SITE, m->actuator_trnid[i*2+1]); + return (s1 == mjS_AWAKE || s2 == mjS_AWAKE) ? mjS_AWAKE : mjS_ASLEEP; + + case mjTRN_TENDON: + return mj_tendonSleepState(m, d, trnid); + + case mjTRN_SITE: + return mj_sleepState(m, d, mjOBJ_SITE, trnid); + + case mjTRN_BODY: + return mj_sleepState(m, d, mjOBJ_BODY, trnid); + + case mjTRN_UNDEFINED: + return mjS_AWAKE; + } + + return mjS_AWAKE; +} + + +// return sleep state of equality i +static mjtSleepState mj_equalitySleepState(const mjModel* m, const mjData* d, int i) { + mjtEq eqtype = m->eq_type[i]; + mjtObj objtype; + + switch (eqtype) { + case mjEQ_CONNECT: + case mjEQ_WELD: + objtype = m->eq_objtype[i]; + break; + case mjEQ_JOINT: + objtype = mjOBJ_JOINT; + break; + case mjEQ_TENDON: + objtype = mjOBJ_TENDON; + break; + case mjEQ_FLEX: + objtype = mjOBJ_FLEX; + break; + default: + return mjS_AWAKE; + } + + int id1 = m->eq_obj1id[i]; + int id2 = m->eq_obj2id[i]; + mjtSleepState s1 = (id1 >= 0) ? mj_sleepState(m, d, objtype, id1) : mjS_STATIC; + mjtSleepState s2 = (id2 >= 0) ? mj_sleepState(m, d, objtype, id2) : mjS_STATIC; + + // return ASLEEP if both objects are asleep or static, AWAKE otherwise + int neither_awake = (s1 != mjS_AWAKE && s2 != mjS_AWAKE); + return neither_awake ? mjS_ASLEEP : mjS_AWAKE; +} + + +// return sleep state of sensor i +static mjtSleepState mj_sensorSleepState(const mjModel* m, const mjData* d, int i) { + mjtSensor type = m->sensor_type[i]; + mjtObj objtype = m->sensor_objtype[i]; + int objid = m->sensor_objid[i]; + mjtObj reftype = m->sensor_reftype[i]; + int refid = m->sensor_refid[i]; + + // get sleep state of the primary and reference objects + mjtSleepState s_obj = mj_sleepState(m, d, objtype, objid); + mjtSleepState s_ref = mj_sleepState(m, d, reftype, refid); + + // special handling for specific sensor types + switch (type) { + + // USER and PLUGIN sensors are always awake + case mjSENS_USER: + case mjSENS_PLUGIN: + return mjS_AWAKE; + + // sensors that use sites to define a volume are always awake + case mjSENS_INSIDESITE: + case mjSENS_TOUCH: + return mjS_AWAKE; + + // contact sensors + case mjSENS_CONTACT: + // site used to define a volume: always awake + if (objtype == mjOBJ_SITE || reftype == mjOBJ_SITE) { + return mjS_AWAKE; + } + + // for contact sensors UNKNOWN means undefined, so the AWAKE returned by mj_sleepState is wrong + + // if both are UNKNOWN (all contacts), return ASLEEP iff everything is alseep + if (objtype == mjOBJ_UNKNOWN && reftype == mjOBJ_UNKNOWN) { + return d->ntree_awake == 0 ? mjS_ASLEEP : mjS_AWAKE; + } + + // if only one is UNKNOWN, return state of other object + if (objtype == mjOBJ_UNKNOWN) { + return s_ref; + } else if (reftype == mjOBJ_UNKNOWN) { + return s_obj; + } + break; + + // sensors whose value depends on objects other than the two they are attached to are always awake + case mjSENS_RANGEFINDER: + return mjS_AWAKE; + + default: + break; + } + + // if either object is awake, return AWAKE + if (s_obj == mjS_AWAKE || s_ref == mjS_AWAKE) { + return mjS_AWAKE; + } + + // otherwise return ASLEEP + return mjS_ASLEEP; +} + + +// return sleep state of object i +mjtSleepState mj_sleepState(const mjModel* m, const mjData* d, mjtObj type, int i) { + const char* typename; + + switch (type) { + + // simple types + case mjOBJ_BODY: + case mjOBJ_XBODY: + return (mjtSleepState) d->body_awake[i]; + case mjOBJ_JOINT: + return (mjtSleepState) d->body_awake[m->jnt_bodyid[i]]; + case mjOBJ_SITE: + return (mjtSleepState) d->body_awake[m->site_bodyid[i]]; + case mjOBJ_DOF: + return (mjtSleepState) d->body_awake[m->dof_bodyid[i]]; + case mjOBJ_GEOM: + return (mjtSleepState) d->body_awake[m->geom_bodyid[i]]; + case mjOBJ_CAMERA: + return (mjtSleepState) d->body_awake[m->cam_bodyid[i]]; + case mjOBJ_LIGHT: + return (mjtSleepState) d->body_awake[m->light_bodyid[i]]; + + // complex types + case mjOBJ_EQUALITY: + return mj_equalitySleepState(m, d, i); + case mjOBJ_TENDON: + return mj_tendonSleepState(m, d, i); + case mjOBJ_ACTUATOR: + return mj_actuatorSleepState(m, d, i); + case mjOBJ_SENSOR: + return mj_sensorSleepState(m, d, i); + + // always awake + case mjOBJ_FLEX: + case mjOBJ_UNKNOWN: + return mjS_AWAKE; + + // unsupported + default: + typename = mju_type2Str(type); + if (typename) { + mjERROR("unsupported object type '%s'", typename); + } else { + mjERROR("unsupported object type %d", type); + } + return mjS_AWAKE; + } +} + + +#ifdef MJ_DEBUG_SLEEP + #undef MJ_DEBUG_SLEEP +#endif diff --git a/src/engine/engine_sleep.h b/src/engine/engine_sleep.h index 92a7a27d..6bb65616 100644 --- a/src/engine/engine_sleep.h +++ b/src/engine/engine_sleep.h @@ -24,11 +24,38 @@ extern "C" { #endif // compute sleeping arrays from tree_asleep, if flg_staticawake is set, treat static bodies as awake -MJAPI void mj_updateSleepInit(const mjModel* m, mjData* d, int flg_staticawake); +void mj_updateSleepInit(const mjModel* m, mjData* d, int flg_staticawake); // compute {ntree,nbody,nv}_awake, {tree,body}_awake, {body,dof}_awake_ind from tree_asleep MJAPI void mj_updateSleep(const mjModel* m, mjData* d); +// return the first tree in the sleep cycle that starts at i, -1 if error +int mj_sleepCycle(const int* tree_asleep, int ntree, int i); + +// return the first tree in the sleep cycle that starts at i, -1 if error +int mj_sleepCycle(const int* tree_asleep, int ntree, int i); + +// wake tree i and its related island cycle, return number of woke trees +MJAPI int mj_wakeTree(int* tree_asleep, int ntree, int i, int wakeval); + +// wake trees with nonzero velocity or external forces, return number of woke trees +int mj_wake(const mjModel* m, mjData* d); + +// wake sleeping trees that touch awake trees, return number of woke trees +int mj_wakeCollision(const mjModel* m, mjData* d); + +// wake sleeping trees with a constrained tendon to a waking tree, return number of woke trees +int mj_wakeTendon(const mjModel* m, mjData* d); + +// wake sleeping trees with an equality to a waking tree, return number of woke trees +int mj_wakeEquality(const mjModel* m, mjData* d); + +// put trees to sleep according to tolerance, return number of slept trees +int mj_sleep(const mjModel* m, mjData* d); + +// return sleep state of object i +mjtSleepState mj_sleepState(const mjModel* m, const mjData* d, mjtObj type, int i); + #ifdef __cplusplus } #endif diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index ec16ee53..4f52626d 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -1063,7 +1063,7 @@ static void CGupdateGradient(mjCGContext* ctx, int flg_Newton) { else { mju_copy(ctx->Mgrad, ctx->grad, nv); mj_solveLD(ctx->Mgrad, ctx->qLD, ctx->qLDiagInv, nv, 1, - ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind); + ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind, NULL); } } diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 4090673f..63678380 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -561,40 +561,51 @@ void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt, } -// integrate qpos with given qvel -void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt) { - // loop over joints - for (int j=0; j < m->njnt; j++) { - // get addresses in qpos and qvel - int padr = m->jnt_qposadr[j]; - int vadr = m->jnt_dofadr[j]; +// integrate qpos with given qvel for given body indices +void mj_integratePosInd(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt, + const int* index, int nbody) { + for (int b=1; b < nbody; b++) { + int k = index ? index[b] : b; + int start = m->body_jntadr[k]; + int end = start + m->body_jntnum[k]; + for (int j=start; j < end; j++) { + // get addresses in qpos and qvel + int padr = m->jnt_qposadr[j]; + int vadr = m->jnt_dofadr[j]; - switch ((mjtJoint) m->jnt_type[j]) { - case mjJNT_FREE: - // position update - for (int i=0; i < 3; i++) { - qpos[padr+i] += dt * qvel[vadr+i]; + switch ((mjtJoint) m->jnt_type[j]) { + case mjJNT_FREE: + // position update + for (int i=0; i < 3; i++) { + qpos[padr+i] += dt * qvel[vadr+i]; + } + padr += 3; + vadr += 3; + + // continue with rotation update + mjFALLTHROUGH; + + case mjJNT_BALL: + // quaternion update + mju_quatIntegrate(qpos+padr, qvel+vadr, dt); + break; + + case mjJNT_HINGE: + case mjJNT_SLIDE: + // scalar update: same for rotation and translation + qpos[padr] += dt * qvel[vadr]; } - padr += 3; - vadr += 3; - - // continue with rotation update - mjFALLTHROUGH; - - case mjJNT_BALL: - // quaternion update - mju_quatIntegrate(qpos+padr, qvel+vadr, dt); - break; - - case mjJNT_HINGE: - case mjJNT_SLIDE: - // scalar update: same for rotation and translation - qpos[padr] += dt * qvel[vadr]; } } } +// integrate qpos with given qvel +void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt) { + mj_integratePosInd(m, qpos, qvel, dt, NULL, m->nbody); +} + + // normalize all quaternions in qpos-type vector void mj_normalizeQuat(const mjModel* m, mjtNum* qpos) { // find quaternion fields and normalize diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index 6a5f9b91..283028d1 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -89,6 +89,10 @@ MJAPI mjtNum mj_geomDistance(const mjModel* m, const mjData* d, int geom1, int g MJAPI void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt, const mjtNum* qpos1, const mjtNum* qpos2); +// integrate qpos with given qvel for given body indices +MJAPI void mj_integratePosInd(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt, + const int* index, int nbody); + // integrate position with given velocity MJAPI void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt); diff --git a/src/engine/engine_util_blas.c b/src/engine/engine_util_blas.c index 24e7b05a..e5232431 100644 --- a/src/engine/engine_util_blas.c +++ b/src/engine/engine_util_blas.c @@ -273,6 +273,14 @@ void mju_zero(mjtNum* res, int n) { } +// res = 0, at given indices +void mju_zeroInd(mjtNum* res, int n, const int* ind) { + for (int i = 0; i < n; i++) { + res[ind[i]] = 0; + } +} + + // res = val void mju_fill(mjtNum* res, mjtNum val, int n) { for (int i=0; i < n; i++) { @@ -287,6 +295,14 @@ void mju_copy(mjtNum* res, const mjtNum* vec, int n) { } +// res = vec, at given indices +void mju_copyInd(mjtNum* res, const mjtNum* vec, const int* ind, int n) { + for (int i = 0; i < n; i++) { + res[ind[i]] = vec[ind[i]]; + } +} + + // sum(vec) mjtNum mju_sum(const mjtNum* vec, int n) { mjtNum res = 0; @@ -397,6 +413,15 @@ void mju_add(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) { } +// res = vec1 + vec2, at selected indices +void mju_addInd(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, const int* ind, int n) { + for (int i = 0; i < n; i++) { + int j = ind[i]; + res[j] = vec1[j] + vec2[j]; + } +} + + // res = vec1 - vec2 void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) { int i = 0; @@ -439,6 +464,15 @@ void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n) { } +// res = vec1 - vec2, at selected indices +void mju_subInd(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, const int* ind, int n) { + for (int i = 0; i < n; i++) { + int j = ind[i]; + res[j] = vec1[j] - vec2[j]; + } +} + + // res += vec void mju_addTo(mjtNum* res, const mjtNum* vec, int n) { int i = 0; @@ -481,6 +515,15 @@ void mju_addTo(mjtNum* res, const mjtNum* vec, int n) { } +// res += vec, at selected indices +void mju_addToInd(mjtNum* res, const mjtNum* vec, const int* ind, int n) { + for (int i = 0; i < n; i++) { + int j = ind[i]; + res[j] += vec[j]; + } +} + + // res -= vec void mju_subFrom(mjtNum* res, const mjtNum* vec, int n) { int i = 0; @@ -568,6 +611,18 @@ void mju_addToScl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n) { #endif } + + +// res += vec*scl, at given indices +void mju_addToSclInd(mjtNum* res, const mjtNum* vec, const int* ind, mjtNum scl, int n) { + for (int i=0; i < n; i++) { + int k = ind[i]; + res[k] += vec[k]*scl; + } +} + + + // res = vec1 + vec2*scl void mju_addScl(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, mjtNum scl, int n) { int i = 0; @@ -706,6 +761,20 @@ mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, int n) { return res; } + + +// vector dot-product, at given indices +mjtNum mju_dotInd(const mjtNum* vec1, const mjtNum* vec2, const int* ind, int n) { + mjtNum res = 0; + for (int i = 0; i < n; i++) { + int k = ind[i]; + res += vec1[k] * vec2[k]; + } + return res; +} + + + //------------------------------ matrix-vector operations ------------------------------------------ // multiply matrix and vector @@ -771,6 +840,14 @@ void mju_eye(mjtNum* mat, int n) { } +// res[ind, :] = mat[ind, :] +void mju_copyRows(mjtNum* res, const mjtNum* mat, const int* ind, int n, int nc) { + for (int i = 0; i < n; i++) { + mju_copy(res + nc*ind[i], mat + nc*ind[i], nc); + } +} + + //------------------------------ matrix-matrix operations ------------------------------------------ // multiply matrices, exploit sparsity of mat1 diff --git a/src/engine/engine_util_blas.h b/src/engine/engine_util_blas.h index 4b2519c9..6f359852 100644 --- a/src/engine/engine_util_blas.h +++ b/src/engine/engine_util_blas.h @@ -141,12 +141,18 @@ MJAPI mjtNum mju_normalize4(mjtNum vec[4]); // res = 0 MJAPI void mju_zero(mjtNum* res, int n); +// res = 0, at given indices +void mju_zeroInd(mjtNum* res, int n, const int* ind); + // res = val MJAPI void mju_fill(mjtNum* res, mjtNum val, int n); // res = vec MJAPI void mju_copy(mjtNum* res, const mjtNum* vec, int n); +// res = vec, at given indices +void mju_copyInd(mjtNum* res, const mjtNum* vec, const int* ind, int n); + // sum(vec) MJAPI mjtNum mju_sum(const mjtNum* vec, int n); @@ -159,18 +165,30 @@ MJAPI void mju_scl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n); // res = vec1 + vec2 MJAPI void mju_add(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n); +// res = vec1 + vec2, at given indices +void mju_addInd(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, const int* ind, int n); + // res = vec1 - vec2 MJAPI void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n); +// res = vec1 - vec2, at selected indices +void mju_subInd(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, const int* ind, int n); + // res += vec MJAPI void mju_addTo(mjtNum* res, const mjtNum* vec, int n); +// res += vec, at selected indices +void mju_addToInd(mjtNum* res, const mjtNum* vec, const int* ind, int n); + // res -= vec MJAPI void mju_subFrom(mjtNum* res, const mjtNum* vec, int n); // res += vec*scl MJAPI void mju_addToScl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n); +// res += vec*scl, at given indices +void mju_addToSclInd(mjtNum* res, const mjtNum* vec, const int* ind, mjtNum scl, int n); + // res = vec1 + vec2*scl MJAPI void mju_addScl(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, mjtNum scl, int n); @@ -183,16 +201,16 @@ MJAPI mjtNum mju_norm(const mjtNum* res, int n); // vector dot-product MJAPI mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, int n); +// vector dot-product, at given indices +mjtNum mju_dotInd(const mjtNum* vec1, const mjtNum* vec2, const int* ind, int n); //------------------------------ matrix-vector operations ------------------------------------------ // multiply matrix and vector -MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, - int nr, int nc); +MJAPI void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc); // multiply transposed matrix and vector -MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, - int nr, int nc); +MJAPI void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr, int nc); // multiply square matrix with vectors on both sides: return vec1'*mat*vec2 MJAPI mjtNum mju_mulVecMatVec(const mjtNum* vec1, const mjtNum* mat, const mjtNum* vec2, int n); @@ -209,6 +227,9 @@ MJAPI void mju_symmetrize(mjtNum* res, const mjtNum* mat, int n); // identity matrix MJAPI void mju_eye(mjtNum* mat, int n); +// copy selected rows: res[ind, :] = mat[ind, :] +void mju_copyRows(mjtNum* res, const mjtNum* mat, const int* ind, int n, int nc); + //------------------------------ matrix-matrix operations ------------------------------------------ // multiply matrices diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 79f1730b..0deb7915 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -612,17 +612,26 @@ void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec, // sparse reverse-order LU factorization, no fill-in (assuming tree topology) // result: LU = L + U; original = (U+I) * L; scratch size is n void mju_factorLUSparse(mjtNum* LU, int n, int* scratch, - const int* rownnz, const int* rowadr, const int* colind) { + const int* rownnz, const int* rowadr, const int* colind, + const int* index) { int* remaining = scratch; // set remaining = rownnz - mju_copyInt(remaining, rownnz, n); + if (index) { + for (int i=0; i < n; i++) { + remaining[i] = rownnz[index[i]]; + } + } else { + mju_copyInt(remaining, rownnz, n); + } // diagonal elements (i,i) - for (int i=n-1; i >= 0; i--) { + for (int r=n-1; r >= 0; r--) { + int i = index ? index[r] : r; + // get address of last remaining element of row i, adjust remaining counter - int ii = rowadr[i] + remaining[i] - 1; - remaining[i]--; + int ii = rowadr[i] + remaining[r] - 1; + remaining[r]--; // make sure ii is on diagonal if (colind[ii] != i) { @@ -635,14 +644,16 @@ void mju_factorLUSparse(mjtNum* LU, int n, int* scratch, } // rows j above i - for (int j=i-1; j >= 0; j--) { + for (int c=r-1; c >= 0; c--) { + int j = index ? index[c] : c; + // get address of last remaining element of row j - int ji = rowadr[j] + remaining[j] - 1; + int ji = rowadr[j] + remaining[c] - 1; // process row j if (j,i) is non-zero if (colind[ji] == i) { // adjust remaining counter - remaining[j]--; + remaining[c]--; // (j,i) = (j,i) / (i,i) LU[ji] = LU[ji] / LU[ii]; @@ -650,7 +661,7 @@ void mju_factorLUSparse(mjtNum* LU, int n, int* scratch, // (j,k) = (j,k) - (i,k) * (j,i) for k= 0; i--) { + for (int k=n-1; k >= 0; k--) { + int i = index ? index[k] : k; + // init: diagonal of (U+I) is 1 res[i] = vec[i]; @@ -703,7 +718,9 @@ void mju_solveLUSparse(mjtNum* res, const mjtNum* LU, const mjtNum* vec, int n, } //------------------ solve L*res(new) = res - for (int i=0; i < n; i++) { + for (int k=0; k < n; k++) { + int i = index ? index[k] : k; + // res[i] -= sum_k= 0 ? kIslandSaturation : 0; - hsv2rgb(rgba, hue, saturation, kIslandValue); +static void islandColor(float rgba[4], int h, int awake) { + // default to gray R = G = B = 0.7; + float hue = 1.0f; + float saturation = 0.0f; + float value = 0.7f; + + // island index given, use Halton sequence to generate pseudo-random color + if (h >= 0) { + // hue in [0, 1] + hue = mju_Halton(h + 1, 7); + + // saturation in [0.5, 1.0] + saturation = .5 + .5*mju_Halton(h + 1, 3); + + // value in [0.6, 1.0] + value = .6 + .4*mju_Halton(h + 1, 5); + } + + // if asleep, decrease saturation and value + if (!awake) { + value *= 0.6; + saturation *= 0.7; + } + + hsv2rgb(rgba, hue, saturation, value); rgba[3] = 1; } @@ -572,7 +591,7 @@ static void addContactGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, if (vopt->flags[mjVIS_ISLAND] && efc_adr >= 0) { // set hue using island's first dof int h = d->nisland > 0 ? d->island_dofadr[d->efc_island[efc_adr]] : -1; - islandColor(thisgeom->rgba, h); + islandColor(thisgeom->rgba, h, /*awake*/1); } // otherwise regular colors (different for included and excluded contacts) @@ -593,8 +612,7 @@ static void addContactGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, const char* geomname = mj_id2name(m, mjOBJ_GEOM, con->geom[k]); if (geomname) { mjSNPRINTF(contactlabel[k], "%s", geomname); - } - else { + } else { mjSNPRINTF(contactlabel[k], "g%d", con->geom[k]); } } @@ -605,16 +623,14 @@ static void addContactGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, if (flexname) { if (con->elem[k] >= 0) { mjSNPRINTF(contactlabel[k], "%s.e%d", flexname, con->elem[k]); - } - else { + } else { mjSNPRINTF(contactlabel[k], "%s.v%d", flexname, con->vert[k]); } } else { if (con->elem[k] >= 0) { mjSNPRINTF(contactlabel[k], "f%d.e%d", con->flex[k], con->elem[k]); - } - else { + } else { mjSNPRINTF(contactlabel[k], "f%d.v%d", con->flex[k], con->vert[k]); } } @@ -865,9 +881,19 @@ static void addGeomGeoms(const mjModel* m, mjData* d, const mjvOption* vopt, thisgeom->matid = -1; // set hue using first island dof, -1 if no island - int island = d->nisland ? d->dof_island[m->body_dofadr[weld_id]] : -1; + int dof = m->body_dofadr[weld_id]; + int island = d->nisland ? d->dof_island[dof] : -1; int h = island >= 0 ? d->island_dofadr[island] : -1; - islandColor(thisgeom->rgba, h); + int awake = d->body_awake[m->geom_bodyid[i]]; + + // if sleep is enabled, color by first tree dof + if (h == -1 && mjENABLED(mjENBL_SLEEP)) { + int tree = m->dof_treeid[dof]; + if (!awake) tree = mj_sleepCycle(d->tree_asleep, m->ntree, tree); + h = m->tree_dofadr[tree]; + } + + islandColor(thisgeom->rgba, h, awake); } } @@ -1102,13 +1128,13 @@ static void addSpatialTendonGeoms(const mjModel* m, mjData* d, const mjvOption* // strip material thisgeom->matid = -1; - // set hue with first island dof, if constrained - int h = -1; - if (d->nisland && d->tendon_efcadr[i] >= 0) { - h = d->island_dofadr[d->efc_island[d->tendon_efcadr[i]]]; + // set hue with first island dof, if constrained + int h = -1; + if (d->nisland && d->tendon_efcadr[i] >= 0) { + h = d->island_dofadr[d->efc_island[d->tendon_efcadr[i]]]; + } + islandColor(thisgeom->rgba, h, 1); } - islandColor(thisgeom->rgba, h); - } // vopt->label: only the first segment if (vopt->label == mjLABEL_TENDON && j == d->ten_wrapadr[i]) { diff --git a/src/engine/engine_vis_visualize.h b/src/engine/engine_vis_visualize.h index daad04a9..046d6634 100644 --- a/src/engine/engine_vis_visualize.h +++ b/src/engine/engine_vis_visualize.h @@ -70,6 +70,9 @@ void mjv_cameraFrustum(float zver[2], float zhor[2], float zclip[2], const mjMo int mjv_catenary(const mjtNum x0[3], const mjtNum x1[3], const mjtNum gravity[3], mjtNum length, mjtNum* catenary, int ncatenary); +// convert HSV to RGB +void hsv2rgb(float *RGB, float H, float S, float V); + #ifdef __cplusplus } #endif diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc index f8cad0f4..afab5dcf 100644 --- a/src/experimental/studio/app.cc +++ b/src/experimental/studio/app.cc @@ -641,10 +641,14 @@ void App::UpdateProfilerData() { sqrt_nnz = mju_sqrt(sqrt_nnz); dim_dof_.erase(dim_dof_.begin()); - dim_dof_.push_back(Model()->nv); + int nv = (Model()->opt.enableflags & mjENBL_SLEEP) ? Data()->nv_awake + : Model()->nv; + dim_dof_.push_back(nv); dim_body_.erase(dim_body_.begin()); - dim_body_.push_back(Model()->nbody); + int nbody = (Model()->opt.enableflags & mjENBL_SLEEP) ? Data()->nbody_awake + : Model()->nbody; + dim_body_.push_back(nbody); dim_constraint_.erase(dim_constraint_.begin()); dim_constraint_.push_back(Data()->nefc); @@ -1719,6 +1723,7 @@ void App::PhysicsGui() { ImGui_Input("Noslip Tol", &opt.noslip_tolerance, {0, 1, 0.01, 0.1, w}); ImGui_Input("CCD Iter", &opt.ccd_iterations, {0, 1000, 1, 100, w}); ImGui_Input("CCD Tol", &opt.ccd_tolerance, {0, 1, 0.01, 0.1, w}); + ImGui_Input("Sleep Tol", &opt.sleep_tolerance, {0, 1, 0.01, 0.1, w}); ImGui_Input("SDF Iter", &opt.sdf_iterations, {1, 20, 1, 10, w}); ImGui_Input("SDF Init", &opt.sdf_initpoints, {1, 100, 1, 10, w}); ImGui::TreePop(); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index c2f11402..ad7af6dd 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -4935,6 +4935,8 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // create data int disableflags = m->opt.disableflags; m->opt.disableflags |= mjDSBL_CONTACT; + int enableflags = m->opt.enableflags; + m->opt.enableflags &= ~mjENBL_SLEEP; mj_makeRawData(&d, m); if (!d) { // m will be deleted by the catch statement in mjCModel::Compile() @@ -4980,18 +4982,39 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // delete partial mjData (no plugins), make a complete one mj_deleteData(d); d = nullptr; + + // if sleep was enabled, check for trees initialized as sleeping + bool asleep_init = false; + if (enableflags & mjENBL_SLEEP) { + for (int i=0; i < m->ntree; i++) { + if (m->tree_sleep_policy[i] == mjSLEEP_INIT) { + asleep_init = true; + break; + } + } + } + + // if any trees initialized as sleeping, restore flags before mj_makeData + if (asleep_init) { + m->opt.disableflags = disableflags; + m->opt.enableflags = enableflags; + } + d = mj_makeData(m); if (!d) { // m will be deleted by the catch statement in mjCModel::Compile() throw mjCError(0, "could not create mjData"); } - // test forward simulation - mj_step(m, d); + // test forward simulation unless asleep_init is true (potentially expensive) + if (!asleep_init) { + mj_step(m, d); + } - // delete data + // delete data, restore flags mj_deleteData(d); m->opt.disableflags = disableflags; + m->opt.enableflags = enableflags; d = nullptr; // pass warning back diff --git a/test/benchmark/factorI_benchmark_test.cc b/test/benchmark/factorI_benchmark_test.cc index 835856b2..73e0af06 100644 --- a/test/benchmark/factorI_benchmark_test.cc +++ b/test/benchmark/factorI_benchmark_test.cc @@ -59,7 +59,7 @@ static void BM_factorI(benchmark::State& state, bool legacy, bool coil) { } else { mju_copy(d->qLD, M, m->nC); mj_factorI(d->qLD, d->qLDiagInv, m->nv, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); } } } diff --git a/test/benchmark/inertia_benchmark_test.cc b/test/benchmark/inertia_benchmark_test.cc index a0fb14c0..cd204663 100644 --- a/test/benchmark/inertia_benchmark_test.cc +++ b/test/benchmark/inertia_benchmark_test.cc @@ -73,9 +73,9 @@ static void BM_solve(benchmark::State& state, SolveType type) { case SolveType::kCsr: mju_copy(d->qLD, M, m->nC); mj_factorI(d->qLD, d->qLDiagInv, m->nv, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); } } } diff --git a/test/benchmark/solveLD_benchmark_test.cc b/test/benchmark/solveLD_benchmark_test.cc index 20646b84..a6c1e157 100644 --- a/test/benchmark/solveLD_benchmark_test.cc +++ b/test/benchmark/solveLD_benchmark_test.cc @@ -64,7 +64,7 @@ static void BM_solveLD(benchmark::State& state, bool featherstone, bool coil) { mj_solveLD_legacy(m, res, 1, LDlegacy, d->qLDiagInv); } else { mj_solveLD(res, d->qLD, d->qLDiagInv, m->nv, 1, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); } } } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 0f4a10ae..017ff64e 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -742,7 +742,7 @@ TEST_F(CoreSmoothTest, SolveLDs) { mj_solveLD_legacy(m, vec.data(), 1, LDlegacy.data(), d->qLDiagInv); mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, 1, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); // expect vectors to match up to floating point precision for (int i=0; i < nv; i++) { @@ -777,7 +777,7 @@ TEST_F(CoreSmoothTest, SolveLDmultipleVectors) { mj_solveLD_legacy(m, vec.data(), n, LDlegacy.data(), d->qLDiagInv); mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); // expect vectors to match up to floating point precision for (int i=0; i < nv*n; i++) { @@ -815,7 +815,7 @@ TEST_F(CoreSmoothTest, SolveM2) { mj_solveM2(m, d, res.data(), vec.data(), sqrtInvD.data(), n); mj_solveLD(vec2.data(), d->qLD, d->qLDiagInv, nv, n, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); // expect equality of dot(v, M^-1 * v) and dot(M^-1/2 * v, M^-1/2 * v) for (int i=0; i < n; i++) { @@ -854,7 +854,7 @@ TEST_F(CoreSmoothTest, FactorIs) { vector qLDiagInv(nv, 0); mj_factorI(qLD.data(), qLDiagInv.data(), nv, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); // expect outputs to match to floating point precision EXPECT_THAT(qLD, Pointwise(DoubleNear(1e-12), qLDexpected)); diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 5c8529fc..c887f4b7 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -438,7 +438,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { Ac[nv*nv + i*nv + i] = -m->dof_damping[i]; } mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2*nv, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); // A = [dt*Ac; Ac] mju_transpose(A, Ac, 2*nv, nv); @@ -466,7 +466,7 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz, d->moment_rowadr, d->moment_colind); mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu, - m->M_rownnz, m->M_rowadr, m->M_colind); + m->M_rownnz, m->M_rowadr, m->M_colind, nullptr); mju_transpose(BcT, Bc, nu, nv); mju_scl(B, BcT, dt*dt, nu*nv); mju_scl(B+nu*nv, BcT, dt, nu*nv); diff --git a/test/engine/engine_sleep_test.cc b/test/engine/engine_sleep_test.cc index da8a3681..2a7c1ad2 100644 --- a/test/engine/engine_sleep_test.cc +++ b/test/engine/engine_sleep_test.cc @@ -14,8 +14,11 @@ // Tests for engine/engine_sleep.c. +#include + #include #include +#include #include #include #include "src/engine/engine_sleep.h" @@ -25,7 +28,10 @@ namespace mujoco { namespace { using ::testing::ElementsAre; +using ::testing::IsNull; +using ::testing::HasSubstr; using ::testing::NotNull; +using ::std::string; using SleepTest = MujocoTest; @@ -177,5 +183,348 @@ TEST_F(SleepTest, MjSleepUpdate) { mj_deleteModel(m); } +TEST_F(SleepTest, MjWakeTree) { + // one awake tree and two cycles + int asleep[] = {kAwake, 2, 1, 3}; + EXPECT_EQ(mj_wakeTree(asleep, 4, 0, kAwake), 0); + EXPECT_THAT(AsVector(asleep, 4), ElementsAre(kAwake, 2, 1, 3)); + EXPECT_EQ(mj_wakeTree(asleep, 4, 1, kAwake), 2); + EXPECT_THAT(AsVector(asleep, 4), + ElementsAre(kAwake, kAwake, kAwake, 3)); + EXPECT_EQ(mj_wakeTree(asleep, 4, 3, kAwake), 1); + EXPECT_THAT(AsVector(asleep, 4), + ElementsAre(kAwake, kAwake, kAwake, kAwake)); +} + +TEST_F(SleepTest, BadWakeTree) { + EXPECT_FATAL_FAILURE( + ([] { + int asleep_bad1[] = {-1, 0}; + mj_wakeTree(asleep_bad1, 2, 1, kAwake); + }()), + "invalid sleep state index -1 when waking tree 1"); + + EXPECT_FATAL_FAILURE( + ([] { + int asleep_bad2[] = {-1, 2}; + mj_wakeTree(asleep_bad2, 2, 1, kAwake); + }()), + "invalid sleep state index 2 when waking tree 1"); + + EXPECT_FATAL_FAILURE( + ([] { + int asleep_bad3[] = {1, 2, 1}; + mj_wakeTree(asleep_bad3, 3, 0, kAwake); + }()), + "tree 0 is not in a cycle"); +} + +static const char* const kStaticModel = "engine/testdata/sleep/static.xml"; +static const char* const kSmoothModel = "engine/testdata/sleep/smooth.xml"; +static const char* const kInitModel = "engine/testdata/sleep/init.xml"; +static const char* const kInitIslandModel = + "engine/testdata/sleep/init_island.xml"; +static const char* const kTendonModel = "engine/testdata/sleep/tendon.xml"; +static const char* const kContactModel = "engine/testdata/sleep/contact.xml"; +static const char* const kPairModel = "engine/testdata/sleep/contactpair.xml"; +static const char* const kSensorModel = "engine/testdata/sleep/sensor.xml"; + +// roll out some models with sleeping enabled, valuable under ASAN and MSAN +TEST_F(SleepTest, KickTires) { + for (const char* path : + {kStaticModel, kInitModel, kInitIslandModel, kSensorModel, kTendonModel, + kContactModel, kPairModel, kSmoothModel}) { + const std::string xml_path = GetTestDataFilePath(path); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), 0, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << error; + + int duration_id = mj_name2id(m, mjOBJ_NUMERIC, "duration"); + ASSERT_GE(duration_id, 0); + mjtNum duration = m->numeric_data[m->numeric_adr[duration_id]]; + + mjData* d = mj_makeData(m); + while (d->time < duration) { + mj_step(m, d); + } + + mj_deleteData(d); + mj_deleteModel(m); + } +} + +// Test that sleeping does not affect the simulation of awake trees: +// Roll out kSmoothModel, where all trees go to sleep within `duration` seconds +// in two mjData's, one with sleeping enabled and one without; expect the same +// values (for selected arrays) in awake trees in both. +TEST_F(SleepTest, WakingUnaffectedBySleeping) { + const std::string xml_path = GetTestDataFilePath(kSmoothModel); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), 0, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << error; + + int duration_id = mj_name2id(m, mjOBJ_NUMERIC, "duration"); + ASSERT_GE(duration_id, 0); + mjtNum duration = m->numeric_data[m->numeric_adr[duration_id]]; + + for (mjtJacobian jacobian : {mjJAC_DENSE, mjJAC_SPARSE}) { + m->opt.jacobian = jacobian; + for (mjtIntegrator integrator : // TODO: b/457674312 - Add support for RK4. + {mjINT_EULER, mjINT_IMPLICITFAST, mjINT_IMPLICIT}) { + m->opt.integrator = integrator; + + // make data with sleeping enabled + m->opt.enableflags |= mjENBL_SLEEP; + mjData* d_sleep = mj_makeData(m); + + // make data with sleeping disabled + m->opt.enableflags &= ~mjENBL_SLEEP; + mjData* d_nosleep = mj_makeData(m); + + // disable constraints, contacts + m->opt.disableflags |= mjDSBL_CONSTRAINT | mjDSBL_CONTACT; + + ASSERT_EQ(d_sleep->nbody_awake, m->nbody); + int nbody_awake = -1; + + while (d_nosleep->time < duration) { + m->opt.enableflags |= mjENBL_SLEEP; + mj_step(m, d_sleep); + m->opt.enableflags &= ~mjENBL_SLEEP; + mj_step(m, d_nosleep); + + // if nbody_awake is not changed, skip + if (d_sleep->nbody_awake == nbody_awake) { + continue; + } + + // compare xpos + for (int i = 0; i < m->nbody; i++) { + if (d_sleep->body_awake[i] == mjS_ASLEEP) continue; + auto xpos1 = AsVector(d_nosleep->xpos + 3 * i, 3); + auto xpos2 = AsVector(d_sleep->xpos + 3 * i, 3); + EXPECT_EQ(xpos1, xpos2) + << " xpos[" << i << "] at time " << d_nosleep->time; + } + + // compare M and qLD + for (int i = 0; i < d_sleep->nv_awake; i++) { + int j = d_sleep->dof_awake_ind[i]; + auto M1 = AsVector(d_nosleep->M + m->M_rowadr[j], m->M_rownnz[j]); + auto M2 = AsVector(d_sleep->M + m->M_rowadr[j], m->M_rownnz[j]); + EXPECT_EQ(M1, M2) << " M[" << j << ",:] at time " << d_nosleep->time; + auto qLD1 = AsVector(d_nosleep->qLD + m->M_rowadr[j], m->M_rownnz[j]); + auto qLD2 = AsVector(d_sleep->qLD + m->M_rowadr[j], m->M_rownnz[j]); + EXPECT_EQ(qLD1, qLD2) + << " qLD[" << j << ",:] at time " << d_nosleep->time; + } + + // compare cvel + for (int i = 0; i < d_sleep->nbody_awake; i++) { + if (d_sleep->body_awake[i] == mjS_ASLEEP) continue; + auto cvel1 = AsVector(d_nosleep->cvel + 6 * i, 6); + auto cvel2 = AsVector(d_sleep->cvel + 6 * i, 6); + EXPECT_EQ(cvel1, cvel2) + << " cvel[" << i << "] at time " << d_nosleep->time; + } + + // compare subtree_angmom, only for dynamic bodies + for (int i = 0; i < d_sleep->nbody_awake; i++) { + if (d_sleep->body_awake[i] != mjS_AWAKE) continue; + auto subtree_angmom1 = AsVector(d_nosleep->subtree_angmom + 3 * i, 3); + auto subtree_angmom2 = AsVector(d_sleep->subtree_angmom + 3 * i, 3); + EXPECT_EQ(subtree_angmom1, subtree_angmom2) + << " subtree_angmom[" << i << "] at time " << d_nosleep->time; + } + + // compare qfrc/qacc arrays + for (int i = 0; i < d_sleep->nv_awake; i++) { + int j = d_sleep->dof_awake_ind[i]; + EXPECT_EQ(d_nosleep->qfrc_smooth[j], d_sleep->qfrc_smooth[j]) + << " qfrc_smooth[" << j << "] at time " << d_nosleep->time; + EXPECT_EQ(d_nosleep->qacc_smooth[j], d_sleep->qacc_smooth[j]) + << " qacc_smooth[" << j << "] at time " << d_nosleep->time; + EXPECT_EQ(d_nosleep->qacc[j], d_sleep->qacc[j]) + << " qacc[" << j << "] at time " << d_nosleep->time; + } + + nbody_awake = d_sleep->nbody_awake; + } + + mj_deleteData(d_sleep); + mj_deleteData(d_nosleep); + } + } + mj_deleteModel(m); +} + + +// Test that waking does not affect sleeping trees for pos/vel-dependent arrays. +// Roll out models where some trees wake and/or sleep. At kCompare intervals, +// copy the state from the mjData with sleeping enabled to another mjData and +// call mj_forward with sleeping disabled. Expect pos/vel-dependent arrays to be +// unchanged for all trees and frc/acc-dependent arrays to be the same for awake +// trees. +TEST_F(SleepTest, SleepingUnaffectedByWaking) { + for (const char* path : {kInitModel, kInitIslandModel, kTendonModel, + kContactModel, kSensorModel, kSmoothModel}) { + const std::string xml_path = GetTestDataFilePath(path); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), 0, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << error; + + const int kCompare = 10; // number of comparisons per rollout + + // TODO: b/457674312 - Add support for RK4. + for (mjtIntegrator integrator : + {mjINT_EULER, mjINT_IMPLICITFAST, mjINT_IMPLICIT}) { + m->opt.integrator = integrator; + + // make data with sleeping enabled + m->opt.enableflags |= mjENBL_SLEEP; + mjData* d_sleep = mj_makeData(m); + + // make data with sleeping disabled + m->opt.enableflags &= ~mjENBL_SLEEP; + mjData* d_nosleep = mj_makeData(m); + + int duration_id = mj_name2id(m, mjOBJ_NUMERIC, "duration"); + ASSERT_GE(duration_id, 0); + mjtNum duration = m->numeric_data[m->numeric_adr[duration_id]]; + + int compare_interval = duration / (m->opt.timestep * kCompare); + int nsteps = 0; + while (d_sleep->time < duration) { + // step d_sleep with sleeping enabled + m->opt.enableflags |= mjENBL_SLEEP; + mj_step(m, d_sleep); + nsteps++; + + // every compare_interval steps, compare with d_nosleep + if (nsteps % compare_interval != 0) { + continue; + } + + // call mj_forward to update d_sleep + mj_forward(m, d_sleep); + + // copy state from d_sleep to d_nosleep + mj_copyData(d_nosleep, m, d_sleep); + + // forward d_nosleep with sleeping disabled + m->opt.enableflags &= ~mjENBL_SLEEP; + mj_forward(m, d_nosleep); + + // ==== compare arrays for all dofs / bodies / sensors ==== + + // compare xpos + for (int i = 0; i < m->nbody; i++) { + auto xpos1 = AsVector(d_sleep->xpos + 3 * i, 3); + auto xpos2 = AsVector(d_nosleep->xpos + 3 * i, 3); + EXPECT_EQ(xpos1, xpos2) + << " xpos[" << i << "] at time " << d_sleep->time; + } + + // compare M and qLD + for (int i = 0; i < m->nv; i++) { + auto M1 = AsVector(d_sleep->M + m->M_rowadr[i], m->M_rownnz[i]); + auto M2 = AsVector(d_nosleep->M + m->M_rowadr[i], m->M_rownnz[i]); + EXPECT_EQ(M1, M2) << " M[" << i << ",:] at time " << d_sleep->time; + auto qLD1 = AsVector(d_sleep->qLD + m->M_rowadr[i], m->M_rownnz[i]); + auto qLD2 = AsVector(d_nosleep->qLD + m->M_rowadr[i], m->M_rownnz[i]); + EXPECT_EQ(qLD1, qLD2) + << " qLD[" << i << ",:] at time " << d_sleep->time; + } + + // compare cvel + for (int i = 0; i < m->nbody; i++) { + auto cvel1 = AsVector(d_sleep->cvel + 6 * i, 6); + auto cvel2 = AsVector(d_nosleep->cvel + 6 * i, 6); + EXPECT_EQ(cvel1, cvel2) + << " cvel[" << i << "] at time " << d_sleep->time; + } + + // compare qfrc arrays + for (int i = 0; i < m->nv; i++) { + EXPECT_EQ(d_sleep->qfrc_fluid[i], d_nosleep->qfrc_fluid[i]) + << " qfrc_fluid[" << i << "] at time " << d_sleep->time; + EXPECT_EQ(d_sleep->qfrc_damper[i], d_nosleep->qfrc_damper[i]) + << " qfrc_damper[" << i << "] at time " << d_sleep->time; + EXPECT_EQ(d_sleep->qfrc_spring[i], d_nosleep->qfrc_spring[i]) + << " qfrc_spring[" << i << "] at time " << d_sleep->time; + EXPECT_EQ(d_sleep->qfrc_gravcomp[i], d_nosleep->qfrc_gravcomp[i]) + << " qfrc_gravcomp[" << i << "] at time " << d_sleep->time; + EXPECT_EQ(d_sleep->qfrc_bias[i], d_nosleep->qfrc_bias[i]) + << " qfrc_bias[" << i << "] at time " << d_sleep->time; + } + + // compare sensordata + for (int i = 0; i < m->nsensor; i++) { + int dim = m->sensor_dim[i]; + int adr = m->sensor_adr[i]; + auto data1 = AsVector(d_sleep->sensordata + adr, dim); + auto data2 = AsVector(d_nosleep->sensordata + adr, dim); + EXPECT_EQ(data1, data2) + << " sensor " << i << " at time " << d_sleep->time; + } + + // ==== compare arrays for awake dofs only ==== + + // compare qacc arrays for awake dofs + for (int j = 0; j < d_sleep->nv_awake; j++) { + int i = d_sleep->dof_awake_ind[j]; + EXPECT_EQ(d_sleep->qacc_smooth[i], d_nosleep->qacc_smooth[i]) + << " qacc_smooth[" << i << "] at time " << d_sleep->time; + EXPECT_EQ(d_sleep->qacc[i], d_nosleep->qacc[i]) + << " qacc[" << i << "] at time " << d_sleep->time; + } + } + + mj_deleteData(d_nosleep); + mj_deleteData(d_sleep); + } + mj_deleteModel(m); + } +} + +static const char* const kEqualityModel = "engine/testdata/sleep/equality.xml"; + +// Activate equality between sleeping and awake trees, useful under ASAN/MSAN. +TEST_F(SleepTest, Equality) { + const std::string xml_path = GetTestDataFilePath(kEqualityModel); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), 0, error, sizeof(error)); + ASSERT_THAT(m, NotNull()) << error; + + mjData* d = mj_makeData(m); + while (d->ntree_awake == m->ntree) { + mj_step(m, d); + } + + int dd = mj_name2id(m, mjOBJ_EQUALITY, "dyn/dyn"); + ASSERT_GE(dd, 0); + + mj_step(m, d); + d->eq_active[dd] = 1; + mj_step(m, d); + + mj_deleteData(d); + mj_deleteModel(m); +} + +static const char* const kInitIslandFailModel = + "engine/testdata/sleep/init_island_fail.xml"; + +TEST_F(SleepTest, InitIslandFail) { + const std::string xml_path = GetTestDataFilePath(kInitIslandFailModel); + char error[1024]; + mjModel* m = mj_loadXML(xml_path.c_str(), 0, error, sizeof(error)); + EXPECT_THAT(m, IsNull()); + EXPECT_THAT(error, + HasSubstr("3 trees were marked as sleep='init' but only 0 could " + "be slept.\nBody 'asleep_init0' (id=1) is the root of " + "the first tree that could not be slept.")); +} + } // namespace } // namespace mujoco diff --git a/test/engine/testdata/sleep/contact.xml b/test/engine/testdata/sleep/contact.xml new file mode 100644 index 00000000..5192285c --- /dev/null +++ b/test/engine/testdata/sleep/contact.xml @@ -0,0 +1,36 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/contactpair.xml b/test/engine/testdata/sleep/contactpair.xml new file mode 100644 index 00000000..22bd76dc --- /dev/null +++ b/test/engine/testdata/sleep/contactpair.xml @@ -0,0 +1,39 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/equality.xml b/test/engine/testdata/sleep/equality.xml new file mode 100644 index 00000000..413d4675 --- /dev/null +++ b/test/engine/testdata/sleep/equality.xml @@ -0,0 +1,60 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/init.xml b/test/engine/testdata/sleep/init.xml new file mode 100644 index 00000000..0040fea3 --- /dev/null +++ b/test/engine/testdata/sleep/init.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/init_island.xml b/test/engine/testdata/sleep/init_island.xml new file mode 100644 index 00000000..e8823201 --- /dev/null +++ b/test/engine/testdata/sleep/init_island.xml @@ -0,0 +1,26 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/init_island_fail.xml b/test/engine/testdata/sleep/init_island_fail.xml new file mode 100644 index 00000000..24ea60c3 --- /dev/null +++ b/test/engine/testdata/sleep/init_island_fail.xml @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/perf/2humanoid100.xml b/test/engine/testdata/sleep/perf/2humanoid100.xml new file mode 100644 index 00000000..d7dd930e --- /dev/null +++ b/test/engine/testdata/sleep/perf/2humanoid100.xml @@ -0,0 +1,123 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/perf/dominos.xml b/test/engine/testdata/sleep/perf/dominos.xml new file mode 100644 index 00000000..805c6541 --- /dev/null +++ b/test/engine/testdata/sleep/perf/dominos.xml @@ -0,0 +1,77 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/perf/humanoid.xml b/test/engine/testdata/sleep/perf/humanoid.xml new file mode 100644 index 00000000..b3192a1b --- /dev/null +++ b/test/engine/testdata/sleep/perf/humanoid.xml @@ -0,0 +1,254 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/perf/smooth.xml b/test/engine/testdata/sleep/perf/smooth.xml new file mode 100644 index 00000000..08ddc78b --- /dev/null +++ b/test/engine/testdata/sleep/perf/smooth.xml @@ -0,0 +1,40 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/perf/static.xml b/test/engine/testdata/sleep/perf/static.xml new file mode 100644 index 00000000..4afc67be --- /dev/null +++ b/test/engine/testdata/sleep/perf/static.xml @@ -0,0 +1,40 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/perf/world.xml b/test/engine/testdata/sleep/perf/world.xml new file mode 100644 index 00000000..32ce26f8 --- /dev/null +++ b/test/engine/testdata/sleep/perf/world.xml @@ -0,0 +1,40 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/sensor.xml b/test/engine/testdata/sleep/sensor.xml new file mode 100644 index 00000000..2d3e9174 --- /dev/null +++ b/test/engine/testdata/sleep/sensor.xml @@ -0,0 +1,40 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/smooth.xml b/test/engine/testdata/sleep/smooth.xml new file mode 100644 index 00000000..96113f8f --- /dev/null +++ b/test/engine/testdata/sleep/smooth.xml @@ -0,0 +1,54 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/static.xml b/test/engine/testdata/sleep/static.xml new file mode 100644 index 00000000..c363fb90 --- /dev/null +++ b/test/engine/testdata/sleep/static.xml @@ -0,0 +1,35 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/tendon.xml b/test/engine/testdata/sleep/tendon.xml new file mode 100644 index 00000000..3482d0a5 --- /dev/null +++ b/test/engine/testdata/sleep/tendon.xml @@ -0,0 +1,33 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/sleep/tendon2.xml b/test/engine/testdata/sleep/tendon2.xml new file mode 100644 index 00000000..89dfdfb3 --- /dev/null +++ b/test/engine/testdata/sleep/tendon2.xml @@ -0,0 +1,40 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/sample/testspeed_test.sh b/test/sample/testspeed_test.sh index c5a76445..79d6c8c5 100755 --- a/test/sample/testspeed_test.sh +++ b/test/sample/testspeed_test.sh @@ -33,7 +33,7 @@ test_model() { "$model" == */replicate/bunnies.xml || "$model" == */replicate/leaves.xml || "$model" == */replicate/particle.xml || - "$model" == */engine/testdata/collision_convex/perf/* + "$model" == */perf/* ]]; then # these tests can take several minutes under ASAN return 0 @@ -76,7 +76,7 @@ for model_dir in ${MODEL_DIRS[@]}; do echo "Skipping $model" >&2 continue fi - if [[ $(basename $model) == malformed* ]]; then + if [[ $(basename $model) == *_fail.xml ]]; then echo "Skipping $model" >&2 continue fi diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 4823aa07..73e36a48 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -531,9 +531,11 @@ TEST_F(PluginTest, RecompileCompare) { if (p.path().extension() == ext) { std::string xml = p.path().string(); - // if file is meant to fail, skip it + // if file is meant to fail or model is too slow to load, skip it if (absl::StrContains(p.path().string(), "malformed_") || + absl::StrContains(p.path().string(), "_fail") || absl::StrContains(p.path().string(), "touch_grid") || + absl::StrContains(p.path().string(), "perf") || absl::StrContains(p.path().string(), "cow")) { continue; } diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc index c6148234..48db4604 100644 --- a/test/xml/xml_native_writer_test.cc +++ b/test/xml/xml_native_writer_test.cc @@ -1378,25 +1378,27 @@ TEST_F(XMLWriterTest, WriteReadCompare) { if (p.path().extension() == ext) { std::string xml = p.path().string(); - // if file is meant to fail, skip it - if (absl::StrContains(p.path().string(), "malformed_") || - // exclude files that are too slow to load - absl::StrContains(p.path().string(), "cow") || - absl::StrContains(p.path().string(), "gmsh_") || - absl::StrContains(p.path().string(), "shark_") || - absl::StrContains(p.path().string(), "spheremesh") || - // exclude files that fail the comparison test - absl::StrContains(p.path().string(), "tactile") || - absl::StrContains(p.path().string(), "makemesh") || - absl::StrContains(p.path().string(), "many_dependencies") || - absl::StrContains(p.path().string(), "usd") || - absl::StrContains(p.path().string(), "torus_maxhull") || - absl::StrContains(p.path().string(), "fitmesh_") || - absl::StrContains(p.path().string(), "lengthrange") || - absl::StrContains(p.path().string(), "hfield_xml") || - absl::StrContains(p.path().string(), "fromto_convex") || - absl::StrContains(p.path().string(), "cube_skin") || - absl::StrContains(p.path().string(), "cube_3x3x3")) { + + if ( // if file is meant to fail, skip it + absl::StrContains(p.path().string(), "malformed_") || + absl::StrContains(p.path().string(), "_fail") || + // exclude files that are too slow to load + absl::StrContains(p.path().string(), "cow") || + absl::StrContains(p.path().string(), "gmsh_") || + absl::StrContains(p.path().string(), "shark_") || + absl::StrContains(p.path().string(), "perf") || + // exclude files that fail the comparison test + absl::StrContains(p.path().string(), "tactile") || + absl::StrContains(p.path().string(), "makemesh") || + absl::StrContains(p.path().string(), "many_dependencies") || + absl::StrContains(p.path().string(), "usd") || + absl::StrContains(p.path().string(), "torus_maxhull") || + absl::StrContains(p.path().string(), "fitmesh_") || + absl::StrContains(p.path().string(), "lengthrange") || + absl::StrContains(p.path().string(), "hfield_xml") || + absl::StrContains(p.path().string(), "fromto_convex") || + absl::StrContains(p.path().string(), "cube_skin") || + absl::StrContains(p.path().string(), "cube_3x3x3")) { continue; } // load model