PaperlayerView exact revision on arXiv

Source-linked AI summary

Guided Riemannian Optimization (GuRO): Bridging Model Predictive Control and Decision Transformers

Hossein Abdi, Satya Prakash Dash, Mingfei Sun

arXiv:2608.23204v1cs.RO

TL;DR

High-dimensional nonlinear robotic control is challenged by inaccurate MPC models, sample-intensive RL, and difficult transformer optimization. The paper combines MPC-guided Decision Transformer learning with efficient curvature-aware Riemannian optimization, and reports higher rewards and faster convergence than strong baselines across quadruped tasks.

  • Problem

    High-dimensional nonlinear robotic decision-making is limited by inaccurate or short-horizon MPC models, high sample complexity in model-free RL, and non-convex optimization of transformer policies.

  • Method

    The framework integrates MPC and model-free RL in a sequence decision-making architecture, using locally optimal MPC trajectories to guide a Decision Transformer and Riemannian curvature-aware policy optimization.

  • Results

    Across three Unitree AlienGo locomotion tasks, the MPC-guided Decision Transformer with Riemannian optimization outperforms TRPO, SAC, Online-DT, SGD, and Adam variants in reward, sample efficiency, and convergence speed.

  • Takeaways & Limitations

    The reported results support combining MPC guidance with curvature-aware optimization for sequence-based quadruped control.

Abstract

from arXiv · show

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.

I. INTRODUCTION

High-dimensional nonlinear robotic decision-making exposes limitations in both MPC and model-free RL. The paper combines MPC-guided sequence modeling with curvature-aware Riemannian optimization and evaluates the framework on quadruped control tasks.

  • MPC offers interpretability and sample efficiency when its dynamics model is accurate, but performance degrades with inaccurate models and long-horizon prediction.
  • Sequence decision-making unifies complementary strengths of RL and MPC for complex robotic tasks involving long-term temporal dependencies.
  • MPC-generated locally near-optimal trajectories guide Decision Transformer learning, reducing reliance on extensive offline pretraining and real-environment interactions.
  • Riemannian optimization accounts for loss-manifold curvature to provide faster and more stable training of transformer-based policies.
  • The framework is evaluated on high-dimensional quadruped control tasks against TRPO, SAC, and Online-DT using reward and convergence-speed metrics.

II. RELATED WORKS

Related work spans model-based RL, MPC-based value learning, transformer sequence decision-making, and curvature-aware optimization. These lines of research provide the conceptual foundations for the paper’s hybrid approach.

  • Model-based RL commonly learns approximate Markov decision process structure, with MPC planning over models refined to capture environment dynamics or rewards.
  • Decision Transformer and Trajectory Transformer reframed policy learning as sequence modeling, motivating extensions such as Online-DT and constrained decision transformers.
  • Euclidean second-order methods use Hessian information to precondition gradients and adapt updates to local loss curvature.
  • Natural-gradient and related Riemannian methods adapt optimization to parameter-space geometry, with TRPO using an approximate natural-gradient step.
  • The paper formulates sequential decision-making within the Markov Decision Process framework, specifying states, actions, transitions, rewards, initial-state distribution, and discounting.

B. Model Predictive Control (MPC)

The MPC formulation models stochastic nonlinear dynamics and selects a finite-horizon control sequence by minimizing a cost over predicted states and actions. A trajectory optimizer produces the control sequence used by the framework.

  • The stochastic discrete-time dynamics model combines a nominal transition function with independently and identically distributed Gaussian process noise.
  • MPC determines a finite-horizon sequence of control actions by minimizing a cost function over predicted system trajectories.
  • Predicted states and controls are conditioned on the observed state at the current time step.
  • A standard trajectory optimization algorithm solves the MPC problem and produces the optimal control sequence for the prediction horizon.

C. Riemannian Optimization

Riemannian optimization treats policy parameters as points on a manifold equipped with a metric tensor. The steepest-descent update uses the manifold gradient and a retraction to remain on that space.

  • The parameter space is represented as an n-dimensional smooth Riemannian manifold with metric tensor G(θ) defined at each parameter point.
  • Riemannian steepest descent updates parameters using the Riemannian gradient, a positive step size, and a retraction from the tangent space back to the manifold.

A. Sequence Modeling Setup

The framework represents sequential decision-making as trajectories of states, actions, and return-to-go values, which a decision transformer uses to generate policy distributions.

  • A. Sequence Modeling Setup: An MDP interaction produces a trajectory containing states, actions, and rewards, represented with return-to-go values instead of immediate rewards.The return-to-go at time i is the discounted sum of rewards from i onward.
  • A. Sequence Modeling Setup: The trajectory τ 0:t is provided to the decision transformer, which outputs the policy distribution πθ(·|τ 0:t).

B. Guiding by MPC

MPC generates online trajectories that guide decision-transformer learning, while entropy and KL constraints support exploration and limit destabilizing policy updates.

  • B. Guiding by MPC: At each rollout step, finite-horizon MPC predicts an optimal control sequence and supplies trajectories for subsequent decision-transformer training.
  • B. Guiding by MPC: The policy is trained by minimizing negative log-likelihood over trajectories sampled from the replay buffer.
  • B. Guiding by MPC: Maximum entropy encourages stochastic exploration, while a KL trust region limits successive policy deviations to stabilize training.
  • B. Guiding by MPC: The constrained objective is solved through primal-dual optimization, with the KL constraint enforced by backtracking line search.
  • B. Guiding by MPC: The entropy constraint uses a positive dual variable α, optimized in log-space so α remains positive.

