Add invdiscrete flag to enable discrete-time inverse dynamics with mj_inverse.

PiperOrigin-RevId: 560161140
Change-Id: If3e30be5e9af4982571dd310172f2fc0d8d1c551

doc/XMLreference.rst

@@ -2130,6 +2130,17 @@ from its default.
    mjData.solver_fwdinv[2]. The first value is the relative norm of the discrepancy in joint space, the next is in
    constraint space.
 
+.. _option-flag-invdiscrete:
+
+:at:`invdiscrete`: :at-val:`[disable, enable], "disable"`
+   This flag enables discrete-time inverse dynamics with :ref:`mj_inverse` for all
+   :ref:`integrators<option-integrator>` other than ``RK4``. Recall from the
+   :ref:`numerical integration<geIntegration>` section that the one-step integrators (``Euler``, ``implicit`` and
+   ``implicitfast``), modify the mass matrix :math:`M \rightarrow M-hD`. This implies that finite-differenced
+   accelerations :math:`(v_{t+h} - v_t)/h` will not correspond to the continuous-time acceleration ``mjData.qacc``.
+   When this flag is enabled, :ref:`mj_inverse` will  interpret ``qacc`` as having been computed from the difference of
+   two sequential velocities, and undo the above modification.
+
 .. _option-flag-sensornoise:
 
 :at:`sensornoise`: :at-val:`[disable, enable], "disable"`

doc/XMLschema.rst

@@ -241,9 +241,9 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`sensor<option-flag-sensor>`                               | :ref:`midphase<option-flag-midphase>`                           | :ref:`eulerdamp<option-flag-eulerdamp>`                         | :ref:`override<option-flag-override>`                           | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`energy<option-flag-energy>`                               | :ref:`fwdinv<option-flag-fwdinv>`                               | :ref:`sensornoise<option-flag-sensornoise>`                     | :ref:`multiccd<option-flag-multiccd>`                           | |
+|                                    |    |    | :ref:`energy<option-flag-energy>`                               | :ref:`fwdinv<option-flag-fwdinv>`                               | :ref:`invdiscrete<option-flag-invdiscrete>`                     | :ref:`sensornoise<option-flag-sensornoise>`                     | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`island<option-flag-island>`                               |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`multiccd<option-flag-multiccd>`                           | :ref:`island<option-flag-island>`                               |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | mujoco |br| |L|                    |    | .. table::                                                                                                                                                                                                                                                                   |

doc/changelog.rst

@@ -36,17 +36,19 @@ General
    :ref:`actuatorfrclimited<body-joint-actuatorfrclimited>`, respectively.
 #. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
    Euler integrator. See the :ref:`Numerical Integration<geIntegration>` section for more details.
+#. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
+   :ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
 
 Python bindings
 ^^^^^^^^^^^^^^^
 
-7. Fixed `#870 <https://github.com/deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
+8. Fixed `#870 <https://github.com/deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
    camera name used the default camera.
 
 Bug fixes
 ^^^^^^^^^
 
-#. Fixed a bug that was causing the geom margins to be ignored during the midphase.
+9. Fixed a bug that was causing the geom margins to be ignored during the midphase.
 
 
 Version 2.3.7 (July 20, 2023)

doc/computation.rst

@@ -1607,6 +1607,8 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
 #. Compute sensor data that depends on velocity, and the kinetic energy if enabled.
 #. Compute all passive forces.
 #. Compute the reference constraint acceleration.
+#. If the :ref:`invdiscrete<option-flag-invdiscrete>` flag is set and the :ref:`integrator<option-integrator>` is not
+   ``RK4``, convert input accelerations from discrete to continuous time.
 #. Compute the constraint force. This is done analytically, without using a numerical solver.
 #. Compute the inverse dynamics for the unconstrained system.
 #. Compute sensor data that depends on force and acceleration if enabled.

doc/includes/references.h

