Add mj_fwdKinematics and a table improving documentation of the MuJoCo computation pipeline structure.

PiperOrigin-RevId: 838768130
Change-Id: I1a22d5cdca7db704ba4cea1667fcbecb16b87005

doc/APIreference/functions.rst

@@ -715,6 +715,15 @@ Components
 These are components of the simulation pipeline, called internally from :ref:`mj_step`, :ref:`mj_forward` and
 :ref:`mj_inverse`. It is unlikely that the user will need to call them.
 
+.. _mj_fwdKinematics:
+
+`mj_fwdKinematics <#mj_fwdKinematics>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_fwdKinematics
+
+Run all kinematics-like computations (kinematics, comPos, camlight, flex, tendon).
+
 .. _mj_fwdPosition:
 
 `mj_fwdPosition <#mj_fwdPosition>`__

doc/changelog.rst

@@ -18,6 +18,8 @@ General
   flex object is similar to the "trilinear" option, but it includes curved deformations.
 - Raise an error if there are name collisions also during parsing.
 - Increase Windows stack size to 16MB to enable models with deep nested body hierarchies.
+- Added a new pipeline component function :ref:`mj_fwdKinematics` that combines all kinematics-like sub-components.
+  Relatedly, added a clarifying table at the top of the :ref:`Simulation Pipeline<Pipeline>` chapter.
 - Added a new :ref:`mj_extractState` function that allows a subset of a state that was previously returned by
   :ref:`mj_getState` to be extracted without having to be written back into ``mjData`` first.
 - Added a new :ref:`mj_copyState` function that copies state components from one ``mjData`` to another.

doc/computation/index.rst

@@ -1741,12 +1741,32 @@ Top level
 - The top-level function :ref:`mj_step` invokes the entire sequence of computations below.
 - :ref:`mj_forward` invokes only stages **2-22**, computing the continuous-time forward dynamics, ending with the
   acceleration ``mjData.qacc``.
-- :ref:`mj_step1` invokes stages **1-18** and :ref:`mj_step2` invokes stages **19-25**, breaking :ref:`mj_step` into two
+- :ref:`mj_step1` invokes stages **1-19** and :ref:`mj_step2` invokes stages **20-26**, breaking :ref:`mj_step` into two
   distinct phases. This allows the user to write controllers that depend on quantities derived from the positions and
   velocities (but not forces, since those have not yet been computed). Note that the :ref:`mj_step1` → :ref:`mj_step2`
   pipeline does not support the Runge Kutta integrator.
 - :ref:`mj_fwdPosition` invokes stages **2-11**, the position-dependent part of the pipeline.
 
+.. the table below was created and is editable in tablesgenerator.com
+
+.. table:: Breakdown of the forward dynamic pipeline
+   :class: small-centered no-stripes
+
+   +-------------------------+----------------------------------------------------------------------------------------------------------------------+
+   | top-level functions     |                                                    :ref:`mj_step`                                                    |
+   |                         +------------------------------------------------------------------------------------+---------------------------------+
+   |                         |                                   :ref:`mj_step1`                                  |         :ref:`mj_step2`         |
+   |                         +------------------------------------------------------------------------------------+-------------+-------+-----------+
+   |                         |                                      :ref:`mj_forward`                                           |       |           |
+   +-------------------------+---+----------------------------------------------+----+-----------------------+----+-------------+-------+-----------+
+   | component / description |   |             :ref:`mj_fwdPosition`            |    | :ref:`mj_fwdVelocity` |    | force / acc |       | integrate |
+   |                         +---+-------------------------+--------------------+----+-----------------------+----+-------------+-------+-----------+
+   |                         |   | :ref:`mj_fwdKinematics` | inertia, collision |    |                       |    |             |       |           |
+   +-------------------------+---+-------------------------+--------------------+----+-----------------------+----+-------------+-------+-----------+
+   | stage                   | 1 |         2,3,4,5         |    6,7,8,9,10,11   | 12 |   13,14,15,16,17,18   | 19 | 20,21,22,23 | 24,25 |     26    |
+   +-------------------------+---+-------------------------+--------------------+----+-----------------------+----+-------------+-------+-----------+
+
+
 .. _piStages:
 
 Stages
@@ -1796,24 +1816,28 @@ dependence structure of the pipeline, the actual dependence is on both ``qpos``
 17. Compute the reference constraint acceleration: :ref:`mj_referenceConstraint`
 18. Compute the vector of Coriolis, centrifugal and gravitational forces: :ref:`mj_rne`
 