C. Efficient Riemannian Optimization

GuRO applies an efficient Riemannian optimization procedure that approximates local curvature with a diagonal, smoothed metric and uses trust-region updates for decision-transformer training.

  • C. Efficient Riemannian Optimization: Riemannian updates exploit loss-landscape curvature to improve per-iteration progress, accelerate convergence, and reduce required environment interactions.
  • C. Efficient Riemannian Optimization: GuRO iteratively rolls out MPC trajectories, updates entropy and policy quantities, estimates curvature, and performs trust-region backtracking.
  • C. Efficient Riemannian Optimization: The Riemannian gradient uses the metric tensor’s Moore–Penrose pseudo-inverse to precondition and reorient the Euclidean gradient.
  • C. Efficient Riemannian Optimization: A mini-batch Gauss–Newton estimator, diagonal representation, and exponential moving average provide a scalable, stable metric approximation for high-dimensional policies.
  • C. Efficient Riemannian Optimization: The diagonal metric is inverted elementwise, and the resulting direction is combined with the Euclidean gradient through an element-wise product.
  • C. Efficient Riemannian Optimization: Backtracking adapts the step size while checking mini-batch KL divergence to satisfy the trust-region constraint and maintain monotonic objective improvement.
  • C. Efficient Riemannian Optimization: Experiments use parallelized GPU simulations of the Unitree AlienGo quadruped in MJX within JAX.

V. EXPERIMENTS AND DISCUSSION

The evaluation examines GuRO on high-dimensional, nonlinear quadruped locomotion tasks in challenging environments.

  • V. EXPERIMENTS AND DISCUSSION: The study evaluates the proposed algorithm on high-dimensional, nonlinear robotic control tasks focused on quadruped locomotion.

A. Validation Setup

Experiments use a Unitree AlienGo quadruped across three challenging locomotion tasks, comparing the proposed method with RL, policy-optimization, and decision-transformer baselines. Figures assess reward versus timesteps and training loss, while reported settings specify implementation and optimization configurations.

  • Environments and Tasks: Experiments evaluate uneven-terrain locomotion, stair climbing, and low-friction incline ascent on the Unitree AlienGo quadruped.The robot is simulated with the MJX physics engine within JAX.
  • Baselines and Metrics: The method is compared against SAC, TRPO, and Online-DT, with mean total reward and convergence speed as evaluation metrics.The comparisons also include Euclidean optimizer ablations in the broader experimental design.
  • Implementation Settings: Reported results use average and standard deviation over 5 runs with different initial random seeds.The implementation uses a parallelized GPU setup on an NVIDIA GeForce RTX 4090.
  • Evaluation Procedure: Figures 2–4 plot mean total reward against timesteps and training loss across iterations for the three locomotion scenarios.The plots are used to assess sample efficiency, final performance, and convergence behavior.
  • Implementation Settings: Hyperparameters include batch size 64, MPC horizon 25, discount factor 0.99, and Online-DT batch size 256.Learning rates and forgetting factors are selected through grid search over predefined ranges.

B. Main Results and Discussions

Across uneven terrain, stair climbing, and low-friction incline tasks, the MPC-guided Decision Transformer with Riemannian optimization achieves higher rewards, better sample efficiency, and faster convergence than the evaluated baselines and Euclidean ablations. Statistical analyses report strong differences in some comparisons, while significance varies by task.

  • Overall Performance: Across three locomotion tasks, the proposed method outperforms TRPO, SAC, Online-DT, Ours-SGD, and Ours-Adam in reward, sample efficiency, and convergence rate.MPC supplies locally optimal predicted trajectories, while Riemannian optimization improves training convergence.
  • Reward and Sample Efficiency: The proposed method achieves higher average total reward with fewer environment steps than TRPO, SAC, and Online-DT across all three tasks.The comparison is described for uneven terrain, stair climbing, and low-friction incline ascent.
  • Convergence: Training-loss curves in Figures 2–4 show faster convergence for the Riemannian optimizer than conventional gradient-based methods.The reported explanation is curvature-aware updates that account for the intrinsic geometry of the loss manifold.
  • Statistical Analysis: In uneven terrain, comparisons with TRPO, SAC, and Online-DT yield t-statistics of 20.32, 7.12, and 4.34, respectively.The corresponding p-values are 3.60 × 10^-8, 1.00 × 10^-4, and 3.43 × 10^-3.

VI. CONCLUSIONS

The paper presents a hybrid MPC–RL sequence-modeling framework using a Decision Transformer and efficient Riemannian optimization for transformer-policy training. On three challenging quadruped locomotion scenarios, it reports higher mean total reward, superior sample efficiency, and faster convergence than strong baselines.

  • Conclusion: The framework integrates model-based MPC with model-free RL through a Decision Transformer sequence-modeling paradigm.It combines the two approaches within a hybrid decision-making framework.
  • Conclusion: Efficient Riemannian optimization provides curvature-aware updates intended to improve numerical stability and convergence rate for non-convex transformer-policy optimization.The approach addresses challenges associated with highly non-convex optimization landscapes.
  • Conclusion: Across uneven terrains, stair climbing, and low-friction inclined surfaces, the method achieves higher mean total reward, superior sample efficiency, and faster convergence than state-of-the-art baselines.The evaluation uses the Unitree AlienGo quadruped robot.
Loading 2608.23204v1…