// Copyright 2021 DeepMind Technologies Limited // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. #include "engine/engine_core_smooth.h" #include #include #include #include #include // IWYU pragma: keep #include "engine/engine_core_constraint.h" #include "engine/engine_core_util.h" #include "engine/engine_crossplatform.h" #include "engine/engine_inline.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 part 1: bodies void mj_kinematics1(const mjModel* m, mjData* d) { int nbody = m->nbody; // set world position and orientation mju_zero3(d->xpos); mju_unit4(d->xquat); mju_zero3(d->xipos); mju_zero(d->xmat, 9); mju_zero(d->ximat, 9); 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]; // free joint if (jntnum == 1 && m->jnt_type[jntadr] == mjJNT_FREE) { // get qpos address int qadr = m->jnt_qposadr[jntadr]; // copy pos and quat from qpos mji_copy3(xpos, d->qpos+qadr); mji_copy4(xquat, d->qpos+qadr+3); mju_normalize4(xquat); // assign xanchor and xaxis mji_copy3(d->xanchor+3*jntadr, xpos); mji_copy3(d->xaxis+3*jntadr, m->jnt_axis+3*jntadr); } // regular or no joint else { int pid = m->body_parentid[i]; // get body pos and quat: from model or mocap mjtNum *bodypos, *bodyquat, quat[4]; if (m->body_mocapid[i] >= 0) { bodypos = d->mocap_pos + 3*m->body_mocapid[i]; mji_copy4(quat, d->mocap_quat + 4*m->body_mocapid[i]); mju_normalize4(quat); bodyquat = quat; } else { bodypos = m->body_pos+3*i; bodyquat = m->body_quat+4*i; } // apply fixed translation and rotation relative to parent if (pid) { mji_mulMatVec3(xpos, d->xmat+9*pid, bodypos); mji_addTo3(xpos, d->xpos+3*pid); mji_mulQuat(xquat, d->xquat+4*pid, bodyquat); } else { // parent is the world mji_copy3(xpos, bodypos); mji_copy4(xquat, bodyquat); } // accumulate joints, compute xpos and xquat for this body mjtNum xanchor[3], xaxis[3]; for (int j=0; j < jntnum; j++) { // get joint id, qpos address, joint type int jid = jntadr + j; int qadr = m->jnt_qposadr[jid]; mjtJoint jtype = m->jnt_type[jid]; // compute axis in global frame; ball jnt_axis is (0,0,1), set by compiler mji_rotVecQuat(xaxis, m->jnt_axis+3*jid, xquat); // compute anchor in global frame mji_rotVecQuat(xanchor, m->jnt_pos+3*jid, xquat); mji_addTo3(xanchor, xpos); // apply joint transformation switch (jtype) { case mjJNT_SLIDE: mji_addToScl3(xpos, xaxis, d->qpos[qadr] - m->qpos0[qadr]); break; case mjJNT_BALL: case mjJNT_HINGE: { // compute local quaternion rotation mjtNum qloc[4]; if (jtype == mjJNT_BALL) { mji_copy4(qloc, d->qpos+qadr); mju_normalize4(qloc); } else { mji_axisAngle2Quat(qloc, m->jnt_axis+3*jid, d->qpos[qadr] - m->qpos0[qadr]); } // apply rotation mju_mulQuat(xquat, xquat, qloc); // correct for off-center rotation mjtNum vec[3]; mji_rotVecQuat(vec, m->jnt_pos+3*jid, xquat); mji_sub3(xpos, xanchor, vec); } break; default: mjERROR("unknown joint type %d", jtype); // SHOULD NOT OCCUR } // assign xanchor and xaxis mji_copy3(d->xanchor+3*jid, xanchor); mji_copy3(d->xaxis+3*jid, xaxis); } } // 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 mji_copy4(d->xquat+4*i, xquat); mji_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 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 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, 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 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 b=0; b < nbody; b++) { int i = sleep_filter ? d->body_awake_ind[b] : b; mji_scl3(d->subtree_com+3*i, d->xipos+3*i, m->body_mass[i]); } // subtree_com: accumulate to parent in backward pass 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]; mji_scl3(child_moment, d->subtree_com+3*i, m->body_subtreemass[i]); mji_addTo3(d->subtree_com+3*parent, child_moment); } else { mji_addTo3(d->subtree_com+3*parent, d->subtree_com+3*i); } } // subtree_com: normalize for (int b=0; b < nbody; b++) { int i = sleep_filter ? d->body_awake_ind[b] : b; if (m->body_subtreemass[i] < mjMINVAL) { mji_copy3(d->subtree_com+3*i, d->xipos+3*i); } else { mju_scl3(d->subtree_com + 3 * i, d->subtree_com + 3 * i, 1.0 / m->body_subtreemass[i]); } } // zero out CoM frame inertia for the world body mju_zero(d->cinert, 10); // map inertias to frame centered at subtree_com for (int b=1; b < nbody; b++) { int i = sleep_filter ? d->body_awake_ind[b] : b; mjtNum offset[3]; mji_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]); } // map motion dofs to global frame centered at subtree_com for (int b=1; b < nbody; b++) { int i = sleep_filter ? d->body_awake_ind[b] : b; int jntnum = m->body_jntnum[i]; if (!jntnum) continue; 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]; mji_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; } } } } // compute camera and light positions and orientations void mj_camlight(const mjModel* m, mjData* d) { 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 < 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: break; case mjCAMLIGHT_TRACK: case mjCAMLIGHT_TRACKCOM: // fixed global orientation mji_copy9(d->cam_xmat+9*i, m->cam_mat0+9*i); // position: track camera body if (m->cam_mode[i] == mjCAMLIGHT_TRACK) { mji_add3(d->cam_xpos+3*i, d->xpos+3*id, m->cam_pos0+3*i); } // position: track subtree com else { mji_add3(d->cam_xpos+3*i, d->subtree_com+3*id, m->cam_poscom0+3*i); } break; case mjCAMLIGHT_TARGETBODY: 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) { mji_copy3(pos, d->xpos+3*id1); } else { mji_copy3(pos, d->subtree_com+3*id1); } // zaxis = -desired camera direction, in global frame mjtNum matT[9]; mji_sub3(matT+6, d->cam_xpos+3*i, pos); mju_normalize3(matT+6); // xaxis: orthogonal to zaxis and to (0,0,1) matT[3] = 0; matT[4] = 0; matT[5] = 1; mji_cross(matT, matT+3, matT+6); mju_normalize3(matT); // yaxis: orthogonal to xaxis and zaxis mji_cross(matT+3, matT+6, matT); mju_normalize3(matT+3); // set camera frame mju_transpose(d->cam_xmat+9*i, matT, 3, 3); } } } // compute Cartesian positions and directions of lights 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); mji_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: break; case mjCAMLIGHT_TRACK: case mjCAMLIGHT_TRACKCOM: // fixed global orientation mji_copy3(d->light_xdir+3*i, m->light_dir0+3*i); // position: track light body if (m->light_mode[i] == mjCAMLIGHT_TRACK) { mji_add3(d->light_xpos+3*i, d->xpos+3*id, m->light_pos0+3*i); } // position: track subtree com else { mji_add3(d->light_xpos+3*i, d->subtree_com+3*id, m->light_poscom0+3*i); } break; case mjCAMLIGHT_TARGETBODY: case mjCAMLIGHT_TARGETBODYCOM: // only if target body is specified if (id1 >= 0) { // get position to look at mjtNum lookat[3]; if (m->light_mode[i] == mjCAMLIGHT_TARGETBODY) { mji_copy3(lookat, d->xpos+3*id1); } else { mji_copy3(lookat, d->subtree_com+3*id1); } // set dir mji_sub3(d->light_xdir+3*i, lookat, d->light_xpos+3*i); } } // normalize dir mju_normalize3(d->light_xdir+3*i); } } // update dynamic BVH; leaf aabbs must be updated before call void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) { mj_markStack(d); int* modified = mjSTACKALLOC(d, bvhnum, int); mju_zeroInt(modified, bvhnum); // mark leafs as modified for (int i=0; i < bvhnum; i++) { if (m->bvh_nodeid[bvhadr+i] >= 0) { modified[i] = 1; } } // update non-leafs in backward pass (parents come before children) for (int i=bvhnum-1; i >= 0; i--) { if (m->bvh_nodeid[bvhadr+i] < 0) { int child1 = m->bvh_child[2*(bvhadr+i)]; int child2 = m->bvh_child[2*(bvhadr+i)+1]; // update if either child is modified if (modified[child1] || modified[child2]) { mjtNum* aabb = d->bvh_aabb_dyn + 6*(bvhadr - m->nbvhstatic + i); const mjtNum* aabb1 = d->bvh_aabb_dyn + 6*(bvhadr - m->nbvhstatic + child1); const mjtNum* aabb2 = d->bvh_aabb_dyn + 6*(bvhadr - m->nbvhstatic + child2); // compute new (min, max) mjtNum xmin[3], xmax[3]; for (int k=0; k < 3; k++) { xmin[k] = mju_min(aabb1[k] - aabb1[k+3], aabb2[k] - aabb2[k+3]); xmax[k] = mju_max(aabb1[k] + aabb1[k+3], aabb2[k] + aabb2[k+3]); } // convert to (center, size) for (int k=0; k < 3; k++) { aabb[k] = 0.5*(xmax[k]+xmin[k]); aabb[k+3] = 0.5*(xmax[k]-xmin[k]); } modified[i] = 1; } } } mj_freeStack(d); } // C(3x2) = A(3x2) * B(2x2) static inline void mju_mulMatMat322(mjtNum* C, const mjtNum* A, const mjtNum* B) { C[0] = A[0]*B[0] + A[1]*B[2]; C[1] = A[0]*B[1] + A[1]*B[3]; C[2] = A[2]*B[0] + A[3]*B[2]; C[3] = A[2]*B[1] + A[3]*B[3]; C[4] = A[4]*B[0] + A[5]*B[2]; C[5] = A[4]*B[1] + A[5]*B[3]; } // compute flex-related quantities void mj_flex(const mjModel* m, mjData* d) { int nv = m->nv; int* rowadr = m->flexedge_J_rowadr; int* vrowadr = m->flexvert_J_rowadr, *vrownnz = m->flexvert_J_rownnz; // skip if no flexes if (!m->nflex) { return; } // compute Cartesian positions of flex vertices for (int f=0; f < m->nflex; f++) { int vstart = m->flex_vertadr[f]; int vend = m->flex_vertadr[f] + m->flex_vertnum[f]; int nstart = m->flex_nodeadr[f]; int nend = m->flex_nodeadr[f] + m->flex_nodenum[f]; // 0: vertices are the mesh vertices, 1: vertices are interpolated from nodal dofs if (m->flex_interp[f] == 0) { for (int i=vstart; i < vend; i++) { if (m->flex_centered[f] || (m->flex_vert[3*i+0] == 0 && m->flex_vert[3*i+1] == 0 && m->flex_vert[3*i+2] == 0)) { mji_copy3(d->flexvert_xpos+3*i, d->xpos+3*m->flex_vertbodyid[i]); } else { mji_mulMatVec3(d->flexvert_xpos+3*i, d->xmat+9*m->flex_vertbodyid[i], m->flex_vert+3*i); mji_addTo3(d->flexvert_xpos+3*i, d->xpos+3*m->flex_vertbodyid[i]); } } } // trilinear/quadratic interpolation else { int nodenum = nend - nstart; mj_markStack(d); mjtNum* nodexpos = mjSTACKALLOC(d, 3*nodenum, mjtNum); for (int i=nstart; i < nend; i++) { int j = i - nstart; if (m->flex_centered[f] || (m->flex_node[3*i+0] == 0 && m->flex_node[3*i+1] == 0 && m->flex_node[3*i+2] == 0)) { mji_copy3(nodexpos + 3*j, d->xpos + 3*m->flex_nodebodyid[i]); } else { mji_mulMatVec3(nodexpos + 3*j, d->xmat + 9*m->flex_nodebodyid[i], m->flex_node + 3*i); mji_addTo3(nodexpos + 3*j, d->xpos + 3*m->flex_nodebodyid[i]); } } int order = m->flex_interp[f]; int cx = m->flex_cellnum[3*f+0]; int cy = m->flex_cellnum[3*f+1]; int cz = m->flex_cellnum[3*f+2]; int nx_g = cx * order + 1; int ny_g = cy * order + 1; int nz_g = cz * order + 1; if (nend - nstart != nx_g * ny_g * nz_g) { mjERROR("flex_interp_order mismatch"); } for (int i=vstart; i < vend; i++) { mju_zero3(d->flexvert_xpos+3*i); // cell lookup: get local coords and node indices mjtNum local[3]; int nodeindices[27]; // max npc for quadratic: 3^3 = 27 mju_cellLookup(m->flex_vert0 + 3*i, m->flex_cellnum+3*f, order, local, nodeindices); mju_interpolate3D(d->flexvert_xpos+3*i, local, nodexpos, order, nodeindices); } mj_freeStack(d); } } // compute flex element aabb for (int f=0; f < m->nflex; f++) { int dim = m->flex_dim[f]; // process elements of this flex int elemnum = m->flex_elemnum[f]; for (int e=0; e < elemnum; e++) { const int* edata = m->flex_elem + m->flex_elemdataadr[f] + e*(dim+1); const mjtNum* vert = d->flexvert_xpos + 3*m->flex_vertadr[f]; // compute min and max along each global axis mjtNum xmin[3], xmax[3]; mji_copy3(xmin, vert+3*edata[0]); mji_copy3(xmax, vert+3*edata[0]); for (int i=1; i <= dim; i++) { for (int j=0; j < 3; j++) { mjtNum value = vert[3*edata[i]+j]; xmin[j] = mju_min(xmin[j], value); xmax[j] = mju_max(xmax[j], value); } } // compute aabb (center, size) int base = m->flex_elemadr[f] + e; d->flexelem_aabb[6*base+0] = 0.5*(xmax[0]+xmin[0]); d->flexelem_aabb[6*base+1] = 0.5*(xmax[1]+xmin[1]); d->flexelem_aabb[6*base+2] = 0.5*(xmax[2]+xmin[2]); d->flexelem_aabb[6*base+3] = 0.5*(xmax[0]-xmin[0]) + m->flex_radius[f]; d->flexelem_aabb[6*base+4] = 0.5*(xmax[1]-xmin[1]) + m->flex_radius[f]; d->flexelem_aabb[6*base+5] = 0.5*(xmax[2]-xmin[2]) + m->flex_radius[f]; } } // update flex bhv_aabb_dyn if needed if (!mjDISABLED(mjDSBL_MIDPHASE)) { for (int f=0; f < m->nflex; f++) { if (m->flex_bvhadr[f] >= 0) { int flex_bvhadr = m->flex_bvhadr[f]; int flex_bvhnum = m->flex_bvhnum[f]; // copy element aabbs to bhv leaf aabbs for (int i=flex_bvhadr; i < flex_bvhadr+flex_bvhnum; i++) { if (m->bvh_nodeid[i] >= 0) { mji_copy6(d->bvh_aabb_dyn + 6*(i - m->nbvhstatic), d->flexelem_aabb + 6*(m->flex_elemadr[f] + m->bvh_nodeid[i])); } } // update dynamic BVH mj_updateDynamicBVH(m, d, m->flex_bvhadr[f], m->flex_bvhnum[f]); } } } // allocate space mj_markStack(d); mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); int* chain = mjSTACKALLOC(d, nv, int); // clear Jacobian mju_zero(d->flexvert_J, 2*m->nJfv); mju_zero(d->flexedge_J, m->nJfe); // compute lengths and Jacobians of edges for (int f=0; f < m->nflex; f++) { // skip if edges cannot generate forces if (m->flex_rigid[f] || m->flex_interp[f]) { continue; } // skip edge Jacobian if no built-in passive force is needed int skipjacobian = m->flex_edgeequality[f] != 1 && !m->flex_edgedamping[f] && !m->flex_edgestiffness[f] && !m->flex_damping[f]; // process edges of this flex int vbase = m->flex_vertadr[f]; int ebase = m->flex_edgeadr[f]; int edgenum = m->flex_edgenum[f]; for (int e=0; e < edgenum; e++) { int v1 = m->flex_edge[2*(ebase+e)]; int v2 = m->flex_edge[2*(ebase+e)+1]; int b1 = m->flex_vertbodyid[vbase+v1]; int b2 = m->flex_vertbodyid[vbase+v2]; mjtNum* pos1 = d->flexvert_xpos + 3*(vbase+v1); mjtNum* pos2 = d->flexvert_xpos + 3*(vbase+v2); // vec = unit vector from v1 to v2, compute edge length mjtNum vec[3]; mji_sub3(vec, pos2, pos1); d->flexedge_length[ebase+e] = mju_normalize3(vec); // skip Jacobian if not needed if (skipjacobian) { continue; } // get endpoint Jacobians, subtract int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2, jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1, /*skipcommon=*/0); // no dofs: skip if (!NV) { continue; } // apply chain rule to compute edge Jacobian mju_mulMatTVec(d->flexedge_J + rowadr[ebase+e], jacdif, vec, 3, NV); } // if dim=2 and constraints are active we use the vertex-based constraint defined in // Chen, Kry, and Vouga, "Locking-free Simulation of Isometric Thin Plates", 2019. if (m->flex_dim[f] == 2 && m->flex_edgeequality[f] == 2) { int nvert = m->flex_vertnum[f]; // use global vertex adjacency list int* v_edge_cnt = m->flex_vertedgenum + vbase; int* v_edge_adr = m->flex_vertedgeadr + vbase; int* adj_edges = m->flex_vertedge; mj_markStack(d); // clear Jacobian and assemble vertex by vertex int* chain1 = mjSTACKALLOC(d, nv, int); int* chain2 = mjSTACKALLOC(d, nv, int); mjtNum* J0_dense = mjSTACKALLOC(d, nv, mjtNum); mjtNum* J1_dense = mjSTACKALLOC(d, nv, mjtNum); mju_zero(J0_dense, nv); mju_zero(J1_dense, nv); // temporary buffer for Jacobian accumulation mjtNum* J_local = mjSTACKALLOC(d, nv, mjtNum); for (int v = 0; v < nvert; v++) { mjtNum A[6] = {0}; int vadr = vbase + v; mjtNum* metric = m->flex_vertmetric + 4 * vadr; for (int k = 0; k < v_edge_cnt[v]; k++) { int e = adj_edges[v_edge_adr[v] + k]; // compute rest configuration edge vector mjtNum dx[3]; int v1 = m->flex_edge[2 * (ebase + e)]; int v2 = m->flex_edge[2 * (ebase + e) + 1]; mju_sub3(dx, m->flex_vert0 + 3 * (vbase + v2), m->flex_vert0 + 3 * (vbase + v1)); // apply scaling since they are half sizes dx[0] *= 2 * m->flex_size[3 * f + 0]; dx[1] *= 2 * m->flex_size[3 * f + 1]; dx[2] *= 2 * m->flex_size[3 * f + 2]; mjtNum dy[3]; mju_sub3(dy, d->flexvert_xpos + 3 * (vbase + v2), d->flexvert_xpos + 3 * (vbase + v1)); // get mass of neighbor vertex mjtNum weight = 1.0; int neighbor_v = (v == v1) ? v2 : v1; int b_neighbor = m->flex_vertbodyid[vbase + neighbor_v]; if (b_neighbor >= 0) { weight = m->body_mass[b_neighbor]; if (weight < mjMINVAL) weight = mjMINVAL; } // accumulate A += w * dy * dx' A[0] += weight * dy[0] * dx[0]; A[1] += weight * dy[0] * dx[1]; A[2] += weight * dy[1] * dx[0]; A[3] += weight * dy[1] * dx[1]; A[4] += weight * dy[2] * dx[0]; A[5] += weight * dy[2] * dx[1]; } mjtNum F[6]; mju_mulMatMat322(F, A, metric); // compute Cauchy strain tensor F^T F mjtNum cauchy[4]; cauchy[0] = F[0] * F[0] + F[2] * F[2] + F[4] * F[4]; // c00 cauchy[1] = F[0] * F[1] + F[2] * F[3] + F[4] * F[5]; // c01 cauchy[2] = F[1] * F[0] + F[3] * F[2] + F[5] * F[4]; // c10 cauchy[3] = F[1] * F[1] + F[3] * F[3] + F[5] * F[5]; // c11 // mass scaling: scale constraint by sqrt(mass) to improve condition number // note: departure from original algorithm in Chen, Kry, and Vouga 2019 mjtNum scale = 1.0; int b = m->flex_vertbodyid[vadr]; if (b >= 0) { mjtNum mass = m->body_mass[b]; if (mass > mjMINVAL) { scale = mju_sqrt(mass); } } // compute tensor invariants d->flexvert_length[2 * vadr + 0] = (cauchy[0] + cauchy[3] - 2) * scale; d->flexvert_length[2 * vadr + 1] = (cauchy[0] * cauchy[3] - cauchy[1] * cauchy[2] - 1) * scale; // Jacobian computation mjtNum FB[6], adj[4], Fadj[6], FadjBinv[6]; mju_mulMatMat322(FB, F, metric); adj[0] = cauchy[3]; adj[1] = -cauchy[1]; adj[2] = -cauchy[2]; adj[3] = cauchy[0]; mju_mulMatMat322(Fadj, F, adj); mju_mulMatMat322(FadjBinv, Fadj, metric); for (int k = 0; k < v_edge_cnt[v]; ++k) { int e = adj_edges[v_edge_adr[v] + k]; mjtNum dx[3]; // rest edge vector int v1 = m->flex_edge[2 * (ebase + e)]; int v2 = m->flex_edge[2 * (ebase + e) + 1]; mju_sub3(dx, m->flex_vert0 + 3 * (vbase + v2), m->flex_vert0 + 3 * (vbase + v1)); dx[0] *= 2 * m->flex_size[3 * f + 0]; dx[1] *= 2 * m->flex_size[3 * f + 1]; dx[2] *= 2 * m->flex_size[3 * f + 2]; mjtNum weight = 1.0; int neighbor_v = (v == v1) ? v2 : v1; int b_neighbor = m->flex_vertbodyid[vbase + neighbor_v]; if (b_neighbor >= 0) { weight = m->body_mass[b_neighbor]; if (weight < mjMINVAL) weight = mjMINVAL; } mjtNum dI1dy1[3], dI1dy2[3], dI2dy[3], dI2dy1[3], dI2dy2[3]; // dI1/dy1, dI1/dy2 (scaled by weight) mju_mulMatVec(dI1dy1, FB, dx, 3, 2); mju_scl3(dI1dy1, dI1dy1, -2 * weight); mju_scl3(dI1dy2, dI1dy1, -1); // dI2/dy1, dI2/dy2 (scaled by weight) mju_mulMatVec(dI2dy, FadjBinv, dx, 3, 2); mju_scl3(dI2dy1, dI2dy, -2 * weight); mju_scl3(dI2dy2, dI2dy1, -1); // get endpoint Jacobians int b1 = m->flex_vertbodyid[vbase+v1]; int b2 = m->flex_vertbodyid[vbase+v2]; int NV1 = mj_bodyChain(m, b1, chain1); mj_jacSparse(m, d, jac1, NULL, d->flexvert_xpos + 3*(vbase+v1), b1, NV1, chain1, /*flg_skipcommon=*/0); int NV2 = mj_bodyChain(m, b2, chain2); mj_jacSparse(m, d, jac2, NULL, d->flexvert_xpos + 3*(vbase+v2), b2, NV2, chain2, /*flg_skipcommon=*/0); // accumulate dense Jacobians for vertex v mju_mulMatTVec(J_local, jac1, dI1dy1, 3, NV1); for (int j=0; jflexvert_J_colind[vrowadr[row0] + j]; d->flexvert_J[vrowadr[row0] + j] += J0_dense[col] * scale; J0_dense[col] = 0; } int row1 = 2 * vadr + 1; int nnz1 = vrownnz[row1]; for (int j = 0; j < nnz1; j++) { int col = m->flexvert_J_colind[vrowadr[row1] + j]; d->flexvert_J[vrowadr[row1] + j] += J1_dense[col] * scale; J1_dense[col] = 0; } } mj_freeStack(d); } } mj_freeStack(d); } // compute tendon lengths and moments void mj_tendon(const mjModel* m, mjData* d) { int nv = m->nv, nten = m->ntendon; const int *rownnz = m->ten_J_rownnz, *rowadr = m->ten_J_rowadr, *colind = m->ten_J_colind; mjtNum *L = d->ten_length, *J = d->ten_J; if (!nten) { return; } // allocate stack arrays int *chain; mjtNum *jac1, *jac2, *jacdif, *tmp; mj_markStack(d); jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); tmp = mjSTACKALLOC(d, nv, mjtNum); chain = mjSTACKALLOC(d, nv, int); // clear results mju_zero(L, nten); // clear Jacobian mju_zero(J, m->nJten); // 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; d->ten_wrapnum[i] = 0; int tendon_num = m->tendon_num[i]; // process fixed tendon if (m->wrap_type[adr] == mjWRAP_JOINT) { // process all defined joints for (int j=0; j < tendon_num; j++) { // get joint id int k = m->wrap_objid[adr+j]; // add to length L[i] += m->wrap_prm[adr+j] * d->qpos[m->jnt_qposadr[k]]; mjtNum coef = 1; int dofadr = m->jnt_dofadr[k]; mju_combineSparseInc(J + rowadr[i], &coef, m->nv, 1, m->wrap_prm[adr+j], rownnz[i], 1, colind + rowadr[i], &dofadr); } continue; } // process spatial tendon mjtNum divisor = 1; int wraptype, j = 0; while (j < tendon_num-1) { // get 1st and 2nd object int type0 = m->wrap_type[adr+j+0]; int type1 = m->wrap_type[adr+j+1]; int id0 = m->wrap_objid[adr+j+0]; int id1 = m->wrap_objid[adr+j+1]; // pulley if (type0 == mjWRAP_PULLEY || type1 == mjWRAP_PULLEY) { // get divisor, insert obj=-2 if (type0 == mjWRAP_PULLEY) { divisor = m->wrap_prm[adr+j]; mju_zero3(d->wrap_xpos+wrapcount*3); d->wrap_obj[wrapcount] = -2; d->ten_wrapnum[i]++; wrapcount++; } // move to next j++; continue; } // init sequence; assume it starts with site mjtNum wlen = -1; int wrapid = -1; mjtNum wpnt[12]; mji_copy3(wpnt, d->site_xpos+3*id0); int wbody[4]; wbody[0] = m->site_bodyid[id0]; // second object is geom: process site-geom-site if (type1 == mjWRAP_SPHERE || type1 == mjWRAP_CYLINDER) { // reassign, get 2nd site info wraptype = type1; wrapid = id1; type1 = m->wrap_type[adr+j+2]; id1 = m->wrap_objid[adr+j+2]; // do wrapping, possibly get 2 extra points (wlen>=0) int sideid = mju_round(m->wrap_prm[adr+j+1]); if (sideid < -1 || sideid >= m->nsite) { mjERROR("invalid sideid %d in wrap_prm", sideid); // SHOULD NOT OCCUR } wlen = mju_wrap(wpnt+3, d->site_xpos+3*id0, d->site_xpos+3*id1, d->geom_xpos+3*wrapid, d->geom_xmat+9*wrapid, m->geom_size[3*wrapid], wraptype, (sideid >= 0 ? d->site_xpos+3*sideid : 0)); } else { wraptype = mjWRAP_NONE; } // complete sequence, accumulate lengths if (wlen < 0) { mji_copy3(wpnt+3, d->site_xpos+3*id1); wbody[1] = m->site_bodyid[id1]; L[i] += mju_dist3(wpnt, wpnt+3) / divisor; } else { mji_copy3(wpnt+9, d->site_xpos+3*id1); wbody[1] = wbody[2] = m->geom_bodyid[wrapid]; wbody[3] = m->site_bodyid[id1]; L[i] += (mju_dist3(wpnt, wpnt+3) + wlen + mju_dist3(wpnt+6, wpnt+9)) / divisor; } // accumulate moments if consecutive points are in different bodies for (int k=0; k < (wlen < 0 ? 1 : 3); k++) { if (wbody[k] != wbody[k+1]) { // get 3D position difference, normalize mjtNum dif[3]; mji_sub3(dif, wpnt+3*k+3, wpnt+3*k); mju_normalize3(dif); // get endpoint Jacobians, subtract int NV = mj_jacDifPair(m, d, chain, wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1, /*skipcommon=*/0); // no dofs: skip if (!NV) { continue; } // apply chain rule to compute tendon Jacobian mju_mulMatTVec(tmp, jacdif, dif, 3, NV); // add to existing mju_combineSparseInc(J+rowadr[i], tmp, nv, 1, 1/divisor, rownnz[i], NV, colind+rowadr[i], chain); } } // assign to wrap if (wlen < 0) { mji_copy3(d->wrap_xpos+wrapcount*3, wpnt); } else { mji_copy9(d->wrap_xpos+wrapcount*3, wpnt); } d->wrap_obj[wrapcount] = -1; if (wlen >= 0) { d->wrap_obj[wrapcount+1] = d->wrap_obj[wrapcount+2] = wrapid; } d->ten_wrapnum[i] += (wlen < 0 ? 1 : 3); wrapcount += (wlen < 0 ? 1 : 3); // advance j += (wraptype != mjWRAP_NONE ? 2 : 1); // assign last site before pulley or tendon end if (j == tendon_num-1 || m->wrap_type[adr+j+1] == mjWRAP_PULLEY) { mji_copy3(d->wrap_xpos+wrapcount*3, d->site_xpos+3*id1); d->wrap_obj[wrapcount] = -1; d->ten_wrapnum[i]++; wrapcount++; } } } mj_freeStack(d); } // return dot product of tendon Jacobian time derivative with vector mjtNum mj_tendonDot(const mjModel* m, mjData* d, int id, const mjtNum* vec) { int nv = m->nv; mjtNum res = 0; // tendon id is invalid: return if (id < 0 || id >= m->ntendon) { return 0; } // fixed tendon has zero Jdot: return int adr = m->tendon_adr[id]; if (m->wrap_type[adr] == mjWRAP_JOINT) { return 0; } // allocate stack arrays mj_markStack(d); int issparse = mj_isSparse(m); int* chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL; mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* tmp = mjSTACKALLOC(d, nv, mjtNum); // process spatial tendon mjtNum divisor = 1; int wraptype, j = 0; int num = m->tendon_num[id]; while (j < num-1) { // get 1st and 2nd object int type0 = m->wrap_type[adr+j+0]; int type1 = m->wrap_type[adr+j+1]; int id0 = m->wrap_objid[adr+j+0]; int id1 = m->wrap_objid[adr+j+1]; // pulley if (type0 == mjWRAP_PULLEY || type1 == mjWRAP_PULLEY) { // get divisor, insert obj=-2 if (type0 == mjWRAP_PULLEY) { divisor = m->wrap_prm[adr+j]; } // move to next j++; continue; } // init sequence; assume it starts with site mjtNum wpnt[6]; mji_copy3(wpnt, d->site_xpos+3*id0); mjtNum vel[6]; mj_objectVelocity(m, d, mjOBJ_SITE, id0, vel, /*flg_local=*/0); mjtNum wvel[6] = {vel[3], vel[4], vel[5], 0, 0, 0}; int wbody[2]; wbody[0] = m->site_bodyid[id0]; // second object is geom: process site-geom-site if (type1 == mjWRAP_SPHERE || type1 == mjWRAP_CYLINDER) { // TODO(tassa) support geom wrapping (requires derivatives of mju_wrap) mjERROR("geom wrapping not supported"); } else { wraptype = mjWRAP_NONE; } // complete sequence wbody[1] = m->site_bodyid[id1]; mji_copy3(wpnt+3, d->site_xpos+3*id1); mj_objectVelocity(m, d, mjOBJ_SITE, id1, vel, /*flg_local=*/0); mji_copy3(wvel+3, vel+3); // accumulate moments if consecutive points are in different bodies if (wbody[0] != wbody[1]) { // dpnt = 3D position difference, normalize mjtNum dpnt[3]; mju_sub3(dpnt, wpnt+3, wpnt); mjtNum norm = mju_normalize3(dpnt); // dvel = d / dt (dpnt) mjtNum dvel[3]; mju_sub3(dvel, wvel+3, wvel); mjtNum dot = mju_dot3(dpnt, dvel); mju_addToScl3(dvel, dpnt, -dot); mju_scl3(dvel, dvel, norm > mjMINVAL ? 1/norm : 0); // sparse if (issparse) { // construct merged chain int NV = mj_mergeChain(m, chain, wbody[0], wbody[1], /*flg_skipcommon=*/0); if (NV) { // get endpoint JacobianDots, subtract mj_jacDotSparse(m, d, jac1, 0, wpnt, wbody[0], NV, chain); mj_jacDotSparse(m, d, jac2, 0, wpnt+3, wbody[1], NV, chain); mju_sub(jacdif, jac2, jac1, 3*NV); // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt mju_mulMatTVec(tmp, jacdif, dpnt, 3, NV); for (int k=0; k < NV; k++) { res += (tmp[k] / divisor) * vec[chain[k]]; } // get endpoint Jacobians, subtract mj_jacSparse(m, d, jac1, 0, wpnt, wbody[0], NV, chain, /*flg_skipcommon=*/0); mj_jacSparse(m, d, jac2, 0, wpnt+3, wbody[1], NV, chain, /*flg_skipcommon=*/0); mju_sub(jacdif, jac2, jac1, 3*NV); // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt) mju_mulMatTVec(tmp, jacdif, dvel, 3, NV); for (int k=0; k < NV; k++) { res += (tmp[k] / divisor) * vec[chain[k]]; } } } // dense else { // get endpoint JacobianDots, subtract mj_jacDot(m, d, jac1, 0, wpnt, wbody[0]); mj_jacDot(m, d, jac2, 0, wpnt+3, wbody[1]); mju_sub(jacdif, jac2, jac1, 3*nv); // chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt mju_mulMatTVec(tmp, jacdif, dpnt, 3, nv); res += mju_dot(tmp, vec, nv) / divisor; // get endpoint Jacobians, subtract mj_jac(m, d, jac1, 0, wpnt, wbody[0]); mj_jac(m, d, jac2, 0, wpnt+3, wbody[1]); mju_sub(jacdif, jac2, jac1, 3*nv); // chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt) mju_mulMatTVec(tmp, jacdif, dvel, 3, nv); res += mju_dot(tmp, vec, nv) / divisor; } } // advance j += (wraptype != mjWRAP_NONE ? 2 : 1); } mj_freeStack(d); return res; } // compute actuator/transmission lengths and moments void mj_transmission(const mjModel* m, mjData* d) { int nv = m->nv, nu = m->nu; // nothing to do if (!nu) { return; } // outputs mjtNum* length = d->actuator_length; mjtNum* moment = d->actuator_moment; int *rownnz = d->moment_rownnz; int *rowadr = d->moment_rowadr; int *colind = d->moment_colind; // allocate Jacbians mj_markStack(d); mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jacA = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jacS = mjSTACKALLOC(d, 3*nv, mjtNum); // define stack variables required for body transmission, don't allocate int issparse = mj_isSparse(m); mjtNum* efc_force = NULL; // used as marker for allocation requirement mjtNum *moment_exclude, *jacdifp, *jac1p, *jac2p; int *chain; // define stack variables required for site transmission, don't allocate mjtNum *jacref = NULL, *moment_row = 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; // process according to transmission type switch ((mjtTrn) m->actuator_trntype[i]) { case mjTRN_JOINT: // joint case mjTRN_JOINTINPARENT: // joint, force in parent frame // slide and hinge joint: scalar gear if (m->jnt_type[id] == mjJNT_SLIDE || m->jnt_type[id] == mjJNT_HINGE) { // sparsity rownnz[i] = 1; colind[adr] = m->jnt_dofadr[id]; length[i] = d->qpos[m->jnt_qposadr[id]]*gear[0]; moment[adr] = gear[0]; } // ball joint: 3D wrench gear else if (m->jnt_type[id] == mjJNT_BALL) { // axis: expmap representation of quaternion mjtNum axis[3], quat[4]; mji_copy4(quat, d->qpos+m->jnt_qposadr[id]); mju_normalize4(quat); mji_quat2Vel(axis, quat, 1); // gearAxis: rotate to parent frame if necessary mjtNum gearAxis[3]; if (m->actuator_trntype[i] == mjTRN_JOINT) { mji_copy3(gearAxis, gear); } else { mju_negQuat(quat, quat); mji_rotVecQuat(gearAxis, gear, quat); } // length: axis*gearAxis length[i] = mju_dot3(axis, gearAxis); // dof start address int jnt_dofadr = m->jnt_dofadr[id]; // sparsity for (int j = 0; j < 3; j++) { colind[adr+j] = jnt_dofadr + j; } rownnz[i] = 3; // moment: gearAxis mji_copy3(moment+adr, gearAxis); } // free joint: 6D wrench gear else { // cannot compute meaningful length, set to 0 length[i] = 0; // gearAxis: rotate to world frame if necessary mjtNum gearAxis[3]; if (m->actuator_trntype[i] == mjTRN_JOINT) { mji_copy3(gearAxis, gear+3); } else { mjtNum quat[4]; mji_copy4(quat, d->qpos+m->jnt_qposadr[id]+3); mju_normalize4(quat); mju_negQuat(quat, quat); mji_rotVecQuat(gearAxis, gear+3, quat); } // dof start address int jnt_dofadr = m->jnt_dofadr[id]; // sparsity for (int j = 0; j < 6; j++) { colind[adr+j] = jnt_dofadr + j; } rownnz[i] = 6; // moment: gear(tran), gearAxis mji_copy3(moment+adr, gear); mji_copy3(moment+adr+3, gearAxis); } break; case mjTRN_SLIDERCRANK: // slider-crank { // get data int idslider = m->actuator_trnid[2*i+1]; mjtNum rod = m->actuator_cranklength[i]; mjtNum axis[3] = {d->site_xmat[9 * idslider + 2], d->site_xmat[9 * idslider + 5], d->site_xmat[9 * idslider + 8]}; mjtNum vec[3]; mju_sub3(vec, d->site_xpos+3*id, d->site_xpos+3*idslider); // compute length and determinant // length = a'*v - sqrt(det); det = (a'*v)^2 + r^2 - v'*v) mjtNum av = mju_dot3(vec, axis); mjtNum sdet, det = av*av + rod*rod - mju_dot3(vec, vec); int ok = 1; if (det <= 0) { ok = 0; sdet = 0; length[i] = av; } else { sdet = mju_sqrt(det); length[i] = av - sdet; } // compute derivatives of length w.r.t. vec and axis mjtNum dlda[3], dldv[3]; if (ok) { mju_scl3(dldv, axis, 1-av/sdet); mju_scl3(dlda, vec, 1/sdet); // use dlda as temp mji_addTo3(dldv, dlda); mju_scl3(dlda, vec, 1-av/sdet); } else { mji_copy3(dlda, vec); mji_copy3(dldv, axis); } // get Jacobians of axis(jacA) and vec(jac) mj_jacPointAxis(m, d, jacS, jacA, d->site_xpos+3*idslider, axis, m->site_bodyid[idslider]); mj_jacSite(m, d, jac, 0, id); mju_subFrom(jac, jacS, 3*nv); moment_row = mjSTACKALLOC(d, nv, mjtNum); // clear moment mju_zero(moment_row, nv); // apply chain rule for (int j=0; j < nv; j++) { for (int k=0; k < 3; k++) { moment_row[j] += dlda[k]*jacA[k*nv+j] + dldv[k]*jac[k*nv+j]; } } // scale by gear ratio length[i] *= gear[0]; // sparsity (compress) nnz = 0; for (int j = 0; j < nv; j++) { if (moment_row[j]) { moment[adr+nnz] = moment_row[j] * gear[0]; colind[adr+nnz] = j; nnz++; } } rownnz[i] = nnz; } break; case mjTRN_TENDON: // tendon length[i] = d->ten_length[id]*gear[0]; // moment { int ten_J_rownnz = m->ten_J_rownnz[id]; int ten_J_rowadr = m->ten_J_rowadr[id]; rownnz[i] = ten_J_rownnz; mju_copyInt(colind + adr, m->ten_J_colind + ten_J_rowadr, ten_J_rownnz); mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz); } break; case mjTRN_SITE: // site // get site translation (jac) and rotation (jacS) Jacobians in global frame mj_jacSite(m, d, jac, jacS, id); // clear length length[i] = 0; if (!moment_row) moment_row = mjSTACKALLOC(d, nv, mjtNum); // reference site undefined if (m->actuator_trnid[2*i+1] == -1) { // wrench: gear expressed in global frame mjtNum wrench[6]; mji_mulMatVec3(wrench, d->site_xmat+9*id, gear); // translation mji_mulMatVec3(wrench+3, d->site_xmat+9*id, gear+3); // rotation // moment: global Jacobian projected on wrench mju_mulMatTVec(moment_row, jac, wrench, 3, nv); // translation mju_mulMatTVec(jac, jacS, wrench+3, 3, nv); // rotation mju_addTo(moment_row, jac, nv); // add the two } // reference site defined else { int refid = m->actuator_trnid[2*i+1]; if (!jacref) jacref = mjSTACKALLOC(d, 3*nv, mjtNum); // initialize last dof address for each body int b0 = m->body_weldid[m->site_bodyid[id]]; int b1 = m->body_weldid[m->site_bodyid[refid]]; int dofadr0 = m->body_dofadr[b0] + m->body_dofnum[b0] - 1; int dofadr1 = m->body_dofadr[b1] + m->body_dofnum[b1] - 1; // find common ancestral dof, if any int dofadr_common = -1; if (dofadr0 >= 0 && dofadr1 >= 0) { // traverse up the tree until common ancestral dof is found while (dofadr0 != dofadr1) { if (dofadr0 < dofadr1) { dofadr1 = m->dof_parentid[dofadr1]; } else { dofadr0 = m->dof_parentid[dofadr0]; } if (dofadr0 == -1 || dofadr1 == -1) { // reached tree root, no common ancestral dof break; } } // found common ancestral dof if (dofadr0 == dofadr1) { dofadr_common = dofadr0; } } // clear moment mju_zero(moment_row, nv); // translational transmission if (!mju_isZero(gear, 3)) { // vec: site position in reference site frame mjtNum vec[3]; mju_sub3(vec, d->site_xpos+3*id, d->site_xpos+3*refid); mju_mulMatTVec3(vec, d->site_xmat+9*refid, vec); // length: dot product with gear length[i] += mju_dot3(vec, gear); // jacref: global Jacobian of reference site mj_jacSite(m, d, jacref, NULL, refid); // subtract jacref from jac mju_subFrom(jac, jacref, 3*nv); // if common ancestral dof was found, clear the columns of its parental chain int da = dofadr_common; while (da >= 0) { jac[nv*0 + da] = 0; jac[nv*1 + da] = 0; jac[nv*2 + da] = 0; da = m->dof_parentid[da]; } // wrench: translational gear expressed in global frame mjtNum wrench[6]; mji_mulMatVec3(wrench, d->site_xmat+9*refid, gear); // moment: global Jacobian projected on wrench mju_mulMatTVec(moment_row, jac, wrench, 3, nv); } // rotational transmission if (!mju_isZero(gear+3, 3)) { mjtNum refquat[4]; // get site and refsite quats from parent bodies (avoiding mju_mat2Quat) mjtNum quat[4]; mji_mulQuat(quat, m->site_quat+4*id, d->xquat+4*m->site_bodyid[id]); mji_mulQuat(refquat, m->site_quat+4*refid, d->xquat+4*m->site_bodyid[refid]); // convert difference to expmap (axis-angle) mjtNum vec[3]; mji_subQuat(vec, quat, refquat); // add length: dot product with gear length[i] += mju_dot3(vec, gear+3); // jacref: global rotational Jacobian of reference site mj_jacSite(m, d, NULL, jacref, refid); // subtract jacref from jacS mju_subFrom(jacS, jacref, 3*nv); // if common ancestral dof was found, clear the columns of its parental chain int da = dofadr_common; while (da >= 0) { jacS[nv*0 + da] = 0; jacS[nv*1 + da] = 0; jacS[nv*2 + da] = 0; da = m->dof_parentid[da]; } // wrench: rotational gear expressed in global frame mjtNum wrench[6]; mji_mulMatVec3(wrench, d->site_xmat+9*refid, gear+3); // global Jacobian projected on wrench, add to moment // reuse jac as temporary storage mju_mulMatTVec(jac, jacS, wrench, 3, nv); mju_addTo(moment_row, jac, nv); } } // sparsity (compress) nnz = 0; for (int j = 0; j < nv; j++) { if (moment_row[j]) { moment[adr+nnz] = moment_row[j]; colind[adr+nnz] = j; nnz++; } } rownnz[i] = nnz; break; case mjTRN_BODY: // body (adhesive contacts) // cannot compute meaningful length, set to 0 length[i] = 0; // clear moment if (!moment_row) moment_row = mjSTACKALLOC(d, nv, mjtNum); mju_zero(moment_row, nv); // moment is average of all contact normal Jacobians { // allocate stack variables for the first mjTRN_BODY if (!efc_force) { efc_force = mjSTACKALLOC(d, d->nefc, mjtNum); moment_exclude = mjSTACKALLOC(d, nv, mjtNum); jacdifp = mjSTACKALLOC(d, 3*nv, mjtNum); jac1p = mjSTACKALLOC(d, 3*nv, mjtNum); jac2p = mjSTACKALLOC(d, 3*nv, mjtNum); chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL; } // clear efc_force and moment_exclude mju_zero(efc_force, d->nefc); mju_zero(moment_exclude, nv); // count all relevant contacts, accumulate Jacobians int counter = 0, ncon = d->ncon; for (int j=0; j < ncon; j++) { const mjContact* con = d->contact+j; // get geom ids int g1 = con->geom[0]; int g2 = con->geom[1]; // contact involving flex, continue if (g1 < 0 || g2 < 0) { continue; } // get body ids int b1 = m->geom_bodyid[g1]; int b2 = m->geom_bodyid[g2]; // irrelevant contact, continue if (b1 != id && b2 != id) { continue; } // mark contact normals in efc_force if (!con->exclude) { counter++; // condim 1 or elliptic cones: normal is in the first row if (con->dim == 1 || m->opt.cone == mjCONE_ELLIPTIC) { efc_force[con->efc_address] = 1; } // pyramidal cones: average all pyramid directions else { int npyramid = con->dim-1; // number of frictional directions for (int k=0; k < 2*npyramid; k++) { efc_force[con->efc_address+k] = 0.5/npyramid; } } } // excluded contact in gap: get sparse or dense Jacobian, accumulate else if (con->exclude == 1) { counter++; // get Jacobian difference int NV = mj_jacDifPair(m, d, chain, b1, b2, con->pos, con->pos, jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse, /*flg_skipcommon=*/0); // project Jacobian along the normal of the contact frame mju_mulMatMat(jac, con->frame, jacdifp, 1, 3, NV); // accumulate in moment_exclude if (issparse) { for (int k=0; k < NV; k++) { moment_exclude[chain[k]] += jac[k]; } } else { mju_addTo(moment_exclude, jac, nv); } } } // moment is average over contact normal Jacobians, make negative for adhesion if (counter) { // accumulate active contact Jacobians into moment mj_mulJacTVec(m, d, moment_row, efc_force); // add Jacobians from excluded contacts mju_addTo(moment_row, moment_exclude, nv); // normalize by total contacts, flip sign mju_scl(moment_row, moment_row, -1.0/counter, nv); } } // sparsity (compress) nnz = 0; for (int j = 0; j < nv; j++) { if (moment_row[j]) { moment[adr+nnz] = moment_row[j]; colind[adr+nnz] = j; nnz++; } } rownnz[i] = nnz; break; default: mjERROR("unknown transmission type %d", m->actuator_trntype[i]); // SHOULD NOT OCCUR } } mj_freeStack(d); } //-------------------------- inertia --------------------------------------------------------------- // add tendon armature to M void mj_tendonArmature(const mjModel* m, mjData* d) { int nv = m->nv, ntendon = m->ntendon; const int* M_rownnz = m->M_rownnz; const int* M_rowadr = m->M_rowadr; const int* M_colind = m->M_colind; // 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] + mj_actuatorArmature(m, mjOBJ_TENDON, k); if (!armature) { continue; } // get sparse info for tendon k int J_rowadr = m->ten_J_rowadr[k]; int J_rownnz = m->ten_J_rownnz[k]; const int* J_colind = m->ten_J_colind + J_rowadr; mjtNum* ten_J = d->ten_J + J_rowadr; // M += armature * ten_J' * ten_J for (int j=0; j < J_rownnz; j++) { mjtNum ten_J_i = ten_J[j]; if (!ten_J_i) { continue; } // M[i,:] += armature * ten_J[i] * ten_J int i = J_colind[j]; int M_adr = M_rowadr[i]; mju_addToSclSparseInc(d->M + M_adr, ten_J, M_rownnz[i], M_colind + M_adr, J_rownnz, J_colind, armature * ten_J_i); } } } // composite rigid body inertia algorithm void mj_crb(const mjModel* m, mjData* d) { // 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* 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 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 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 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 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]) { M[adr] = dof_M0[i]; continue; } // init M(i,i) with armature inertia int Madr_ij = adr + rownnz[i] - 1; M[Madr_ij] = dof_armature[i] + mj_actuatorArmature(m, mjOBJ_JOINT, m->dof_jntid[i]); // precompute buf = crb_body_i * cdof_i mjtNum buf[6]; mju_mulInertVec(buf, crb+10*dof_bodyid[i], cdof+6*i); // sparse backward pass over ancestors for (int j=i; j >= 0; j = dof_parentid[j]) { // M(i,j) += cdof_j * (crb_body_i * cdof_i) M[Madr_ij--] += mji_dot6(cdof+6*j, buf); } } } 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); // TODO(tassa): scatter only awake dofs TM_END(mjTIMER_POS_INERTIA); } // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd // (legacy implementation) void mj_factorI_legacy(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv) { int cnt; int Madr_kk, Madr_ki; mjtNum tmp; // local copies of key variables int* dof_Madr = m->dof_Madr; int* dof_parentid = m->dof_parentid; int nv = m->nv; // copy M into LD mju_copy(qLD, M, m->nM); // dense backward loop over dofs (regular only, simple diagonal already copied) for (int k=nv-1; k >= 0; k--) { // get address of M(k,k) Madr_kk = dof_Madr[k]; // check for small/negative numbers on diagonal if (qLD[Madr_kk] < mjMINVAL) { mj_warning(d, mjWARN_INERTIA, k); qLD[Madr_kk] = mjMINVAL; } // skip the rest if simple if (m->dof_simplenum[k]) { continue; } // sparse backward loop over ancestors of k (excluding k) Madr_ki = Madr_kk + 1; int i = dof_parentid[k]; while (i >= 0) { tmp = qLD[Madr_ki] / qLD[Madr_kk]; // tmp = M(k,i) / M(k,k) // get number of ancestors of i (including i) if (i < nv-1) { cnt = dof_Madr[i+1] - dof_Madr[i]; } else { cnt = m->nM - dof_Madr[i+1]; } // M(i,j) -= M(k,j) * tmp mju_addToScl(qLD+dof_Madr[i], qLD+Madr_ki, -tmp, cnt); qLD[Madr_ki] = tmp; // M(k,i) = tmp // advance to i's parent i = dof_parentid[i]; Madr_ki++; } } // compute 1/diag(D) for (int i=0; i < nv; i++) { qLDiagInv[i] = 1.0 / qLD[dof_Madr[i]]; } } // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd void mj_factorM(const mjModel* m, mjData* d) { TM_START; // 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 (with dof skipping) void mj_factorI(mjtNum* mat, mjtNum* diaginv, int nv, const int* rownnz, const int* rowadr, const int* colind, const int* index) { // backward loop over rows 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; int end = start + diag; mjtNum invD = 1 / mat[end]; if (diaginv) diaginv[k] = invD; // update triangle above row k for (int adr=end - 1; adr >= start; adr--) { // update row i < k: L(i, 0..i) -= L(i, 0..i) * L(k, i) / L(k, k) int i = colind[adr]; mju_addToScl(mat + rowadr[i], mat + start, -mat[adr] * invD, rownnz[i]); } // update row k: L(k, :) /= L(k, k) mju_scl(mat + start, mat + start, invD, diag); } } // in-place sparse backsubstitution: x = inv(L'*D*L)*x // (legacy implementation) void mj_solveLD_legacy(const mjModel* m, mjtNum* restrict x, int n, const mjtNum* qLD, const mjtNum* qLDiagInv) { // local copies of key variables int* dof_Madr = m->dof_Madr; int* dof_parentid = m->dof_parentid; int nv = m->nv; // single vector if (n == 1) { // x <- inv(L') * x; skip simple, exploit sparsity of input vector for (int i=nv-1; i >= 0; i--) { if (!m->dof_simplenum[i] && x[i]) { // init int Madr_ij = dof_Madr[i]+1; int j = dof_parentid[i]; // traverse ancestors backwards // read directly from x[i] since i cannot be a parent of itself while (j >= 0) { x[j] -= qLD[Madr_ij++]*x[i]; // x(j) -= L(i,j) * x(i) // advance to parent j = dof_parentid[j]; } } } // x <- inv(D) * x for (int i=0; i < nv; i++) { x[i] *= qLDiagInv[i]; // x(i) /= L(i,i) } // x <- inv(L) * x; skip simple for (int i=0; i < nv; i++) { if (!m->dof_simplenum[i]) { // init int Madr_ij = dof_Madr[i]+1; int j = dof_parentid[i]; // traverse ancestors backwards // write directly in x[i] since i cannot be a parent of itself while (j >= 0) { x[i] -= qLD[Madr_ij++]*x[j]; // x(i) -= L(i,j) * x(j) // advance to parent j = dof_parentid[j]; } } } } // multiple vectors else { int offset; mjtNum tmp; // x <- inv(L') * x; skip simple for (int i=nv-1; i >= 0; i--) { if (!m->dof_simplenum[i]) { // init int Madr_ij = dof_Madr[i]+1; int j = dof_parentid[i]; // traverse ancestors backwards while (j >= 0) { // process all vectors, exploit sparsity for (offset=0; offset < n*nv; offset+=nv) if ((tmp = x[i+offset])) { x[j+offset] -= qLD[Madr_ij]*tmp; // x(j) -= L(i,j) * x(i) } // advance to parent Madr_ij++; j = dof_parentid[j]; } } } // x <- inv(D) * x for (int i=0; i < nv; i++) { for (offset=0; offset < n*nv; offset+=nv) { x[i+offset] *= qLDiagInv[i]; // x(i) /= L(i,i) } } // x <- inv(L) * x; skip simple for (int i=0; i < nv; i++) { if (!m->dof_simplenum[i]) { // init int Madr_ij = dof_Madr[i]+1; int j = dof_parentid[i]; // traverse ancestors backwards tmp = x[i+offset]; while (j >= 0) { // process all vectors for (offset=0; offset < n*nv; offset+=nv) { x[i+offset] -= qLD[Madr_ij]*x[j+offset]; // x(i) -= L(i,j) * x(j) } // advance to parent Madr_ij++; j = dof_parentid[j]; } } } } } // 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* index) { // x <- L^-T x for (int k = nv - 1; k >= 0; k--) { int i = index ? index[k] : k; // skip diagonal rows if (rownnz[i] == 1) { continue; } // one vector if (n == 1) { mjtNum x_i; if ((x_i = x[i])) { int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int adr=start; adr < end; adr++) { x[colind[adr]] -= qLD[adr] * x_i; } } } // multiple vectors else { int start = rowadr[i]; int end = start + rownnz[i] - 1; for (int offset=0; offset < n*nv; offset+=nv) { mjtNum x_i; if ((x_i = x[i+offset])) { for (int adr=start; adr < end; adr++) { x[offset + colind[adr]] -= qLD[adr] * x_i; } } } } } // x <- D^-1 x for (int k = 0; k < nv; k++) { int i = index ? index[k] : k; mjtNum invD_i = qLDiagInv[i]; // one vector if (n == 1) { x[i] *= invD_i; } // multiple vectors else { for (int offset=0; offset < n*nv; offset+=nv) { x[i+offset] *= invD_i; } } } // x <- L^-1 x for (int k = 0; k < nv; k++) { int i = index ? index[k] : k; // skip diagonal rows if (rownnz[i] == 1) { continue; } int d; if ((d = rownnz[i] - 1) > 0) { int adr = rowadr[i]; // one vector if (n == 1) { x[i] -= mju_dotSparse(qLD+adr, x, d, colind+adr); } // multiple vectors else { for (int offset=0; offset < n*nv; offset+=nv) { x[i+offset] -= mju_dotSparse(qLD+adr, x+offset, d, colind+adr); } } } } } // sparse backsubstitution: x = inv(L'*D*L)*y // use factorization in d 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, NULL); } // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, const mjtNum* sqrtInvD, int n) { int nv = m->nv; // local copies of key variables const int* rownnz = m->M_rownnz; const int* rowadr = m->M_rowadr; const int* colind = m->M_colind; const int* diagnum = m->dof_simplenum; const mjtNum* qLD = d->qLD; // x = y mju_copy(x, y, n * nv); // x <- L^-T x for (int i=nv-1; i > 0; i--) { // skip diagonal rows if (diagnum[i]) { continue; } // prepare row i column address range int start = rowadr[i]; int end = start + rownnz[i] - 1; // process all vectors for (int offset=0; offset < n*nv; offset+=nv) { mjtNum x_i; if ((x_i = x[i+offset])) { for (int adr=start; adr < end; adr++) { x[offset + colind[adr]] -= qLD[adr] * x_i; } } } } // x <- D^-1/2 x for (int i=0; i < nv; i++) { mjtNum invD_i = sqrtInvD[i]; for (int offset=0; offset < n*nv; offset+=nv) { x[i+offset] *= invD_i; } } } //---------------------------------- velocity ------------------------------------------------------ // compute cvel, cdof_dot void mj_comVel(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; // set world vel to 0 mju_zero(d->cvel, 6); // forward pass over bodies for (int b=1; b < nbody; b++) { int i = sleep_filter ? d->body_awake_ind[b] : b; // cvel = cvel_parent mjtNum cvel[6]; mji_copy6(cvel, d->cvel+6*m->body_parentid[i]); // 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]; // compute cvel and cdofdot switch ((mjtJoint) m->jnt_type[m->dof_jntid[bda+j]]) { case mjJNT_FREE: // cdofdot = 0 mju_zero(cdofdot, 18); // update velocity mju_mulDofVec(tmp, d->cdof+6*bda, d->qvel+bda, 3); mju_addTo(cvel, tmp, 6); // continue with rotations j += 3; mjFALLTHROUGH; case mjJNT_BALL: // compute all 3 cdofdots using parent velocity mji_crossMotion(cdofdot+6*(j+0), cvel, d->cdof+6*(bda+j+0)); mji_crossMotion(cdofdot+6*(j+1), cvel, d->cdof+6*(bda+j+1)); mji_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); mju_addTo(cvel, tmp, 6); // adjust for 3-dof joint j += 2; break; default: // in principle we should use the new velocity to compute cdofdot, // but it makes no difference because crossMotion(cdof, cdof) = 0, // and using the old velocity may be more accurate numerically mji_crossMotion(cdofdot+6*j, cvel, d->cdof+6*(bda+j)); // update velocity mju_mulDofVec(tmp, d->cdof+6*(bda+j), d->qvel+bda+j, 1); mju_addTo(cvel, tmp, 6); } } // assign cvel, cdofdot mji_copy6(d->cvel+6*i, cvel); mju_copy(d->cdof_dot+6*bda, cdofdot, 6*dofnum); } } // subtree linear velocity and angular momentum void mj_subtreeVel(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; mj_markStack(d); mjtNum* body_vel = mjSTACKALLOC(d, 6*m->nbody, mjtNum); // bodywise quantities 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); // body linear momentum 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]; dv[2] *= m->body_inertia[3*i+2]; mji_mulMatVec3(d->subtree_angmom+3*i, d->ximat+9*i, dv); } // 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) { mji_addTo3(d->subtree_linvel+3*m->body_parentid[i], d->subtree_linvel+3*i); } // convert linear momentum to linear velocity mju_scl3(d->subtree_linvel+3*i, d->subtree_linvel+3*i, 1/mju_max(mjMINVAL, m->body_subtreemass[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]); mji_cross(dL, dx, dp); // add to subtree i mji_addTo3(d->subtree_angmom+3*i, dL); // add to parent mji_addTo3(d->subtree_angmom+3*parent, d->subtree_angmom+3*i); // momentum wrt parent mju_sub3(dx, d->subtree_com+3*i, d->subtree_com+3*parent); mju_sub3(dv, d->subtree_linvel+3*i, d->subtree_linvel+3*parent); mju_scl3(dv, dv, m->body_subtreemass[i]); mji_cross(dL, dx, dv); // add to parent mji_addTo3(d->subtree_angmom+3*parent, dL); } mj_freeStack(d); // mark as computed d->flg_subtreevel = 1; } //---------------------------------- RNE ----------------------------------------------------------- // 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 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); // set world acceleration to -gravity mju_zero(loc_cacc, 6); if (!mjDISABLED(mjDSBL_GRAVITY)) { mju_scl3(loc_cacc+3, m->opt.gravity, -1); } // forward pass over bodies: accumulate cacc, set cfrc_body 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); // cacc += cdof * qacc if (flg_acc) { mju_mulDofVec(tmp, d->cdof+6*bda, d->qacc+bda, m->body_dofnum[i]); mju_addTo(loc_cacc+6*i, tmp, 6); } // 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]; mji_crossForce(tmp1, d->cvel+6*i, tmp); mju_addTo(loc_cfrc_body+6*i, tmp1, 6); } // clear world cfrc_body mju_zero(loc_cfrc_body, 6); // backward pass over bodies: accumulate cfrc_body from children 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 v=0; v < nv; v++) { int i = sleep_filter ? d->dof_awake_ind[v] : v; result[i] = mji_dot6(d->cdof+6*i, loc_cfrc_body+6*m->dof_bodyid[i]); } mj_freeStack(d); } // RNE with complete data: compute cacc, cfrc_ext, cfrc_int void mj_rnePostConstraint(const mjModel* m, mjData* d) { int nbody = m->nbody; mjtNum cfrc_com[6], cfrc[6], lfrc[6]; mjContact* con; // clear cacc, set world acceleration to -gravity mju_zero(d->cacc, 6); if (!mjDISABLED(mjDSBL_GRAVITY)) { mju_scl3(d->cacc+3, m->opt.gravity, -1); } // cfrc_ext = perturb mju_zero(d->cfrc_ext, 6*nbody); for (int i=1; i < nbody; i++) { if (!mju_isZero(d->xfrc_applied+6*i, 6)) { // rearrange as torque:force mji_copy3(cfrc, d->xfrc_applied+6*i+3); mji_copy3(cfrc+3, d->xfrc_applied+6*i); // map force from application point to com; both world-oriented mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[i], d->xipos+3*i, 0); // accumulate mju_addTo(d->cfrc_ext+6*i, cfrc_com, 6); } } // cfrc_ext += contacts int ncon = d->ncon; for (int i=0; i < ncon; i++) { // get contact pointer con = d->contact+i; // skip excluded contacts if (con->efc_address < 0) { continue; } // skip contact involving flex if (con->geom[0] < 0 || con->geom[1] < 0) { continue; } // tmp = contact-local force:torque vector mj_contactForce(m, d, i, lfrc); // cfrc = world-oriented torque:force vector (swap in the process) mju_mulMatTVec3(cfrc, con->frame, lfrc+3); mju_mulMatTVec3(cfrc+3, con->frame, lfrc); // body 1 int k; if ((k = m->geom_bodyid[con->geom[0]])) { // tmp = subtree CoM-based torque_force vector mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[k], con->pos, 0); // apply (opposite for body 1) mju_subFrom(d->cfrc_ext+6*k, cfrc_com, 6); } // body 2 if ((k = m->geom_bodyid[con->geom[1]])) { // tmp = subtree CoM-based torque_force vector mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[k], con->pos, 0); // apply mju_addTo(d->cfrc_ext+6*k, cfrc_com, 6); } } // cfrc_ext += connect, weld, flex constraints int i = 0, ne = d->ne; while (i < ne) { if (d->efc_type[i] != mjCNSTR_EQUALITY) { mjERROR("row %d of efc is not an equality constraint", i); // SHOULD NOT OCCUR } int id = d->efc_id[i]; mjtNum* eq_data = m->eq_data + mjNEQDATA*id; mjtNum pos[3], *offset; int k, obj1, obj2, body_semantic; switch ((mjtEq) m->eq_type[id]) { case mjEQ_CONNECT: case mjEQ_WELD: // cfrc = world-oriented torque:force vector mji_copy3(cfrc + 3, d->efc_force + i); if (m->eq_type[id] == mjEQ_WELD) { mji_copy3(cfrc, d->efc_force + i + 3); } else { mju_zero3(cfrc); // no torque from connect } body_semantic = m->eq_objtype[id] == mjOBJ_BODY; // body 1 obj1 = m->eq_obj1id[id]; k = body_semantic ? obj1 : m->site_bodyid[obj1]; if (k) { offset = body_semantic ? eq_data + 3 * (m->eq_type[id] == mjEQ_WELD) : m->site_pos + 3 * obj1; // transform point on body1: local -> global mj_local2Global(d, pos, 0, offset, 0, k, 0); // tmp = subtree CoM-based torque_force vector mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[k], pos, 0); // apply (opposite for body 1) mju_addTo(d->cfrc_ext+6*k, cfrc_com, 6); } // body 2 obj2 = m->eq_obj2id[id]; k = body_semantic ? obj2 : m->site_bodyid[obj2]; if (k) { offset = body_semantic ? eq_data + 3 * (m->eq_type[id] == mjEQ_CONNECT) : m->site_pos + 3 * obj2; // transform point on body2: local -> global mj_local2Global(d, pos, 0, offset, 0, k, 0); // tmp = subtree CoM-based torque_force vector mju_transformSpatial(cfrc_com, cfrc, 1, d->subtree_com+3*m->body_rootid[k], pos, 0); // apply mju_subFrom(d->cfrc_ext+6*k, cfrc_com, 6); } // increment rows i += m->eq_type[id] == mjEQ_WELD ? 6 : 3; break; case mjEQ_JOINT: case mjEQ_TENDON: // increment 1 row i++; break; case mjEQ_FLEX: // increment with number of non-rigid edges k = m->eq_obj1id[id]; int flex_edgeadr = m->flex_edgeadr[k]; int flex_edgenum = m->flex_edgenum[k]; for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) { if (!m->flexedge_rigid[e]) { i++; } } break; case mjEQ_FLEXVERT: k = m->eq_obj1id[id]; i += 2*m->flex_vertnum[k]; break; case mjEQ_FLEXSTRAIN: { // increment: trilinear uses 2 center (I1,J-1) + 3*ngauss shear, quadratic uses 6*ngauss k = m->eq_obj1id[id]; int order = m->flex_interp[k]; int nodenum = m->flex_nodenum[k]; if (order && nodenum) { int nquad = order + 1; int ngauss = nquad * nquad * nquad; int ncells = m->flex_cellnum[3*k+0] * m->flex_cellnum[3*k+1] * m->flex_cellnum[3*k+2]; i += ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss)); } break; } default: mjERROR("unknown constraint type type %d", m->eq_type[id]); // SHOULD NOT OCCUR } } // forward pass over bodies: compute cacc, cfrc_int mjtNum cacc[6], cfrc_body[6], cfrc_corr[6]; mju_zero(d->cfrc_int, 6); for (int j=1; j < nbody; j++) { // get body's first dof address int bda = m->body_dofadr[j]; // cacc = cacc_parent + cdofdot * qvel + cdof * qacc mju_mulDofVec(cacc, d->cdof_dot+6*bda, d->qvel+bda, m->body_dofnum[j]); mju_add(d->cacc+6*j, d->cacc+6*m->body_parentid[j], cacc, 6); mju_mulDofVec(cacc, d->cdof+6*bda, d->qacc+bda, m->body_dofnum[j]); mju_addTo(d->cacc+6*j, cacc, 6); // cfrc_body = cinert * cacc + cvel x (cinert * cvel) mju_mulInertVec(cfrc_body, d->cinert+10*j, d->cacc+6*j); mju_mulInertVec(cfrc_corr, d->cinert+10*j, d->cvel+6*j); mji_crossForce(cfrc, d->cvel+6*j, cfrc_corr); mju_addTo(cfrc_body, cfrc, 6); // set cfrc_int = cfrc_body - cfrc_ext mju_sub(d->cfrc_int+6*j, cfrc_body, d->cfrc_ext+6*j, 6); } // backward pass over bodies: accumulate cfrc_int from children for (int j=nbody-1; j > 0; j--) { mju_addTo(d->cfrc_int+6*m->body_parentid[j], d->cfrc_int+6*j, 6); } // mark as computed d->flg_rnepost = 1; } // 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; // 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] + mj_actuatorArmature(m, mjOBJ_TENDON, i); // no armature: skip if (!armature) { continue; } // get d/dt(tendon Jacobian) dotted with qvel for tendon i mjtNum dot = mj_tendonDot(m, d, i, d->qvel); // add bias term: qfrc += ten_J * armature * dot(ten_Jdot, qvel) mjtNum coef = armature * dot; if (coef) { // sparse int nnz = m->ten_J_rownnz[i]; int adr = m->ten_J_rowadr[i]; const int* colind = m->ten_J_colind + adr; const mjtNum* ten_J = d->ten_J + adr; for (int j=0; j < nnz; j++) { qfrc[colind[j]] += coef * ten_J[j]; } } } }