fastwam-dmd-1step

FastWAM-AR · DMD Self-Forcing · 1-step LoRA

LoRA adapter that distills an autoregressive dual-stream FastWAM (video + action world-action model) down to 1-step generation via DMD (Distribution Matching Distillation) + Self-Forcing. Evaluated on LIBERO (spatial / object / goal / long).

• Streams: video (5B DiT) + action (1B DiT), joint MoT, group-internal bidirectional, group-to-group AR.
• Distillation: video → DMD (real/fake score = the bidirectional FastWAM release); action → GT behaviour-cloning on the Self-Forcing rollout. LoRA on both experts.
• LoRA config: r=64, alpha=64, dropout=0, targets self_attn.{q,k,v,o}, cross_attn.{q,k,v,o}, ffn.0, ffn.2.
• Result: ~96% LIBERO success at 1 step; ~0.3–0.4 s model inference per control chunk.

⚠️ This is a LoRA delta, not a full model

It must be folded into its base checkpoint to be usable: base = the DF-trained AR checkpoint (FastWAMAR), itself warm-started from the FastWAM release libero_uncond_2cam224.pt. You need that base separately.

Usage

# 1) download the LoRA adapter
hf download ElysiaTrue/fastwam-dmd-1step dmd_lora_step005500.pt --local-dir ./lora

# 2) fold it into the AR base -> a standalone eval-loadable checkpoint
PYTHONPATH=src python scripts/distill/merge_ar_lora.py \
    task=libero_ar_2cam224_dmd_sf_4step method=dmd \
    init_ckpt=<your_AR_base>.pt \
    +lora_ckpt=./lora/dmd_lora_step005500.pt \
    +merged_out=checkpoints/distilled/ar_dmd_sf_1step_merged.pt

# 3) eval at 1 step (streaming AR); 1-step deploys fastest with the naive path
CKPT=checkpoints/distilled/ar_dmd_sf_1step_merged.pt \
STEPS=1 NUM_GPUS=4 NUM_TRIALS=5 USE_CACHE=0 CONCURRENT=0 \
RUN_NAME=DMD_SF_1step bash eval_DF_full.sh

The adapter key names carry tag gen (video: gen_video, action: gen_action); the merge script applies the same spec and merge_and_unloads them into the base.

Files

• dmd_lora_step005500.pt — {"student_lora": {name: tensor}, "step": 5500}; the trainable LoRA params of the AR generator (video + action experts).