Make actuator kp implicit (∂τ/∂q) to lift the stiffness cap

[qiayuanl]

The feature, motivation and pitch

I propose linearizing actuator forces in position as well as velocity inside the implicit integrators, so that kp receives the treatment kv already gets. This removes the timestep from the definition of a valid gain and replaces a hard divergence boundary with a bounded, documented approximation. A working implementation exists (C, plus a MuJoCo Warp prototype); the questions at the end are the design decisions I would like input on.

Problem statement

Notation: throughout, $f_n = \omega_n/2\pi$ with $\omega_n = \sqrt{k_p/I}$ is the undamped natural frequency of the closed position loop, where $I$ is the reflected inertia at the joint. It is determined by kp and $I$ alone; it is not the control rate of the motor driver, and not a −3 dB bandwidth.

Stiff robot joints close their position loop on the motor driver at 1–2 kHz and commonly reach $f_n$ of 100 Hz or more. MuJoCo cannot represent such a joint at reinforcement-learning timesteps. Because qDeriv contains only $\partial\tau/\partial v$, the position gain is integrated explicitly and the closed loop is subject to the stability bound

$$\omega_n < \frac{2\zeta + 2\sqrt{\zeta^2 + 1}}{h}$$

(The circles are bisected against mjWARN_BADQACC on a one-hinge model and match the formula to four digits.)

At $h = 5$ ms, a servo with $f_n = 100$ Hz at the damping a real motor has already exceeds the bound, and $f_n = 160$ Hz exceeds it at $h = 2$ ms. The available workarounds are to add damping the motor does not have, to add armature the rotor does not have, or to reduce the timestep at the cost of training throughput. Each modifies the model to accommodate the integrator, and afterwards the kp in the model file no longer corresponds to the gain deployed on the robot.

A further consequence is that admissibility depends on $h$, so whether a gain is valid is not a property of the model: raising the timestep from 2 ms to 5 ms shrinks the admissible set by a factor of 6.25, and the failure mode is NaN rather than a warning. The same bound applies to jnt_stiffness.

Proposed discretization

Backward Euler in both position and velocity, with $q^+ = q_0 + h v^+$, gives

$$\left(M - h\frac{\partial\tau}{\partial v} - h^2\frac{\partial\tau}{\partial q}\right)\ddot{q} = \tau_0 + h\frac{\partial\tau}{\partial q}v_0$$

The scheme has two parts: a matrix term $-h^2 \partial\tau/\partial q$ added to $\hat{M}$, and a force term $h(\partial\tau/\partial q)v_0$ added to the right-hand side. Both are required; the matrix term alone is a different, non-dissipative scheme. The right-hand side carries a single power of $h$ because the equation is solved for $\ddot{q}$: dividing the momentum-form update by $h$ leaves the matrix unchanged and removes one factor of $h$ from the shift.

For a position-like actuator, $\partial\tau/\partial q = b_1 m^\mathsf{T} m$, where $m$ is the transmission moment and $b_1$ is biasprm[1] $= -k_p$. The matrix addition $h^2 k_p m^\mathsf{T} m$ is therefore positive semidefinite for $k_p &gt; 0$ and $\hat{M}$ remains positive definite. The term has exactly the sparsity of $M$ and accumulates through the moment CSR that mjd_actuator_vel already uses (engine_derivative.c:2182), so it introduces no new fill-in and no additional factorization.

