diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 2c484096..e86ad896 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -3068,10 +3068,13 @@ mjd_transitionFD .. mujoco-include:: mjd_transitionFD -Finite differenced transition matrices. Letting :math:`x, u` denote the current :ref:`state` and control -vectors and letting :math:`y, s` denote the next state and sensor values, the top-level :ref:`mj_step` function computes -:math:`(x,u) \rightarrow (y,s)`. :ref:`mjd_transitionFD` computes the four associated Jacobians using -finite-differencing. These matrices and their dimensions are: +Finite-differenced discrete-time transition matrices. + +Letting :math:`x, u` denote the current :ref:`state` and :ref:`control` +vector in an :ref:`mjData` instance, and letting :math:`y, s` denote the next state and sensor +values, the top-level :ref:`mj_step` function computes :math:`(x,u) \rightarrow (y,s)`. +:ref:`mjd_transitionFD` computes the four associated Jacobians using finite-differencing. +These matrices and their dimensions are: .. csv-table:: :header: "matrix", "Jacobian", "dimension" @@ -3083,10 +3086,47 @@ finite-differencing. These matrices and their dimensions are: ``C``, :math:`\partial s / \partial x`, ``nsensordata x 2*nv+na`` ``D``, :math:`\partial s / \partial u`, ``nsensordata x nu`` -- All four matrix outputs are optional (can be NULL). +- All outputs are optional (can be NULL). - ``eps`` is the finite-differencing epsilon. - ``flg_centered`` denotes whether to use forward (0) or centered (1) differences. +.. _mjd_inverseFD: + +mjd_inverseFD +~~~~~~~~~~~~~ + +.. mujoco-include:: mjd_inverseFD + +Finite differenced continuous-time inverse-dynamics Jacobians. + +Letting :math:`x, a` denote the current :ref:`state` and acceleration vectors in an :ref:`mjData` instance, and +letting :math:`f, s` denote the forces computed by the inverse dynamics (``qfrc_inverse``), the function +:ref:`mj_inverse` computes :math:`(x,a) \rightarrow (f,s)`. :ref:`mjd_inverseFD` computes seven associated Jacobians +using finite-differencing. These matrices and their dimensions are: + +.. csv-table:: + :header: "matrix", "Jacobian", "dimension" + :widths: auto + :align: left + + ``DfDq``, :math:`\partial f / \partial q`, ``nv x nv`` + ``DfDv``, :math:`\partial f / \partial v`, ``nv x nv`` + ``DfDa``, :math:`\partial f / \partial a`, ``na x nv`` + ``DsDq``, :math:`\partial s / \partial q`, ``nv x nsensordata`` + ``DsDv``, :math:`\partial s / \partial v`, ``nv x nsensordata`` + ``DsDa``, :math:`\partial s / \partial a`, ``nv x nsensordata`` + ``DmDq``, :math:`\partial M / \partial q`, ``nv x nM`` + +- All outputs are optional (can be NULL). +- All outputs are transposed relative to Control Theory convention (i.e., column major). +- ``DmDq``, which contains a sparse representation of the ``nv x nv x nv`` tensor :math:`\partial M / \partial q`, is + not strictly an inverse dynamics Jacobian but is useful in related applications. It is provided as a convenience to + the user, since the required values are already computed if either of the other two :math:`\partial / \partial q` + Jacobians are requested. +- ``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. + .. _Plugins-api: Plugins diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index abd1184d..9a766d1e 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -371,10 +371,13 @@ Symmetrize square matrix :math:`R = \frac{1}{2}(M + M^T)`. .. _mjd_transitionFD: -Finite differenced transition matrices. Letting :math:`x, u` denote the current :ref:`state` and control -vectors and letting :math:`y, s` denote the next state and sensor values, the top-level :ref:`mj_step` function computes -:math:`(x,u) \rightarrow (y,s)`. :ref:`mjd_transitionFD` computes the four associated Jacobians using -finite-differencing. These matrices and their dimensions are: +Finite-differenced discrete-time transition matrices. + +Letting :math:`x, u` denote the current :ref:`state` and :ref:`control` +vector in an mjData instance, and letting :math:`y, s` denote the next state and sensor +values, the top-level :ref:`mj_step` function computes :math:`(x,u) \rightarrow (y,s)`. +:ref:`mjd_transitionFD` computes the four associated Jacobians using finite-differencing. +These matrices and their dimensions are: .. csv-table:: :header: "matrix", "Jacobian", "dimension" @@ -386,7 +389,38 @@ finite-differencing. These matrices and their dimensions are: ``C``, :math:`\partial s / \partial x`, ``nsensordata x 2*nv+na`` ``D``, :math:`\partial s / \partial u`, ``nsensordata x nu`` -- All four matrix outputs are optional (can be NULL). +- All outputs are optional (can be NULL). - ``eps`` is the finite-differencing epsilon. - ``flg_centered`` denotes whether to use forward (0) or centered (1) differences. +.. _mjd_inverseFD: + +Finite differenced continuous-time inverse-dynamics Jacobians. + +Letting :math:`x, a` denote the current :ref:`state` and acceleration vectors in an mjData instance, and +letting :math:`f, s` denote the forces computed by the inverse dynamics (``qfrc_inverse``), the function +:ref:`mj_inverse` computes :math:`(x,a) \rightarrow (f,s)`. :ref:`mjd_inverseFD` computes seven associated Jacobians +using finite-differencing. These matrices and their dimensions are: + +.. csv-table:: + :header: "matrix", "Jacobian", "dimension" + :widths: auto + :align: left + + ``DfDq``, :math:`\partial f / \partial q`, ``nv x nv`` + ``DfDv``, :math:`\partial f / \partial v`, ``nv x nv`` + ``DfDa``, :math:`\partial f / \partial a`, ``na x nv`` + ``DsDq``, :math:`\partial s / \partial q`, ``nv x nsensordata`` + ``DsDv``, :math:`\partial s / \partial v`, ``nv x nsensordata`` + ``DsDa``, :math:`\partial s / \partial a`, ``nv x nsensordata`` + ``DmDq``, :math:`\partial M / \partial q`, ``nv x nM`` + +- All outputs are optional (can be NULL). +- All outputs are transposed relative to Control Theory convention (i.e., column major). +- ``DmDq``, which contains a sparse representation of the ``nv x nv x nv`` tensor :math:`\partial M / \partial q`, is + not strictly an inverse dynamics Jacobian but is useful in related applications. It is provided as a convenience to + the user, since the required values are already computed if either of the other two :math:`\partial / \partial q` + Jacobians are requested. +- ``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. diff --git a/doc/computation.rst b/doc/computation.rst index bd47c07a..15bcb48c 100644 --- a/doc/computation.rst +++ b/doc/computation.rst @@ -1590,24 +1590,20 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu Derivatives ----------- -MuJoCo's entire computational pipline and uniquely -- its constraint solver -- are analytically differentiable. Writing +MuJoCo's entire computational pipline including its constraint solver are analytically differentiable. Writing efficient implementations of these derivatives is a long term goal of the development team. Analytic derivatives of the -smooth dynamics with respect to velocity are already in place and power the :ref:`implicit integrator`. +smooth dynamics (excluding constraints) with respect to velocity are already computed and enable the two +:ref:`implicit integrators`. -The function :ref:`mjd_transitionFD` computes state-transition and control-transition Jacobians. Given any valid MuJoCo -model ``mjModel* m`` with an initial :ref:`simulation state` in ``mjData* d``, +Two functions are currently available which use efficient finite-differencing in order to compute dynamics Jacobians: -- Let :math:`x` denote the :ref:`physics state` of the simulation at time :math:`t` -- the concatenation - of positions, velocities and actuator states ``[d->qpos; d->qvel; d->act]``. -- Let :math:`u` denote the vector of controls at time :math:`t`, corresponding to ``d->ctrl``. -- Let :math:`y` denote the physical state of the simulation at time :math:`t+h`, where :math:`h` corresponds to - ``m->opt.timstep``. -- Let :math:`s` denote the values of the sensors defined in the model. -- The high level function :ref:`mj_step` computes :math:`(x,u) \rightarrow (y,s)`: the next state and - sensor values as a function of the current state and control. -- ``mjd_transitionFD`` computes the Jacobians :math:`A = \frac{\partial y}{\partial x}`, - :math:`B = \frac{\partial y}{\partial u}`, :math:`C = \frac{\partial s}{\partial x}` and - :math:`D = \frac{\partial s}{\partial u}` using efficient finite-differencing of :ref:`mj_step`. +:ref:`mjd_transitionFD`: + Computes state-transition and control-transition Jacobians for the discrete-time forward dynamics (:ref:`mj_step`). + See :ref:`API documentation`. + +:ref:`mjd_inverseFD`: + Computes Jacobians for the continuous-time inverse dynamics (:ref:`mj_inverse`). + See :ref:`API documentation`. These derivatives are made efficient by exploiting MuJoCo's configurable computation pipeline so that quantities are not recomputed when not required. For example when differencing with respect to controls, quantities which depend only on diff --git a/doc/includes/references.h b/doc/includes/references.h index 07500615..a7d748f0 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -2441,6 +2441,10 @@ char* mju_strncpy(char *dst, const char *src, int n); mjtNum mju_sigmoid(mjtNum x); void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered, mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D); +void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation, + mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa, + mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa, + mjtNum *DmDq); 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 2ea4d19c..4102096e 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1200,7 +1200,25 @@ MJAPI mjtNum mju_sigmoid(mjtNum x); MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered, mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D); - +// 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: (na 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 +MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation, + mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa, + mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa, + mjtNum *DmDq); //---------------------- Plugins ------------------------------------------------------------------- diff --git a/introspect/functions.py b/introspect/functions.py index c0ca9320..0a12a860 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -7533,6 +7533,76 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Finite differenced transition matrices (control theory notation) d(x_next) = A*dx + B*du d(sensor) = C*dx + D*du required output matrix dimensions: A: (2*nv+na x 2*nv+na) B: (2*nv+na x nu) D: (nsensordata x 2*nv+na) C: (nsensordata x nu)', # pylint: disable=line-too-long )), + ('mjd_inverseFD', + FunctionDecl( + name='mjd_inverseFD', + 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'), + ), + ), + FunctionParameterDecl( + name='eps', + type=ValueType(name='mjtNum'), + ), + FunctionParameterDecl( + name='flg_actuation', + type=ValueType(name='mjtByte'), + ), + FunctionParameterDecl( + name='DfDq', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='DfDv', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='DfDa', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='DsDq', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='DsDv', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='DsDa', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='DmDq', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + ), + 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: (na 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 + )), ('mjp_defaultPlugin', FunctionDecl( name='mjp_defaultPlugin', diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index d5d90a90..e8a1ee45 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -41,7 +41,7 @@ TEST_XML = r""" - + @@ -54,6 +54,10 @@ TEST_XML = r""" + + + + """ @@ -1046,6 +1050,31 @@ Euler integrator, semi-implicit in velocity. bodyexclude=0, geomid=geomid) + def test_inverse_fd_none(self): + eps = 1e-6 + flg_centered = 0 + mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered, + None, None, None, None, None, None, None) + + def test_inverse_fd(self): + eps = 1e-6 + flg_centered = 0 + df_dq = np.zeros((self.model.nv, self.model.nv)) + df_dv = np.zeros((self.model.nv, self.model.nv)) + df_da = np.zeros((self.model.nv, self.model.nv)) + ds_dq = np.zeros((self.model.nv, self.model.nsensordata)) + ds_dv = np.zeros((self.model.nv, self.model.nsensordata)) + ds_da = np.zeros((self.model.nv, self.model.nsensordata)) + dm_dq = np.zeros((self.model.nv, self.model.nM)) + mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered, + df_dq, df_dv, df_da, ds_dq, ds_dv, ds_da, dm_dq) + self.assertGreater(np.linalg.norm(df_dq), eps) + self.assertGreater(np.linalg.norm(df_dv), eps) + self.assertGreater(np.linalg.norm(df_da), eps) + self.assertGreater(np.linalg.norm(ds_dq), eps) + self.assertGreater(np.linalg.norm(ds_dv), eps) + self.assertGreater(np.linalg.norm(ds_da), eps) + def test_mju_box_qp(self): n = 5 res = np.zeros(n) diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index f033d700..32407590 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -1131,7 +1131,8 @@ PYBIND11_MODULE(_functions, pymodule) { }); Def( pymodule, - [](const raw::MjModel* m, raw::MjData* d, mjtNum eps, mjtByte centered, + [](const raw::MjModel* m, raw::MjData* d, + mjtNum eps, mjtByte flg_centered, std::optional> A, std::optional> B, std::optional> C, @@ -1153,12 +1154,61 @@ PYBIND11_MODULE(_functions, pymodule) { throw py::type_error("D should be of shape (nsensordata, nu)"); } return InterceptMjErrors(::mjd_transitionFD)( - m, d, eps, centered, + m, d, eps, flg_centered, A.has_value() ? A->data() : nullptr, B.has_value() ? B->data() : nullptr, C.has_value() ? C->data() : nullptr, D.has_value() ? D->data() : nullptr); }); + Def( + pymodule, + [](const raw::MjModel* m, raw::MjData* d, + mjtNum eps, mjtByte flg_actuation, + std::optional> DfDq, + std::optional> DfDv, + std::optional> DfDa, + std::optional> DsDq, + std::optional> DsDv, + std::optional> DsDa, + std::optional> DmDq) { + if (DfDq.has_value() && + (DfDq->rows() != m->nv || DfDq->cols() != m->nv)) { + throw py::type_error("DfDq should be of shape (nv, nv)"); + } + if (DfDv.has_value() && + (DfDv->rows() != m->nv || DfDv->cols() != m->nv)) { + throw py::type_error("DfDv should be of shape (nv, nv)"); + } + if (DfDa.has_value() && + (DfDa->rows() != m->nv || DfDa->cols() != m->nv)) { + throw py::type_error("DfDa should be of shape (nv, nv)"); + } + if (DsDq.has_value() && + (DsDq->rows() != m->nv || DsDq->cols() != m->nsensordata)) { + throw py::type_error("DsDq should be of shape (nv, nsensordata)"); + } + if (DsDv.has_value() && + (DsDv->rows() != m->nv || DsDv->cols() != m->nsensordata)) { + throw py::type_error("DsDv should be of shape (nv, nsensordata)"); + } + if (DsDa.has_value() && + (DsDa->rows() != m->nv || DsDa->cols() != m->nsensordata)) { + throw py::type_error("DsDa should be of shape (nv, nsensordata)"); + } + if (DmDq.has_value() && + (DmDq->rows() != m->nv || DmDq->cols() != m->nM)) { + throw py::type_error("DmDq should be of shape (nv, nM)"); + } + return InterceptMjErrors(::mjd_inverseFD)( + m, d, eps, flg_actuation, + DfDq.has_value() ? DfDq->data() : nullptr, + DfDv.has_value() ? DfDv->data() : nullptr, + DfDa.has_value() ? DfDa->data() : nullptr, + DsDq.has_value() ? DsDq->data() : nullptr, + DsDv.has_value() ? DsDv->data() : nullptr, + 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); diff --git a/src/engine/engine_derivative_fd.c b/src/engine/engine_derivative_fd.c index 5cca1a57..de7bf14a 100644 --- a/src/engine/engine_derivative_fd.c +++ b/src/engine/engine_derivative_fd.c @@ -176,6 +176,19 @@ void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) { +// compute qfrc_inverse, optionally subtracting qfrc_actuator +static void inverseSkip(const mjModel* m, mjData* d, mjtStage stage, int skipsensor, + int flg_actuation, mjtNum* force) { + mj_inverseSkip(m, d, stage, skipsensor); + mju_copy(force, d->qfrc_inverse, m->nv); + if (flg_actuation) { + mj_fwdActuation(m, d); + mju_subFrom(force, d->qfrc_actuator, m->nv); + } +} + + + //------------------------- derivatives of passive forces ------------------------------------------ // add forward fin-diff approximation of (d qfrc_passive / d qvel) to qDeriv @@ -598,3 +611,123 @@ void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_cente mjFREESTACK; } + +// 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: (na 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 compute mass matrix Jacobian DmDq +// flg_actuation specifies whether to subtract qfrc_actuator from qfrc_inverse +void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation, + mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa, + mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa, + mjtNum *DmDq) { + int nq = m->nq, nv = m->nv, nM = m->nM, ns = m->nsensordata; + mjMARKSTACK; + + if (m->opt.integrator == mjINT_RK4) { + mju_error("RK4 integrator is not supported by mjd_inverseFD"); + } + + if (m->opt.noslip_iterations) { + mju_error("The noslip solver is not supported by mjd_inverseFD"); + } + + // skip sensor computations if no sensor Jacobians requested + int skipsensor = !DsDq && !DsDv && !DsDa; + + // local vectors + mjtNum *pos = mj_stackAlloc(d, nq); // position + mjtNum *force = mj_stackAlloc(d, nv); // force + mjtNum *force_plus = mj_stackAlloc(d, nv); // nudged force + mjtNum *sensor = skipsensor ? NULL : mj_stackAlloc(d, ns); // sensor values + mjtNum *mass = DmDq ? mj_stackAlloc(d, nM) : NULL; // mass matrix + + // save current positions + mju_copy(pos, d->qpos, nq); + + // center point outputs + inverseSkip(m, d, mjSTAGE_NONE, skipsensor, flg_actuation, force); + if (sensor) mju_copy(sensor, d->sensordata, ns); + if (mass) mju_copy(mass, d->qM, nM); + + // acceleration: skip = mjSTAGE_VEL + if (DfDa || DsDa) { + for (int i=0; i < nv; i++) { + // nudge acceleration + mjtNum tmp = d->qacc[i]; + d->qacc[i] += eps; + + // inverse dynamics, get force output + inverseSkip(m, d, mjSTAGE_VEL, skipsensor, flg_actuation, force_plus); + + // restore + d->qacc[i] = tmp; + + // row of force Jacobian + if (DfDa) diff(DfDa + i*nv, force, force_plus, eps, nv); + + // row of sensor Jacobian + if (DsDa) diff(DsDa + i*ns, sensor, d->sensordata, eps, ns); + } + } + + // velocity: skip = mjSTAGE_POS + if (DfDv || DsDv) { + for (int i=0; i < nv; i++) { + // nudge velocity + mjtNum tmp = d->qvel[i]; + d->qvel[i] += eps; + + // inverse dynamics, get force output + inverseSkip(m, d, mjSTAGE_POS, skipsensor, flg_actuation, force_plus); + + // restore + d->qvel[i] = tmp; + + // row of force Jacobian + if (DfDv) diff(DfDv + i*nv, force, force_plus, eps, nv); + + // row of sensor Jacobian + if (DsDv) diff(DsDv + i*ns, sensor, d->sensordata, eps, ns); + } + } + + // position: skip = mjSTAGE_NONE + if (DfDq || DsDq || DmDq) { + mjtNum *dpos = mj_stackAlloc(d, nv); // allocate position perturbation + for (int i=0; i < nv; i++) { + // nudge + mju_zero(dpos, nv); + dpos[i] = 1.0; + mj_integratePos(m, d->qpos, dpos, eps); + + // inverse dynamics, get force output + inverseSkip(m, d, mjSTAGE_NONE, skipsensor, flg_actuation, force_plus); + + // restore + mju_copy(d->qpos, pos, nq); + + // row of force Jacobian + if (DfDq) diff(DfDq + i*nv, force, force_plus, eps, nv); + + // row of sensor Jacobian + if (DsDq) diff(DsDq + i*ns, sensor, d->sensordata, eps, ns); + + // row of inertia Jacobian + if (DmDq) diff(DmDq + i*nM, mass, d->qM, eps, nM); + } + } + + mjFREESTACK; +} diff --git a/src/engine/engine_derivative_fd.h b/src/engine/engine_derivative_fd.h index d19667e0..00ad6ccc 100644 --- a/src/engine/engine_derivative_fd.h +++ b/src/engine/engine_derivative_fd.h @@ -36,6 +36,12 @@ MJAPI void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsenso MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D); +// finite differenced Jacobian of (force, sensors) = mj_inverse(state, acceleration) +MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation, + mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa, + mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa, + mjtNum *DmDq); + #ifdef __cplusplus } #endif diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 1d251e15..68a80c9d 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -23,6 +23,7 @@ #include "src/engine/engine_core_smooth.h" #include "src/engine/engine_derivative.h" #include "src/engine/engine_derivative_fd.h" +#include "src/engine/engine_forward.h" #include "src/engine/engine_io.h" #include "src/engine/engine_util_blas.h" #include "src/engine/engine_util_errmem.h" @@ -435,7 +436,7 @@ TEST_F(DerivativeTest, ClampedCtrlDerivatives) { // expect derivatives to be 0 EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0))); - // expect ctrl to remain unchanged (despite intenal clamping) + // expect ctrl to remain unchanged (despite internal clamping) EXPECT_EQ(data->ctrl[0], 2.0); EXPECT_EQ(data->ctrl[1], -2.0); @@ -646,5 +647,95 @@ TEST_F(DerivativeTest, DenseSparseRneEquivalent) { } } +// compare FD inverse derivatives to analytic derivatives of linear system +TEST_F(DerivativeTest, LinearSystemInverse) { + const std::string xml_path = GetTestDataFilePath(kLinearPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + + static const int nv = 3; + EXPECT_EQ(nv, model->nv); + static const int nu = 2; + EXPECT_EQ(nu, model->nu); + static const int ns = 5; + EXPECT_EQ(ns, model->nsensordata); + static const int nM = 6; + EXPECT_EQ(nM, model->nM); + + mjtNum DfDq[nv*nv]; + mjtNum DfDv[nv*nv]; + mjtNum DfDa[nv*nv]; + mjtNum DsDq[nv*ns]; + mjtNum DsDv[nv*ns]; + mjtNum DsDa[nv*ns]; + mjtNum DmDq[nv*nM]; + + // call mj_forward to get accelerations at initial state + mj_forward(model, data); + + // get derivatives + mjtNum eps = 1e-6; + mjtByte flg_actuation = 0; + mjd_inverseFD(model, data, eps, flg_actuation, + DfDq, DfDv, DfDa, + DsDq, DsDv, DsDa, + DmDq); + + // expect that position derivatives are the stiffnesses + mjtNum DfDq_expect[3*3] = {model->jnt_stiffness[0], 0, 0, + 0, model->jnt_stiffness[1], 0, + 0, 0, model->jnt_stiffness[2]}; + EXPECT_THAT(AsVector(DfDq, nv*nv), + Pointwise(DoubleNear(eps), AsVector(DfDq_expect, nv*nv))); + + // expect that velocity derivatives are the dampings + mjtNum DfDv_expect[3*3] = {model->dof_damping[0], 0, 0, + 0, model->dof_damping[1], 0, + 0, 0, model->dof_damping[2]}; + EXPECT_THAT(AsVector(DfDv, nv*nv), + Pointwise(DoubleNear(eps), AsVector(DfDv_expect, nv*nv))); + + // expect that acceleration derivatives are the mass matrix + mjtNum DfDa_expect[3*3]; + mj_fullM(model, DfDa_expect, data->qM); + EXPECT_THAT(AsVector(DfDa, nv*nv), + Pointwise(DoubleNear(eps), AsVector(DfDa_expect, nv*nv))); + + // expect that sensor derivatives w.r.t position only see sensor 1 at dof 0 + mjtNum DsDq_expect[3*5] = {0}; + int dof_index = 0; + int sensordata_index = model->sensor_adr[1]; + DsDq_expect[dof_index*ns + sensordata_index] = 1; + EXPECT_THAT(AsVector(DsDq, nv*ns), + Pointwise(DoubleNear(eps), AsVector(DsDq_expect, nv*ns))); + + // expect that sensor derivatives w.r.t velocity only see sensor 0 at dof 1 + mjtNum DsDv_expect[3*5] = {0}; + dof_index = 1; + sensordata_index = model->sensor_adr[0]; + DsDv_expect[dof_index*ns + sensordata_index] = 1; + EXPECT_THAT(AsVector(DsDv, nv*ns), + Pointwise(DoubleNear(eps), AsVector(DsDv_expect, nv*ns))); + + // expect that sensor derivatives w.r.t acceleration see the accelerometer + // in the y-axis, affected by both dof 0 and dof 1 + mjtNum DsDa_expect[3*5] = {0}; + dof_index = 0; + sensordata_index = model->sensor_adr[2] + 1; + DsDa_expect[dof_index*ns + sensordata_index] = 1; + dof_index = 1; + DsDa_expect[dof_index*ns + sensordata_index] = 1; + EXPECT_THAT(AsVector(DsDa, nv*ns), + Pointwise(DoubleNear(eps), AsVector(DsDa_expect, nv*ns))); + + // expect that mass matrix derivatives are zero + mjtNum DmDq_expect[nv*nM] = {0}; + EXPECT_THAT(AsVector(DmDq, nv*nM), + Pointwise(DoubleNear(eps), AsVector(DmDq_expect, nv*nM))); + + mj_deleteData(data); + mj_deleteModel(model); +} + } // namespace } // namespace mujoco diff --git a/test/engine/testdata/derivative/linear.xml b/test/engine/testdata/derivative/linear.xml index d136c9d2..c151bea1 100644 --- a/test/engine/testdata/derivative/linear.xml +++ b/test/engine/testdata/derivative/linear.xml @@ -5,6 +5,10 @@ + + @@ -12,6 +16,7 @@ + @@ -25,4 +30,10 @@ + + + + + + diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 15327f47..f1708192 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -3826,5 +3826,8 @@ public static unsafe extern double mju_sigmoid(double x); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mjd_transitionFD(mjModel_* m, mjData_* d, double eps, byte flg_centered, double* A, double* B, double* C, double* D); + +[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); } }