You can deploy this model and start earning money today!

Parv-09/
patchpolicy-so101-expert16w216-flow-bidir

Name: patchpolicy-so101-expert16w216-flow-bidir
Availability: InStock

Views

No views yet

patchpolicy-so101-expert16w216-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, 16 layers, width 216 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, 16 layers, width 216
Objective	flow matching
Action attention	bidirectional
Trainable params	11,765,190
Final train loss	0.010001
Final held-out loss	0.052711

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.

python

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.

Parv-09/patchpolicy-so101-expert16w216-flow-bidir