@@ -397,12 +397,13 @@ typedef enum mjtEnableBit_ {      // enable optional feature bitflags
   mjENBL_OVERRIDE     = 1<<0,     // override contact parameters
   mjENBL_ENERGY       = 1<<1,     // energy computation
   mjENBL_FWDINV       = 1<<2,     // record solver statistics
-  mjENBL_SENSORNOISE  = 1<<3,     // add noise to sensor data
+  mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
+  mjENBL_SENSORNOISE  = 1<<4,     // add noise to sensor data
                                   // experimental features:
-  mjENBL_MULTICCD     = 1<<4,     // multi-point convex collision detection
+  mjENBL_MULTICCD     = 1<<5,     // multi-point convex collision detection
-  mjENBL_ISLAND       = 1<<5,     // constraint island discovery
+  mjENBL_ISLAND       = 1<<6,     // constraint island discovery
 
-  mjNENABLE           = 6         // number of enable flags
+  mjNENABLE           = 7         // number of enable flags
 } mjtEnableBit;
 typedef enum mjtJoint_ {          // type of degree of freedom
   mjJNT_FREE          = 0,        // global position and orientation (quat)       (7)

include/mujoco/mjmodel.h

@@ -68,12 +68,13 @@ typedef enum mjtEnableBit_ {      // enable optional feature bitflags
   mjENBL_OVERRIDE     = 1<<0,     // override contact parameters
   mjENBL_ENERGY       = 1<<1,     // energy computation
   mjENBL_FWDINV       = 1<<2,     // record solver statistics
-  mjENBL_SENSORNOISE  = 1<<3,     // add noise to sensor data
+  mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
+  mjENBL_SENSORNOISE  = 1<<4,     // add noise to sensor data
                                   // experimental features:
-  mjENBL_MULTICCD     = 1<<4,     // multi-point convex collision detection
+  mjENBL_MULTICCD     = 1<<5,     // multi-point convex collision detection
-  mjENBL_ISLAND       = 1<<5,     // constraint island discovery
+  mjENBL_ISLAND       = 1<<6,     // constraint island discovery
 
-  mjNENABLE           = 6         // number of enable flags
+  mjNENABLE           = 7         // number of enable flags
 } mjtEnableBit;
 
 

introspect/enums.py

@@ -53,10 +53,11 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjENBL_OVERRIDE', 1),
              ('mjENBL_ENERGY', 2),
              ('mjENBL_FWDINV', 4),
-             ('mjENBL_SENSORNOISE', 8),
+             ('mjENBL_INVDISCRETE', 8),
-             ('mjENBL_MULTICCD', 16),
+             ('mjENBL_SENSORNOISE', 16),
-             ('mjENBL_ISLAND', 32),
+             ('mjENBL_MULTICCD', 32),
-             ('mjNENABLE', 6),
+             ('mjENBL_ISLAND', 64),
+             ('mjNENABLE', 7),
          ]),
      )),
     ('mjtJoint',

introspect/enums_test.py

@@ -41,10 +41,11 @@ class EnumsTest(absltest.TestCase):
         tuple(enum_decl.values.items()), (('mjENBL_OVERRIDE', 1<<0),
                                           ('mjENBL_ENERGY', 1<<1),
                                           ('mjENBL_FWDINV', 1<<2),
-                                          ('mjENBL_SENSORNOISE', 1<<3),
+                                          ('mjENBL_INVDISCRETE', 1<<3),
-                                          ('mjENBL_MULTICCD', 1<<4),
+                                          ('mjENBL_SENSORNOISE', 1<<4),
-                                          ('mjENBL_ISLAND', 1<<5),
+                                          ('mjENBL_MULTICCD', 1<<5),
-                                          ('mjNENABLE', 6)))
+                                          ('mjENBL_ISLAND', 1<<6),
+                                          ('mjNENABLE', 7)))
 
   # values mostly increment by one with occasional overrides
   def test_mjtGeom(self):  # pylint: disable=invalid-name

python/mujoco/bindings_test.py

