Source-linked AI summary

WAM-OPD: On-Policy Distillation for World Action Models

Liuhaichen Yang, Zhuang Jiang, Chenchao Sheng, Zezhi Tang

arXiv:2608.22364v1cs.AIcs.RO

TL;DR

Accelerated video-first WAM students may lose task capabilities and encounter Student-induced histories not covered by offline data, motivating OPD without sparse-reward RL. WAM-OPD uses frozen-Teacher labels on Student histories while training actions under the Student’s own video plan. In two preliminary RoboTwin 2.0 tasks, Flash-WAM success rose substantially, but the evidence is limited to a narrow two-task study.

  • Problem

    Accelerated WAM distillation can lose capabilities, while offline data may poorly represent histories generated by the deployed Student.

  • Method

    WAM-OPD trains lightweight shared-backbone adapters with joint video–action and flow-matching supervision, using frozen-Teacher labels on Student histories and the Student’s own video plan.

  • Results

    58.3% versus 0.0% success on HANDOVER MIC and 33.3% versus 16.7% on PUT OBJECT CABINET were reported for released one-video/one-action-step Flash-WAM.

  • Takeaways & Limitations

    Dense Teacher supervision on Student-induced histories is a promising post-training interface for video-first WAMs within the tested scope.

  • Takeaways & Limitations

    Evidence covers only two clean simulation tasks and a small held-out set, with no evidence of real-robot transfer or broad generalization.

Abstract

from arXiv · show

World action models (WAMs) couple visual future prediction with robot action generation, but accelerated students can lose task capabilities during distillation and later encounter states that are poorly represented by offline data. We study whether on-policy distillation (OPD) can repair such a student without requiring sparse-reward reinforcement learning. We introduce WAM-OPD, a deployment-consistent post-training recipe for a video-first WAM. The student acts in the environment and therefore determines the history distribution. A frozen teacher labels those student histories with coherent video and action targets, while the student action branch is trained under its own generated video plan, as it is at deployment. Joint video and action losses update lightweight adapters in the shared backbone, together with an action flow-matching regularizer. In preliminary RoboTwin 2.0 studies on two tasks, the released one-video/one-action-step Flash-WAM improves from 0.0% to 58.3% success on HANDOVER MIC, and from 16.7% to 33.3% on PUT OBJECT CABINET. These task-specific results are an initial capability proof rather than evidence of broad or uniform generalization. They nevertheless suggest that dense teacher supervision on student-induced histories is a promising post-training interface for video-first WAMs.

1. Introduction

WAM-OPD addresses capability loss and distribution mismatch in accelerated video-first WAMs by supervising Student-induced histories while training the deployed Student conditioning path. Preliminary two-task RoboTwin results show improved held-out success, but remain a narrow capability proof.

  • Motivation: Accelerated distillation can leave WAM students with capability gaps and expose them to deployment histories poorly represented by offline data.Offline imitation and distillation use fixed data distributions, whereas Student decisions determine later closed-loop states.
  • Motivation: Video-first WAM OPD faces both Student-induced history shift and a Student-plan/Teacher-plan conditioning mismatch.The deployed action branch consumes the Student-generated video plan, not the Teacher plan.
  • Method: WAM-OPD labels Student-generated histories with a frozen Teacher while updating a shared WAM through joint video–action supervision and flow matching.The method is designed as deployment-consistent post-training rather than sparse-reward reinforcement learning.
  • Results: 58.3% success on HANDOVER MIC, up from 0.0%, and 33.3% on PUT OBJECT CABINET, up from 16.7%, were obtained in preliminary RoboTwin 2.0 studies.The evaluation used released one-video/one-action-step Flash-WAM and exact-paired held-out scene-seed runs.
  • Scope: The contribution is an initial capability proof that leaves broad or uniform generalization unestablished.The two tasks show different improvement magnitudes, so the evidence still required for a general claim remains substantial.

2. Related Work

