Add documentation on pipeline stages and mjData self-consistency. Fixes #1667

PiperOrigin-RevId: 638619787 Change-Id: Ic8dd0fbd178c4946757a07d74fca3b2c905f9080
2024-05-30 05:37:03 -07:00
parent c511d02265
commit 0d1c92b20a
4 changed files with 91 additions and 47 deletions
@@ -57,18 +57,18 @@ Main simulation

 These are the main entry points to the simulator. Most users will only need to call :ref:`mj_step`, which computes
 everything and advanced the simulation state by one time step. Controls and applied forces must either be set in advance
-(in mjData.{ctrl, qfrc_applied, xfrc_applied}), or a control callback :ref:`mjcb_control` must be installed which will
-be called just before the controls and applied forces are needed. Alternatively, one can use :ref:`mj_step1` and
+(in ``mjData.{ctrl, qfrc_applied, xfrc_applied}``), or a control callback :ref:`mjcb_control` must be installed which
+will be called just before the controls and applied forces are needed. Alternatively, one can use :ref:`mj_step1` and
 :ref:`mj_step2` which break down the simulation pipeline into computations that are executed before and after the
 controls are needed; in this way one can set controls that depend on the results from :ref:`mj_step1`. Keep in mind
-though that the RK4 solver does not work with mj_step1/2.
+though that the RK4 solver does not work with mj_step1/2. See :ref:`Pipeline` for a more detailed description.

 mj_forward performs the same computations as :ref:`mj_step` but without the integration. It is useful after loading or
 resetting a model (to put the entire mjData in a valid state), and also for out-of-order computations that involve
 sampling or finite-difference approximations.

-mj_inverse runs the inverse dynamics, and writes its output in ``mjData.qfrc_inverse``. Note that ``mjData.qacc`` must
-be set before calling this function. Given the state (qpos, qvel, act), mj_forward maps from force to acceleration,
+:ref:`mj_inverse` runs the inverse dynamics, and writes its output in ``mjData.qfrc_inverse``. Note that ``mjData.qacc``
+must be set before calling this function. Given the state (qpos, qvel, act), mj_forward maps from force to acceleration,
 while mj_inverse maps from acceleration to force. Mathematically these functions are inverse of each other, but
 numerically this may not always be the case because the forward dynamics rely on a constraint optimization algorithm
 which is usually terminated early. The difference between the results of forward and inverse dynamics can be computed
@@ -77,7 +77,7 @@ general sanity check).

 The skip version of :ref:`mj_forward` and :ref:`mj_inverse` are useful for example when qpos was unchanged but qvel was
 changed (usually in the context of finite differencing). Then there is no point repeating the computations that only
-depend on qpos. Calling the dynamics with skipstage = mjSTAGE_POS will achieve these savings.
+depend on qpos. Calling the dynamics with skipstage = :ref:`mjSTAGE_POS<mjtStage>` will achieve these savings.

 .. _mj_step:

@@ -249,8 +249,8 @@ degrees-of-freedom and a given point. Given a body specified by its integer id (
 frame (``point``) treated as attached to the body, the Jacobian has both translational (``jacp``) and rotational
 (``jacr``) components. Passing ``NULL`` for either pointer will skip that part of the computation. Each component is a
 3-by-nv matrix. Each row of this matrix is the gradient of the corresponding coordinate of the specified point with
-respect to the degrees-of-freedom. The ability to compute end-effector Jacobians efficiently and analytically is one of
-the advantages of working in minimal coordinates.
+respect to the degrees-of-freedom. The :ref:`pipeline stages<piStages>` required for Jacobian computations to be
+consistenst with the current generalized positions ``mjData.qpos`` are :ref:`mj_kinematics` and :ref:`mj_comPos`.

 .. _mj_jacBody:

@@ -33,18 +33,18 @@ The model and all files referenced in it can be loaded from disk or from a VFS w

 These are the main entry points to the simulator. Most users will only need to call :ref:`mj_step`, which computes
 everything and advanced the simulation state by one time step. Controls and applied forces must either be set in advance
-(in mjData.{ctrl, qfrc_applied, xfrc_applied}), or a control callback :ref:`mjcb_control` must be installed which will
-be called just before the controls and applied forces are needed. Alternatively, one can use :ref:`mj_step1` and
+(in ``mjData.{ctrl, qfrc_applied, xfrc_applied}``), or a control callback :ref:`mjcb_control` must be installed which
+will be called just before the controls and applied forces are needed. Alternatively, one can use :ref:`mj_step1` and
 :ref:`mj_step2` which break down the simulation pipeline into computations that are executed before and after the
 controls are needed; in this way one can set controls that depend on the results from :ref:`mj_step1`. Keep in mind
-though that the RK4 solver does not work with mj_step1/2.
+though that the RK4 solver does not work with mj_step1/2. See :ref:`Pipeline` for a more detailed description.

 mj_forward performs the same computations as :ref:`mj_step` but without the integration. It is useful after loading or
 resetting a model (to put the entire mjData in a valid state), and also for out-of-order computations that involve
 sampling or finite-difference approximations.

-mj_inverse runs the inverse dynamics, and writes its output in ``mjData.qfrc_inverse``. Note that ``mjData.qacc`` must
-be set before calling this function. Given the state (qpos, qvel, act), mj_forward maps from force to acceleration,
+:ref:`mj_inverse` runs the inverse dynamics, and writes its output in ``mjData.qfrc_inverse``. Note that ``mjData.qacc``
+must be set before calling this function. Given the state (qpos, qvel, act), mj_forward maps from force to acceleration,
 while mj_inverse maps from acceleration to force. Mathematically these functions are inverse of each other, but
 numerically this may not always be the case because the forward dynamics rely on a constraint optimization algorithm
 which is usually terminated early. The difference between the results of forward and inverse dynamics can be computed
@@ -53,7 +53,7 @@ general sanity check).

 The skip version of :ref:`mj_forward` and :ref:`mj_inverse` are useful for example when qpos was unchanged but qvel was
 changed (usually in the context of finite differencing). Then there is no point repeating the computations that only
