Add missing term in mj_jacDot

yuvaltassa · copybara-github · commit de48f4178f97 · 2025-03-11T03:01:29.000-07:00
PiperOrigin-RevId: 735686836
Change-Id: I813ca46f71e368ddc97a6d2e540102137b31dbcf
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+Bug fixes
+^^^^^^^^^
+- :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to
+  which the Jacobian is computed, now fixed.
+
 Version 3.3.0 (Feb 26, 2025)
 ----------------------------
 
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
@@ -1893,7 +1893,7 @@ void mj_comVel(const mjModel* m, mjData* d) {
 
     // assign cvel, cdofdot
     mju_copy(d->cvel+6*i, cvel, 6);
-    mju_copy(d->cdof_dot+6*bda, cdofdot, 6*m->body_dofnum[i]);
+    mju_copy(d->cdof_dot+6*bda, cdofdot, 6*dofnum);
   }
 }
 
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
@@ -811,11 +811,14 @@ void mj_jacDot(const mjModel* m, const mjData* d,
                mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) {
   int nv = m->nv;
   mjtNum offset[3];
+  mjtNum pvel[6];  // point velocity (rot:lin order)
 
-  // clear jacobians, compute offset if required
+  // clear jacobians, compute offset and pvel if required
   if (jacp) {
     mju_zero(jacp, 3*nv);
-    mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]);
+    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);
@@ -838,6 +841,7 @@ void mj_jacDot(const mjModel* m, const mjData* d,
   while (i >= 0) {
     mjtNum cdof_dot[6];
     mju_copy(cdof_dot, d->cdof_dot+6*i, 6);
+    mjtNum* cdof = d->cdof+6*i;
 
     // check for quaternion
     mjtJoint type = m->jnt_type[m->dof_jntid[i]];
@@ -846,7 +850,7 @@ void mj_jacDot(const mjModel* m, const mjData* d,
 
     // compute cdof_dot for quaternion (use current body cvel)
     if (is_quat) {
-      mju_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[i], d->cdof+6*i);
+      mju_crossMotion(cdof_dot, d->cvel+6*m->dof_bodyid[i], cdof);
     }
 
     // construct rotation jacobian
@@ -858,11 +862,17 @@ void mj_jacDot(const mjModel* m, const mjData* d,
 
     // construct translation jacobian (correct for rotation)
     if (jacp) {
-      mjtNum tmp[3] = {0};
-      mju_cross(tmp, cdof_dot, offset);
-      jacp[i+0*nv] += cdof_dot[3] + tmp[0];
-      jacp[i+1*nv] += cdof_dot[4] + tmp[1];
-      jacp[i+2*nv] += cdof_dot[5] + tmp[2];
+      // first correction term, account for varying cdof
+      mjtNum tmp1[3];
+      mju_cross(tmp1, cdof_dot, offset);
+
+      // second correction term, account for point translational velocity
+      mjtNum tmp2[3];
+      mju_cross(tmp2, cdof, pvel + 3);
+
+      jacp[i+0*nv] += cdof_dot[3] + tmp1[0] + tmp2[0];
+      jacp[i+1*nv] += cdof_dot[4] + tmp1[1] + tmp2[1];
+      jacp[i+2*nv] += cdof_dot[5] + tmp1[2] + tmp2[2];
     }
 
     // advance to parent dof
diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc
@@ -288,6 +288,7 @@ static constexpr char kQuat[] = R"(
     <body name="query">
       <joint type="ball"/>
       <geom size="1"/>
+      <site name="query" pos=".1 .2 .3"/>
     </body>
   </worldbody>
   <keyframe>
@@ -315,6 +316,7 @@ static constexpr char kFreeBall[] = R"(
           <body name="query" pos="0 .2 0">
             <joint type="slide" axis="1 1 1"/>
             <geom size=".05"/>
+            <site name="query" pos=".1 .2 .3"/>
           </body>
         </body>
       </body>
@@ -350,6 +352,7 @@ static constexpr char kQuatlessPendulum[] = R"(
           <body name="query" pos="0 .1 0">
             <joint axis="1 0 0"/>
             <geom type="capsule" size="0.02" fromto="0 0 0 0 .1 0"/>
+            <site name="query" pos=".1 0 0"/>
           </body>
         </body>
       </body>
@@ -373,6 +376,7 @@ static constexpr char kTelescope[] = R"(
           <body pos=".1 .02 0" name="query">
             <joint type="slide" axis="1 0 0"/>
             <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+            <site name="query" pos=".1 0 0"/>
           </body>
         </body>
       </body>
@@ -384,12 +388,27 @@ static constexpr char kTelescope[] = R"(
 </mujoco>
 )";
 
+static constexpr char kHinge[] = R"(
+<mujoco>
+  <worldbody>
+    <body name="query">
+      <joint name="link1" axis="0 1 0"/>
+      <geom type="capsule" size=".02" fromto="0 0 0 0 0 -1"/>
+      <site name="query" pos="0 0 -1"/>
+    </body>
+  </worldbody>
+
+  <keyframe>
+    <key qpos="1" qvel="1"/>
+  </keyframe>
+</mujoco>
+)";
+
 // compare mj_jacDot with finite-differenced mj_jac
 TEST_F(JacobianTest, JacDot) {
-  for (auto xml : {kQuat, kFreeBall, kQuatlessPendulum, kTelescope}) {
+  for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
     mjModel* model = LoadModelFromString(xml);
     int nv = model->nv;
-    mjtNum point[3] = {.01, .02, .03};
     mjData* data = mj_makeData(model);
 
     // load keyframe if present, step for a bit
@@ -407,34 +426,43 @@ TEST_F(JacobianTest, JacDot) {
     int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
     EXPECT_GT(bodyid, 0);
 
+    // get site position
+    int siteid = mj_name2id(model, mjOBJ_SITE, "query");
+    EXPECT_GT(siteid, -1);
+    mjtNum point[3];
+    mju_copy3(point, data->site_xpos+3*siteid);
+
     // jac, jac_dot
-    mj_markStack(data);
-    mjtNum* jac = mj_stackAllocNum(data, 6*nv);
-    mj_jac(model, data, jac, jac+3*nv, point, bodyid);
-    mjtNum* jac_dot =  mj_stackAllocNum(data, 6*nv);
-    mj_jacDot(model, data, jac_dot, jac_dot+3*nv, point, bodyid);
+    vector<mjtNum> jacp(3*nv);
+    vector<mjtNum> jacr(3*nv);
+    mj_jac(model, data, jacp.data(), jacr.data(), point, bodyid);
+    vector<mjtNum> jacp_dot(3*nv);
+    vector<mjtNum> jacr_dot(3*nv);
+    mj_jacDot(model, data, jacp_dot.data(), jacr_dot.data(), point, bodyid);
 
     // jac_h: jacobian after integrating qpos with a timestep of h
     mjtNum h = 1e-7;
     mj_integratePos(model, data->qpos, data->qvel, h);
     mj_kinematics(model, data);
     mj_comPos(model, data);
-    mjtNum* jac_h =  mj_stackAllocNum(data, 6*nv);;
-    mj_jac(model, data, jac_h, jac_h+3*nv, point, bodyid);
+    vector<mjtNum> jacp_h(3*nv);
+    vector<mjtNum> jacr_h(3*nv);
+    mju_copy3(point, data->site_xpos+3*siteid);  // get updated site position
+    mj_jac(model, data, jacp_h.data(), jacr_h.data(), point, bodyid);
 
     // jac_dot_h finite-difference approximation
-    mjtNum* jac_dot_h = mj_stackAllocNum(data, 6*nv);;
-    mju_sub(jac_dot_h, jac_h, jac, 6*nv);
-    mju_scl(jac_dot_h, jac_dot_h, 1/h, 6*nv);
+    vector<mjtNum> jacp_dot_h(3*nv);
+    mju_sub(jacp_dot_h.data(), jacp_h.data(), jacp.data(), 3*nv);
+    mju_scl(jacp_dot_h.data(), jacp_dot_h.data(), 1/h, 3*nv);
+    vector<mjtNum> jacr_dot_h(3*nv);
+    mju_sub(jacr_dot_h.data(), jacr_h.data(), jacr.data(), 3*nv);
+    mju_scl(jacr_dot_h.data(), jacr_dot_h.data(), 1/h, 3*nv);
 
     // compare finite-differenced and analytic
     mjtNum tol = 1e-5;
-    for (int j=0; j < 6; j++) {
-      EXPECT_THAT(AsVector(jac_dot_h + j*nv, nv),
-                  Pointwise(DoubleNear(tol), AsVector(jac_dot + j*nv, nv)));
-    }
+    EXPECT_THAT(jacp_dot, Pointwise(DoubleNear(tol), jacp_dot_h));
+    EXPECT_THAT(jacr_dot, Pointwise(DoubleNear(tol), jacr_dot_h));
 
-    mj_freeStack(data);
     mj_deleteData(data);
     mj_deleteModel(model);
   }
@@ -654,19 +682,19 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
   int size = mj_stateSize(model, spec);
 
   // save the initial state and step
-  std::vector<mjtNum> state0a(size);
+  vector<mjtNum> state0a(size);
   mj_getState(model, data, state0a.data(), spec);
 
   // get the initial state, expect equality
-  std::vector<mjtNum> state0b(size);
+  vector<mjtNum> state0b(size);
   mj_getState(model, data, state0b.data(), spec);
   EXPECT_EQ(state0a, state0b);
 
   // take one step
   mj_step(model, data);
 
   // save the resulting state
-  std::vector<mjtNum> state1a(size);
+  vector<mjtNum> state1a(size);
   mj_getState(model, data, state1a.data(), spec);
 
   // expect the state to be different after stepping
@@ -675,7 +703,7 @@ TEST_F(SupportTest, GetSetStateStepEqual) {
   // reset to the saved state, step again, get the resulting state
   mj_setState(model, data, state0a.data(), spec);
   mj_step(model, data);
-  std::vector<mjtNum> state1b(size);
+  vector<mjtNum> state1b(size);
   mj_getState(model, data, state1b.data(), spec);
 
   // expect the state to be the same after re-stepping
@@ -701,13 +729,13 @@ TEST_F(InertiaTest, DenseSameAsSparse) {
   }
 
   // dense zero matrix
-  std::vector<mjtNum> dst_sparse(nv * nv, 0.0);
+  vector<mjtNum> dst_sparse(nv * nv, 0.0);
 
   // sparse zero matrix
-  std::vector<mjtNum> dst_dense(nv * nv, 0.0);
-  std::vector<int> rownnz(nv, nv);
-  std::vector<int> rowadr(nv, 0);
-  std::vector<int> colind(nv * nv, 0);
+  vector<mjtNum> dst_dense(nv * nv, 0.0);
+  vector<int> rownnz(nv, nv);
+  vector<int> rowadr(nv, 0);
+  vector<int> colind(nv * nv, 0);
 
   // set sparse structure
   for (int i = 0; i < nv; i++) {
@@ -902,44 +930,44 @@ TEST_F(SupportTest, GeomDistance) {
   EXPECT_EQ(mj_geomDistance(model, data, 0, 1, distmax, nullptr), 0.5);
   mjtNum fromto[6];
   EXPECT_EQ(mj_geomDistance(model, data, 0, 1, distmax, fromto), 0.5);
-  EXPECT_THAT(fromto, Pointwise(Eq(), std::vector<mjtNum>{0, 0, 0, 0, 0, 0}));
+  EXPECT_THAT(fromto, Pointwise(Eq(), vector<mjtNum>{0, 0, 0, 0, 0, 0}));
 
   // plane-sphere
   distmax = 1.0;
   EXPECT_DOUBLE_EQ(mj_geomDistance(model, data, 0, 1, 1.0, fromto), 0.8);
   mjtNum eps = 1e-12;
   EXPECT_THAT(fromto, Pointwise(DoubleNear(eps),
-                                std::vector<mjtNum>{0, 0, 0, 0, 0, 0.8}));
+                                vector<mjtNum>{0, 0, 0, 0, 0, 0.8}));
 
   // sphere-plane
   EXPECT_DOUBLE_EQ(mj_geomDistance(model, data, 1, 0, 1.0, fromto), 0.8);
   EXPECT_THAT(fromto, Pointwise(DoubleNear(eps),
-                                std::vector<mjtNum>{0, 0, 0.8, 0, 0, 0}));
+                                vector<mjtNum>{0, 0, 0.8, 0, 0, 0}));
 
   // sphere-sphere
   EXPECT_DOUBLE_EQ(mj_geomDistance(model, data, 1, 2, 1.0, fromto), 0.5);
   EXPECT_THAT(fromto, Pointwise(DoubleNear(eps),
-                                std::vector<mjtNum>{.2, 0, 1, .7, 0, 1}));
+                                vector<mjtNum>{.2, 0, 1, .7, 0, 1}));
 
   // sphere-sphere, flipped order
   EXPECT_DOUBLE_EQ(mj_geomDistance(model, data, 2, 1, 1.0, fromto), 0.5);
   EXPECT_THAT(fromto, Pointwise(DoubleNear(eps),
-                                std::vector<mjtNum>{.7, 0, 1, .2, 0, 1}));
+                                vector<mjtNum>{.7, 0, 1, .2, 0, 1}));
 
   // mesh-sphere (close distmax)
   distmax = 0.701;
   eps = model->opt.ccd_tolerance;
   EXPECT_THAT(mj_geomDistance(model, data, 3, 1, distmax, fromto),
               DoubleNear(0.7, eps));
   EXPECT_THAT(fromto, Pointwise(DoubleNear(eps),
-                                std::vector<mjtNum>{0, 0, .8, 0, 0, .1}));
+                                vector<mjtNum>{0, 0, .8, 0, 0, .1}));
 
   // mesh-sphere (far distmax)
   distmax = 1.0;
   EXPECT_THAT(mj_geomDistance(model, data, 3, 1, distmax, fromto),
               DoubleNear(0.7, eps));
   EXPECT_THAT(fromto, Pointwise(DoubleNear(eps),
-                                std::vector<mjtNum>{0, 0, .8, 0, 0, .1}));
+                                vector<mjtNum>{0, 0, .8, 0, 0, .1}));
 
   mj_deleteData(data);
   mj_deleteModel(model);

Original file line number	Diff line number	Diff line change
`@@ -1893,7 +1893,7 @@ void mj_comVel(const mjModel* m, mjData* d) {`
`1893`	`1893`
`1894`	`1894`	`// assign cvel, cdofdot`
`1895`	`1895`	`mju_copy(d->cvel+6*i, cvel, 6);`
`1896`		`- mju_copy(d->cdof_dot+6bda, cdofdot, 6m->body_dofnum[i]);`
	`1896`	`+ mju_copy(d->cdof_dot+6bda, cdofdot, 6dofnum);`
`1897`	`1897`	`}`
`1898`	`1898`	`}`
`1899`	`1899`