patchpolicy-so101-expert8w288-flow-bidir

Visuomotor policy for the SO-101 arm. Frozen DINOv2 ViT-S/14 trunk with dense patch tokens (Patch Policy), plus a SmolVLA action expert, 8 layers, width 288 action head trained with flow matching, using bidirectional attention among action tokens.

One of 8 arms from a controlled ablation. All 8 share the trunk, dataset, split, seed and optimizer. They differ only in the action head and the action-attention pattern.

Trunk	DINOv2 ViT-S/14, frozen (22,056,576 params, 0 trainable)
Head	SmolVLA action expert, 8 layers, width 288
Objective	flow matching
Action attention	bidirectional
Trainable params	8,551,046
Final train loss	0.010678
Final held-out loss	0.051919

Data

phi_so101_cubes_cylinder_recovery_v1: 143 episodes (120 teleop + 23 learning-from-failure recovery), 81,943 frames, 30 fps, 3 cameras (wrist, front, top). 113 train / 30 held-out episodes.

Shapes

input   3 cameras x 2 timesteps, 240x320 RGB
        -> crop 210x280 -> DINOv2 patch 14 -> 15x20 = 300 tokens x 384 per camera
        -> 900 patch tokens + 1 state token = 901 per frame
        -> memory (B, 1802, 384)
output  action chunk (B, 24, 6)   24 steps x 6 joints, degrees

Predicts 48 steps, returns the first 24 for execution.

Training

seed 1000, batch 64, 29 epochs = 29,493 steps, Adam(0.9, 0.95), lr 1e-4 cosine with 500 warmup steps, weight decay 1e-6, grad clip 10.0, EMA. Random crop at train, center crop at eval.

Usage

The checkpoint is a whole pickled nn.Module, so the model classes must be importable. Code and a runnable infer.py are in the companion repo.

import torch, einops
model = torch.load("model_final.pt", map_location="cpu", weights_only=False).eval()
# obs: (B, 2, 3, 3, 210, 280) float in [0,1], cameras in wrist/front/top order
tokens = encoder(obs)                                    # (B,2,3,300,384)
flat = einops.rearrange(tokens, "N T V P E -> N T (V P) E")
state = torch.zeros(B, 2, 1, 6)                          # normalized joint positions
action, _, _ = model(flat, state, None)                  # (B, 24, 6) in degrees

Camera order is fixed. Each camera owns its own 300-token slice of the memory; swapping them feeds the model the wrong slice.

Results across all 8 arms

Bidirectional action attention won every pair:

pair	causal	bidirectional
ppformer + DDPM	0.037980	0.016703
ppformer + flow	0.067604	0.042783
expert8w288 + flow	0.079359	0.051919
expert16w216 + flow	0.081903	0.052711

Among the three flow-matching arms (comparable, same loss target) the Patch Policy transformer beats both expert configurations, and doubling expert depth did not help.

Limitations

• Single seed (1000). No error bars on any comparison.
• Held-out loss is not task success. SO-101 has no simulator, so a real ranking needs physical rollouts on the arm.
• Do not compare DDPM to flow-matching losses. They minimize different targets; the DDPM arm's lower number does not mean a better policy.