From 5cfbb6ac8b8a6322b2c695734e87fbda6b8e5970 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 7 Jul 2023 05:32:18 -0700 Subject: [PATCH] Add `mjd_quatIntegrate`, expose `mjd_subQuat`. PiperOrigin-RevId: 546252926 Change-Id: I7d6b2bd6536e7e8afc368f3c46037c20f6a1cb5e --- doc/APIreference/APIfunctions.rst | 2 +- doc/APIreference/functions.rst | 40 ++++++++ doc/APIreference/functions_override.rst | 30 ++++++ doc/changelog.rst | 5 +- doc/includes/references.h | 3 + include/mujoco/mujoco.h | 7 ++ introspect/functions.py | 76 ++++++++++++++++ python/mujoco/bindings_test.py | 15 +++ python/mujoco/functions.cc | 10 +- src/engine/engine_derivative.c | 58 ++++++++++++ src/engine/engine_derivative.h | 6 +- test/engine/engine_derivative_test.cc | 116 +++++++++++++++++++++++- unity/Runtime/Bindings/MjBindings.cs | 6 ++ 13 files changed, 366 insertions(+), 8 deletions(-) diff --git a/doc/APIreference/APIfunctions.rst b/doc/APIreference/APIfunctions.rst index 96aab950..acf5f274 100644 --- a/doc/APIreference/APIfunctions.rst +++ b/doc/APIreference/APIfunctions.rst @@ -29,7 +29,7 @@ API function can be classified as: - :ref:`Poses transformations`. - :ref:`Matrix decompositions and solvers`. - :ref:`Miscellaneous` functions. -- :ref:`Dynamics derivatives`. +- :ref:`Derivatives`. - :ref:`Plugin` related functions. - :ref:`Macros`. diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 0e302c5c..caa62507 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -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 diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 21027e3f..c749668a 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -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. diff --git a/doc/changelog.rst b/doc/changelog.rst index 78a076e7..1b0ebc66 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -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` and :ref:`integration` + (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) ----------------------------- diff --git a/doc/includes/references.h b/doc/includes/references.h index ca0bc529..b2c37fb2 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -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(); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index c5894980..0b5dc24a 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -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. diff --git a/introspect/functions.py b/introspect/functions.py index e309785e..91feff67 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -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', diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index d016c355..c079c814 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -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 diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 9db25866..3b52a937 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -1239,6 +1239,11 @@ PYBIND11_MODULE(_functions, pymodule) { return InterceptMjErrors(::mju_insertionSortInt)( res.data(), res.size()); }); + Def(pymodule); + // Skipped: mju_strncpy (doesn't make sense in Python) + Def(pymodule); + + // Derivatives Def( 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(pymodule); - // Skipped: mju_strncpy (doesn't make sense in Python) - Def(pymodule); + Def(pymodule); + Def(pymodule); } // PYBIND11_MODULE NOLINT(readability/fn_size) } // namespace } // namespace mujoco::python diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index cb14f667..fdc67b91 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -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 diff --git a/src/engine/engine_derivative.h b/src/engine/engine_derivative.h index 52fdf5c8..8f316e6c 100644 --- a/src/engine/engine_derivative.h +++ b/src/engine/engine_derivative.h @@ -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); diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 5c55f7ff..a55d7f77 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -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 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 diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 6a6224c3..14474b2f 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -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); } }