-depend on qpos. Calling the dynamics with skipstage = mjSTAGE_POS will achieve these savings.
+depend on qpos. Calling the dynamics with skipstage = :ref:`mjSTAGE_POS<mjtStage>` will achieve these savings.

 .. _Initialization:

@@ -165,8 +165,8 @@ degrees-of-freedom and a given point. Given a body specified by its integer id (
 frame (``point``) treated as attached to the body, the Jacobian has both translational (``jacp``) and rotational
 (``jacr``) components. Passing ``NULL`` for either pointer will skip that part of the computation. Each component is a
 3-by-nv matrix. Each row of this matrix is the gradient of the corresponding coordinate of the specified point with
-respect to the degrees-of-freedom. The ability to compute end-effector Jacobians efficiently and analytically is one of
-the advantages of working in minimal coordinates.
+respect to the degrees-of-freedom. The :ref:`pipeline stages<piStages>` required for Jacobian computations to be
+consistenst with the current generalized positions ``mjData.qpos`` are :ref:`mj_kinematics` and :ref:`mj_comPos`.

 .. _mj_jacBody:

@@ -1575,11 +1575,23 @@ Top level
  pipeline does not support the Runge Kutta integrator.
 - :ref:`mj_fwdPosition` invokes stages **2-11**, the position-dependent part of the pipeline.

+.. _piStages:
+
 Stages
 ^^^^^^
+Below we describe the pipeline stages and API functions corresponding to each stage. All functions write their outputs
+into attributes of :ref:`mjData`. It is informative to compare the list below with the :ref:`mjData` struct definition,
+wherein comments above blocks of attributes specify the function that computes them. Note that each stage depends on
+the values computed in some or all of the previous stages.
+
 1. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence
   is detected, the state is automatically reset and the corresponding warning is raised:
   :ref:`mj_checkPos`, :ref:`mj_checkVel`
+
+Position
+''''''''
+The stages below compute quantities that depend on the generalized positions ``mjData.qpos``.
+
 2. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites,
   cameras and lights. It also normalizes all quaternions: :ref:`mj_kinematics`, :ref:`mj_camLight`
 3. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding
@@ -1596,6 +1608,12 @@ Stages
 11. Compute the matrices and vectors needed by the constraint solvers: :ref:`mj_projectConstraint`
 12. Compute sensor data that only depends on position, and the potential energy if enabled: :ref:`mj_sensorPos`,
    :ref:`mj_energyPos`
+
+Velocity
+''''''''
+The stages below compute quantities that depend on the generalized velocity ``mjData.qvel``. Due to the sequential
+dependence structure of the pipeline, the actual dependence is on both ``qpos`` and ``qvel``.
+
 13. Compute the tendon, flex edge and actuator velocities: :ref:`mj_fwdVelocity`
 14. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at
    the subtree centers of mass: :ref:`mj_comVel`
@@ -1604,6 +1622,12 @@ Stages
    (if required by sensors, call :ref:`mj_subtreeVel`): :ref:`mj_sensorVel`
 17. Compute the reference constraint acceleration: :ref:`mj_referenceConstraint`
 18. Compute the vector of Coriolis, centrifugal and gravitational forces: :ref:`mj_rne`
+
+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 joint acceleration resulting from all forces except for the (still unknown) constraint forces:
    :ref:`mj_fwdAcceleration`
@@ -1620,12 +1644,53 @@ Stages
    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`

+.. _piConsistency:
+
+Consistency in ``mjData``
+~~~~~~~~~~~~~~~~~~~~~~~~~
+The MuJoCo computation pipeline is entirely imperative, nothing happens automatically. This leads to behavior which
+can seem unexpected to users more familiar with other paradigms. Here are two examples of intended behaviors which can
+sometimes be surprising:
+
+- After setting :ref:`the state<geState>`, state-derived quantities do not automatically correspond to the new state.
+  The required stage or stages must be manually invoked. For example after setting the generalized positions
+  ``mjData.qpos``, Cartesian positions and orientations will not be consistent with ``qpos`` without first calling
+  :ref:`mj_kinematics`.
+- After an :ref:`mj_step`, which terminates immediately after updating the state, quantities in ``mjData`` correspond
+  to the *previous* state (or more precisely, the *transition* between the previous and current state).
+  In particular, all position-dependent sensor values and position-dependent computations like kinematic
+  :ref:`Jacobians<mj_jac>`, will be with respect to the *previous positions*.
+
+
+
+.. _piReproducibility:
+
+Reproducibility
+~~~~~~~~~~~~~~~
+
+MuJoCo's simulation pipeline is entirely deterministic and reproducible -- if a :ref:`state<geState>` in a trajectory is
+saved and reloaded and :ref:`mj_step` called again, the resulting next state will be identical. However, there are some
+important caveats:
+
+- Save all the required :ref:`integration state<geIntegrationState>` components. In particular :ref:`warmstart
+  accelerations<geWarmstart>` have only a very small effect on the next state, but should be saved if bit-wise equality
+  is required.
+- Any numerical difference between states, no matter how small, will become significant upon integration, especially for
+  systems with contact. Contact events have high `Lyapunov exponents
+  <https://en.wikipedia.org/wiki/Lyapunov_exponent>`__; this is a property of any rigid-body simulator (and indeed of
+  `real-world physics <https://en.wikipedia.org/wiki/Roulette>`__) and is not MuJoCo-specific.
+- Exact reproducibility is only guaranteed within a **single version**, on the **same architecture**. Small numerical
+  differences are quite common between versioned releases, for example due to code optimizations. This means that when
+  saving an initial state and an open-loop control sequence, the resulting rolled-out trajectory will be identical
+  within the same version, but will likely be different between MuJoCo versions or different operating systems.
+
 .. _piInverse:

 Inverse dynamics
 ~~~~~~~~~~~~~~~~

