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

Implement sparse Jacobian time derivative

Browse Source 
PiperOrigin-RevId: 897609541
Change-Id: Ic86b6026cfc369c584d3edfe12a9a2fd14cc4357
This commit is contained in:
Yuval Tassa
2026-04-10 04:06:20 -07:00
committed by Copybara-Service
parent 6b724616c0
commit 025ba59fab
4 changed files with 220 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+54 -18
View File
@@ -1113,6 +1113,8 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
// 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);
@@ -1178,30 +1180,64 @@ void mj_tendonDot(const mjModel* m, mjData* d, int id, mjtNum* Jdot) {
mju_addToScl3(dvel, dpnt, -dot);
mju_scl3(dvel, dvel, norm > mjMINVAL ? 1/norm : 0);
// TODO(tassa ) write sparse branch, requires mj_jacDotSparse
// if (mj_isSparse(m)) { ... }
// sparse
if (issparse) {
// construct merged chain
int NV = mj_mergeChain(m, chain, wbody[0], wbody[1], /*flg_skipcommon=*/0);
// 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);
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);
// chain rule, first term: Jdot += d/dt(jac2 - jac1) * dpnt
mju_mulMatTVec(tmp, jacdif, dpnt, 3, NV);
// add to existing
mju_addToScl(Jdot, tmp, 1/divisor, nv);
// scatter into dense output
for (int k=0; k < NV; k++) {
Jdot[chain[k]] += tmp[k] / 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);
// 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);
// chain rule, second term: Jdot += (jac2 - jac1) * d/dt(dpnt)
mju_mulMatTVec(tmp, jacdif, dvel, 3, NV);
// add to existing
mju_addToScl(Jdot, tmp, 1/divisor, nv);
// scatter into dense output
for (int k=0; k < NV; k++) {
Jdot[chain[k]] += tmp[k] / divisor;
}
}
}
// 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);
// add to existing
mju_addToScl(Jdot, tmp, 1/divisor, nv);
// 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);
// add to existing
mju_addToScl(Jdot, tmp, 1/divisor, nv);
}
}
// advance
src/engine/engine_core_util.c
+86
View File
@@ -660,6 +660,92 @@ void mj_jacDot(const mjModel* m, const mjData* d,
}
// compute 3/6-by-NV sparse Jacobian time derivative of global point attached to given body
void mj_jacDotSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain) {
mjtNum offset[3];
mjtNum pvel[6];
// clear jacobians, compute offset and pvel if required
if (jacp) {
mju_zero(jacp, 3*NV);
const mjtNum* com = d->subtree_com+3*m->body_rootid[body];
mju_sub3(offset, point, com);
mju_transformSpatial(pvel, d->cvel+6*body, 0, point, com, 0);
}
if (jacr) {
mju_zero(jacr, 3*NV);
}
// skip fixed bodies
body = m->body_weldid[body];
// no movable body found: nothing to do
if (!body) {
return;
}
// get last dof that affects this body
int da = m->body_dofadr[body] + m->body_dofnum[body] - 1;
// start at end of chain (chain is in increasing order)
int ci = NV-1;
// backward pass over dof ancestor chain
while (da >= 0) {
// find chain index for this dof
while (ci >= 0 && chain[ci] > da) {
ci--;
}
// dof not in chain: SHOULD NOT OCCUR
if (ci < 0 || chain[ci] != da) {
mjERROR("dof index %d not found in chain", da);
}
mjtNum cdof_dot[6];
mji_copy6(cdof_dot, d->cdof_dot+6*da);
mjtNum* cdof = d->cdof+6*da;
// check for quaternion
mjtJoint type = m->jnt_type[m->dof_jntid[da]];
int dofadr = m->jnt_dofadr[m->dof_jntid[da]];
int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && da >= dofadr + 3);
// compute cdof_dot for quaternion (use current body cvel)
if (is_quat) {
mji_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[da], cdof);
}
// construct rotation jacobian
if (jacr) {
jacr[ci+0*NV] += cdof_dot[0];
jacr[ci+1*NV] += cdof_dot[1];
jacr[ci+2*NV] += cdof_dot[2];
}
// construct translation jacobian (correct for rotation)
if (jacp) {
// first correction term, account for varying cdof
mjtNum tmp1[3];
mji_cross(tmp1, cdof_dot, offset);
// second correction term, account for point translational velocity
mjtNum tmp2[3];
mji_cross(tmp2, cdof, pvel + 3);
jacp[ci+0*NV] += cdof_dot[3] + tmp1[0] + tmp2[0];
jacp[ci+1*NV] += cdof_dot[4] + tmp1[1] + tmp2[1];
jacp[ci+2*NV] += cdof_dot[5] + tmp1[2] + tmp2[2];
}
// advance to parent dof
da = m->dof_parentid[da];
}
}
// compute subtree angular momentum matrix
void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
int nv = m->nv;
src/engine/engine_core_util.h
+6
View File
@@ -18,6 +18,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#ifdef __cplusplus
extern "C" {
@@ -85,6 +86,11 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d,
mjtNum* jacdifp, mjtNum* jacdifr, const mjtNum* point,
int body, int flg_second, int NV, int start);
// compute 3/6-by-NV sparse Jacobian time derivative of global point attached to given body
MJAPI void mj_jacDotSparse(const mjModel* m, const mjData* d,
mjtNum* jacp, mjtNum* jacr, const mjtNum* point, int body,
int NV, const int* chain);
// dense or sparse Jacobian difference for two body points: pos2 - pos1, global
MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain,
int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3],
test/engine/engine_support_test.cc
+74
View File
@@ -17,6 +17,7 @@
#include "src/engine/engine_core_util.h"
#include "src/engine/engine_support.h"
#include <algorithm>
#include <cstring>
#include <limits>
#include <random>
@@ -483,6 +484,79 @@ TEST_F(JacobianTest, JacDot) {
}
}
// compare mj_jacDotSparse with dense mj_jacDot
TEST_F(JacobianTest, JacDotSparse) {
for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
int nv = model->nv;
mjData* data = mj_makeData(model);
// load keyframe if present, step for a bit
if (model->nkey) mj_resetDataKeyframe(model, data, 0);
while (data->time < 0.1) {
mj_step(model, data);
}
// minimal call required for mj_jacDot outputs to be valid
mj_kinematics(model, data);
mj_comPos(model, data);
mj_comVel(model, data);
// get bodyid and site position
int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
EXPECT_GT(bodyid, 0);
int siteid = mj_name2id(model, mjOBJ_SITE, "query");
EXPECT_GT(siteid, -1);
mjtNum point[3];
mju_copy3(point, data->site_xpos+3*siteid);
// dense jacDot
vector<mjtNum> jacp_dense(3*nv);
vector<mjtNum> jacr_dense(3*nv);
mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid);
// compute body chain using public mjModel fields
vector<int> chain(nv);
int NV = 0;
int weldbody = model->body_weldid[bodyid];
if (weldbody) {
int da = model->body_dofadr[weldbody] + model->body_dofnum[weldbody] - 1;
while (da >= 0) {
chain[NV++] = da;
da = model->dof_parentid[da];
}
std::reverse(chain.begin(), chain.begin() + NV);
}
EXPECT_GT(NV, 0);
// sparse jacDot
vector<mjtNum> jacp_sparse(3*NV);
vector<mjtNum> jacr_sparse(3*NV);
mj_jacDotSparse(model, data, jacp_sparse.data(), jacr_sparse.data(),
point, bodyid, NV, chain.data());
// expand sparse to dense and compare
vector<mjtNum> jacp_expanded(3*nv, 0);
vector<mjtNum> jacr_expanded(3*nv, 0);
for (int ci = 0; ci < NV; ci++) {
int di = chain[ci];
for (int r = 0; r < 3; r++) {
jacp_expanded[di+r*nv] = jacp_sparse[ci+r*NV];
jacr_expanded[di+r*nv] = jacr_sparse[ci+r*NV];
}
}
// expect bitwise equality
EXPECT_EQ(jacp_expanded, jacp_dense);
EXPECT_EQ(jacr_expanded, jacr_dense);
mj_deleteData(data);
mj_deleteModel(model);
}
}
using Name2idTest = MujocoTest;
static constexpr char name2idTestingModel[] = R"(