This is the actuator instance of the discretization that shipped for flex in 3.11.0 (ea230a95, #3386): mjd_effBuild assembles $\tilde{M} = M + (h^2 + hd)K$ and mjd_effShift computes $c = hK\dot{q}$. The body of #3386 lists "the general effective-metric mode (all solvers, all PSD-safe force classes, behind an enable flag)" as follow-up work, and the present proposal is one named force class from that list. It is also the discretization PhysX documents for its PGS force drives.

Measured behavior

The one-step update matches the analytic backward-Euler solution

$$\ddot{q} = \frac{k_p(q^* - q_0) - (k_v + h k_p)v_0}{I + h k_v + h^2 k_p}$$

to between 0 and 3.4e-16 for kp from 1e2 to 1e8, so the discretization itself is exact.

What it delivers above the old cap is bounded. Delivered closed-loop poles at $h = 5$ ms, $I = 0.01576$, $k_v = 1$:

The delivered damped frequency saturates at $\pi/(2h)$, half the Nyquist rate: 50 Hz at $h = 5$ ms and 125 Hz at $h = 2$ ms, measured at 0.99993 of $\pi/(2h)$ for $k_p = 10^{12}$. Beyond that point, raising kp drives $\zeta \to 1$ and increases the decay rate roughly as $\ln(h^2 k_p/M)/(2h)$ while the oscillation frequency stays fixed. The behavior above the cap therefore degrades gracefully, toward a correctly damped and faster-settling but no faster-oscillating response, instead of diverging. It is not high fidelity.

Step response at $h = 5$ ms against the exact continuous PD response (matrix exponential; RK4 at 1e-6 agrees to 2.5e-13):

Below the cap the flag is less accurate than the current integrator. At $h\omega_n = 1$ and $\zeta = 0.1$, stock MuJoCo delivers $\omega_n$ to −0.5% with $\zeta = 0.092$; the flag gives −20% and $\zeta = 0.379$. The current treatment is nearly exact wherever it is stable, so the flag is for the regime where the alternative is divergence, not a general accuracy improvement.

Interaction with constraints

On model/humanoid/humanoid.xml with position servos on all 21 actuators at $h = 5$ ms, the flag is stable to $k_p = 10^{7}$ with contact disabled but diverges at $k_p = 10^{5}$ with contact enabled, three decades lower. This is limitation 1 below acting as a hard divergence rather than as increased penetration, and it currently prevents use on contact-rich legged models at the gains the flag targets. Question 1 is therefore the substantive one.

The equivalence to a userland gain fold holds only while no constraint is active. Once a contact or joint limit engages, the fold changes qfrc_actuator and the solver is handed a different force. That is the reason to implement the scheme in the engine rather than in user code: the fold obtains stability by misreporting the applied force.

Limitations

1. Constraint forces are computed against the unmodified $M$. The constraint solve precedes the implicit update, so impulses are sized for an effective inertia smaller than the one the drive presents. This is pre-existing for $h k_v$ and amplified here.
2. The scheme changes the damping, by $\Delta\zeta \approx (h/2)\sqrt{k_p/m}$ for small $h\omega_n$ (the exact value at $h\omega_n = 2$ is 0.588, not 1). The sign is relevant: the existing implicit kv treatment under-damps by $-\zeta^2 W^2$, and the added term partly cancels that error. The weight that would make the damping exact exceeds 1 for every $\zeta \ge 0.7$, so at $\zeta = 2$ the flag improves the delivered damping; genuine over-damping occurs only for lightly damped actuators.
3. Scope is hinge and slide joints reached through a joint transmission, with an affine gain or bias. Spatial tendons are excluded because the neglected $(\partial m/\partial q)^\mathsf{T} f$ term is roughly 1.2 times the retained derivative with the opposite sign, so the matrix addition would not be PSD. Ball and free joints are excluded for the same reason (for a ball joint the length is expmap(q)·gearAxis against a constant moment), and under implicitfast a multi-dof block would additionally lose its off-diagonal entries in the gather into the reduced mass matrix. Body transmissions and site transmissions without a refsite have actuator_length identically zero. Fixed tendons and site transmissions with a refsite are exactly supportable and are arguably in scope.
4. An actuator on its forcerange clamp, or a joint on its actuatorfrcrange clamp, contributes no stiffness, since the applied force no longer varies with position. The term is added only when the position slope is restoring, which keeps $\hat{M}$ positive definite.
5. The flag has no effect under euler or RK4, nor on the mj_flexCG path. On every excluded configuration it is bit-identical to flag-off.

Questions

1. Mechanism. The implementation adds a separate qDeriv contribution rather than extending mjd_effBuild/efm_*, for three reasons: the exactness test at engine_derivative.c:3383 requires body_simple == 2, which actuator dofs on a serial chain never satisfy, so an actuator term would force the iterative path permanently; mj_flexCG disables islands; and the metric is selected under !flg_Newton, while Newton is the default solver. The virtue of the actuator term is precisely that it has $M$'s sparsity and rides the existing factorization. @quagla — rather than asking you to freeze an interface while Replace the flex metric factorization with a block preconditioner #3439 is in flight: what shape should this take so that it can later be absorbed into the general effective-metric mode? Is efm_K intended to remain a plain nv-row dof CSR that a second assembler could union into?
2. Exposure. The nearest precedent is mjDSBL_EULERDAMP, though that is a disable flag and this change wants the opposite polarity for a first landing. I propose mjENBL_IMPLICITPOS, off by default, with flag-off behavior bit-identical. I would avoid a per-actuator attribute: the choice is a property of the discretization, and per-actuator granularity would allow two actuators on the same dof to be integrated by different schemes. Is this acceptable?
3. Scope. Should jnt_stiffness be part of the first change? It is the easier case, diagonal and exact with no transmission moment involved, and it is subject to the identical bound. The order could reasonably be inverted, validating the mechanism there first, which would also let Euler benefit.

Alternatives

Folding $h k_p$ into kv, by hand or as a compile-time transform, produces the same dynamics on the smooth path. However, actuator_force and qfrc_actuator then report the proportional term as damping, forcerange and jnt_actfrcrange clamp a fictitious quantity, and the constraint solver is handed the wrong force. The compile-time variant additionally cannot track a gain written into mjModel after mj_setConst, as domain randomization does, and becomes stale if m->opt.timestep is changed at runtime.

Reducing the timestep is the correct answer whenever it is affordable; for reinforcement learning it is precisely the resource being economized.

Additional context

Implemented against main behind mjENBL_IMPLICITPOS, roughly 750 changed lines including tests and documentation. mj_discreteAcc mirrors both terms, so mjENBL_INVDISCRETE is consistent with the flag, and the Python, WASM and Unity binding tables are regenerated. The remaining gaps I know of: no test/benchmark/ entry yet, and MJX and Warp reject the unknown bit but do not implement it.

I expect the cost to be near zero — one extra accumulation into an existing sparse structure and one sparse matrix–vector product, no new factorization — and I will measure ms/step flag on and off on the standard benchmark models in the 2.3.3 changelog format rather than assert it. The work splits naturally as (1) the term with tests, (2) benchmark, (3) Warp.

[yuvaltassa]

Yes! Yes of course we should have this term, it's an obvious and strict improvement.
@qiayuanl , you're a legend ❤️
I will follow up soon.

[yuvaltassa]

Okay I've thought this through and I was wrong, it's not a strict improvement, we gain unconditional stability at the cost of accuracy, which is why you propose to put this behind a flag.

Giving the solver the effective (modified) inertia is a whole different can of worms since it invalidates the contract of mj_forward (output becomes dt-dependent). @quagla is indeed exploring this in the context of flex and we've been discussing whether we should make this a global option. But I think it's out of scope here; let me know if you disagree.

I'll keep investigating.

[qiayuanl]

Agreed on keeping the solver inertia out of scope — that belongs to the flex/effective-metric discussion.

On "unconditional stability at the cost of accuracy": I measured where the trade applies. There are three regimes relative to the cap $h\omega_n &lt; 2\zeta + 2\sqrt{\zeta^2 + 1}$.

1. Well below the cap the sentence is exactly right. The flag adds damping the model does not have, while the current integrator is nearly exact. This is why the flag defaults to off.
2. Above the cap there is no accuracy to trade. The current integrator does not give a less accurate trajectory, it gives NaN. The flag's error is bounded; the alternative's is not.
3. Near the cap the accuracy comparison flips before the stability one. Approaching the cap, the integrator's poles approach the unit circle and the response rings without decaying: stable, but wrong.

The curves are the RMS error of a unit step against the exact continuous PD, over the settling window (one-step maps, verified against MuJoCo to 12 digits). The crossover where flag-on becomes the more accurate one sits near $h\omega_n \approx 1$ at every damping: 1.26, 1.36, 1.12, 0.73 for $\zeta$ = 0.1, 0.5, 1, 2. As a fraction of the stable range that is 57% down to 9%, mainly because the cap grows like $4\zeta$ while the crossover stays put. For a critically damped or overdamped servo — the realistic case here — the flag is more accurate over three quarters or more of the stable range, and the only stable option beyond it.

So the narrower summary I would defend: the flag costs accuracy only for lightly damped servos well below the cap, and is equal or better everywhere else. "Unconditional" still carries the smooth-path caveat from the body: with contacts, limitation 1 can still diverge.

[yuvaltassa]

PTAL at yuvaltassa#4, WIP on a new integrator the will include this change (but also uses solver effective inertia)

In my mind this change belongs there though we could still do the more minimal thing you propose. Let me know if this makes sense to you, feel free to comment.