-The top-level function :ref:`mj_inverse` invokes the following sequence of computations.
+The top-level function :ref:`mj_inverse` invokes the following sequence of computations. The notes above regarding
+:ref:`consistency<piConsistency>` and :ref:`reproducibility<piReproducibility>` apply here as well.

 #. Compute the forward kinematics.
 #. Compute the body inertias and joint axes.
@@ -1652,27 +1717,6 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
   equals the sum of external and actuation forces.


-.. _piReproducibility:
-
-Reproducibility
-~~~~~~~~~~~~~~~
-
-MuJoCo's simulation pipeline is entirely deterministic and reproducible -- if a :ref:`state<geState>` in a trajectory is
-saved and reloaded and :ref:`mj_step` called again, the resulting next state will be identical. However, there are some
-important caveats:
-
- Save all the required :ref:`integration state<geIntegrationState>` components. In particular :ref:`warmstart
-  accelerations<geWarmstart>` have only a very small effect on the next state, but should be saved if bit-wise equality
-  is required.
- Any numerical difference between states, no matter how small, will become significant upon integration, especially for
-  systems with contact. Contact events have high `Lyapunov exponents
-  <https://en.wikipedia.org/wiki/Lyapunov_exponent>`__; this is a property of any rigid-body simulator (and indeed of
-  `real-world physics <https://en.wikipedia.org/wiki/Roulette>`__) and is not MuJoCo-specific.
- Exact reproducibility is only guaranteed within a **single version**, on the **same architecture**. Small numerical
-  differences are quite common between versioned releases, for example due to code optimizations. This means that when
-  saving an initial state and an open-loop control sequence, the resulting rolled-out trajectory will be identical
-  within the same version, but will likely be different between MuJoCo versions or different operating systems.
-
 .. _derivatives:

 Derivatives
@@ -150,10 +150,10 @@ continuous time, and then integrate over a time period specified by ``mjModel.op
 the acceleration ``mjData.qacc`` at time ``mjData.time``, given the :ref:`state and control <siStateControl>` at time
 ``mjData.time``. The numerical integrator then advances the state and time to ``mjData.time + mjModel.opt.timestep``.
 Now, the control is required to be a function of the state at time ``mjData.time``. However a general feedback
-controller can be a very complex function, depending on various features of the state - in particular all the features
+controller can be a very complex function, depending on various features of the state -- in particular all the features
 computed by MuJoCo as intermediate results of the simulation. These may include contacts, Jacobians, passive forces.
