WAM-OPD: Sharpening World Action Models via On-Policy Distillation

2026-09-28

AI Summary

This study addresses the challenges of skill degradation during task adaptation and the prohibitive data collection costs associated with online distillation for pretrained world-action models. To this end, we propose Prefix-Weighted Trajectory Replay (PWTR), a method that integrates denoising path generation with a surrogate importance-weight-based loss reweighting mechanism. By leveraging only an initial student policy and a small set of teacher trajectories, PWTR enables efficient online policy distillation over a fixed trajectory buffer without requiring additional environment interactions. Experiments in both simulation and real-world robotic settings demonstrate that the proposed approach significantly enhances performance on target tasks while effectively preserving the model’s original capabilities on unadapted tasks, thereby achieving low-cost knowledge transfer and task adaptation.

Abstract

Pretrained world action models (WAMs) provide generalist capabilities across diverse robotic manipulation tasks, yet improving target-task performance to an expert level without degrading pretrained skills remains challenging. We explore on-policy distillation (OPD) for WAMs and introduce WAM-OPD. WAM-OPD inherits the advantage of OPD methods that transfer task-specific teacher knowledge under the student's own induced distribution, rather than directly fitting the student to a narrow task-specific data distribution. However, in closed-loop manipulation, the observation histories change as the student policy evolves, requiring fresh environment rollouts to remain on-policy. Applying OPD to WAMs entails repeated data collection, which is costly even in simulation and often impractical on real robots. To avoid repeated environment rollouts during distillation, we introduce prefix-weighted trajectory replay (PWTR). PWTR uses a fixed trajectory pool composed primarily of initial-student rollouts, supplemented with task-specific teacher rollouts to broaden trajectory coverage. For each trajectory replayed from this pool, PWTR conditions the current policy on successive stored histories to generate fresh denoising paths, along which the task-specific teacher provides supervision. Although these denoising paths are refreshed as the policy evolves, the replayed environment trajectories remain fixed. PWTR therefore reweights per-decision distillation losses using proxy importance weights derived from path scores accumulated over the trajectory prefix preceding each decision to mitigate the resulting shift in the history distribution. Simulated and real-world experiments demonstrate task adaptation without additional environment interaction during distillation. In both settings, WAM-OPD improves target-task performance while retaining near-initial performance on tasks excluded from adaptation.

Problem

• World Action Models
• On-Policy Distillation
• Robotic Manipulation
• Task Adaptation
• Catastrophic Forgetting

Innovation

• On-Policy Distillation
• World Action Models
• Prefix-Weighted Trajectory Replay
• Robotic Manipulation
• Importance Weighting