+Control callback
+''''''''''''''''
+19. Invoke the user-defined control callback if defined: :ref:`mjcb_control`
+
 Force/acceleration
 ''''''''''''''''''
 The stages below compute quantities that depend on :ref:`user inputs<geInput>`. Due to the sequential nature
 of the pipeline, the actual dependence is on the entire :ref:`integration state<geIntegrationState>`.
 
-19. Compute the actuator forces and activation dynamics if defined: :ref:`mj_fwdActuation`
+20. Compute the actuator forces and activation dynamics if defined: :ref:`mj_fwdActuation`
-20. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces:
+21. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces:
     :ref:`mj_fwdAcceleration`
-21. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
+22. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
     constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics:
     :ref:`mj_fwdConstraint`
-22. Compute sensor data that depends on force and acceleration if enabled
+23. Compute sensor data that depends on force and acceleration if enabled
     (if required by sensors, call :ref:`mj_rnePostConstraint`): :ref:`mj_sensorAcc`
-23. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
+24. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
     automatically reset and the corresponding warning is raised: :ref:`mj_checkAcc`
-24. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
+25. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
     dynamics. This is an optional step, and is performed only when enabled: :ref:`mj_compareFwdInv`
-25. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator
+26. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator
     repeats the above sequence three more times, except for the optional computations which are performed only once:
     one of :ref:`mj_Euler`, :ref:`mj_RungeKutta`, :ref:`mj_implicit`
 

doc/includes/references.h

@@ -3132,6 +3132,7 @@ int mj_printSchema(const char* filename, char* buffer, int buffer_sz,
 void mj_printScene(const mjvScene* s, const char* filename);
 void mj_printFormattedScene(const mjvScene* s, const char* filename,
                             const char* float_format);
+void mj_fwdKinematics(const mjModel* m, mjData* d);
 void mj_fwdPosition(const mjModel* m, mjData* d);
 void mj_fwdVelocity(const mjModel* m, mjData* d);
 void mj_fwdActuation(const mjModel* m, mjData* d);

include/mujoco/mujoco.h

@@ -328,6 +328,9 @@ MJAPI void mj_printFormattedScene(const mjvScene* s, const char* filename,
 
 //---------------------------------- Components ----------------------------------------------------
 
+// Run all kinematics-like computations (kinematics, comPos, camlight, flex, tendon).
+MJAPI void mj_fwdKinematics(const mjModel* m, mjData* d);
+
 // Run position-dependent computations.
 MJAPI void mj_fwdPosition(const mjModel* m, mjData* d);
 

python/mujoco/functions.cc

@@ -214,6 +214,7 @@ PYBIND11_MODULE(_functions, pymodule) {
       });
 
   // Components
+  Def<traits::mj_fwdKinematics>(pymodule);
   Def<traits::mj_fwdPosition>(pymodule);
   Def<traits::mj_fwdVelocity>(pymodule);
   Def<traits::mj_fwdActuation>(pymodule);

python/mujoco/introspect/functions.py

@@ -1498,6 +1498,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Print scene to text file, specifying format. float_format must be a valid printf-style format string for a single float value.',  # pylint: disable=line-too-long
      )),
+    ('mj_fwdKinematics',
+     FunctionDecl(
+         name='mj_fwdKinematics',
+         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'),
+                 ),
+             ),
+         ),
+         doc='Run all kinematics-like computations (kinematics, comPos, camlight, flex, tendon).',  # pylint: disable=line-too-long
+     )),
     ('mj_fwdPosition',
      FunctionDecl(
          name='mj_fwdPosition',

src/engine/engine_forward.c

@@ -136,12 +136,8 @@ void* mj_collisionThreaded(void* args) {
   return NULL;
 }
 
-
+// kinematics-related computations
-// position-dependent computations
+void mj_fwdKinematics(const mjModel* m, mjData* d) {
-void mj_fwdPosition(const mjModel* m, mjData* d) {
-  TM_START1;
-
-  TM_START;
   mj_kinematics(m, d);
   mj_comPos(m, d);
   mj_camlight(m, d);
@@ -150,6 +146,14 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
   if (mj_wakeTendon(m, d)) {
     mj_updateSleep(m, d);
   }
+}
+
+// position-dependent computations
+void mj_fwdPosition(const mjModel* m, mjData* d) {
+  TM_START1;
+
+  TM_START;
+  mj_fwdKinematics(m, d);
 
   TM_END(mjTIMER_POS_KINEMATICS);
 

src/engine/engine_forward.h

@@ -67,6 +67,9 @@ MJAPI void mj_implicitSkip(const mjModel *m, mjData *d, int skipfactor);
 
 //-------------------------------- solver components -----------------------------------------------
 
+// all kinematics-like computations
+MJAPI void mj_fwdKinematics(const mjModel* m, mjData* d);
+
 // computations that depend only on qpos
 MJAPI void mj_fwdPosition(const mjModel* m, mjData* d);
 

unity/Runtime/Bindings/MjBindings.cs

@@ -6522,6 +6522,9 @@ public static unsafe extern void mj_printScene(mjvScene_* s, [MarshalAs(Unmanage
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_printFormattedScene(mjvScene_* s, [MarshalAs(UnmanagedType.LPStr)]string filename, [MarshalAs(UnmanagedType.LPStr)]string float_format);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_fwdKinematics(mjModel_* m, mjData_* d);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_fwdPosition(mjModel_* m, mjData_* d);
 

wasm/codegen/generated/bindings.cc

@@ -8021,6 +8021,10 @@ void mj_fwdConstraint_wrapper(const MjModel& m, MjData& d) {
   mj_fwdConstraint(m.get(), d.get());
 }
 
+void mj_fwdKinematics_wrapper(const MjModel& m, MjData& d) {
+  mj_fwdKinematics(m.get(), d.get());
+}
+
 void mj_fwdPosition_wrapper(const MjModel& m, MjData& d) {
   mj_fwdPosition(m.get(), d.get());
 }
@@ -12325,6 +12329,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   function("mj_fwdAcceleration", &mj_fwdAcceleration_wrapper);
   function("mj_fwdActuation", &mj_fwdActuation_wrapper);
   function("mj_fwdConstraint", &mj_fwdConstraint_wrapper);
+  function("mj_fwdKinematics", &mj_fwdKinematics_wrapper);
   function("mj_fwdPosition", &mj_fwdPosition_wrapper);
   function("mj_fwdVelocity", &mj_fwdVelocity_wrapper);
   function("mj_geomDistance", &mj_geomDistance_wrapper);