ur5e_pi05_real_oa_3ep
π₀.₅ with obstacle-attention only (no target/destination roles), fine-tuned
for 3 epochs on annotated UR5e real-world data. Warm-started from
mahgoobi/ur5e_pi05_all_3cam_20k.
- Supervise layers 0, 1, 2; inject into 15, 16, 17
- Contact-mode OA, 20 cm beta threshold
- Five tasks: cup-in-bowl, bowl-on-rack, mug-on-coaster, stack-two-cubes, place-three-cups. Book is excluded (no annotated sidecars).
- Cameras: countertop / right (wrist) / left (side)
- 7-D UR5e state/action (6 joints + gripper). Joints are deltas; gripper is absolute.
- 3 × H200 FSDP, global batch 48, 20,082 steps
The exact TrainConfig used for this run is in
config_snippet.py and copied below.
TrainConfig(
name="pi05_obs_only_firstandlast3_injection_ur5e_real",
checkpoint_base_dir="/work/shared/outputs/pi05_obj/checkpoints",
model=pi0_config.Pi0Config(
pi05_obj=True,
obstacle_attention=pi0_config.ObstacleAttentionConfig(
enabled=True,
supervised_layers=(0, 1, 2),
supervised_layer_lr=(0.01, 0.01, 0.01),
heatmap_sigma=15.0,
target_attn=False,
dest_attn=False,
inject_attention=True,
inject_layers=(15, 16, 17),
),
),
data=GroundingDataConfig(
repo_id="ur5e_real_grounding",
assets=AssetsConfig(assets_dir="./assets/pi05_obs_only_firstandlast3_injection_ur5e_real"),
grounding_root="/work/data/real_world_data",
grounding_annotation_root="/work/data/real_world_data_annotated",
grounding_exclude_tasks=("put_book_in_box", "put_book_in_box_amir", "put_book_on_shelf"),
grounding_cameras=("countertop_camera", "right_camera", "left_camera"),
grounding_frame_stride=1,
grounding_beta_cache_root="/work/data/real_world_beta",
grounding_attention_mask_mode="contact",
adapt_to_pi=False,
physical_action_dim=7,
),
weight_loader=weight_loaders.CheckpointWeightLoader(
"/work/richard/safevla-home/checkpoints/ur5e_pi05_all_3cam_20k/params",
missing_regex=".*(lora|attn_query_head|oa_inject).*",
),
num_workers=6,
num_epochs=3,
num_train_steps=21_000,
batch_size=48,
viz_interval=500,
viz_num_samples=2,
save_interval=2500,
keep_period=2500,
fsdp_devices=3,
),
Layout
Each checkpoint is under <step>/ (params/, assets/, _CHECKPOINT_METADATA).
Included steps: 2500, 5000, 6694, 7500, 10000, 12500, 13388,
15000, 17500, 20000, 20082 (final). Optimizer train_state/ is
not published.
Load like any openpi Orbax checkpoint: point the weight loader at
20082/params/ and use 20082/assets/ur5e_real_grounding/norm_stats.json.