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

Implement sleeping in engine

Browse Source 
PiperOrigin-RevId: 829361787
Change-Id: I6f64d8e25c4248cf32c18cd94d37ff5def78946e
This commit is contained in:
Yuval Tassa
2025-11-07 03:32:03 -08:00
committed by Copybara-Service
parent 1e0226d360
commit 769f37b653
55 changed files with 3602 additions and 677 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+331 -127
View File
@@ -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