Learning Legged MPC with Smooth Neural Surrogates

TL;DR

The paper tackles the challenge of integrating learned dynamics with online planning for legged robots by identifying stiffness, nonsmoothness, and heavy-tailed residuals as key failure modes in gradient-based MPC. It introduces the Smooth Neural Surrogate (SNS), a Lipschitz-controlled MLP, paired with heavy-tailed (Cauchy) maximum-likelihood estimation, to provide well-behaved predictions and derivatives through contact. Coupled with a predictor–corrector state estimator and a gray-box Generalized Gauss–Newton MPC, the approach achieves robust zero-shot generalization and efficient planning, both in simulation and hardware, under domain randomization and partial observability. The results show substantial robustness and scalability gains, enabling reliable whole-body control across unseen terrains and gaits without task-specific retraining. Overall, the work demonstrates that carefully designed smoothness constraints and robust losses can dramatically improve the compatibility between learned dynamics and model-based planners in legged robotics.

Abstract

Deep learning and model predictive control (MPC) can play complementary roles in legged robotics. However, integrating learned models with online planning remains challenging. When dynamics are learned with neural networks, three key difficulties arise: (1) stiff transitions from contact events may be inherited from the data; (2) additional non-physical local nonsmoothness can occur; and (3) training datasets can induce non-Gaussian model errors due to rapid state changes. We address (1) and (2) by introducing the smooth neural surrogate, a neural network with tunable smoothness designed to provide informative predictions and derivatives for trajectory optimization through contact. To address (3), we train these models using a heavy-tailed likelihood that better matches the empirical error distributions observed in legged-robot dynamics. Together, these design choices substantially improve the reliability, scalability, and generalizability of learned legged MPC. Across zero-shot locomotion tasks of increasing difficulty, smooth neural surrogates with robust learning yield consistent reductions in cumulative cost on simple, well-conditioned behaviors (typically 10-50%), while providing substantially larger gains in regimes where standard neural dynamics often fail outright. In these regimes, smoothing enables reliable execution (from 0/5 to 5/5 success) and produces about 2-50x lower cumulative cost, reflecting orders-of-magnitude absolute improvements in robustness rather than incremental performance gains.

Figures (20)

• Figure 1: Learning model predictive control and state estimation for legged robots. Stiff legged-robot dynamics, when modeled with standard neural networks, introduce nonsmooth behavior that destabilizes learning and MPC. Motivated by smooth approximations of nonsmooth functions and by robust optimization, we introduce a neural architecture and a heavy-tailed likelihood that provides stable learning and informative gradients through contact events. Our smooth neural surrogates learn terrain-varying dynamics and state estimation while generalizing to new tasks at test time with MPC. Together, they combine the representational flexibility of deep learning with the task-level adaptability of MPC, enabling reliable whole-body control across diverse behaviors and environments.
• Figure 2: Our framework for learning generalizable MPC and state estimation. The backbone of our method consists of dynamics and estimator modules with an architecture we introduce called the smooth neural surrogate. The smooth neural surrogate is an MLP provides informative derivatives for planning and state estimation, even in the presence of nonsmooth dynamics and out-of-distribution data. To combat impulse-like residuals in the learned dynamics and estimation, we train both modules via Cauchy maximum-likelihood estimation (MLE). Finally, we introduce a gray-box model predictive controller and model-based state-estimation strategy that exploit the smooth derivatives provided by our learned dynamics.
• Figure 3: Smooth neural surrogates converge under strict Lipschitz constraints and enable zero-shot generalization. Top: Lowest-MSE models for the 2-D shape interpolation task. Only the SNS converges under tight Lipschitz budgets; the Lipschitz MLP liu2022learning collapses. Middle: Evolution of Lipschitz upper bounds and training MSE.
• Figure 4: First- and second-order smooth neural surrogates for the ReLU function. First-order surrogates eliminate large jumps but can form cusp-like shapes if trained for longer periods. Second-order surrogates bound curvature throughout training, yielding smooth transitions similar to analytical smooth surrogates like $\text{softplus}$ or $\text{mish}$.
• Figure 5: Smooth neural surrogates provide informative derivatives throughout training.Top: Learned surrogates for a nonsmooth piecewise function. Standard MLPs develop extremely stiff gradients unsuitable for gradient-based optimization, and Lipschitz MLPs briefly collapse under smoothness constraints. Smooth neural surrogates yield informative derivatives for downstream optimization throughout training. Bottom: Evolution of Lipschitz bounds and MSE during training.