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

Add mj_jacDot: Time derivative of kinematic Jacobian.

Browse Source 
PiperOrigin-RevId: 670505417
Change-Id: Icf30cff5dc955e002bd2f5d5fb3a04f8388d0780
This commit is contained in:
Yuval Tassa
2024-09-03 04:48:44 -07:00
committed by Copybara-Service
parent fb38727303
commit 2d3d5415b7
12 changed files with 363 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+14 -1
View File
@@ -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<piStages>` 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<piStages>` 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
doc/APIreference/functions_override.rst
+9 -1
View File
@@ -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<piStages>` 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<piStages>` 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
doc/changelog.rst
+1
View File
@@ -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
^^^
doc/includes/references.h
+2
View File
@@ -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);
include/mujoco/mujoco.h
+4
View File
@@ -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);
introspect/functions.py
+43
View File
@@ -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',
python/mujoco/functions.cc
+20
View File
@@ -449,6 +449,26 @@ PYBIND11_MODULE(_functions, pymodule) {
jacr.has_value() ? jacr->data() : nullptr,
&(*point)[0], &(*axis)[0], body);
});
Def<traits::mj_jacDot>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
std::optional<Eigen::Ref<EigenArrayXX>> jacp,
std::optional<Eigen::Ref<EigenArrayXX>> 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<traits::mj_angmomMat>(
pymodule, [](const raw::MjModel* m, raw::MjData* d,
Eigen::Ref<EigenArrayXX> mat, int body) {
src/engine/engine_core_smooth.c
+4 -1
View File
@@ -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:
src/engine/engine_support.c
+96 -30
View File
@@ -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;
src/engine/engine_support.h
+4
View File
@@ -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);
test/engine/engine_support_test.cc
+163 -4
View File
@@ -29,16 +29,17 @@
namespace mujoco {
namespace {
std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(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<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(array, array + n);
}
using AngMomMatTest = MujocoTest;
@@ -285,6 +286,164 @@ TEST_F(JacobianTest, SubtreeJacNoInternalAcc) {
mj_deleteModel(model);
}
static constexpr char kQuat[] = R"(
<mujoco>
<worldbody>
<body name="query">
<joint type="ball"/>
<geom size="1"/>
</body>
</worldbody>
<keyframe>
<key qvel="2 3 5"/>
</keyframe>
</mujoco>
)";
static constexpr char kFreeBall[] = R"(
<mujoco>
<worldbody>
<body name="distractor1" pos="0 0 .3">
<freejoint/>
<geom size=".1"/>
</body>
<body name="main">
<freejoint/>
<geom size=".1"/>
<body pos=".1 0 0">
<joint axis="0 1 0"/>
<geom type="capsule" size=".03" fromto="0 0 0 .2 0 0"/>
<body pos=".2 0 0">
<joint type="ball" stiffness="20"/>
<geom type="capsule" size=".03" fromto="0 0 0 0 .2 0"/>
<body name="query" pos="0 .2 0">
<joint type="slide" axis="1 1 1"/>
<geom size=".05"/>
</body>
</body>
</body>
</body>
<body name="distractor2" pos="0 0 -.3">
<freejoint/>
<geom size=".1"/>
</body>
</worldbody>
<keyframe>
<key qvel="1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1"/>
</keyframe>
</mujoco>
)";
static constexpr char kQuatlessPendulum[] = R"(
<mujoco>
<option integrator="implicit">
<flag constraint="disable"/>
</option>
<worldbody>
<body pos="0.15 0 0">
<joint type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<body pos="0.1 0 0">
<joint type="slide" axis="1 0 0" stiffness="200"/>
<geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
<body pos=".1 0 0">
<joint axis="1 0 0"/>
<joint axis="0 1 0"/>
<joint axis="0 0 1"/>
<geom type="box" size=".02" fromto="0 0 0 0 .1 0"/>
<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"/>
</body>
</body>
</body>
</body>
</worldbody>
</mujoco>
)";
static constexpr char kTelescope[] = R"(
<mujoco>
<worldbody>
<body>
<joint type="ball"/>
<geom type="capsule" size="0.02" fromto="0 .02 0 .1 .02 0"/>
<body pos=".1 .02 0">
<joint type="slide" axis="1 0 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<body pos=".1 .02 0">
<joint type="slide" axis="1 0 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<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"/>
</body>
</body>
</body>
</body>
</worldbody>
<keyframe>
<key qvel="1 1 1 1 1 1"/>
</keyframe>
</mujoco>
)";
// 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"(
unity/Runtime/Bindings/MjBindings.cs
+3
View File
@@ -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);