@@ -821,8 +821,8 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1<<0)
     self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1<<1)
     self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1<<2)
-    self.assertEqual(mujoco.mjtEnableBit.mjENBL_SENSORNOISE, 1<<3)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_SENSORNOISE, 1<<4)
-    self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 6)
+    self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 7)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_PLANE, 0)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_HFIELD, 1)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_SPHERE, 2)

src/engine/engine_inverse.c

@@ -22,6 +22,7 @@
 #include "engine/engine_collision_driver.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_smooth.h"
+#include "engine/engine_derivative.h"
 #include "engine/engine_io.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_passive.h"
@@ -88,6 +89,88 @@ void mj_invVelocity(const mjModel* m, mjData* d) {
 
 
 
+// convert discrete-time qacc to continuous-time qacc
+static void mj_discreteAcc(const mjModel* m, mjData* d) {
+  int nv = m->nv, dof_damping;
+  mjtNum *qacc = d->qacc;
+
+  mjMARKSTACK;
+  mjtNum* qfrc = mj_stackAlloc(d, nv);
+
+  // use selected integrator
+  switch ((mjtIntegrator) m->opt.integrator) {
+  case mjINT_RK4:
+    // not supported by RK4
+    return;
+
+  case mjINT_EULER:
+    // check for dof damping if disable flag is not set
+    dof_damping = 0;
+    if (!mjDISABLED(mjDSBL_EULERDAMP)) {
+      for (int i=0; i < nv; i++) {
+        if (m->dof_damping[i] > 0) {
+          dof_damping = 1;
+          break;
+        }
+      }
+    }
+
+    // if disabled or no dof damping, nothing to do
+    if (!dof_damping) {
+      return;
+    }
+
+    // set qfrc = (M + h*diag(B)) * qacc
+    mj_mulM(m, d, qfrc, qacc);
+    for (int i=0; i < nv; i++) {
+      qfrc[i] += m->opt.timestep * m->dof_damping[i] * d->qacc[i];
+    }
+    break;
+
+  case mjINT_IMPLICIT:
+    // compute qDeriv
+    mjd_smooth_vel(m, d, /* flg_bias = */ 1);
+
+    // set qLU = qM
+    mj_copyM2DSparse(m, d, d->qLU, d->qM);
+
+    // set qLU = qM - dt*qDeriv
+    mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD);
+
+    // set qfrc = qLU * qacc
+    mju_mulMatVecSparse(qfrc, d->qLU, qacc, nv,
+                        d->D_rownnz, d->D_rowadr, d->D_colind, /*rowsuper=*/NULL);
+    break;
+
+  case mjINT_IMPLICITFAST:
+    // compute analytical derivative qDeriv; skip rne derivative
+    mjd_smooth_vel(m, d, /* flg_bias = */ 0);
+
+    // save mass matrix
+    mjtNum* qMsave = mj_stackAlloc(d, m->nM);
+    mju_copy(qMsave, d->qM, m->nM);
+
+    // set M = M - dt*qDeriv (reduced to M nonzeros)
+    mjtNum* qDerivReduced = mj_stackAlloc(d, m->nM);
+    mj_copyD2MSparse(m, d, qDerivReduced, d->qDeriv);
+    mju_addToScl(d->qM, qDerivReduced, -m->opt.timestep, m->nM);
+
+    // set qfrc = (M - dt*qDeriv) * qacc
+    mj_mulM(m, d, qfrc, qacc);
+
+    // restore mass matrix
+    mju_copy(d->qM, qMsave, m->nM);
+    break;
+  }
+
+  // solve for qacc: qfrc = M * qacc
+  mj_solveM(m, d, qacc, qfrc, 1);
+
+  mjFREESTACK;
+}
+
+
+
 // inverse constraint solver
 void mj_invConstraint(const mjModel* m, mjData* d) {
   TM_START;
@@ -144,6 +227,10 @@ void mj_inverseSkip(const mjModel* m, mjData* d,
     }
   }
 
