Skip to content

README.md

Latest commit

 

History

History
389 lines (297 loc) · 13.3 KB

README.md

File metadata and controls

389 lines (297 loc) · 13.3 KB
Raw

PI0.5: Training, Inference, and Serving

This guide covers how to train, run inference, and serve PI0.5 models using FlagScale.

Installation

Clone Repository

git clone https://github.com/FlagOpen/FlagScale.git
cd FlagScale/

Setup Conda Environment

Create a new conda environment for robotics training:

conda create -n flagos-robo python=3.12
conda activate flagos-robo

Install FlagScale and robotics dependencies:

cd FlagScale/
# replace "[cuda]" with "[ascend]" on Huawei Ascend, or "[musa]" on Moore Threads MUSA
pip install ".[cuda]" --verbose

pip install git+https://github.com/huggingface/transformers.git@fix/lerobot_openpi

Install additional dependencies for downloading models/datasets:

# For HuggingFace Hub
pip install huggingface_hub

# For ModelScope (optional)
pip install modelscope

Download Models and Tokenizers

Download models and tokenizers using the provided script. Choose either HuggingFace Hub or ModelScope based on your preference:

Using HuggingFace Hub:

cd FlagScale/
python examples/pi0/download.py \
    --repo_id lerobot/pi05_base \
    --output_dir /workspace/models \
    --source huggingface

python examples/pi0/download.py \
    --repo_id google/paligemma-3b-pt-224 \
    --output_dir /workspace/models \
    --source huggingface

Using ModelScope:

cd FlagScale/
python examples/pi0/download.py \
    --repo_id lerobot/pi05_base \
    --output_dir /workspace/models \
    --source modelscope

python examples/pi0/download.py \
    --repo_id google/paligemma-3b-pt-224 \
    --output_dir /workspace/models \
    --source modelscope

The models will be downloaded to (example with /workspace/models):

  • /workspace/models/lerobot/pi05_base
  • /workspace/models/google/paligemma-3b-pt-224

Training

Prepare Dataset

FlagScale uses the LeRobotDataset v3.0 format. For detailed information about the format structure, see the LeRobotDataset v3.0 documentation.

For example, to download the aloha_mobile_cabinet dataset:

Using HuggingFace Hub:

cd FlagScale/
python examples/pi0/download.py \
    --repo_id lerobot/aloha_mobile_cabinet \
    --output_dir /workspace/datasets \
    --repo_type dataset \
    --source huggingface

Using ModelScope:

cd FlagScale/
python examples/pi0/download.py \
    --repo_id lerobot/aloha_mobile_cabinet \
    --output_dir /workspace/datasets \
    --repo_type dataset \
    --source modelscope

The dataset will be downloaded to (example with /workspace/datasets):

  • /workspace/datasets/lerobot/aloha_mobile_cabinet

Note: PI0.5 supports two normalization methods:

  • Quantile normalization (default when use_quantiles: true): Requires quantile statistics in the dataset
  • MEAN_STD normalization (when use_quantiles: false): Uses mean and standard deviation statistics

If your dataset doesn't have quantile statistics and you want to use quantile normalization, you can augment the dataset with quantile stats. For more details, see the PI0.5 documentation.

Alternatively, you can train PI0.5 with MEAN_STD normalization by setting use_quantiles: false in the config (which is the default for PI0.5 in FlagScale).

Edit Config

FlagScale uses a two-level configuration system:

  1. Experiment-level config (examples/pi0_5/conf/train.yaml): Defines experiment settings, environment variables, and resource allocation
  2. Task-level config (examples/pi0_5/conf/train/pi0_5.yaml): Defines model, dataset, and training hyperparameters

Experiment-Level Config

Edit the experiment-level config for multi-GPU training:

cd FlagScale/
vim examples/pi0_5/conf/train.yaml

Configure the following fields:

  • experiment.envs.CUDA_VISIBLE_DEVICES - GPU devices to use (e.g., "0,1,2,3" for 4 GPUs, "0,1" for 2 GPUs)，Use ASCEND_RT_VISIBLE_DEVICES for Huawei Ascend, MUSA_VISIBLE_DEVICES for Moore Threads MUSA
  • experiment.envs.CUDA_DEVICE_MAX_CONNECTIONS - Connection limit (typically 1)，Use MUSA_DEVICE_MAX_CONNECTIONS for Moore Threads MUSA
  • experiment.exp_name - Experiment name
  • experiment.exp_dir - Output directory for checkpoints and logs
  • runner.nproc_per_node - Number of processes per node for multi-GPU training (required for Huawei Ascend)

Task-Level Config

Edit the task-level config for model and training settings:

