Source-linked AI summary
Accelerating Nonlinear MPC by Behavior Cloning with Separated Timescales: Application to Nuclear Reactor Load-Following
Perceval Beja-Battais, Alain Grosset{ê}te, Nicolas Vayatis
TL;DR
Behavior cloning for multiscale systems faces a representation challenge because slow and fast dynamics may be conflated in a shared latent space. SRLBC separates slow, fast, and exogenous inputs; across the reported benchmark, it improves performance, and its safety-augmented policies are faster than NMPC with high feasibility.
Problem
A shared latent space may conflate slow and fast dynamics in multiscale systems, degrading closed-loop generalization for behavior-cloned controllers.
Source: §I · p. 1Method
SRLBC uses separate latent encoders for slow states, fast states, and exogenous inputs, with a controller acting on their product latent space.
Source: §I · p. 2Results
Separated embeddings improve closed-loop performance over shared embeddings; safety augmentation yields a 6.8-fold speedup over NMPC and at least 99.8% feasibility, with SRLBC's lower closed-loop cost than baselines.
Source: §I · p. 2Takeaways & Limitations
On the PWR benchmark, separated embeddings yield sparser, less redundant, more interpretable representations aligned with system physics.
Source: §I · p. 2The timescale partition must be known a priori, and encoder latent dimensions require expert knowledge or prior hyperparameter optimization.
Source: §V · p. 7
Abstract
from arXiv · showhide
Behavior Cloning (BC) of expert controllers for multiscale dynamical systems is challenging: state variables evolving at widely separated timescales induce highly nonlinear couplings that a shared latent space may fail to adequately represent, degrading closed-loop generalization. We propose Structured Representation Learning for Behavior Cloning (SRLBC), an architecture designed to leverage the physical knowledge of multiscale systems in which slow state variables, fast (stiff) state variables, and predictable exogenous inputs evolve at different timescales. On 500 long-horizon roll-outs of a pressurized water reactor (PWR) load-following benchmark -a high-dimensional, stiff industrial system with challenging safety constraints -and across 20 training seeds, separated embeddings systematically outperform shared embedding baselines on cumulative cost, imitation accuracy and constraint satisfaction. The resulting standalone policies approach NMPC closed-loop performance at a computational cost 15 times lower on average. Since standalone policies still leave a few percent of trajectories infeasible, we embed the learned policies in a safety-augmented approximate NMPC framework, in which the optimizer is only called when the predicted input sequence is not certified as safe. This yields controllers still 6.8 times faster than the NMPC controller on average, with a feasibility (i.e., constraint satisfaction) of at least 99.8% for all architectures, SRLBC retaining the lowest closed-loop cost. Under distribution shift, standalone BC policies degrade sharply: SRLBC degrades the least, and once safety-augmented, it retains a feasibility above 98% with the lowest closed-loop cost. Beyond the quantitative results, Jacobian-based sensitivity analysis shows that the resulting controller relies on a sparse and less redundant set of latent features, each attributable to a single group of physical variables, and aligned with the physical timescale decomposition of the system.
I. INTRODUCTION
A shared latent space may conflate slow and fast dynamics, limiting behavior cloning for multiscale systems. SRLBC separates timescales and, with safety augmentation, retains lower closed-loop cost while achieving high feasibility and faster computation than NMPC.
I. INTRODUCTION
- This paper extends a preliminary short version with 20 random seeds, in- and out-of-distribution experiments, and a safety-augmented pipeline.The preliminary version was accepted for a non-archival workshop presentation.Source: §I · p. 1
- Behavior cloning can avoid online optimization, but nonlinear MPC's computational burden limits deployment in resource-constrained industrial pipelines.BC mimics an expert through supervised learning and can retain near-optimal performance when its policy generalizes well.Source: §I · p. 1
Contributions
- The decomposition is prescribed by known physical timescale separation, and is evaluated by closed-loop performance and expert imitation rather than reconstruction or disentanglement metrics.The approach differs from task-agnostic factor discovery and modality-based encoder separation.Source: §I · p. 2
- SRLBC encodes slow states, fast states, and exogenous inputs separately, then trains a structured controller on their product latent space.The architecture's separated embeddings improve closed-loop performance over a shared embedding and reduce latent-feature redundancy across 20 training seeds.Source: Figure 1 · p. 1, §I · p. 2
- Within safety-augmented approximate NMPC, learned policies call the optimizer at less than 1.6% of timesteps, achieving a 6.8-fold speedup over NMPC and at least 99.8% feasibility; SRLBC lowers closed-loop cost versus baselines, including on out-of-distribution transients.Source: §I · p. 2
- On a full-scale PWR load-following benchmark, separated embeddings yield sparser, less redundant, more interpretable representations aligned with system physics.The paper evaluates optimal load-following control and reports the interpretability analysis in Fig. 4.Source: §I · p. 2
II. PROBLEM STATEMENT & RELATED WORK
The section formulates a fast–slow control problem with a known exogenous signal and an expert NMPC policy to imitate. It motivates separate representations for slow, fast, and exogenous variables, while discussing BC and safety options.
Multiscale control problem
- The system state comprises slow x and fast z components separated by stiffness parameter ε, with an exogenous signal w known in advance.The control objective is to design explicit state feedback that imitates the expert input without online optimization, subject to system dynamics and constraints.Source: §II · p. 2
- The assumed NMPC scheme may be too slow for embedded industrial deployment, motivating imitation of its control policy.Source: §II · p. 2
Representations for control
- Separate encoders for slow, fast, and exogenous variables preserve the multiscale structure instead of imposing approximately linear dynamics over a shared latent space.The latent dimensions should be large enough to retain complex control-relevant couplings, but need not be unnecessarily large.Source: §II · p. 2
BC and safety
- Behavior cloning learns a policy from expert demonstrations, but generating representative demonstrations can be computationally costly for complex systems.The section describes safety pipelines that call NMPC when a learned policy is unsafe, alongside robust MPC approaches with statistical safety guarantees.Source: §II · p. 2
- For safety-critical applications, the authors prioritize interpretability of the cloned policy over network complexity and its statistical guarantees.Source: §II · p. 3
III. PROPOSED FRAMEWORK
The framework uses separate encoders for slow states, fast states, and exogenous variables, feeding a nonlinear controller trained to imitate NMPC. It also investigates decorrelating the embeddings and includes one DAgger round in every policy’s training pipeline.
Designing the controller and training the ensemble
- Three encoders map slow states, fast states, and exogenous variables into separate embeddings for a latent controller trained to approximate NMPC.Source: §III · p. 3
- The architectural prior separates variable groups during encoding while allowing the nonlinear controller to represent arbitrary couplings between their latent blocks.This constrains the space of admissible representations, although a sufficiently expressive shared embedding could in principle recover the same decomposition.Source: §III · p. 3
Enforcing decorrelation
- The framework investigates training the structured encoder to decorrelate its embedding components by minimizing their covariance deviation from identity matrices.Source: §III · p. 3
Finetuning with DAgger
- Every policy, including the baselines, receives one DAgger round that adds expert actions along trajectories generated by the current policy.Source: §III · p. 3
IV. EXPERIMENTS
Experiments evaluate the separated-timescale controller on PWR load-following, including in-distribution and shifted transients. SRLBC improves standalone performance over shared-embedding baselines, while safety augmentation raises feasibility and preserves its cost advantage.
System and evaluation setup
- The benchmark is a high-dimensional, stiff PWR load-following system with slow states, fast states, and turbine power known 30 minutes ahead.Slow dynamics evolve over minutes to hours, while fast dynamics typically reach equilibrium in less than one second; the OOD generator introduces abrupt reversals in the known-ahead setpoint during power variations.Source: §IV · pp. 3, 5
In-distribution results
- SRLBC reduces excess closed-loop cost to 7.4% from 16.6% for the shared baseline and raises standalone feasibility from 94.2% to 97.4%.Across 20 training seeds and 500 feasible 48-hour test transients, it also reduces the control-input gap by 30%; the reported differences from both baselines are significant at p < 10^-4.Source: §IV · p. 5
Out-of-distribution results
- On 175 feasible out-of-distribution transients, standalone policies reach 3–5 times NMPC cost and at most 25% feasibility, while SRLBC degrades less than the baselines.Its feasibility is twice that of the baselines and its excess cost is almost halved, but no learned policy is suitable for standalone deployment under this shift.Source: §IV · p. 5
Safety-augmented results
- Under safety augmentation, SRLBC has cost ratios of 1.053 in distribution and 1.138 out of distribution, the lowest among the compared policies.Feasibility is 99.8–99.9% in distribution and 98.5–98.7% out of distribution for all architectures; average speedup is about 7× in distribution and 4.8× out of distribution.Source: §IV · p. 6
Computational time and hardware
- Standalone BC roll-outs take about 20 seconds versus 5–6 minutes for NMPC, an acceleration of approximately 15×.For the safety-augmented pipeline, optimizer calls are avoided at up to 98.4% of timesteps; absolute timings were measured on a workstation CPU.Source: §IV · p. 6
Analysis of the learned representations and interpretability of the controller
- Jacobian analysis finds SRLBC sensitivity concentrated in fewer latent components, while the shared embedding spreads contributions across a more redundant representation.The separated design assigns each component to one physical-variable group; controller output is mainly driven by boron concentration and the fast state, while other state and exogenous variables mostly regularize.Source: §IV · pp. 6–7
Behavior on infeasible transients
- On transients infeasible even for NMPC, standalone SRLBC exceeds NMPC’s constraint violation by about 20%, versus 34–44% for the baselines.Safety augmentation matches NMPC’s violation when demand leaves the admissible domain; on infeasible OOD transients, standalone violations reach 10–24 times NMPC’s, while augmentation again matches it.Source: §IV · p. 7
V. DISCUSSION
Safety augmentation makes feasibility, computation time, and rejection rate nearly architecture-independent, but closed-loop cost remains tied to policy quality. The approach still lacks recursive-feasibility guarantees and assumes an accurate simulation model.
V. DISCUSSION
- Safety filtering nearly equalizes feasibility, computation time, and rejection rates across architectures, but closed-loop performance remains closely related to policy quality.The structured representation's cost advantage is more visible for transients outside the training distribution.Source: §V · p. 7
- Without terminal ingredients, the safety-augmented pipeline provides no recursive-feasibility guarantee, and certification relies on an accurate simulation model.Feasibility remains slightly below that of the NMPC controller.Source: §V · p. 7
- New online control strategies with different cost or constraint functions remain computationally expensive.The authors identify this as an industrial motivation for the approach.Source: §V · p. 8
VI. CONCLUSION
SRLBC uses structured embeddings for behavior cloning in multiscale nonlinear control, with consistent performance and interpretability gains on a PWR-simulator benchmark. Safety augmentation retains fast execution and high feasibility, including under distribution shift.
VI. CONCLUSION
- Across 20 training seeds, SRLBC improves integral cost, constraint violations, and expert imitation while maintaining a structured embedding on the PWR-simulator benchmark.The safety-augmented policies run up to 7 times faster than NMPC with at least 99.8% feasibility, and SRLBC retains the lowest closed-loop cost.Source: §VI · p. 8
- Under distribution shift, standalone policies fail, while the certified combination remains feasible on more than 98% of trajectories, with SRLBC achieving the lowest cost.Source: §VI · p. 8
- The authors conclude that physics-based structured embeddings can improve controller performance and interpretability when the network is expressive enough to approximate NMPC.Source: §VI · p. 8