Related work spans generalist robot policies, predictive world-action models, few-step generative acceleration, and interactive post-training. WAM-OPD targets an unaddressed video-first interface combining Student occupancy with Student-plan-conditioned action learning.

  • Predictive robot policies: Predictive robot models make future visual states explicit and increasingly couple visual prediction with robot control inside unified generative models.WAMs use visual prediction as an intended state-change representation while inverse dynamics grounds it in control.
  • Few-step acceleration: Progressive distillation and consistency methods reduce iterative generative inference cost, while Flash-WAM adapts consistency distillation to asymmetric video and action noise regimes.These methods accelerate the solver but do not determine which closed-loop histories the accelerated Student encounters.
  • Interactive post-training: Interactive post-training methods address occupancy mismatch by querying supervision on learner- or policy-induced states, including environment rollouts and continuous generator states.This shared occupancy principle differs from sparse-reward policy-gradient reinforcement learning.
  • Gap and contribution: Existing OPD methods do not define the video-first WAM case where actions condition on a separately generated Student video plan and modalities share trainable blocks.WAM-OPD therefore uses a WAM-specific deployment-consistent objective rather than token reverse-KL or a reproduction of full pathwise matching.

3. Preliminaries

A video-first WAM hierarchically generates a future video plan before an action chunk, with both modalities interacting through a shared Transformer. Flow matching and on-policy distillation provide the generative and occupancy foundations for WAM-OPD.

  • World Action Models: A WAM policy factorizes as pθ(zt, at | ht) = pθ(zt | ht) pθ(at | ht, zt), generating a video plan before an action chunk from closed-loop history.Here ht contains observation–action history and language, zt is the future video-plan latent, and at is the action chunk.
  • World Action Models: Although outputs are factorized, LingBot-VA and Flash-WAM use modality-specific modules around one shared 30-block Transformer.Flash-WAM distills the generative solver without changing the backbone architecture.
  • Flow matching: Flow matching learns a time-dependent neural vector field by regressing to the velocity of a prescribed probability path between clean data and noise.The Flash-WAM convention uses xσ = (1 −σ)x0 + σϵ and integrates generation from σ = 1 toward σ = 0.
  • Few-step inference: Consistency distillation accelerates the solver by mapping different points on a Teacher trajectory to a common clean endpoint, with video and action losses optimized jointly.Flash-WAM uses different consistency parameterizations for high-noise video and low-noise action streams.
  • On-policy distillation: On-policy distillation evaluates a frozen Teacher on contexts generated by the Student and minimizes a dense discrepancy on the Student-induced occupancy.The sampled trajectory is normally treated as data rather than differentiated through.
  • WAM-OPD: WAM-OPD applies this principle to environment histories while supervising the Student’s video and action computation under deployment-consistent conditioning.The overview combines Student action, frozen-Teacher labels, Student video plans, and shared-block adapter updates.

4. Method

WAM-OPD collects histories from the released Student, labels them with a frozen Teacher, and trains the Student under its deployment-time video-to-action computation. Joint endpoint losses and action flow matching update shared adapters, while the fixed trajectory package becomes stale as the Student changes.

  • Deployment-consistent distillation: The trainable Student is run in inference order on each released-Student history, generating a video plan before predicting actions from that plan.The action branch receives the Student plan through stop-gradient.
  • Teacher supervision: The frozen Teacher labels the same Student histories with coherent video and action targets, although the action target currently uses the Teacher plan.Video supervision also reduces the conditional gap between Teacher-plan targets and the deployed Student-plan action branch.
  • Joint objective: The objective combines video loss, action loss, and an action flow-matching term weighted (λz, λa, λFM) = (1, 1, 0.2).The flow-matching term retains a high-noise-boundary action signal; the objective is not exact reverse-KL or full trajectory matching.
  • Parameter-efficient update: Rank-8 JointLoRA adapters are inserted across 30 shared Transformer blocks while released weights remain frozen.This shared scope lets both modalities modify representations used by the video-first policy, without guaranteeing nonconflicting gradients.
  • Training protocol: The proof of concept reuses teacher-labeled trajectories for three epochs, remaining on-policy only for the checkpoint that collected them.A complete iterative version would alternate fresh Student collection, Teacher labeling, and bounded optimization.