cd FlagScale/
vim examples/pi0_5/conf/train/pi0_5.yaml

Configure the following fields:

System settings (training hyperparameters):

  • system.batch_size - Batch size per GPU
  • system.train_steps - Total training steps
  • system.checkpoint.save_checkpoint - Whether to save checkpoints (default: true)
  • system.checkpoint.save_freq - Steps between checkpoints (default: 1000)
  • system.checkpoint.output_directory - Checkpoint output directory (default: ${experiment.exp_dir})

Model settings:

  • model.model_name - Model name: "pi0.5"
  • model.checkpoint_dir - Path to pretrained model (e.g., /workspace/models/lerobot/pi05_base)
  • model.tokenizer_path - Path to tokenizer (e.g., /workspace/models/google/paligemma-3b-pt-224)
  • model.tokenizer_max_length - Maximum tokenizer sequence length (default: 200 for pi0.5)
  • model.action_steps - Number of action steps to predict
  • model.optimizer.name - Optimizer name (for example: "AdamW")
  • model.optimizer.lr - Learning rate (for example: 2.5e-5)
  • model.optimizer.betas - Optimizer betas (for example: [0.9, 0.95])
  • model.optimizer.eps - Optimizer epsilon (for example: 1.0e-8)
  • model.optimizer.weight_decay - Weight decay (for example: 0.01)
  • model.optimizer.scheduler.warmup_steps - Warmup steps (for example: 1000)
  • model.optimizer.scheduler.decay_steps - Decay steps (for example: 30000)
  • model.optimizer.scheduler.decay_lr - Final learning rate after decay (for example: 2.5e-6)

Data settings:

  • data.data_path - Path to LeRobot dataset root (e.g., /workspace/datasets/lerobot/aloha_mobile_cabinet)
  • data.use_imagenet_stats - Whether to use ImageNet normalization stats (default: true)
  • data.rename_map - Dictionary mapping dataset keys to policy keys (optional). Check the features key in your dataset's meta/info.json file to determine the correct mapping: 
    rename_map:
  observation.images.cam_high: observation.images.base_0_rgb
  observation.images.cam_left_wrist: observation.images.left_wrist_0_rgb
  observation.images.cam_right_wrist: observation.images.right_wrist_0_rgb
  • data.use_quantiles - Whether to use quantile normalization (default: false for pi0.5, uses MEAN_STD normalization)

Start Training

cd FlagScale/
flagscale train pi0_5 --config ./examples/pi0_5/conf/train.yaml
# or
flagscale train pi0_5 -c ./examples/pi0_5/conf/train.yaml

Training logs are saved to outputs/pi0_5_train/logs/host_0_localhost.output by default.

Checkpoints are saved to ${experiment.exp_dir}/checkpoints (default: outputs/pi0_5_train/checkpoints).

Stop Training

cd FlagScale/
flagscale train pi0_5 --stop

Inference

Prepare Inference Inputs

You can extract inference inputs (images, state, task) from a dataset using the provided script:

cd FlagScale/
python examples/pi0/dump_dataset_inputs.py \
    --dataset_root /workspace/datasets/lerobot/aloha_mobile_cabinet \
    --output_dir ./inference_inputs \
    --frame_index 100

This will create:

  • frame_100_observation_images_*.jpg - Image files
  • frame_100_state.pt - State tensor
  • frame_100_task.txt - Task prompt
  • extraction_summary.json - Summary of extracted files

Alternatively, you can extract from a specific episode and frame:

python examples/pi0/dump_dataset_inputs.py \
    --dataset_root /workspace/datasets/lerobot/aloha_mobile_cabinet \
    --output_dir ./inference_inputs \
    --episode_index 0 \
    --frame_in_episode 50

Or extract multiple samples at once:

python examples/pi0/dump_dataset_inputs.py \
    --dataset_root /workspace/datasets/lerobot/aloha_mobile_cabinet \
    --output_dir ./inference_inputs \
    --frame_indices 100 200 300

Edit Config

cd FlagScale/
vim examples/pi0_5/conf/inference/pi0_5.yaml

Configure the following fields:

Engine settings:

  • engine.model_variant - Model variant: "pi0.5"
  • engine.model - Path to pretrained model (e.g., /workspace/models/lerobot/pi05_base)
  • engine.tokenizer - Path to tokenizer (e.g., /workspace/models/google/paligemma-3b-pt-224)
  • engine.stat_path - Path to dataset statistics (e.g., /workspace/datasets/lerobot/aloha_mobile_cabinet/meta/stats.json)
  • engine.device - Device to use (e.g., "cuda", "npu", "musa")
  • engine.use_quantiles - Whether to use quantile normalization (default: false for pi0.5)

