from openpi.policies import policy_config
from openpi.training import config

# Load the model configuration
config_name = "pi0_trossen_toolbox" 
train_config = config.get_config(config_name)

# Create policy from your local checkpoint
policy = policy_config.create_trained_policy(
    train_config, 
    "path/to/checkpoint",
    default_prompt="manipulate the toolbox"
)

# Use for inference
observation = {
    "images": {
        "cam_high": image_array,  # [H, W, 3] uint8
        "cam_left_wrist": left_wrist_image,  # [H, W, 3] uint8  
        "cam_right_wrist": right_wrist_image,  # [H, W, 3] uint8
    },
    "state": joint_positions,  # [14] float32
    "prompt": "manipulate the toolbox"
}

actions = policy.infer(observation)["actions"]  # [50, 14]

With Policy Server

bash

# Start the policy server
uv run scripts/serve_policy.py policy:checkpoint \
    --policy.config=pi0_trossen_toolbox \
    --policy.dir=path/to/checkpoint

# Use with client
from openpi_client import websocket_client_policy
client = websocket_client_policy.WebsocketClientPolicy("localhost", 8000)
actions = client.infer(observation)

Model Architecture

Vision Encoder: SigLIP-So400m/14
Language Model: Gemma 2B + Gemma 300M (action expert)
Training: Diffusion-based action prediction
Input: Multi-camera RGB + proprioception + language instruction
Output: Future action sequence (50 timesteps)

Training Details

Framework: JAX/Flax with OpenPI
Optimizer: AdamW
Base Checkpoint: Pi-0 base model from Google
Fine-tuning: Task-specific fine-tuning on trossen toolbox data
Normalization: Dataset-specific state/action normalization

License

MIT License

Citation

If you use this model, please cite:

bibtex

@article{pi0,
  title={Pi-Zero: A Diffusion-Based Policy for Robot Manipulation},
  author={TODO: Add authors},
  year={2024}
}

Acknowledgments

Built using the OpenPI framework
Based on the Pi-0 architecture
Training data from Trossen toolbox manipulation demonstrations