Flexible Locomotion Learning with Diffusion Model Predictive Control

TL;DR

Diffusion-MPC addresses the need for adaptable legged locomotion by treating a diffusion model as a trajectory prior and integrating it into a receding-horizon planner. Trajectories are sampled from the prior and steered with reward-based planning and constraint projection, with a compositional reward form R(tau)=alpha_nn R_nn(tau)+alpha_an R_an(tau), and updates applied via pi(tau) ~ p_theta(tau) exp(lambda R(tau)). The planner is trained online with reward-weighted denoising to adapt beyond demonstrations, and test-time compositional control enables flexible behavior without retraining. Real-time performance is achieved through asynchronous planning and early-step caching, and deployment on a Unitree Go2 demonstrates robust adaptation across terrains, energy efficiency, and joint constraints. Overall, diffusion priors offer a practical, data-driven approach to combining learning and planning for adaptable, general-purpose embodied control.

Abstract

Legged locomotion demands controllers that are both robust and adaptable, while remaining compatible with task and safety considerations. However, model-free reinforcement learning (RL) methods often yield a fixed policy that can be difficult to adapt to new behaviors at test time. In contrast, Model Predictive Control (MPC) provides a natural approach to flexible behavior synthesis by incorporating different objectives and constraints directly into its optimization process. However, classical MPC relies on accurate dynamics models, which are often difficult to obtain in complex environments and typically require simplifying assumptions. We present Diffusion-MPC, which leverages a learned generative diffusion model as an approximate dynamics prior for planning, enabling flexible test-time adaptation through reward and constraint based optimization. Diffusion-MPC jointly predicts future states and actions; at each reverse step, we incorporate reward planning and impose constraint projection, yielding trajectories that satisfy task objectives while remaining within physical limits. To obtain a planning model that adapts beyond imitation pretraining, we introduce an interactive training algorithm for diffusion based planner: we execute our reward-and-constraint planner in environment, then filter and reweight the collected trajectories by their realized returns before updating the denoiser. Our design enables strong test-time adaptability, allowing the planner to adjust to new reward specifications without retraining. We validate Diffusion-MPC on real world, demonstrating strong locomotion and flexible adaptation.

Figures (8)

• Figure 2: Left: Calf position under different planned joint targets. Right: SoC during the long-distance real-world evaluation.
• Figure 3: Comparison of behavior patterns with different joint position targets. From left to right: negative calf target, original planner, and positive calf target.
• Figure 4: Height transition experiment.
• Figure 5: Balancing under external disturbance. The robot is subjected to a lateral push at the trunk and subsequently recovers from the perturbed posture, reestablishing balance.
• Figure 6: Zero-shot walking. Left: grass. Right: grassy slope.