5. Experiments

The study asks whether deployment-consistent distillation can convert dense Teacher supervision into closed-loop success for an accelerated WAM Student. It evaluates two clean RoboTwin 2.0 tasks using a fixed training and exact-paired held-out protocol, while reserving broader comparisons for future work.

  • Research question and tasks: The evaluation targets HANDOVER MIC and PUT OBJECT CABINET in RoboTwin 2.0's clean setting.HANDOVER MIC tests microphone transfer between robot hands, while PUT OBJECT CABINET tests placing and releasing an object in a designated drawer.
  • Checkpoint selection: Checkpoints are selected on a disjoint split using predeclared ordering by task success, semantic progress, and calibration loss.The rule selects epoch 3 for both tasks, and screening data are excluded from held-out results.
  • Exact-paired evaluation: Held-out evaluation uses six scene seeds and two fixed noise banks per task, with Released and WAM-OPD exactly paired within each unit.The primary outcome is RoboTwin's latched eval_success; the 12 units per task are not independent scene samples.
  • Results: The same WAM-OPD recipe improves held-out success on both tasks, with a larger gain on HANDOVER MIC than PUT OBJECT CABINET.The task-specific effect sizes motivate broader evidence rather than pooled reporting.
  • Planned evidence: Future evaluation will test multi-task performance, matched baselines, controlled ablations, multiple horizons, randomized settings, and fresh on-policy recollection.Table 2 entries marked TBD are planned evidence, not experimental results.

6. Discussion and Limitations

The two-task studies indicate that joint video-action supervision on Student-induced histories can recover capability in an accelerated WAM, but the evidence remains task-dependent and bounded by several methodological limitations.

  • Supported conclusion: Joint video-action supervision on Student-induced histories can recover task capability in an accelerated WAM.The conclusion is supported by two task studies and is explicitly limited in scope.
  • Evaluation scope: The evidence covers only two clean simulation tasks and a small held-out set, with no representative multi-task average or real-robot transfer evidence.The current sample is therefore insufficient for a general performance claim.
  • On-policy assumption: A fixed trajectory package is on-policy only for the released collector and becomes stale after Student updates.Formal on-policy training after updates requires iterative recollection.
  • Plan mismatch: Teacher action targets use the Teacher video plan, so video alignment reduces but does not eliminate the Teacher-plan/Student-plan mismatch.This leaves a deployment-relevant conditional interface gap.
  • Missing controls: The study does not measure shared video/action gradient interactions or compare compute-matched SFT, RL, action-only distillation, and full pathwise flow supervision.These omissions constrain attribution of the observed gains.
  • Success interpretation: Official latched eval_success is claimed, but restored snapshots do not reliably preserve contact manifolds for auxiliary contact diagnostics.The supported claim is task success and paired semantic progress, not verified stable contact mechanics.

7. Conclusion

WAM-OPD applies on-policy distillation to video-first World Action Models by training actions under the Student's own video plan and using a frozen Teacher on Student-induced histories. Preliminary exact-paired gains on two tasks form a promising vertical slice, not a finished general method.

  • Contribution: WAM-OPD is a preliminary post-training framework for video-first World Action Models.Its central design couples Student-induced histories with deployment-consistent video-to-action computation.
  • Core design: The action branch is trained under the Student's own video plan while a frozen Teacher supplies coherent labels on Student-induced histories.A parameter-efficient joint update adapts the shared computation.
  • Results: 0.0% to 58.3% exact-paired success is reported on HANDOVER MIC, and 16.7% to 33.3% on PUT OBJECT CABINET.These are the two-task results of the released one-video/one-action-step Flash-WAM after WAM-OPD.
  • Next steps: The next stage is to broaden tasks, refresh Student occupancy between updates, and separate video, action, flow-matching, and shared-adaptation contributions.The conclusion explicitly frames the current evidence as a two-task vertical slice.
Loading 2608.22364v1…