+  if (mjENABLED(mjENBL_INVDISCRETE)) {
+    mj_discreteAcc(m, d);
+  }
+
   // acceleration-dependent
   mj_invConstraint(m, d);
   mj_rne(m, d, 1, d->qfrc_inverse);

src/engine/engine_support.c

@@ -67,6 +67,7 @@ const char* mjENABLESTRING[mjNENABLE] = {
   "Energy",
   "Fwdinv",
   "Sensornoise",
+  "InvDiscrete",
   "MultiCCD",
   "Island"
 };
@@ -1122,7 +1123,7 @@ void mj_copyM2DSparse(const mjModel* m, mjData* d, mjtNum* dst, const mjtNum* sr
     }
   }
 
-  mjFREESTACK
+  mjFREESTACK;
 }
 
 

src/xml/xml_native_reader.cc

@@ -103,10 +103,10 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = {
         "solver", "iterations", "noslip_iterations", "mpr_iterations",
         "sdf_iterations", "sdf_initpoints"},
     {"<"},
-        {"flag", "?", "21", "constraint", "equality", "frictionloss", "limit", "contact",
+        {"flag", "?", "22", "constraint", "equality", "frictionloss", "limit", "contact",
             "passive", "gravity", "clampctrl", "warmstart",
             "filterparent", "actuation", "refsafe", "sensor", "midphase", "eulerdamp",
-            "override", "energy", "fwdinv", "sensornoise", "multiccd", "island"},
+            "override", "energy", "fwdinv", "invdiscrete", "sensornoise", "multiccd", "island"},
     {">"},
 
     {"size", "*", "14", "memory", "njmax", "nconmax", "nstack", "nuserdata", "nkey",
@@ -1014,6 +1014,7 @@ void mjXReader::Option(XMLElement* section, mjOption* opt) {
     READENBL("override",    mjENBL_OVERRIDE)
     READENBL("energy",      mjENBL_ENERGY)
     READENBL("fwdinv",      mjENBL_FWDINV)
+    READENBL("invdiscrete", mjENBL_INVDISCRETE)
     READENBL("sensornoise", mjENBL_SENSORNOISE)
     READENBL("multiccd",    mjENBL_MULTICCD)
     READENBL("island",      mjENBL_ISLAND)

src/xml/xml_native_writer.cc

@@ -845,6 +845,7 @@ void mjXWriter::Option(XMLElement* root) {
     WRITEENBL("override",       mjENBL_OVERRIDE)
     WRITEENBL("energy",         mjENBL_ENERGY)
     WRITEENBL("fwdinv",         mjENBL_FWDINV)
+    WRITEENBL("invdiscrete",    mjENBL_INVDISCRETE)
     WRITEENBL("sensornoise",    mjENBL_SENSORNOISE)
     WRITEENBL("multiccd",       mjENBL_MULTICCD)
     WRITEENBL("island",         mjENBL_ISLAND)

test/engine/CMakeLists.txt

@@ -30,6 +30,9 @@ target_link_libraries(engine_derivative_test fixture gmock)
 mujoco_test(engine_forward_test)
 target_link_libraries(engine_forward_test fixture gmock)
 
+mujoco_test(engine_inverse_test)
+target_link_libraries(engine_inverse_test fixture gmock)
+
 mujoco_test(engine_island_test)
 target_link_libraries(engine_island_test fixture gmock)
 

test/engine/engine_inverse_test.cc

@@ -0,0 +1,126 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Tests for engine/engine_inverse.c.
+
+#include "src/engine/engine_inverse.h"
+
+#include <string>
+
+#include <gtest/gtest.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+using InverseTest = MujocoTest;
+
+const int kSteps = 70;
+static const char* const kModelPath = "testdata/model.xml";
+
+// test standard continuous-time inverse dynamics
+TEST_F(InverseTest, ForwardInverseMatch) {
+  const std::string xml_path = GetTestDataFilePath(kModelPath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  mjData* data = mj_makeData(model);
+
+  // simulate, call mj_forward
+  for (int i = 0; i < kSteps; ++i) {
+    mj_step(model, data);
+  }
+  mj_forward(model, data);
+
+  // call built-in testing function
+  mj_compareFwdInv(model, data);
+
+  // expect mismatch to be small
+  mjtNum epsilon = 1e-10;
+  EXPECT_LT(data->solver_fwdinv[0], epsilon);
+  EXPECT_LT(data->solver_fwdinv[1], epsilon);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+// test discrete-time inverse dynamics
+TEST_F(InverseTest, DiscreteInverseMatch) {
+  // load and allocate
+  const std::string xml_path = GetTestDataFilePath(kModelPath);
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0);
+  int nv = model->nv;
+  mjData* data = mj_makeData(model);
+  int nstate = mj_stateSize(model, mjSTATE_INTEGRATION);
+  mjtNum* state = (mjtNum*)mju_malloc(nstate * sizeof(mjtNum));
+  mjtNum* qvel_next = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
+  mjtNum* qacc_fd = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
+
+  for (auto integrator : {mjINT_EULER, mjINT_IMPLICIT, mjINT_IMPLICITFAST}) {
+    model->opt.integrator = integrator;
+    for (bool invdiscrete : {false, true}) {
+      // simulate
+      mj_resetData(model, data);
+      for (int i = 0; i < kSteps; ++i) {
+        mj_step(model, data);
+      }
+
+      // save state
+      mj_getState(model, data, state, mjSTATE_INTEGRATION);
+
+      // call step, save new qvel
+      mj_step(model, data);
+      mju_copy(qvel_next, data->qvel, nv);
+
+      // reset the state, compute discrete-time (finite-differenced) qacc
+      mj_setState(model, data, state, mjSTATE_INTEGRATION);
+      mju_sub(qacc_fd, qvel_next, data->qvel, nv);
+      mju_scl(qacc_fd, qacc_fd, 1/model->opt.timestep, nv);
+
+      // call mj_forward, overwrite qacc with qacc_fd
+      mj_forward(model, data);
+      mju_copy(data->qacc, qacc_fd, nv);
+
+      // set/unset mjENBL_INVDISCRETE flag
+      if (invdiscrete) {
+        model->opt.enableflags |= mjENBL_INVDISCRETE;
+      } else {
+        model->opt.enableflags &= ~mjENBL_INVDISCRETE;
+      }
+
+      // call built-in testing function
+      mj_compareFwdInv(model, data);
+
+      // depending on mjENBL_INVDISCRETE flag, expect mismatch to be small/large
+      if (invdiscrete) {
+        mjtNum epsilon = 1e-10;
+        EXPECT_LT(data->solver_fwdinv[0], epsilon);
+        EXPECT_LT(data->solver_fwdinv[1], epsilon);
+      } else {
+        EXPECT_GT(data->solver_fwdinv[0], 1.0);
+        EXPECT_GT(data->solver_fwdinv[1], 1.0);
+      }
+    }
+  }
+
+  // deallocate
+  mju_free(qacc_fd);
+  mju_free(qvel_next);
+  mju_free(state);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+}  // namespace
+}  // namespace mujoco

unity/Runtime/Bindings/MjBindings.cs

@@ -158,10 +158,11 @@ public enum mjtEnableBit : int{
   mjENBL_OVERRIDE = 1,
   mjENBL_ENERGY = 2,
   mjENBL_FWDINV = 4,
-  mjENBL_SENSORNOISE = 8,
+  mjENBL_INVDISCRETE = 8,
-  mjENBL_MULTICCD = 16,
+  mjENBL_SENSORNOISE = 16,
-  mjENBL_ISLAND = 32,
+  mjENBL_MULTICCD = 32,
-  mjNENABLE = 6,
+  mjENBL_ISLAND = 64,
+  mjNENABLE = 7,
 }
 public enum mjtJoint : int{
   mjJNT_FREE = 0,