README.md

---
library_name: openpi
pipeline_tag: robotics
tags:
  - pi0.5
  - openpi
  - robotics
  - imitation-learning
  - 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

```text
.
+-- 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:

```text
checkpoint/2000
```

Repository ID:

```text
qmddn/pi05-insert-mouse-battery
```