Push-T FlashWAM (from scratch, 30 epochs)

Public archival bundle for a FROM-SCRATCH run on the real-hardware Push-T task: push an orange T-shaped block across a table until it aligns with a pink T outline, on a Franka Emika Panda. The gripper is clamped shut on a marker pen that acts as a single-point pusher, so this is the classic Push-T task run on real hardware rather than in simulation.

Trained on [SleepMastger/pusht-manipulation] (https://huggingface.co/datasets/SleepMastger/pusht-manipulation) — 100 human teleoperation demonstrations, 32,131 frames at 10 Hz, converted to LeRobot v2.1. Companion model: SleepMastger/pusht-fastwam — identical data and recipe, different architecture, so the two are directly comparable.

Architecture: FlashWAM M1 decoupled MoT: fused-KV KV source, fixed action-aligned video RoPE, 1-layer action expert (randomly initialised — fused-KV modes cannot use layer-selective init), 30-layer Wan2.2 video expert.

Training

init from scratch (resume: null)
GPUs 4 x H200
batch 8 per GPU x 4 x accum 1 = global 32
schedule cosine, lr 1e-4, weight decay 1e-2, bf16
epochs 30
steps 1,005/epoch -> 30,150 total
wall clock 8h47m
validation none (val_set_proportion: 0.0) — all 100 episodes are train

Completed all 30 epochs. Checkpoints were saved every 5 epochs:

  • checkpoints/weights/step_005025.pt (epoch 5) — SHA-256 7732ce6db1e2199f67b6b587f027c359c8c0bebefb9bf6f34394e5e76e5a5720
  • checkpoints/weights/step_010050.pt (epoch 10) — SHA-256 4ddef751a683df6b073b3c8e2818f17a03d4e4c502b80d3e2ba167a587ba23b7
  • checkpoints/weights/step_015075.pt (epoch 15) — SHA-256 2a9f762196f66ece1cb2fe6d399912428c49a32d9cd99e234f7a079dfff630c9
  • checkpoints/weights/step_020100.pt (epoch 20) — SHA-256 f13b996c9ce49f1abf6d48db7a6f678c658866a7caf84592ca66bbfb1878f2b5
  • checkpoints/weights/step_025125.pt (epoch 25) — SHA-256 a76af418ff3bdd83547694f8551c30be20d1c43a6f5ac5cdd75d88a5f5598815
  • checkpoints/weights/step_030150.pt (epoch 30) — SHA-256 61829d8bc1cda5ad3e59ed4281d8f76495d322fc928c4786d61491a2410ad440

The model has not been evaluated on a robot.

Conditioning

Exact task text (must byte-match at deployment):

push the T block to the target outline

No trailing period. conditioning/text_embedding.pt is the exact cached T5 embedding consumed during training (Wan text stack, context length 128); its filename hash in the training cache is 71cb088954da46da4e3cb6c6f73ac812690e9bccf886c7858377c61918b4e3e8. The resolved 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 with raw LeRobot keys image / wrist_image, resized to 224x224 and concatenated horizontally to 224x448.
  • Dataset is LeRobot v2.1 (per-episode parquet + AV1 mp4), staged node-locally by training_code/stage_pusht_local.sh.
  • Horizon: 33 observations; 32 action transitions at 10 Hz.
  • State is 8-D: eef_pos(3) + quat2axisangle(eef_quat)(3) + [gripper_width/2, -gripper_width/2].
  • Action is 7-D: delta XYZ (metres per 100 ms step), delta rotation XYZ, grasp. Raw grasp uses the robosuite {-1 open, +1 close} convention and is remapped g -> (1-g)/2 to the LIBERO/RLDS {1 open, 0 close} convention at conversion.
  • dataset_stats.json holds the exact min/max statistics used by this run.

⚠️ Six of the fifteen channels are constant — read before deploying

The teleoperation rig commanded translation only, and the pen stayed clamped for the whole task. Measured across all 100 episodes / 32,131 frames:

field dims raw range
action dx, dy, dz 3.5e-02 … 3.6e-02 carries the signal
action drx, dry, drz 0.0 constant
action grasp 0.0 (raw +1 -> 0.0) constant, gripper never opens
state x, y, z, rx, ry, rz 8.9e-02 … 4.4e-01 carries the signal
state gripL, gripR 3.2e-05 constant

This does not produce NaNs. SingleFieldLinearNormalizer (training_code/normalizer.py) has ignore_dim = input_range < range_tol with range_tol = 1e-4, which catches all six — the tightest, the gripper state dims at 3.2e-05, clears it by about 3x. An ignored dim gets scale = 1.0 and offset = -min, i.e. it normalizes to x - min: exactly 0 for the four action dims whose raw range is exactly 0.0, and [0, 3.2e-05] for the two gripper state dims. Finite and negligible beside the ±1 live channels either way.

Consequences for deployment:

  1. Only dx, dy, dz are meaningful outputs. Dims 3-6 were constant in training, so the model has learned to emit their constant value; denormalize as usual and they come back as the constants. Do not interpret them as commanded rotation or grasp.
  2. The gripper must stay closed. The policy never saw it open and cannot command it to. Clamp the pusher before handing control to the policy.
  3. Do not lower range_tol below 3.2e-05 — the gripper state dims would then be treated as live and divide by ~0.

The full 7-D action / 8-D proprio space was kept deliberately, so the architecture is byte-identical to the dish-utensil and fruit-picking runs and the results stay comparable.

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, dataset converter, staging and launch scripts are included under training_code/; the resolved config is config.yaml. training_code/preflight.py is the CPU-only check that verifies the tensor contract and the degenerate-channel behaviour described above. 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
6
Video Preview
loading