Decision-making in high-dimensional, nonlinear systems remains a central challenge in robotics. While model-based methods like Model Predictive Control (MPC) offer sample efficiency and interpretability, their performance degrades when the dynamics model is inaccurate or long-horizon predictions are required. Conversely, model-free reinforcement learning (RL) learns policies directly from interaction but suffers from high sample complexity and unstable optimization. Recent advances in sequence modeling have inspired transformer-based decision-making frameworks that can unify MPC and RL, but their training typically faces significant optimization challenges due to highly non-convex loss landscapes. In this work, we propose a novel framework that integrates MPC with RL in a sequence decision-making framework and leverages a curvature-aware optimization to efficiently tackle non-convex loss landscapes. MPC provides predictions of locally optimal trajectories that guide the decision transformer, removing the need for extensive offline pretraining. To address the slow and unstable convergence of traditional optimizers, we train the policy in a Riemannian parameter space using an efficient Riemannian (curvature-aware) method, leading to faster and more robust optimization. We evaluate our framework on high-dimensional quadruped control tasks and demonstrate consistent improvements over strong baselines, including TRPO, SAC, and Online Decision Transformer, achieving higher returns and faster convergence.
Figures & tables
Fig. 1: Experimental evaluation on the Unitree AlienGo quadruped robot, simulated using the MuJoCo XLA (MJX) physics engine within the JAX framework. All experiments leverage a parallelized GPU implementation on an NVIDIA GeForce RTX 4090. The translucent robots illustrate the parallel simulations running simultaneously. The tasks include locomotion on uneven terrains (left), stair climbing (center), and ascending a low-friction inclined surface (right).
Fig. 2: Locomotion on uneven terrains. Left : mean total reward versus timesteps, highlighting improved sample efficiency and final performance. Right : training loss across iterations, illustrating faster convergence with the proposed Riemannian optimization method.
Fig. 3: Stair climbing performance. Left : mean total reward over timesteps, showing enhanced sample efficiency and final reward with our approach. Right : corresponding training loss, demonstrating accelerated convergence relative to SGD and Adam.
Fig. 4: Ascending low-friction incline. Left : mean total reward, emphasizing the performance improvement and sample efficiency of our method. Right : training loss across iterations, highlighting the faster convergence of the proposed Riemannian optimizer compared to conventional methods.
Task Scenario
TRPO
SAC
Online-DT
Ours-SGD
Ours-Adam
Ours-Riemannian
Locomotion on uneven terrains
−5.52±0.2
−3.85±0.2
−3.65±0.3
−3.42±0.4
−3.25±0.3
−2.95±0.2
Stair climbing performance
−3.78±0.4
−3.68±0.3
−3.90±0.4
−3.52±0.4
−3.42±0.3
−3.36±0.2
Ascending low-friction inclines
−3.52±0.4
−3.58±0.4
−3.50±0.4
−3.62±0.5
−3.65±0.6
−3.25±0.3
TABLE I: Average total reward (mean ± standard deviation) for quadruped robot across challenging scenarios. Tasks include uneven terrain locomotion, stair climbing, and low-friction slope ascent. Higher values correspond to better performance.
Reinforcement learning control uses reward-driven learning to make decisions in dynamic systems, from robots and vehicles to energy and infrastructure.
State-of-the-art model-based Reinforcement Learning (RL) approaches either use gradient-free, population-based methods for planning, learned policy networks, or a combination of policy networks and planning. Hybrid approaches that combine Model Predictive Control (MPC) with a learned model and a policy prior to leverage the advantages of both paradigms have shown promising results. However, these approaches typically rely on gradient-free optimization methods, which can be computationally expensive for high-dimensional control tasks. While gradient-based methods are a promising alternative, recent works have empirically shown that gradient-based methods often perform worse than their gradient-free counterparts. We propose Dream-MPC, a novel approach that generates few candidate trajectories from a rolled-out policy and optimizes each trajectory by gradient ascent using a learned world model, uncertainty regularization and amortization of optimization iterations over time by reusing previously optimized actions. Our results on 24 continuous control tasks show that Dream-MPC can significantly improve the performance of the underlying policy and can outperform gradient-free MPC and state-of-the-art baselines. Code and videos are available at https://dream-mpc.github.io.
Jonathan Spieler, Sven Behnke
Autonomous Intelligent Systems, Computer Science Institute VI - Intelligent Systems and Robotics, Center for Robotics and the Lamarr Institute for Machine Learning and Artificial Intelligence, University of Bonn, Germany.
Actor-Critic Model Predictive Control (MPC) effectively addresses complex, non-convex control problems, but guaranteeing the closed-loop stability of sequence-based learning models within these pipelines remains challenging. This paper introduces a novel Transformer-Actor-Critic MPC architecture with formal robustness guarantees. First, we prove that Transformer networks can satisfy global incremental Input-to-State Stability (δISS). We then leverage Riemannian contraction theory to analyze the interconnected dynamics between the physical plant and the predictive neural network. Finally, we integrate these theoretical bounds as a training regularizer to yield a certifiably robust policy. The framework is validated on a nonlinear 3D drone model executing target-reaching and obstacle-avoidance maneuvers.
Model Predictive Control (MPC) provides a structured and constraint-aware mechanism for decision-making, but its reliance on optimization-friendly analytical dynamics models limits its use in tasks with contacts and other hard-to-model state dependencies. Model-free reinforcement learning avoids explicit modeling assumptions but typically requires large amounts of interaction data. We present a learning-based MPC framework that combines the data efficiency and structure of local model-based planning with learned components that compensate for incomplete dynamics and finite-horizon myopia. The method augments a nominal analytical model with a residual dynamics network that learns missing state-dependent effects from data and combines the resulting planner with a learned action-value critic that injects long-horizon MDP structure into the local iLQR optimization. To make this practical at reinforcement-learning scale, we develop a GPU-accelerated batched iLQR solver that evaluates learned dynamics and critic networks inside the optimal-control loop and solves thousands of trajectory-optimization problems in parallel. The complete system is integrated into a robotics simulator, enabling scalable model-based reinforcement learning under incomplete dynamics. Experiments on biased and incompletely modeled control tasks show that the approach improves closed-loop control performance while preserving the model-based structure needed for efficient constrained trajectory optimization.
Philipp Schoch, Markus Ryll
Autonomous Aerial Systems Lab Technical University of Munich