armanakbari4/fastwam-robotwin-idm-50k

FastWAM-IDM RoboTwin — clean 50k reproduction (step_050000)

Fast-WAM IDM variant (create_fastwam_idm: FastWAMIDM(FastWAMJoint), inverse-dynamics style — action attends to ALL video latent tokens and teacher-forces the video conditioning when denoising actions), built on Wan2.2-TI2V-5B video DiT + ActionDiT (MoT, ~6B params), trained on RoboTwin 2.0.

Clean 50,000-step run on 8 GPUs (effective batch 128), lr 1e-4 cosine (warmup 5%), AdamW(0.9,0.95), weight_decay 1e-2, bf16, ZeRO-1. 3 cams (cam_high + L/R wrist, 240x320 each, tiled), num_frames 33, action_video_freq_ratio 4 (32 actions / 9 video frames), action&state dim 14.

This matches the BadWAM (arXiv:2607.15207) reproduction recipe (8xH100 / 50k / batch16 / lr1e-4 cosine), whose reported clean RoboTwin success rates are idm 91.4% / joint 90.9% / action-only 92.1%.

Notes:

• Final: step 50,000/50,000, loss 0.0695, LR at the 1e-6 cosine floor; resume-safe closed-form LambdaLR schedule (LR is a pure function of global_step).
• Run resumed once (exact / bit-exact) from step_015000 mid-way — full DeepSpeed state restore (Adam moments + per-rank RNG + scheduler), verified no loss spike and LR back on the cosine curve.
• Companion joint 50k checkpoint: https://huggingface.co/armanakbari4/fastwam-robotwin-joint-50k
• Files: fastwam_idm_step50000.pt (12 GB bf16 weights) + robotwin_idm_dataset_stats.json (normalization stats, REQUIRED for inference/eval).