Source-linked AI summary
Differentiable MPC for End-to-end Planning and Control
Brandon Amos, Ivan Dario Jimenez Rodriguez, Jacob Sacks, Byron Boots, J. Zico Kolter
TL;DR
The paper addresses how to make MPC a differentiable policy class for reinforcement learning in continuous state and action spaces while retaining model-based structure. It differentiates through a convex approximation at the controller's fixed point using KKT conditions, enabling end-to-end learning of costs and dynamics. In imitation learning, the resulting MPC policies are more data-efficient than a generic neural network and outperform traditional system identification when the expert is unrealizable.
Problem
Model-free policies are differentiable but often have poor sample complexity and generalization, while model-based MPC has typically lacked full end-to-end learning through its optimization procedure.
Method
The method treats MPC as a policy parameterized by cost and dynamics, differentiating through a convex approximation at a fixed point using KKT conditions.
Results
In imitation learning, the method recovers MPC cost and dynamics from action-only losses, significantly outperforms a general neural-network policy, and surpasses vanilla system identification.
Takeaways & Limitations
Differentiable MPC provides an end-to-end policy class that combines explicit model-based cost and dynamics with task-loss optimization and efficient planning.
Takeaways & Limitations
Fixed-point differentiation can fail when the controller has not reached a fixed point or when neural-network dynamics prevent one from existing.
Abstract
from arXiv · showhide
We present foundations for using Model Predictive Control (MPC) as a differentiable policy class for reinforcement learning in continuous state and action spaces. This provides one way of leveraging and combining the advantages of model-free and model-based approaches. Specifically, we differentiate through MPC by using the KKT conditions of the convex approximation at a fixed point of the controller. Using this strategy, we are able to learn the cost and dynamics of a controller via end-to-end learning. Our experiments focus on imitation learning in the pendulum and cartpole domains, where we learn the cost and dynamics terms of an MPC policy class. We show that our MPC policies are significantly more data-efficient than a generic neural network and that our method is superior to traditional system identification in a setting where the expert is unrealizable.
1 Introduction
The paper proposes treating MPC as a differentiable policy class whose cost and dynamics can be learned end-to-end from task loss. This addresses the tension between expressive model-based control and the differentiability of simpler policy classes.
- Model-free policies are differentiable but often suffer from poor sample complexity and generalization, whereas MPC uses predictive models and online receding-horizon optimization.MPC applies the first planned action, then resolves the optimization at the next time step.
- The proposed policy is u = π(xinit; C, f), with cost C and dynamics f learned by differentiating through the MPC optimization problem.This enables direct task-loss optimization instead of learning dynamics or trajectory components in isolation.
- The method analytically differentiates through a box-constrained iterative LQR solver using one additional backward pass of a modified solver.The derivative is computed for an iterative non-convex optimization procedure based on box-constrained iterative LQR.
- In imitation learning, the method recovers cost and dynamics from action-only losses and directly optimizing imitation loss outperforms vanilla system identification.The comparison is motivated by learning controller components jointly rather than regressing on collected dynamics or rollout data.
2 Background and Related Work
Related work combines model-based planning structure with model-free learning, including explicit planning modules and model-generated training signals. The paper instead targets efficient differentiation through MPC while preserving explicit cost and dynamics components.
- Model-free methods have motivated hybrid approaches that use learned models or planning structure to improve policy learning.Examples include synthetic model-based training data, model-based priors, and algorithms that learn a model while planning with it.
- Recent architectures incorporate planning, rollouts, value iteration, or goal-conditioned models into end-to-end trainable policies for continuous control.These approaches approximate planning procedures within neural architectures or use abstract models to predict values, actions, and rewards.
- Unlike methods that explicitly unroll optimization, this paper analytically differentiates through an MPC solver's fixed point using a single LQR step in the backward pass.The approach is designed to reduce the memory and computational cost of backward passes while retaining continuous state and action planning.
- The resulting framework keeps explicit cost and dynamics components analyzable while allowing them to be learned entirely end-to-end.This contrasts with model-free policies, which do not expose equivalent cost and dynamics components for separate analysis.
3 Differentiable LQR
The paper formulates finite-horizon LQR as a constrained convex quadratic optimization problem with affine dynamics and parameters that can be embedded in an end-to-end learning system. Its gradients are obtained by implicitly differentiating the KKT conditions and solving another LQR problem.
- Finite-horizon LQR optimizes a convex quadratic objective over affine state-transition dynamics from an initial state xinit.The optimal nominal trajectory is computed over a finite horizon.
- The LQR problem uses trajectory variables τ1:T subject to x1 = xinit and affine dynamics xt+1 = Ftτt + ft.The formulation represents the initial condition and transition constraints explicitly.
- LQR can be treated as a module with parameters θ = {C, c, F, f}, whose outputs and derivatives support end-to-end parameter updates from a loss ℓ.The module takes xinit and its parameters, then computes gradients with respect to C, c, F, f, and xinit.
- The forward pass solves the LQR problem for the optimal trajectory τ⋆1:T, while the backward pass solves a second LQR problem using ∇τ⋆ℓ.The second solve can ideally reuse factorizations from the forward pass.
- Because LQR is a constrained convex quadratic argmin, implicit differentiation of its KKT conditions yields gradients with respect to the LQR parameters.The dynamic Riccati recursion can be interpreted as an efficient way to solve the KKT system.
4 Differentiable MPC
The paper differentiates box-constrained MPC by differentiating the convex approximation reached at an iLQR fixed point, using linear systems that account for active constraints. This fixed-point approach requires a fixed point; otherwise, unrolling provides surrogate gradients with iteration-dependent cost.
- Differentiable MPC: MPC handles nonlinear costs and dynamics with box-constrained controls by iteratively optimizing convex approximations.The cost and dynamics are locally approximated around the current trajectory before solving the resulting convex problem.
- Differentiable MPC: At a fixed point, derivatives of the non-convex MPC problem can be computed from its final convex approximation.The fixed-point approximation uses the corresponding objective and dynamics derivatives rather than differentiating every forward iteration.
- Differentiating with Box Constraints: Active box constraints are treated as equalities, and the derivative computation becomes a zero-constrained LQR problem.Controls at tight bounds receive zero perturbations in the derivative solve.
- Differentiable MPC: The differentiable MPC module solves the control problem with box-DDP, then differentiates the fixed-point solution using an LQR solver that zeros constrained controls.The resulting parameter gradients are obtained by differentiating through the approximated objective and dynamics and applying the chain rule.
- Limitations and Runtime: If no fixed point exists or the controller stops before reaching one, fixed-point differentiation cannot be used reliably.Unrolling can provide surrogate gradients in this case, but its backward-pass cost scales linearly with the number of iLQR iterations, whereas fixed-point differentiation uses one LQR solve.
5 Experimental Results
The experiments evaluate differentiable MPC for solver efficiency and imitation learning, including LQR, pendulum, cartpole, and non-realizable expert settings. Results compare MPC-based learning with neural-network policies and system identification.
- 5.1 MPC Solver Performance: The differentiable MPC solver is more computationally and memory efficient than individually unrolling box-constrained iLQR iterations.Its backward pass reuses forward-pass factorizations and avoids multiple iterations.
- 5.2 Imitation Learning: Linear-Dynamics Quadratic-Cost (LQR): In LQR imitation learning, the model converges to the true dynamics and achieves perfect imitation in half of eight randomly initialized trials.The remaining trials reach a local minimum of imitation loss and diverge from the true model.
- 5.3 Imitation Learning: Non-Convex Continuous Control: The pendulum and cartpole experiments learn dynamics, costs, or both by directly optimizing imitation loss through MPC.Compared methods include an LSTM policy, SysId, mpc.dx, mpc.cost, and mpc.cost.dx.
- 5.3 Imitation Learning: Non-Convex Continuous Control: Differentiable MPC significantly outperforms a general neural-network policy in nearly every pendulum and cartpole setting.In many cases, it also recovers the expert’s true cost function and dynamics.
- 5.3 Imitation Learning: Non-Convex Continuous Control: mpc.dx matches SysId performance while using only control information and no state information.The comparison concerns learned dynamics in the pendulum and cartpole experiments.
- 5.4 Imitation Learning: SysId with a non-realizable expert: With a non-realizable pendulum expert, direct imitation-loss optimization finds a better imitation model than vanilla SysId despite SysId’s better next-state-transition fit.The expert includes damping and wind terms absent from the learned dynamics model class.
6 Conclusion
The paper establishes a framework for differentiating and learning MPC-based controllers, demonstrated in imitation learning. It highlights potential advantages over generic imitation learning and system identification, with future extensions to task-specific model tuning and stochastic control.
- The approach provides foundations for differentiating and learning MPC-based controllers in reinforcement learning and imitation learning.
- The method differentiates through MPC more efficiently than policy unrolling, with a backward pass that is essentially free after iLQR convergence.
- The experiments demonstrate potential advantages over generic imitation learning and system identification.
- The framework is intended to integrate MPC policies into deep network architectures and support task-specific model tuning, joint model-based and policy-based losses, and stochastic control.
A LQR and MPC Algorithms
The algorithms represent MPC costs and dynamics locally around a trajectory, solve the resulting constrained quadratic problem, and iterate forward and backward recursions to update the control sequence.
- LQR Algorithm: The LQR solver uses quadratic cost terms and affine terms over a horizon, with each quadratic cost matrix required to be positive semidefinite.
- Recursions: Forward recursion initializes the state at xinit and propagates the trajectory, while backward recursion proceeds from T to 1.
- MPC Algorithm: Algorithm 2 models the cost as non-convex and twice-differentiable and the dynamics as non-convex and once-differentiable, with bounded controls over a finite horizon.
- MPC Algorithm: At each iteration, the algorithm forms a second-order Taylor expansion of the cost and a first-order Taylor expansion of the dynamics around the current trajectory.
- MPC Algorithm: The local approximations are passed to MPCstep, which produces the next trajectory iterate before the algorithm repeats.
- LQR Algorithm: The constrained control update can be solved with a Projected-Newton method, separating free and clamped dimensions.
- Recursions: The algorithm returns the optimized state and control trajectories over the horizon.
B Imitation learning experiment losses
Figure 6 reports learning results for the simple pendulum and cartpole environments using the best validation loss from each training run.
- Figure 6 reports the corresponding train and test loss at the best validation loss observed during each training run.Each datapoint is averaged over four trials.