Add mjENBL_DIAGEXACT for exact constraint diagonal. Fixes #2472

PiperOrigin-RevId: 916932908
Change-Id: Id23ac39b5cd996afc52990719a4e07c0cc7de600

doc/XMLreference.rst

@@ -697,6 +697,19 @@ from its default.
       order for the :ref:`sleep-init<body-sleep>` policy to take effect. Second, it must be set in order for static
       quantities to be computed. See :ref:`implementation notes<siSleep>` for more details.
 
+.. _option-flag-diagexact:
+
+:at:`diagexact`: :at-val:`[disable, enable], "disable"`
+   This flag enables computation of the exact diagonal of the constraint-space inertia matrix :math:`A = J M^{-1} J^T`,
+   replacing the body-based approximation normally used. The exact diagonal is computed from the whitened Jacobian
+   :math:`Y = J M^{-1/2}` as :math:`A_{ii} = \|Y_i\|^2`. This provides a more accurate
+   :ref:`impedance<soParameters>` computation, which can improve solver quality for models with complex kinematic
+   coupling. See :ref:`Diagonal approximation <soExactDiag>` for details on the approximation errors that this flag
+   eliminates. The cost is one back-substitution with the Cholesky factor of the mass matrix per active constraint row;
+   if dual solvers are used (:ref:`PGS<option-solver>` or :ref:`NoSlip<option-noslip_iterations>`), the cost is
+   negligible since :math:`Y` is computed anyway. Consider enabling this flag when observing divergence or poor
+   constraint quality, particularly in models with highly anisotropic body inertias or bodies operating far from the
+   initial configuration ``qpos0``.
 
 .. _compiler:
 

doc/XMLschema.rst

@@ -175,6 +175,9 @@
             .. grid-item::
                :ref:`sleep<option-flag-sleep>`
 
+            .. grid-item::
+               :ref:`diagexact<option-flag-diagexact>`
+
 
    .. dropdown:: :ref:`compiler<compiler>` |*|
 

doc/changelog.rst

@@ -7,6 +7,12 @@ Upcoming version (not yet released)
 
 General
 ^^^^^^^
+- Added ``mjData.efc_Y``, the whitened constraint Jacobian :math:`Y = J M^{-1/2}`, allocated in the arena when
+  dual solvers (PGS or NoSlip) are used or when :ref:`diagexact<option-flag-diagexact>` is enabled.
+- Added the :ref:`diagexact<option-flag-diagexact>` enable flag, which computes the exact diagonal of the
+  constraint-space inertia matrix at the current configuration, replacing the default compile-time approximation.
+  This improves solver quality for models with anisotropic inertias or complex kinematic coupling. See
+  :ref:`Exact diagonal <soExactDiag>` for details.
 - The pseudo-random constraint visitation order in the :ref:`PGS solver<soAlgorithms>`, introduced in the previous
   release, now uses a fixed seed. The previous implementation seeded with ``mjData.time``, which introduced subtle yet
   undesirable time dependence.

doc/computation/index.rst

@@ -1514,6 +1514,20 @@ constraint would satisfy
 and so we would achieve the desired interpolation effect. This of course does not hold exactly in general, but the goal
 here is to construct a sensible and intuitive parameterization of the constraint model and get the scaling right.
 
+.. _soExactDiag:
+
+**Diagonal approximation:** The approximation has three sources of error: (i) it is frozen at ``qpos0`` rather than
+evaluated at the current configuration; (ii) it averages the directional inverse inertia into a scalar, assuming
+isotropy; and (iii) it treats the contributions of different bodies as independent, ignoring kinematic coupling through
+shared DOFs. These errors are usually modest, but can become significant for models with highly anisotropic inertias or
+long kinematic chains that operate far from ``qpos0``. In severe cases — particularly when the averaged inertia becomes
+near-zero despite finite directional inertia — the regularizer :math:`R` becomes near-zero, making constraints
+infinitely hard and causing divergence. The :ref:`diagexact<option-flag-diagexact>` flag replaces the approximation with
+the exact diagonal :math:`A_{ii} = \|Y_i\|^2`, where :math:`Y = J M^{-1/2}` is the whitened Jacobian, computed at the
+current configuration. This eliminates all three sources of error at a modest runtime cost: computing :math:`Y` requires
+a back-substitution with the Cholesky factor of the mass matrix for each active constraint row; if
+:ref:`dual solvers<soAlgorithms>` are used (PGS or NoSlip), the cost is negligible since :math:`Y` is computed anyway.
+
 Next we explain how the reference acceleration is computed. As already mentioned, we use a spring-damper model
 parameterized by *damping* and *stiffness* coefficients element-wise:
 

doc/includes/references.h

@@ -508,10 +508,10 @@ typedef enum mjtEnableBit_ {      // enable optional feature bitflags
   mjENBL_ENERGY       = 1<<1,     // energy computation
   mjENBL_FWDINV       = 1<<2,     // record solver statistics
   mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
-                                  // experimental features:
   mjENBL_SLEEP        = 1<<4,     // sleeping
+  mjENBL_DIAGEXACT    = 1<<5,     // exact diagonal of constraint inertia
 
-  mjNENABLE           = 5         // number of enable flags
+  mjNENABLE           = 6         // number of enable flags
 } mjtEnableBit;
 typedef enum mjtJoint_ {          // type of degree of freedom
   mjJNT_FREE          = 0,        // global position and orientation (quat)       (7)

doc/modeling.rst

@@ -1855,6 +1855,25 @@ in a visible way, and the energy fluctuates around the initial value instead of
    </worldbody>
 
 
+.. _CConstraintImpedance:
+
+Constraint accuracy
+~~~~~~~~~~~~~~~~~~~
+MuJoCo's :ref:`constraint impedance<soParameters>` computation relies on an approximate diagonal of the constraint-space
+inertia matrix, computed once at compile time from the initial configuration ``qpos0``.
+
+In the vast majority of models this approximation is entirely adequate. However, in certain situations—such as models
+with highly anisotropic inertias, complex kinematic chains, or bodies operating far from ``qpos0``—the approximation
+may become inaccurate. This can occasionally manifest as unexplained solver divergence (``badqacc`` warnings),
+excessive penetration, unrealistic slip, or poor solver convergence. A useful diagnostic is the
+:ref:`fwdinv<option-flag-fwdinv>` flag: if the forward-inverse discrepancy is large, inaccurate constraint scaling may
+be a contributing factor.
+
+If you suspect that the compile-time approximation is insufficient for your model, you can enable the
+:ref:`diagexact<option-flag-diagexact>` flag to compute the exact diagonal at runtime. See :ref:`Diagonal approximation
+<soExactDiag>` for details on the underlying mechanics and performance implications.
+
+
 .. |image3| image:: images/modeling/tendonwraps.png
    :width: 500px
 .. |image4| image:: images/modeling/particle.png