qmddn/pi05-insert-mouse-battery

pi0.5 Insert Mouse Battery Base Policy

This repository contains the selected step2000 Orbax checkpoint for a pi0.5 base policy trained on the mixed-success Insert Mouse Battery dataset.

• Source repo: OpenDriveLab/kai0
• Training config: pi05_insert_mouse_battery_normal
• Selected checkpoint: step2000
• Validation split: mixed_success_seed42_val_lerobot
• Validation mean loss over 100 batches: 0.012273
• Parameters for evaluation: ema_params

Input Format

The policy expects the kai0/openpi pi0.5 observation format after the repository transforms:

• RGB images, float tensors resized to 224 x 224 x 3:
  • base_0_rgb
  • left_wrist_0_rgb
  • right_wrist_0_rgb
• Robot state: float32[32]
• Prompt: Insert mouse battery.

During training/evaluation, batches had this shape:

• Images: batch x 224 x 224 x 3
• State: batch x 32
• Tokenized prompt: batch x 200

Output Format

The policy predicts chunked robot actions:

• Action tensor: batch x 50 x 32
• Action horizon: 50
• Action dimension: 32

Use the same kai0/openpi config and normalization stats in norm_stats.json when loading this checkpoint.

Files

.
+-- README.md
+-- norm_stats.json
+-- checkpoint/
    +-- 2000/
        +-- _CHECKPOINT_METADATA
        +-- assets/
        +-- params/
        |   +-- _METADATA
        |   +-- _sharding
        |   +-- array_metadatas/
        |   +-- d/
        |   +-- manifest.ocdbt
        |   +-- ocdbt.process_0/
        +-- train_state/
            +-- _METADATA
            +-- _sharding
            +-- array_metadatas/
            +-- d/
            +-- manifest.ocdbt
            +-- ocdbt.process_0/

Loading

Download this repository and point the kai0/openpi checkpoint loader at:

checkpoint/2000

Repository ID:

qmddn/pi05-insert-mouse-battery