Computationally efficient Gauss-Newton reinforcement learning for model predictive control

1. Problem Statement

The paper addresses the challenge of applying Reinforcement Learning (RL) to Model Predictive Control (MPC) policies. While MPC offers interpretability, constraint handling, and strong initial performance compared to black-box Neural Network (NN) policies, optimizing MPC parameters via RL is computationally demanding.

• The Bottleneck: Most RL methods rely on first-order optimization (e.g., Gradient Descent/Adam), which converges linearly. This is inefficient when each policy update requires solving an optimal control problem (OCP).
• The Second-Order Dilemma: Second-order methods (e.g., Newton, Quasi-Newton) offer superlinear convergence, which is highly desirable for MPC. However, exact second-order methods require computing the Hessian of the policy. For MPC, this involves calculating second-order Nonlinear Programming (NLP) sensitivities, which depend on third-order derivatives of the system's Lagrangian. This is computationally intractable and memory-intensive, especially for high-dimensional or highly nonlinear systems.
• Stability Issues: Existing second-order approximations often suffer from instability due to noisy gradient/Hessian estimates and poor parameter scaling, leading to erratic training.

2. Methodology

The authors propose a framework that combines a Gauss-Newton approximation of the Hessian with momentum-based averaging and an adaptive trust-region strategy.

A. Gauss-Newton Hessian Approximation

The core theoretical contribution is a novel approximation of the deterministic policy Hessian that eliminates the need for second-order NLP sensitivities.

• Standard Approximation: Previous work (Kordabad et al., 2022) approximated the Hessian as $M = M_1 + M_2$. The term $M_1$ requires second-order policy derivatives (computationally expensive).
• The Innovation: The authors prove that as the policy parameters $\theta$ approach the optimal $\theta^*$, the term $M_1$ vanishes because the gradient of the action-value function with respect to actions becomes zero at the optimum (Lemma 1).
• Result: They propose using only the $M_2$ term (the Gauss-Newton approximation):
  $$M_2(\theta) = \mathbb{E}_s \left[ \nabla_\theta \pi_\theta(s)^\top \nabla_a^2 Q_{\pi_\theta}(s, a)|_{a=\pi_\theta(s)} \nabla_\theta \pi_\theta(s) \right]$$
  This term relies only on first-order NLP sensitivities (Jacobian of the policy), which are computationally tractable, yet the authors rigorously prove this approximation still guarantees superlinear convergence.

B. Momentum-Based Hessian Averaging

To handle the noise inherent in stochastic RL estimates (finite sampling of states and actions), the authors introduce a momentum scheme for the Hessian estimate.

• They define a biased exponentially moving average of the local Hessian estimates.
• Crucially, they propose a specific initialization strategy (using the square root of the first Adam momentum estimate) to prevent the Hessian from being dominated by inaccurate early estimates.
• A bias-correction factor is applied to ensure the estimate is unbiased.

C. Adaptive Trust-Region Optimization

To ensure stability during training, the parameter update is constrained within an adaptive trust region.

• Instead of comparing predicted vs. actual reward (which is noisy in RL), the trust-region radius $\delta_k$ is dynamically scaled based on the 2-norm of the Adam update step.
• This allows large steps when gradients are large (far from optimum) and conservative steps near the optimum, without requiring expensive objective function comparisons.

3. Key Contributions

1. Computationally Efficient Hessian: A Gauss-Newton approximation that removes the need for third-order derivatives and second-order NLP sensitivities, reducing computational complexity while maintaining superlinear convergence guarantees.
2. Robust Training Scheme: A momentum-based Hessian averaging technique combined with an adaptive trust-region constraint, ensuring stable training even with noisy estimates and varying parameter scales.
3. Theoretical Proof: Rigorous mathematical proof that the proposed approximation converges superlinearly to the optimal policy.
4. Empirical Validation: Demonstration on a nonlinear Continuously Stirred Tank Reactor (CSTR) showing superior data efficiency and convergence speed compared to first-order methods (Adam) and deep RL (TD3).

4. Results

The proposed method was evaluated on a nonlinear CSTR benchmark with varying levels of parameterization (2, 13, and 33 trainable parameters).

• Convergence Speed: The Gauss-Newton approach converged significantly faster than the Adam optimizer (first-order) and required fewer episodes to reach optimal performance. It outperformed the "Approximate Newton" method (which uses the full $M_1+M_2$ Hessian) in terms of wall-clock time for highly parameterized policies.
• Computational Efficiency:
  • Adam: Fastest per iteration but slowest convergence.
  • Approximate Newton: Extremely slow per iteration due to the quadratic scaling of second-order sensitivity calculations; became intractable for high parameter counts.
  • Proposed Gauss-Newton: Iteration time was roughly double that of Adam but significantly lower than the full Approximate Newton. The reduction in total iterations resulted in a net gain in total training time for complex policies.
• Robustness to Scaling: The method demonstrated high resilience to poor parameter scaling (varying the sensitivity of the objective function). While Adam failed to converge or got stuck in local optima with poorly scaled parameters, the Gauss-Newton method converged robustly.
• Comparison with Deep RL: Compared to a Neural Network policy trained with TD3:
  • Initialization: MPC-based RL started with a much better initial policy (solving a model-based OCP) compared to the random initialization of NNs.
  • Performance: The trained MPC policy achieved a higher final cumulative reward and better setpoint tracking with fewer constraint violations than the NN policy.
  • Data Efficiency: The MPC approach required significantly fewer interactions with the environment to reach high performance.

5. Significance

This work bridges the gap between the interpretability and safety of Model Predictive Control and the adaptability of Reinforcement Learning. By making second-order optimization computationally feasible for MPC policies, the paper enables:

• Sample Efficiency: Drastic reduction in the number of agent-environment interactions required for training, which is critical for industrial processes where data collection is expensive or risky.
• Scalability: The ability to tune complex, highly parameterized MPC policies without the computational explosion associated with full second-order derivatives.
• Stability: A robust framework that handles the noise and scaling issues typical in industrial control systems, making RL a more viable option for real-world process control applications.

In summary, the paper provides a mathematically sound and practically efficient pathway to optimize MPC policies using advanced second-order RL techniques, overcoming the primary computational barriers that have previously limited their adoption.