diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 09619b39..234cbb8b 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -323,7 +323,8 @@ frame (``point``) treated as attached to the body, the Jacobian has both transla (``jacr``) components. Passing ``NULL`` for either pointer will skip that part of the computation. Each component is a 3-by-nv matrix. Each row of this matrix is the gradient of the corresponding coordinate of the specified point with respect to the degrees-of-freedom. The :ref:`pipeline stages` required for Jacobian computations to be -consistent with the current generalized positions ``mjData.qpos`` are :ref:`mj_kinematics` and :ref:`mj_comPos`. +consistent with the current generalized positions ``mjData.qpos`` are :ref:`mj_kinematics` and :ref:`mj_comPos` (in +that order). .. _mj_jacBody: @@ -380,6 +381,18 @@ mj_jacPointAxis Compute translation end-effector Jacobian of point, and rotation Jacobian of axis. +.. _mj_jacDot: + +mj_jacDot +~~~~~~~~~ + +.. mujoco-include:: mj_jacDot + +This function computes the time-derivative of an end-effector kinematic Jacobian computed by :ref:`mj_jac`. +The minimal :ref:`pipeline stages` required for computation to be +consistent with the current generalized positions and velocities ``mjData.{qpos, qvel}`` are +:ref:`mj_kinematics`, :ref:`mj_comPos`, :ref:`mj_comVel` (in that order). + .. _mj_angmomMat: mj_angmomMat diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index d703cf94..093a21ae 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -184,13 +184,21 @@ frame (``point``) treated as attached to the body, the Jacobian has both transla (``jacr``) components. Passing ``NULL`` for either pointer will skip that part of the computation. Each component is a 3-by-nv matrix. Each row of this matrix is the gradient of the corresponding coordinate of the specified point with respect to the degrees-of-freedom. The :ref:`pipeline stages` required for Jacobian computations to be -consistent with the current generalized positions ``mjData.qpos`` are :ref:`mj_kinematics` and :ref:`mj_comPos`. +consistent with the current generalized positions ``mjData.qpos`` are :ref:`mj_kinematics` and :ref:`mj_comPos` (in +that order). .. _mj_jacBody: This and the remaining variants of the Jacobian function call mj_jac internally, with the center of the body, geom or site. They are just shortcuts; the same can be achieved by calling mj_jac directly. +.. _mj_jacDot: + +This function computes the time-derivative of an end-effector kinematic Jacobian computed by :ref:`mj_jac`. +The minimal :ref:`pipeline stages` required for computation to be +consistent with the current generalized positions and velocities ``mjData.{qpos, qvel}`` are +:ref:`mj_kinematics`, :ref:`mj_comPos`, :ref:`mj_comVel` (in that order). + .. _mj_angmomMat: This function computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear mapping from diff --git a/doc/changelog.rst b/doc/changelog.rst index ca946999..113eae38 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -29,6 +29,7 @@ General parent model, but only on the first attachment. - Added the :ref:`mjtSameFrame` enum which contains the possible frame alignments of bodies and their children. These alignments are used for computation shortcuts in :ref:`mj_kinematics`. +- Added :ref:`mj_jacDot` for computing time-derivatives of kinematic Jacobians. Fixes :github:issue:`411`. MJX ^^^ diff --git a/doc/includes/references.h b/doc/includes/references.h index 8f55f852..5769ba30 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3217,6 +3217,8 @@ void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, i void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site); void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis, const mjtNum point[3], const mjtNum axis[3], int body); +void mj_jacDot(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, + const mjtNum point[3], int body); void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body); int mj_name2id(const mjModel* m, int type, const char* name); const char* mj_id2name(const mjModel* m, int type, int id); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 0d115c9f..98c415e1 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -456,6 +456,10 @@ MJAPI void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* j MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis, const mjtNum point[3], const mjtNum axis[3], int body); +// Compute 3/6-by-nv Jacobian time derivative of global point attached to given body. +MJAPI void mj_jacDot(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, + const mjtNum point[3], int body); + // Compute subtree angular momentum matrix. MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body); diff --git a/introspect/functions.py b/introspect/functions.py index 90e744a0..c90ddb7e 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -2623,6 +2623,49 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Compute translation end-effector Jacobian of point, and rotation Jacobian of axis.', # pylint: disable=line-too-long )), + ('mj_jacDot', + FunctionDecl( + name='mj_jacDot', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData', is_const=True), + ), + ), + FunctionParameterDecl( + name='jacp', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='jacr', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='point', + type=ArrayType( + inner_type=ValueType(name='mjtNum', is_const=True), + extents=(3,), + ), + ), + FunctionParameterDecl( + name='body', + type=ValueType(name='int'), + ), + ), + doc='Compute 3/6-by-nv Jacobian time derivative of global point attached to given body.', # pylint: disable=line-too-long + )), ('mj_angmomMat', FunctionDecl( name='mj_angmomMat', diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 86beb5cd..a82005b2 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -449,6 +449,26 @@ PYBIND11_MODULE(_functions, pymodule) { jacr.has_value() ? jacr->data() : nullptr, &(*point)[0], &(*axis)[0], body); }); + Def( + pymodule, + [](const raw::MjModel* m, raw::MjData* d, + std::optional> jacp, + std::optional> jacr, + const mjtNum (*point)[3], int body) { + if (jacp.has_value() && + (jacp->rows() != 3 || jacp->cols() != m->nv)) { + throw py::type_error("jacp should be of shape (3, nv)"); + } + if (jacr.has_value() && + (jacr->rows() != 3 || jacr->cols() != m->nv)) { + throw py::type_error("jacr should be of shape (3, nv)"); + } + return InterceptMjErrors(::mj_jacDot)( + m, d, + jacp.has_value() ? jacp->data() : nullptr, + jacr.has_value() ? jacr->data() : nullptr, + &(*point)[0], body); + }); Def( pymodule, [](const raw::MjModel* m, raw::MjData* d, Eigen::Ref mat, int body) { diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 2fec9331..5602eece 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -1624,7 +1624,6 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) // compute cvel, cdof_dot void mj_comVel(const mjModel* m, mjData* d) { int nbody = m->nbody; - mjtNum tmp[6], cvel[6], cdofdot[36]; // set world vel to 0 mju_zero(d->cvel, 6); @@ -1635,11 +1634,15 @@ void mj_comVel(const mjModel* m, mjData* d) { int bda = m->body_dofadr[i]; // cvel = cvel_parent + mjtNum cvel[6]; mju_copy(cvel, d->cvel+6*m->body_parentid[i], 6); // cvel = cvel_parent + cdof * qvel, cdofdot = cvel x cdof int dofnum = m->body_dofnum[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: diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 1a343268..06c4380d 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -386,19 +386,17 @@ int mj_bodyChain(const mjModel* m, int body, int* chain) { void mj_jac(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body) { int nv = m->nv; + mjtNum offset[3]; - // clear jacobians + // clear jacobians, compute offset if required if (jacp) { mju_zero(jacp, 3*nv); + mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]); } if (jacr) { mju_zero(jacr, 3*nv); } - // compute point-com offset - mjtNum offset[3]; - mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]); - // skip fixed bodies while (body && !m->body_dofnum[body]) { body = m->body_parentid[body]; @@ -410,30 +408,30 @@ void mj_jac(const mjModel* m, const mjData* d, } // get last dof that affects this (as well as the original) body - int da = m->body_dofadr[body] + m->body_dofnum[body] - 1; + int i = m->body_dofadr[body] + m->body_dofnum[body] - 1; // backward pass over dof ancestor chain - while (da >= 0) { - mjtNum *cdof = d->cdof; + while (i >= 0) { + mjtNum* cdof = d->cdof+6*i; // construct rotation jacobian if (jacr) { - jacr[da] = cdof[6*da]; - jacr[da+nv] = cdof[6*da+1]; - jacr[da+2*nv] = cdof[6*da+2]; + jacr[i+0*nv] = cdof[0]; + jacr[i+1*nv] = cdof[1]; + jacr[i+2*nv] = cdof[2]; } // construct translation jacobian (correct for rotation) if (jacp) { mjtNum tmp[3]; - mju_cross(tmp, cdof+6*da, offset); - jacp[da] = cdof[6*da+3] + tmp[0]; - jacp[da+nv] = cdof[6*da+4] + tmp[1]; - jacp[da+2*nv] = cdof[6*da+5] + tmp[2]; + mju_cross(tmp, cdof, offset); + jacp[i+0*nv] = cdof[3] + tmp[0]; + jacp[i+1*nv] = cdof[4] + tmp[1]; + jacp[i+2*nv] = cdof[5] + tmp[2]; } // advance to parent dof - da = m->dof_parentid[da]; + i = m->dof_parentid[i]; } } @@ -594,7 +592,7 @@ void mj_jacSparse(const mjModel* m, const mjData* d, 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) { - mjtNum offset[3], tmp[3], *cdof = d->cdof; + mjtNum offset[3], tmp[3]; // compute point-com offset mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]); @@ -608,39 +606,41 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d, int ci = start; int end = m->body_dofadr[body] + m->body_dofnum[body]; for (int da=m->body_dofadr[body]; da < end; da++) { + mjtNum *cdof = d->cdof+6*da; + // construct rotation jacobian if (jacdifr) { // plus sign if (flg_second) { - jacdifr[ci] = cdof[6*da]; - jacdifr[ci+NV] = cdof[6*da+1]; - jacdifr[ci+2*NV] = cdof[6*da+2]; + jacdifr[ci+0*NV] = cdof[0]; + jacdifr[ci+1*NV] = cdof[1]; + jacdifr[ci+2*NV] = cdof[2]; } // minus sign else { - jacdifr[ci] = -cdof[6*da]; - jacdifr[ci+NV] = -cdof[6*da+1]; - jacdifr[ci+2*NV] = -cdof[6*da+2]; + jacdifr[ci+0*NV] = -cdof[0]; + jacdifr[ci+1*NV] = -cdof[1]; + jacdifr[ci+2*NV] = -cdof[2]; } } // construct translation jacobian (correct for rotation) if (jacdifp) { - mju_cross(tmp, cdof+6*da, offset); + mju_cross(tmp, cdof, offset); // plus sign if (flg_second) { - jacdifp[ci] = (cdof[6*da+3] + tmp[0]); - jacdifp[ci+NV] = (cdof[6*da+4] + tmp[1]); - jacdifp[ci+2*NV] = (cdof[6*da+5] + tmp[2]); + jacdifp[ci+0*NV] = (cdof[3] + tmp[0]); + jacdifp[ci+1*NV] = (cdof[4] + tmp[1]); + jacdifp[ci+2*NV] = (cdof[5] + tmp[2]); } // plus sign else { - jacdifp[ci] = -(cdof[6*da+3] + tmp[0]); - jacdifp[ci+NV] = -(cdof[6*da+4] + tmp[1]); - jacdifp[ci+2*NV] = -(cdof[6*da+5] + tmp[2]); + jacdifp[ci+0*NV] = -(cdof[3] + tmp[0]); + jacdifp[ci+1*NV] = -(cdof[4] + tmp[1]); + jacdifp[ci+2*NV] = -(cdof[5] + tmp[2]); } } @@ -802,6 +802,72 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain, +// compute 3/6-by-nv Jacobian time derivative of global point attached to given body +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]; + + // clear jacobians, compute offset if required + if (jacp) { + mju_zero(jacp, 3*nv); + mju_sub3(offset, point, d->subtree_com+3*m->body_rootid[body]); + } + if (jacr) { + mju_zero(jacr, 3*nv); + } + + // skip fixed bodies + while (body && !m->body_dofnum[body]) { + body = m->body_parentid[body]; + } + + // no movable body found: nothing to do + if (!body) { + return; + } + + // get last dof that affects this (as well as the original) body + int i = m->body_dofadr[body] + m->body_dofnum[body] - 1; + + // backward pass over dof ancestor chain + while (i >= 0) { + mjtNum cdof_dot[6]; + mju_copy(cdof_dot, d->cdof_dot+6*i, 6); + + // check for quaternion + mjtJoint type = m->jnt_type[m->dof_jntid[i]]; + int dofadr = m->jnt_dofadr[m->dof_jntid[i]]; + int is_quat = type == mjJNT_BALL || (type == mjJNT_FREE && i >= dofadr + 3); + + // 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); + } + + // construct rotation jacobian + if (jacr) { + jacr[i+0*nv] += cdof_dot[0]; + jacr[i+1*nv] += cdof_dot[1]; + jacr[i+2*nv] += cdof_dot[2]; + } + + // 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]; + } + + // advance to parent dof + i = m->dof_parentid[i]; + } +} + + + // compute subtree angular momentum matrix void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) { int nv = m->nv; diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index df1c8e02..f7cdf9b8 100644 --- a/src/engine/engine_support.h +++ b/src/engine/engine_support.h @@ -107,6 +107,10 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain, int n, const int* body, const mjtNum* weight, const mjtNum point[3], mjtNum* jac, int flg_rot); +// compute 3/6-by-nv Jacobian time derivative of global point attached to given body +MJAPI void mj_jacDot(const mjModel* m, const mjData* d, + mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body); + // compute subtree angular momentum matrix MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body); diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index 93e69bfa..67e41787 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -29,16 +29,17 @@ namespace mujoco { namespace { -std::vector AsVector(const mjtNum* array, int n) { - return std::vector(array, array + n); -} - using ::testing::DoubleNear; using ::testing::Eq; using ::testing::ContainsRegex; // NOLINT using ::testing::MatchesRegex; using ::testing::Pointwise; using ::testing::ElementsAreArray; +using ::testing::Pointwise; + +std::vector AsVector(const mjtNum* array, int n) { + return std::vector(array, array + n); +} using AngMomMatTest = MujocoTest; @@ -285,6 +286,164 @@ TEST_F(JacobianTest, SubtreeJacNoInternalAcc) { mj_deleteModel(model); } +static constexpr char kQuat[] = R"( + + + + + + + + + + + +)"; + +static constexpr char kFreeBall[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +static constexpr char kQuatlessPendulum[] = R"( + + + + + + + + + + + + + + + + + + + + + + + +)"; + +static constexpr char kTelescope[] = R"( + + + + + + + + + + + + + + + + + + + + + + + +)"; + +// compare mj_jacDot with finite-differenced mj_jac +TEST_F(JacobianTest, JacDot) { + for (auto xml : {kQuat, kFreeBall, kQuatlessPendulum, kTelescope}) { + 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 + 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 + int bodyid = mj_name2id(model, mjOBJ_BODY, "query"); + EXPECT_GT(bodyid, 0); + + // 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); + + // 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); + + // 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); + + // 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))); + } + + mj_freeStack(data); + mj_deleteData(data); + mj_deleteModel(model); + } +} + using Name2idTest = MujocoTest; static constexpr char name2idTestingModel[] = R"( diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 176f0307..f0cc7180 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -6715,6 +6715,9 @@ public static unsafe extern void mj_jacSite(mjModel_* m, mjData_* d, double* jac [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_jacPointAxis(mjModel_* m, mjData_* d, double* jacPoint, double* jacAxis, double* point, double* axis, int body); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mj_jacDot(mjModel_* m, mjData_* d, double* jacp, double* jacr, double* point, int body); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_angmomMat(mjModel_* m, mjData_* d, double* mat, int body);