End-to-End Multi-Task Policy Learning from NMPC for Quadruped Locomotion

TL;DR

This work addresses the challenge of efficient quadruped locomotion by replacing computation-heavy NMPC with an end-to-end multi-task imitation policy learned from NMPC demonstrations. A single neural network with a shared trunk and gait-specific heads maps raw proprioceptive inputs to joint targets under a PD control framework, enabling multiple gaits without explicit state estimation. Experiments in PyBullet and on the Go1 robot show the MTL policy closely reproduces NMPC trajectories, supports seamless gait switching, and improves predictive accuracy relative to a single-task baseline. While effective for the trained gaits, generalization to unseen locomotion patterns remains open, motivating future work in meta-learning and perception-enabled planning.

Abstract

Quadruped robots excel in traversing complex, unstructured environments where wheeled robots often fail. However, enabling efficient and adaptable locomotion remains challenging due to the quadrupeds' nonlinear dynamics, high degrees of freedom, and the computational demands of real-time control. Optimization-based controllers, such as Nonlinear Model Predictive Control (NMPC), have shown strong performance, but their reliance on accurate state estimation and high computational overhead makes deployment in real-world settings challenging. In this work, we present a Multi-Task Learning (MTL) framework in which expert NMPC demonstrations are used to train a single neural network to predict actions for multiple locomotion behaviors directly from raw proprioceptive sensor inputs. We evaluate our approach extensively on the quadruped robot Go1, both in simulation and on real hardware, demonstrating that it accurately reproduces expert behavior, allows smooth gait switching, and simplifies the control pipeline for real-time deployment. Our MTL architecture enables learning diverse gaits within a unified policy, achieving high $R^{2}$ scores for predicted joint targets across all tasks.

Figures (8)

• Figure 1: Overview of our end-to-end multi-task learning framework for quadrupedal locomotion. A single neural network learns multiple gait-specific policies from an expert NMPC via imitation learning. The shared policy maps proprioceptive observations to joint targets. Red markers indicate foot contact points, highlighting the contact sequence for each gait.
• Figure 2: Multi-Task End-to-End Neural Network Architecture. Raw proprioceptive inputs are processed through two shared hidden layers to learn common locomotion features, which are then branched into three task-specific heads—trot, jump, and bound, each predicting desired joint position targets.
• Figure 3: Simulation snapshots of NMPC expert behaviors in PyBullet. Red markers indicate foot contact points, highlighting the contact sequence for each gait.
• Figure 4: Baseline single‐task neural network architecture. Raw proprioceptive inputs are processed through three fully connected hidden layers to predict the desired joint position targets.
• Figure 5: Validation loss curves comparing the baseline approach with our MTL method across locomotion tasks: trot, jump, and bound. Each subplot shows the validation loss over training epochs.