Add mjd_quatIntegrate, expose mjd_subQuat.

PiperOrigin-RevId: 546252926
Change-Id: I7d6b2bd6536e7e8afc368f3c46037c20f6a1cb5e
This commit is contained in:
Yuval Tassa
2023-07-07 05:32:18 -07:00
committed by Copybara-Service
parent 7401d6786a
commit 5cfbb6ac8b
13 changed files with 366 additions and 8 deletions
+1 -1
View File
@@ -29,7 +29,7 @@ API function can be classified as:
- :ref:`Poses transformations<Poses>`.
- :ref:`Matrix decompositions and solvers<Decompositions>`.
- :ref:`Miscellaneous<Miscellaneous>` functions.
- :ref:`Dynamics derivatives<Derivatives-api>`.
- :ref:`Derivatives<Derivatives-api>`.
- :ref:`Plugin<Plugins-api>` related functions.
- :ref:`Macros<Macros>`.
+40
View File
@@ -3185,6 +3185,10 @@ Sigmoid function over 0<=x<=1 using quintic polynomial.
Derivatives
^^^^^^^^^^^
The functions below provide useful derivatives of various functions, both analytic and
finite-differenced. The latter have names with the suffix ``FD``. Note that unlike much of the API,
outputs of derivative functions are the trailing rather than leading arguments.
.. _mjd_transitionFD:
mjd_transitionFD
@@ -3251,6 +3255,42 @@ using finite-differencing. These matrices and their dimensions are:
- ``flg_actuation`` denotes whether to subtract actuation forces (``qfrc_actuator``) from the output of the inverse
dynamics. If this flag is positive, actuator forces are not considered as external.
.. _mjd_subQuat:
mjd_subQuat
~~~~~~~~~~~
.. mujoco-include:: mjd_subQuat
Derivatives of :ref:`mju_subQuat` (quaternion difference).
.. _mjd_quatIntegrate:
mjd_quatIntegrate
~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjd_quatIntegrate
Derivatives of :ref:`mju_quatIntegrate`.
:math:`{\tt \small mju\_quatIntegrate}(q, v, h)` performs the in-place rotation :math:`q \leftarrow q + v h`,
where :math:`q \in \mathbf{S}^3` is a unit quaternion, :math:`v \in \mathbf{R}^3` is a 3D angular velocity and
:math:`h \in \mathbf{R^+}` is a timestep. This is equivalent to :math:`{\tt \small mju\_quatIntegrate}(q, s, 1.0)`,
where :math:`s` is the scaled velocity :math:`s = h v`.
:math:`{\tt \small mjd\_quatIntegrate}(v, h, D_q, D_v, D_h)` computes the Jacobians of the output :math:`q` with respect
to the inputs. Below, :math:`\bar q` denotes the pre-modified quaternion:
.. math::
\begin{aligned}
D_q &= \partial q / \partial \bar q \\
D_v &= \partial q / \partial v \\
D_h &= \partial q / \partial h
\end{aligned}
Note that derivatives depend only on :math:`h` and :math:`v` (in fact, on :math:`s = h v`).
All outputs are optional.
.. _Plugins-api:
Plugins
+30
View File
@@ -435,6 +435,10 @@ Symmetrize square matrix :math:`R = \frac{1}{2}(M + M^T)`.
.. _Derivatives-api:
The functions below provide useful derivatives of various functions, both analytic and
finite-differenced. The latter have names with the suffix ``FD``. Note that unlike much of the API,
outputs of derivative functions are the trailing rather than leading arguments.
.. _mjd_transitionFD:
Finite-differenced discrete-time transition matrices.
@@ -490,3 +494,29 @@ using finite-differencing. These matrices and their dimensions are:
- ``eps`` is the (forward) finite-differencing epsilon.
- ``flg_actuation`` denotes whether to subtract actuation forces (``qfrc_actuator``) from the output of the inverse
dynamics. If this flag is positive, actuator forces are not considered as external.
.. _mjd_subQuat:
Derivatives of :ref:`mju_subQuat` (quaternion difference).
.. _mjd_quatIntegrate:
Derivatives of :ref:`mju_quatIntegrate`.
:math:`{\tt \small mju\_quatIntegrate}(q, v, h)` performs the in-place rotation :math:`q \leftarrow q + v h`,
where :math:`q \in \mathbf{S}^3` is a unit quaternion, :math:`v \in \mathbf{R}^3` is a 3D angular velocity and
:math:`h \in \mathbf{R^+}` is a timestep. This is equivalent to :math:`{\tt \small mju\_quatIntegrate}(q, s, 1.0)`,
where :math:`s` is the scaled velocity :math:`s = h v`.
:math:`{\tt \small mjd\_quatIntegrate}(v, h, D_q, D_v, D_h)` computes the Jacobians of the output :math:`q` with respect
to the inputs. Below, :math:`\bar q` denotes the pre-modified quaternion:
.. math::
\begin{aligned}
D_q &= \partial q / \partial \bar q \\
D_v &= \partial q / \partial v \\
D_h &= \partial q / \partial h
\end{aligned}
Note that derivatives depend only on :math:`h` and :math:`v` (in fact, on :math:`s = h v`).
All outputs are optional.
+4 -1
View File
@@ -8,9 +8,12 @@ Upcoming version (not yet released)
General
^^^^^^^
- Added primitive collider for sphere-cylinder contacts, previously this pair used the generic convex-convex collider.
- Added analytic derivatives for quaternion :ref:`subtraction<mjd_subQuat>` and :ref:`integration<mjd_quatIntegrate>`
(rotation with an angular velocity). Derivatives are in the 3D tangent space.
- Added :ref:`mjv_connector` which has identical functionality to :ref:`mjv_makeConnector`, but with more convenient
"from-to" argument parametrization. :ref:`mjv_makeConnector` is now deprecated.
- Added primitive collider for sphere-cylinder contacts, previously this pair used the generic convex-convex collider.
Version 2.3.6 (June 20, 2023)
-----------------------------
+3
View File
@@ -2494,6 +2494,9 @@ void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuatio
mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
mjtNum *DmDq);
void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db[9]);
void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]);
void mjp_defaultPlugin(mjpPlugin* plugin);
int mjp_registerPlugin(const mjpPlugin* plugin);
int mjp_pluginCount();
+7
View File
@@ -1246,6 +1246,13 @@ MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_ac
mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
mjtNum *DmDq);
// Derivatives of mju_subQuat.
MJAPI void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db[9]);
// Derivatives of mju_quatIntegrate.
MJAPI void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]);
//---------------------- Plugins -------------------------------------------------------------------
// Set default plugin definition.
+76
View File
@@ -7973,6 +7973,82 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Finite differenced Jacobians of (force, sensors) = mj_inverse(state, acceleration) All outputs are optional. Output dimensions (transposed w.r.t Control Theory convention): DfDq: (nv x nv) DfDv: (nv x nv) DfDa: (nv x nv) DsDq: (nv x nsensordata) DsDv: (nv x nsensordata) DsDa: (nv x nsensordata) DmDq: (nv x nM) single-letter shortcuts: inputs: q=qpos, v=qvel, a=qacc outputs: f=qfrc_inverse, s=sensordata, m=qM notes: optionally computes mass matrix Jacobian DmDq flg_actuation specifies whether to subtract qfrc_actuator from qfrc_inverse', # pylint: disable=line-too-long
)),
('mjd_subQuat',
FunctionDecl(
name='mjd_subQuat',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='qa',
type=ArrayType(
inner_type=ValueType(name='mjtNum', is_const=True),
extents=(4,),
),
),
FunctionParameterDecl(
name='qb',
type=ArrayType(
inner_type=ValueType(name='mjtNum', is_const=True),
extents=(4,),
),
),
FunctionParameterDecl(
name='Da',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(9,),
),
),
FunctionParameterDecl(
name='Db',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(9,),
),
),
),
doc='Derivatives of mju_subQuat.',
)),
('mjd_quatIntegrate',
FunctionDecl(
name='mjd_quatIntegrate',
return_type=ValueType(name='void'),
parameters=(
FunctionParameterDecl(
name='vel',
type=ArrayType(
inner_type=ValueType(name='mjtNum', is_const=True),
extents=(3,),
),
),
FunctionParameterDecl(
name='scale',
type=ValueType(name='mjtNum'),
),
FunctionParameterDecl(
name='Dquat',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(9,),
),
),
FunctionParameterDecl(
name='Dvel',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(9,),
),
),
FunctionParameterDecl(
name='Dscale',
type=ArrayType(
inner_type=ValueType(name='mjtNum'),
extents=(3,),
),
),
),
doc='Derivatives of mju_quatIntegrate.',
)),
('mjp_defaultPlugin',
FunctionDecl(
name='mjp_defaultPlugin',
+15
View File
@@ -1113,6 +1113,21 @@ Euler integrator, semi-implicit in velocity.
self.assertGreater(np.linalg.norm(ds_dv), eps)
self.assertGreater(np.linalg.norm(ds_da), eps)
def test_mjd_sub_quat(self):
quat1 = np.array((0.2, 0.3, 0.3, 0.4))
quat2 = np.array((0.2, 0.3, 0.3, 0.4))
d1 = np.empty(9, np.float64)
d2 = np.empty(9, np.float64)
mujoco.mjd_subQuat(quat1, quat2, d1, d2)
def test_mjd_quat_intergrate(self):
scale = 0.1
vel = np.array((0.2, 0.3, 0.3))
d_quat = np.empty(9, np.float64)
d_vel = np.empty(9, np.float64)
d_h = np.empty(3, np.float64)
mujoco.mjd_quatIntegrate(vel, scale, d_quat, d_vel, d_h)
def test_banded(self):
n_total = 4
n_band = 1
+7 -3
View File
@@ -1239,6 +1239,11 @@ PYBIND11_MODULE(_functions, pymodule) {
return InterceptMjErrors(::mju_insertionSortInt)(
res.data(), res.size());
});
Def<traits::mju_Halton>(pymodule);
// Skipped: mju_strncpy (doesn't make sense in Python)
Def<traits::mju_sigmoid>(pymodule);
// Derivatives
Def<traits::mjd_transitionFD>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d,
@@ -1319,9 +1324,8 @@ PYBIND11_MODULE(_functions, pymodule) {
DsDa.has_value() ? DsDa->data() : nullptr,
DmDq.has_value() ? DmDq->data() : nullptr);
});
Def<traits::mju_Halton>(pymodule);
// Skipped: mju_strncpy (doesn't make sense in Python)
Def<traits::mju_sigmoid>(pymodule);
Def<traits::mjd_subQuat>(pymodule);
Def<traits::mjd_quatIntegrate>(pymodule);
} // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace
} // namespace mujoco::python
+58
View File
@@ -267,6 +267,64 @@ void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db
// derivative of mju_quatIntegrate w.r.t scaled velocity
// reference: https://arxiv.org/abs/1711.02508, Eq. 183
void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]) {
// scaled velocity
mjtNum s[3] = {scale*vel[0], scale*vel[1], scale*vel[2]};
// 3 basis matrices
mjtNum eye[9] = {
1, 0, 0,
0, 1, 0,
0, 0, 1
};
mjtNum cross[9] = {
0, s[2], -s[1],
-s[2], 0, s[0],
s[1], -s[0], 0
};
mjtNum outer[9] = {
s[0]*s[0], s[0]*s[1], s[0]*s[2],
s[1]*s[0], s[1]*s[1], s[1]*s[2],
s[2]*s[0], s[2]*s[1], s[2]*s[2]
};
// squared norm, norm of s
mjtNum xx = mju_dot3(s, s);
mjtNum x = mju_sqrt(xx);
// 4 coefficients: a=cos(x), b=sin(x)/x, c=(1-cos(x))/x^2, d=(x-sin(x))/x^3}
mjtNum a = mju_cos(x);
mjtNum b, c, d;
// x is not small: use full expressions
if (mju_abs(x) > 1.0/32) {
b = mju_sin(x) / x;
c = (1.0 - a) / xx;
d = (1.0 - b) / xx;
}
// |x| <= 1/32: use 6th order Taylor expansion (Horner form)
else {
b = 1 + xx/6 * (xx/20 * (1 - xx/42) - 1);
c = (1 + xx/12 * (xx/30 * (1 - xx/56) - 1)) / 2;
d = (1 + xx/20 * (xx/42 * (1 - xx/72) - 1)) / 6;
}
// derivatives
mjtNum Dvel_[9];
for (int i=0; i < 9; i++) {
if (Dquat) Dquat[i] = a*eye[i] + b*cross[i] + c*outer[i];
if (Dvel || Dscale) Dvel_[i] = b*eye[i] + c*cross[i] + d*outer[i];
}
if (Dvel) mju_copy(Dvel, Dvel_, 9);
if (Dscale) mju_rotVecMat(Dscale, vel, Dvel_);
}
//------------------------- dense derivatives of component functions -------------------------------
// no longer used, except in tests
+5 -1
View File
@@ -23,9 +23,13 @@
extern "C" {
#endif
// derivative of mju_subQuat w.r.t inputs
// derivatives of mju_subQuat w.r.t inputs
MJAPI void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db[9]);
// derivatives of mju_quatIntegrate w.r.t inputs
MJAPI void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]);
// analytical derivative of smooth forces w.r.t velocities:
// d->qDeriv = d (qfrc_actuator + qfrc_passive - [qfrc_bias]) / d qvel
MJAPI void mjd_smooth_vel(const mjModel* m, mjData* d, int flg_bias);
+114 -2
View File
@@ -70,9 +70,9 @@ static mjtNum CompareMatrices(mjtNum* Actual, mjtNum* Expected,
return max_error;
}
// utility function for matrix printing
// utility function for matrix printing (debug)
// NOLINTNEXTLINE(clang-diagnostic-unused-function)
static void PrintMatrix(mjtNum* mat, int nrow, int ncol) {
// NOLINT(clang-diagnostic-unused-function)
std::cerr.precision(5);
std::cerr << "\n";
for (int r=0; r < nrow; r++) {
@@ -836,5 +836,117 @@ TEST_F(DerivativeTest, SubQuat) {
}
}
// utility: random quaternion, 3D velocity
void randomQuatVel(mjtNum quat[4], mjtNum vel[3], int seed) {
// make distribution using seed
std::mt19937_64 rng;
rng.seed(seed);
std::normal_distribution<double> dist(0, 1);
// sample quat
for (int i=0; i < 4; i++) {
quat[i] = dist(rng);
}
mju_normalize4(quat);
// sample vel
for (int i=0; i < 3; i++) {
vel[i] = dist(rng);
}
}
// utility: finite-difference Jacobians of mju_quatIntegrate
void mjd_quatIntegrateFD(mjtNum Dquat[9], mjtNum Ds[9],
mjtNum Dvel[9], mjtNum Dh[3],
const mjtNum quat[4], const mjtNum vel[3],
mjtNum h, mjtNum eps) {
// compute y, output of mju_quatIntegrate(quat, vel, h)
mjtNum y[4] = {quat[0], quat[1], quat[2], quat[3]};
mju_quatIntegrate(y, vel, h);
mjtNum dx[3]; // nudged tangent-space input
mjtNum dq[4]; // quat output
mjtNum dy[3]; // nudged tangent-space output
mjtNum DquatT[9]; // Dquat transposed
mjtNum DsT[9]; // Ds transposed
mjtNum DvelT[9]; // Dvel transposed
for (int i = 0; i < 3; i++) {
// perturbation
mju_zero3(dx);
dx[i] = 1.0;
// d_y / d_quat
mju_copy4(dq, quat);
mju_quatIntegrate(dq, dx, eps); // nudge dq
mju_quatIntegrate(dq, vel, h); // compute nudged
mju_subQuat(dy, dq, y); // subtract
mju_scl3(DquatT + i * 3, dy, 1.0 / eps);
// d_y / d_sv (scaled velocity)
mju_copy4(dq, quat);
mjtNum dsv[3] = {vel[0]*h, vel[1]*h, vel[2]*h};
mju_addToScl3(dsv, dx, eps); // nudge dsv
mju_quatIntegrate(dq, dsv, 1.0); // compute nudged
mju_subQuat(dy, dq, y); // subtract
mju_scl3(DsT + i * 3, dy, 1.0 / eps);
// d_y / d_v (unscaled velocity)
mju_copy4(dq, quat);
mjtNum dv[3] = {vel[0], vel[1], vel[2]};
mju_addToScl3(dv, dx, eps); // nudge dv
mju_quatIntegrate(dq, dv, h); // compute nudged
mju_subQuat(dy, dq, y); // subtract
mju_scl3(DvelT + i * 3, dy, 1.0 / eps);
}
// d_y / d_h (unscaled velocity)
mju_copy4(dq, quat);
mju_quatIntegrate(dq, vel, h + eps); // compute nudged
mju_subQuat(dy, dq, y); // subtract
mju_scl3(Dh, dy, 1.0 / eps);
// transpose
mju_transpose(Dquat, DquatT, 3, 3);
mju_transpose(Ds, DsT, 3, 3);
mju_transpose(Dvel, DsT, 3, 3);
}
TEST_F(DerivativeTest, quatIntegrate) {
const int nrepeats = 10; // number of repeats
const mjtNum eps = 1e-7; // epsilon for finite-differencing and comparison
int seed = 1;
for (int i = 0; i < nrepeats; i++) {
for (mjtNum h : {0.0, 1e-9, 1e-5, 1e-2, 1.0, 4.0}) {
// make random quaternion and velocity
mjtNum quat[4];
mjtNum vel[3];
randomQuatVel(quat, vel, seed++);
// analytic Jacobians
mjtNum Dquat[9]; // d_quatIntegrate(quat, vel, h) / d_quat
mjtNum Dvel[9]; // d_quatIntegrate(quat, vel, h) / d_vel
mjtNum Dh[3]; // d_quatIntegrate(quat, vel, h) / d_h
mjd_quatIntegrate(vel, h, Dquat, Dvel, Dh);
// finite-differenced Jacobians
mjtNum DquatFD[9];
mjtNum DsFD[9];
mjtNum DvelFD[9];
mjtNum DhFD[3];
mjd_quatIntegrateFD(DquatFD, DsFD, DvelFD, DhFD, quat, vel, h, eps);
// expect numerical equality of un/scaled velocity derivatives
EXPECT_THAT(AsVector(DvelFD, 9), Pointwise(DoubleNear(eps), DsFD));
// expect numerical equality of analytic and FD derivatives
EXPECT_THAT(AsVector(DquatFD, 9), Pointwise(DoubleNear(eps), Dquat));
EXPECT_THAT(AsVector(DvelFD, 9), Pointwise(DoubleNear(eps), Dvel));
EXPECT_THAT(AsVector(DhFD, 3), Pointwise(DoubleNear(eps), Dh));
}
}
}
} // namespace
} // namespace mujoco
+6
View File
@@ -3865,5 +3865,11 @@ public static unsafe extern void mjd_transitionFD(mjModel_* m, mjData_* d, doubl
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mjd_inverseFD(mjModel_* m, mjData_* d, double eps, byte flg_actuation, double* DfDq, double* DfDv, double* DfDa, double* DsDq, double* DsDv, double* DsDa, double* DmDq);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mjd_subQuat(double* qa, double* qb, double* Da, double* Db);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mjd_quatIntegrate(double* vel, double scale, double* Dquat, double* Dvel, double* Dscale);
}
}