StevenZhou90/openvla-traj-head-dagger-r2

OpenVLA Trajectory Head DAgger R2

This repository contains the DAgger round-2 trajectory draft head used for the OpenVLA + SimplerEnv fast trajectory policy.

Files

• best.pt: TinyTrajectoryHead checkpoint
• metrics.json: training metrics for the selected checkpoint

Benchmark Snapshot

Matched SimplerEnv benchmark matrix:

• tasks: vertical/horizontal/standing coke-can
• x positions: -0.3500, -0.2925, -0.2350
• obj_init_y = -0.02
• 3 episodes per setting, 27 episodes total
• hardware: 1x A100-SXM4-40GB, bfloat16

Decoder	Success	Avg ms/step	Speedup
Baseline OpenVLA	14/27 (51.9%)	295.2	1.00x
Adaptive fast policy (gate v=5, h=6, s=6)	16/27 (59.3%)	105.0	2.81x

Usage

Download best.pt to:

checkpoints/traj_head_dagger_r2/best.pt

Then run:

uv run python scripts/run_openvla_sim.py \
  --decoder trajectory-spec \
  --trajectory-head-checkpoint checkpoints/traj_head_dagger_r2/best.pt \
  --trajectory-fast-draft-only \
  --trajectory-head-threshold 0.2 \
  --trajectory-fast-min-confident-tokens 5 \
  --task google_robot_pick_vertical_coke_can \
  --published-eval-setup \
  --episodes 3 \
  --steps 80

For full reproduction from data generation through benchmarking, use:

./scripts/reproduce_readme_speedup_local.sh