Generate settings:

  • generate.images - Dictionary mapping image keys to file paths: 
    images:
  observation.images.cam_high: /path/to/image1.jpg
  observation.images.cam_left_wrist: /path/to/image2.jpg
  observation.images.cam_right_wrist: /path/to/image3.jpg
  • generate.state_path - Path to state tensor file (.pt file)
  • generate.task_path - Path to task prompt file (.txt file)
  • generate.rename_map (optional) - Map input keys to policy expected keys. Check the features key in your dataset's meta/info.json file to determine the correct mapping: 
    rename_map:
  observation.images.cam_high: observation.images.base_0_rgb
  observation.images.cam_left_wrist: observation.images.left_wrist_0_rgb
  observation.images.cam_right_wrist: observation.images.right_wrist_0_rgb

Run Inference

cd FlagScale/
flagscale inference pi0_5 --config ./examples/pi0_5/conf/inference.yaml
# or
flagscale inference pi0_5 -c ./examples/pi0_5/conf/inference.yaml

Inference logs are saved to outputs/pi0_5_inference/inference_logs/host_0_localhost.output by default.

The predicted action tensor is printed to the console and saved in the log file.

Serving

Edit Config

cd FlagScale/
vim examples/pi0_5/conf/serve/pi0_5.yaml

Configure the following fields:

Engine arguments:

  • engine_args.host - Server host (default: "0.0.0.0")
  • engine_args.port - Server port (default: 5000)
  • engine_args.model_variant - Model variant: "pi0.5"
  • engine_args.model - Path to pretrained model (e.g., /workspace/models/lerobot/pi05_base)
  • engine_args.tokenizer - Path to tokenizer (e.g., /workspace/models/google/paligemma-3b-pt-224)
  • engine_args.stat_path - Path to dataset statistics (e.g., /workspace/datasets/lerobot/aloha_mobile_cabinet/meta/stats.json)
  • engine_args.device - Device to use (e.g., "cuda", "npu", "musa")
  • engine_args.use_quantiles - Whether to use quantile normalization (default: false for pi0.5)
  • engine_args.images_keys - List of image keys expected by the model (do not change): 
    images_keys:
  - observation.images.base_0_rgb
  - observation.images.left_wrist_0_rgb
  - observation.images.right_wrist_0_rgb
  • engine_args.images_shape - Image shape [C, H, W] for warmup (e.g., [3, 480, 640])
  • engine_args.state_key - Key for state in the batch (e.g., "observation.state")

Run Serving

cd FlagScale/
flagscale serve pi0_5 --config ./examples/pi0_5/conf/serve.yaml
# or
flagscale serve pi0_5 -c ./examples/pi0_5/conf/serve.yaml

Serving logs are saved to outputs/pi0_5_serve/logs/host_0_localhost.output by default.

Stop Serving

cd FlagScale/
flagscale serve pi0_5 --stop

Evaluation

FlagScale supports online evaluation via FlagEval. You will need a FLAGEVAL_SECRET — contact the team to obtain one before proceeding.

Run Evaluation

cd FlagScale/
FLAGEVAL_SECRET=<your_secret> flagscale eval robo \
    --model-name pi0_5 \
    --datasets libero_10 \
    --server-host <your_model_server_host> \
    --server-port <your_model_server_port> \
    --poll-interval 30 \
    --model-id <your_model_id>

Configure the following fields:

  • FLAGEVAL_SECRET - Authentication secret (contact the team to obtain)
  • --model-name - Model name: "pi0_5"
  • --datasets - Evaluation dataset(s), e.g. libero_10
  • --server-host - Host/IP of your model server (FlagEval will connect back to this address)
  • --server-port - Port of your model server
  • --poll-interval - How often (in seconds) to poll for results
  • --model-id - A unique identifier for this evaluation run (e.g., pi0_5_exp1)

By default the script automatically starts the model server before submitting the evaluation. If your model server is already running, add --attach to skip the startup step.

Test Server with Client

The client should send images using keys that match the images_keys in the config. For example, if using the default config:

cd FlagScale/
python examples/pi0/client_pi0.py \
  --host 127.0.0.1 \
  --port 5000 \
  --img1 ./inference_inputs/frame_100_observation_images_cam_high.jpg \
  --img2 ./inference_inputs/frame_100_observation_images_cam_left_wrist.jpg \
  --img3 ./inference_inputs/frame_100_observation_images_cam_right_wrist.jpg \
  --state-path ./inference_inputs/frame_100_state.pt \
  --instruction "Grab the orange and put it into the basket."

Note: The client must send image keys that match the engine_args.images_keys in the config.