Fruit-picking FlashWAM (100-epoch run)

Public archival bundle for the fruit-picking FlashWAM model trained by Amin — the 100-epoch FROM-SCRATCH rerun (run fruit_pick_flashwam_scratch_100ep, 2026-08-17), which completed all 100 epochs (59,200 optimizer steps). Unlike the earlier partial 30-epoch run published at SleepMastger/fruit-picking-flashwam, this run trained on a 95-episode train split with a 5% validation holdout (592 steps/epoch — the 30-epoch runs saw 630 because they used all 100 episodes; do not mix the step arithmetic).

Architecture: FlashWAM M1 fused-KV, fixed-RoPE, 1-layer action expert, from scratch (Wan2.2 base video expert + randomly initialized action expert). Recipe: 4 GPUs x batch 8 x accum 1 = global batch 32, cosine lr 1e-4, weight decay 1e-2, bf16, checkpoint every 10 epochs.

This bundle contains a selected subset of the run's checkpoints:

  • checkpoints/weights/step_035520.pt (epoch 60) — SHA-256 8c21b694e81461a7de83c567cbc9fbb11842216d8a46a6d06894e00d987478e4
  • checkpoints/weights/step_053280.pt (epoch 90) — SHA-256 fae96dee3c2771697a4563ebbd7003cd00888d3c83951a10ab48dcfc8c611758

The remaining every-10-epoch checkpoints exist on the training cluster; this repo intentionally publishes the pre-overfit candidates chosen from the validation curve. The model has not been evaluated on a robot, and the physical meaning of the -1/+1 gripper polarity still needs confirmation.

Conditioning

Exact task text:

Lift the lid, put it aside, and pick the black plum.

conditioning/text_embedding.pt is the exact cached T5 embedding consumed during training (Wan text stack, context length 128, Wan-AI/Wan2.1-T2V-1.3B tokenizer reference). The resolved training config sets load_text_encoder: false, so this cached tensor is part of the required inference bundle.

Input processing and normalization

  • Two 256x256 RGB cameras (agentview, then wrist), resized to 224x224 and concatenated horizontally to 224x448.
  • Horizon: 33 observations; 32 action transitions at 10 Hz.
  • Original 15-D state was converted to 8-D: eef_xyz(3) + quat-to-axis-angle(3) + [gripper_width/2, -gripper_width/2].
  • Action is 7-D: delta XYZ, delta rotation XYZ, and gripper; the gripper channel is absolute rather than delta.
  • dataset_stats.json contains the exact min/max normalization statistics used by this run (computed on the 95-episode train split).

Attention masks

  • video_attention_mask_mode: first_frame_causal
  • action_group_causal_mask_mode: group_diagonal
  • Text cross-attention is enabled for the action expert.

The exact implementations, configs, split arithmetic, and launch script are included under training_code/; the resolved config is config.yaml. No license is asserted here for the bundled upstream code; its original terms continue to apply.

Base components

This weights-only checkpoint is not standalone. It references Wan-AI/Wan2.2-TI2V-5B and requires the matching Wan VAE plus the included FastWAM code/configuration. PyTorch .pt files may contain pickled objects; load only in a trusted environment.

Downloads last month
4
Video Preview
loading