ur5e_pi05_real_oa_3ep

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.