dk1-pretrain-40d-30hz-cart6d-synthetic

Fast-WAM (world–action model) pretrain for the DK-1 bimanual robot. A Wan2.2-TI2V-5B video backbone (sparsely sliced) plus a Mixture-of-Transformers (MoT) action expert, trained jointly with flow matching to predict both future video and a chunk of Cartesian end-effector actions.

Two things distinguish this model:

1. 6-D rotation action representation (Zhou et al. 2019 continuous rotation, as used by NVIDIA Cosmos) instead of 3-D axis-angle.
2. Synthetic position-control targets. The training data is recorded under impedance control, where the commanded action is offset from the measured observation.state by the controller's compliance (the state never catches up — in contact the command sits "inside" the object to make force). Training a position policy on those commands and replaying them on a stiff position controller overshoots. Instead the action is relabeled as the future achieved state — synthetic_action[t] := observation.state[t+1] — i.e. "go to the next achieved pose", which a stiff controller reproduces without the compliance offset.

Checkpoint

File	Contents
step_100000/model.pt	final model weights (bf16)
step_100000/ema.pt	EMA weights (bf16, ema_s=0.10) — usually the better choice for inference

Each file is self-describing: it embeds config, norm_stats (action-normalization statistics) and norm_config. ema.pt carries both model_state_dict (the live weights) and ema_state_dict (the EMA weights). Converted to bf16 from the fp32 training checkpoint.

Action chunk representation

• Head: cartesian_rot6d_multistep_rel — bimanual end-effector relative pose deltas with a 6-D rotation encoding. Trained isolated (single head).
• Width / horizon: action_dim = 20, action_horizon H = 30 (≈ 1.0 s chunks at 30 Hz).
• Per-step 20-D layout: [ L_Δpos(3), L_Δrot6d(6), L_gripper(1), R_Δpos(3), R_Δrot6d(6), R_gripper(1) ] (left arm, then right arm).
• 6-D rotation: the first two rows of the relative rotation matrix R = R_state⁻¹ R_action; decode with Gram-Schmidt. The identity rotation maps to [1, 0, 0, 0, 1, 0].
• Reference frame: every step's target is relative to the current EE frame (chunk start): Δpos = R_stateᵀ·(p_target − p_state), with the rotation taken in that frame. End-effector = tool0 flange (link6-7 origin, zero lever), from dk1-dual-arm-urdf.
• Synthetic target source: the chunk targets are the future achieved states observation.state[t+1 … t+H] (each pose via forward kinematics), expressed relative to the current EE frame — not the recorded impedance command. This removes the controller's steady-state compliance offset.
• Normalization (norm_stats, embedded): the 18 pos+rot6d delta dims are identity-centered — the rotation block has its [1,0,0,0,1,0] offset subtracted so small rotations map to ≈ 0 — then percentile-scaled (q01/q99) and clipped to ±5. The 2 gripper dims are min-max mapped to [−1, 1].

  Note: because of identity-centering plus the small relative deltas, the normalized rot6d targets cluster near 0 (decode adds the identity back).

Backbone & model

• Backbone: Wan2.2-TI2V-5B, sparse 15-of-30 layer slice (keep_layers = [0,1,2,4,7,10,12,15,18,21,23,26,27,28,29]), crossattn_dim 4096.
• Action expert: MoT, sharing the backbone transformer blocks.
• State encoder: conv1d over [pos, vel, torque], history_len 3, state_dim 14. Anti-proprio-shortcut: whole-sample state_dropout 0.5, state_noise_std 0.02 (no action channel as input).
• Auxiliary: progress head (predict_progress).
• Vision: cameras head, left_wrist, right_wrist (top_split grid), image 384×320; 5 observation + 8 future frames, temporal_downsample 4.

Training

• From scratch, 100,000 steps. DDP on 2× H100, bf16, gradient checkpointing, torch.compile.
• Optimizer: 8-bit AdamW, lr 5e-5, 2000-step warmup + cosine decay, weight_decay 0.01, grad_clip 1.0.
• Effective batch 64 (batch_size 16 × 2 ranks × 2 grad-accum).
• Flow matching: shifted_uniform timestep schedule, timestep_shift 4.0, timestep weighting on.
• RTC (real-time chunking) training: rtc_prob 0.1, rtc_max_prefix 16, rtc_decay 0.3.
• Loss weights: video 1.0, gripper 2.0, progress 0.1. ema_s 0.10.
• Final eval: cart6d action-eval L1 0.032 (L2 0.056), RTC suffix-L1 ≈ 0.03–0.05 across prefix lengths.

Data

20-dataset blend, ≈ 5.23 M frames @ 30 Hz:

• DK-1 contact-rich teleop: black/white swan, cutlery basket, duplo (in-box / sorting-by-size / disassembly / simple stacking).
• dk1-merge-2026-03 (large pos-only teleop merge).
• robotwin_dk1_blocks (RoboTwin sim, blocks stacking/sorting, 30 Hz).

Action-normalization stats were precomputed over exactly this blend at H=30, with the synthetic (state-shifted) target distribution.

Code

open-thought/fastwam — config configs/wan2_5b_dk1_pretrain_40d_30hz_cart6d_synthetic_h30.yaml.