-None of these quantities are available before ``mj_step`` is called (or rather, they are available but outdated by one
-time step). In contrast, when ``mj_step`` calls our control callback, it does so as late in the computation as possible
+None of these quantities are available before ``mj_step`` is called (or rather, they are available but *outdated by one
+time step*). In contrast, when ``mj_step`` calls our control callback, it does so as late in the computation as possible
 - namely after all the intermediate results dependent on the state but not on the control have been computed.

 The same effect can be achieved without using a control callback. This is done by breaking ``mj_step`` in two parts:
@@ -253,7 +253,7 @@ notion of state of a dynamical system. Dynamical systems are usually described i
 where ``t`` is the time, ``x`` is the state vector, ``u`` is the control vector, and ``f`` is the function that
 computes the time-derivative of the state. This is a continuous-time formulation, and indeed the physics model
 simulated by MuJoCo is defined in continuous time. Even though the numerical integrator operates in discrete time, the
-main part of the computation - namely the function :ref:`mj_forward` - corresponds to the
+main part of the computation---namely the function :ref:`mj_forward`---corresponds to the
 continuous-time dynamics function ``f(t,x,u)`` above. Here we explain this correspondence.

 The state vector in MuJoCo is:
@@ -426,7 +426,7 @@ the first and second sections (position and velocity dependent) have already bee
 When can we use the above machinery and skip some of the computations? In a regular simulation this is not possible.
 However, MuJoCo is designed not only for simulation but also for more advanced applications such as model-based
 optimization, machine learning etc. In such settings one often needs to sample the dynamics at a cloud of nearby
-states, or approximate derivatives via finite differences - which is another form of sampling. If the samples are
+states, or approximate derivatives via finite differences -- which is another form of sampling. If the samples are
 arranged on a grid, where only the position or only the velocity or only the control is different from the center
 point, then the above mechanism can improve performance by about a factor of 2.

@@ -591,7 +591,7 @@ The only way to guarantee that all changes are saved is to save the model as a b
 changes need to be made programmatically, as in system identification for example, and this can only be done with the
 compiled model. So in summary, we have reasonable but not perfect mechanisms for saving model changes. The reason for
 this lack of perfection is that we are working with a compiled model, so this is like changing a binary executable and
-asking a "decompiler" to make corresponding changes to the C code - it is just not possible in general.
+asking a "decompiler" to make corresponding changes to the C code -- it is just not possible in general.

 .. _siLayout:

@@ -642,7 +642,7 @@ spread out the data so that we do not have to perform many memory moves when ins
 layout "uncompressed". It is still a valid layout, but instead of A_rowadr[r] = A_rowadr[r-1] + A_rownnz[r] which is
 the standard convention, we set A_rowadr[r] = r*n. MuJoCo uses sparse matrices internally

-To represent 3D orientations and rotations, MuJoCo uses unit quaternions - namely 4D unit vectors arranged as q = (w,
+To represent 3D orientations and rotations, MuJoCo uses unit quaternions -- namely 4D unit vectors arranged as q = (w,
 x, y, z). Here (x, y, z) is the rotation axis unit vector scaled by sin(a/2), where a is the rotation angle in
 radians, and w = cos(a/2). Thus the quaternion corresponding to a null rotation is (1, 0, 0, 0). This is the default
 setting of all quaternions in MJCF.
@@ -841,7 +841,7 @@ Jacobians
 The derivative of any vector function with respect to its vector argument is called Jacobian. When this term is used
 in multi-joint kinematics and dynamics, it refers to the derivative of some spatial quantity as a function of the
 system configuration. In that case the Jacobian is also a linear map that operates on vectors in the (co)tangent space
-to the configuration manifold - such as velocities, momenta, accelerations, forces. One caveat here is that the system
+to the configuration manifold -- such as velocities, momenta, accelerations, forces. One caveat here is that the system
 configuration encoded in ``mjData.qpos`` has dimensionality ``mjModel.nq``, while the tangent space has dimensionality
 ``mjModel.nv``, and the latter is smaller when quaternion joints are present. So the size of the Jacobian matrix is
 N-by-``mjModel.nv`` where N is the dimensionality of the spatial quantity being differentiated.
@@ -861,7 +861,7 @@ Jacobians, which tell us how a spatial frame anchored at the given point will tr
 change to the kinematic configuration. More precisely, the Jacobian maps joint velocities to end-effector velocities,
 while the transpose of the Jacobian maps end-effector forces to joint forces. There are also several other
 ``mj_jacXXX`` functions; these are convenience functions that call the main ``mj_jac`` function with different points
-of interest - such as a body center of mass, geom center etc.
+of interest -- such as a body center of mass, geom center etc.

 The ability to compute end-effector Jacobians exactly and efficiently is a key advantage of working in joint
 coordinates. Such Jacobians are the foundation of many control schemes that map end-effector errors to actuator