OpenPI Supervised Fine-Tuning#

https://raw.githubusercontent.com/RLinf/misc/main/pic/pi0_icon.jpg

OpenPI π₀ / π₀.₅ vision-language-action models.#

Run full-parameter or LoRA SFT on π₀ / π₀.₅ in RLinf. Policies use model_type: openpi, a PyTorch port aligned with the JAX reference (the official OpenPI PyTorch code is not). SFT is the usual RL cold start: imitate demonstrations first, then continue with RL from a stronger prior.

Recipes#

Each experiment YAML imports a path-free model template through Hydra defaults, then overrides paths and task fields. Replace /path/to/... before launch.

Task

Experiment config (examples/sft/config/)

Model template

Launch name

LIBERO π₀

libero_sft_openpi.yaml

model/pi0

libero_sft_openpi

Real-world Franka bin-relocation

realworld_bin_relocation_sft_openpi.yaml

model/pi0

realworld_bin_relocation_sft_openpi

Custom LeRobot

custom_sft_openpi.yaml

model/pi0

custom_sft_openpi

BEHAVIOR-1K π₀.₅

behavior_sft_openpi_pi05.yaml

model/pi0_5

behavior_sft_openpi_pi05

RoboTwin adjust_bottle π₀

robotwin_adjust_bottle_sft_openpi.yaml

model/pi0

robotwin_adjust_bottle_sft_openpi

For dual-Franka SFT and deployment, see Using Dual Franka with OpenPI PyTorch.

Models

π₀ · π₀.₅

Methods

Full SFT · LoRA

Data

LeRobot · BEHAVIOR

Hardware

1+ nodes · GPUs

You’ll do: install → prepare data and norm stats → bash examples/sft/run_vla_sft.sh <launch-name> → watch the training loss.
Prerequisites: Installation · a matching dataset.

Shared conventions#

openpi.task must be sft during training. Eval YAMLs use eval; do not copy the training value.

num_action_chunks is the env-executed chunk, not the network horizon:

  • num_action_chunks / openpi.action_chunk: how many actions the env executes per step.

  • Network action_horizon: openpi.action_horizon when set; otherwise the official TrainConfig.model.action_horizon for openpi.config_name. Default templates only interpolate num_action_chunks into action_chunk.

Official LeRobot SFT also sizes its data window from TrainConfig.model.action_horizon; it does not read num_action_chunks. Typical official horizons: BEHAVIOR pi05_behavior 32, RoboTwin pi0_aloha_robotwin / LIBERO pi0_libero 50, real-world pi0_realworld 10. Override openpi.action_horizon only when the checkpoint horizon differs from that TrainConfig.

The BEHAVIOR / dual-Franka streaming SFT loaders currently still use num_action_chunks as the data window, so those recipes keep it equal to the official horizon.

Keep original parameters and optimizer updates in FP32 with actor.model.precision: fp32. Set actor.model.openpi.embed_dtype: bf16 for BF16 backbone activations and residuals. If omitted or null, embed_dtype follows model.precision; precision: null selects BF16 activations. Configure FSDP separately: BF16 parameters for forward/backward computation, FP32 gradient reduction, and FP32 buffers. The shared model templates use precision: null for OpenPI’s selective BF16/FP32 parameter dtypes. BEHAVIOR, RoboTwin, dual-Franka, and RLT Stage 1 SFT recipes set precision: fp32.

Select hybrid_engines/openpi_fsdp to keep projections, RMSNorm weights, and vision embeddings in FP32. The profile preserves FP32 action/time inputs and adaptive-norm conditioning while the backbone computes in BF16. RLT Stage 1 recipes also keep rlt_module in FP32.

Use no_shard and use_orig_params: true for OpenPI. Each rank holds complete parameters, gradients, and optimizer states. BEHAVIOR / RoboTwin recipes use the openpi_cosine learning rate schedule.

defaults:
  - model/pi0_5@actor.model
  - hybrid_engines/openpi_fsdp@actor.fsdp_config
  - _self_

actor:
  model:
    precision: fp32
    openpi:
      task: sft
      embed_dtype: bf16
  fsdp_config:
    sharding_strategy: no_shard
    use_orig_params: true
    mixed_precision:
      param_dtype: bf16
      reduce_dtype: fp32
      buffer_dtype: fp32

The tokenizer is loaded by OpenPI’s ModelTransformFactory from the base model; the SFT YAML does not need a SentencePiece path.

Datasets#

Select the data format with actor.model.openpi.config_name:

config_name

Dataset / environment

pi0_libero · pi05_libero

LIBERO

pi0_realworld

Real-world Franka

pi05_yam_joint

Real-world Dual YAM

pi0_aloha_robotwin

RoboTwin (ALOHA)

pi05_behavior

BEHAVIOR-1K (streaming loader)

pi0_maniskill · pi05_maniskill

ManiSkill

pi05_metaworld

MetaWorld

pi05_calvin

CALVIN

pi0_custom

Custom LeRobot

Custom LeRobot dataset#

  1. Set config_name in examples/sft/config/custom_sft_openpi.yaml:

    actor:
  model:
    openpi:
      task: sft
      config_name: "pi0_custom"

  2. Register pi0_custom in rlinf/models/embodiment/openpi/dataconfig/__init__.py (a template is already there):

    TrainConfig(
    name="pi0_custom",
    model=pi0_config.Pi0Config(),
    data=CustomDataConfig(
        repo_id="physical-intelligence/custom_dataset",
        base_config=DataConfig(prompt_from_task=True),
        assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_base/assets"),
        extra_delta_transform=False,
        action_train_with_rotation_6d=False,
    ),
    pytorch_weight_path="checkpoints/torch/pi0_base",
),

  3. Adapt CustomDataConfig in rlinf/models/embodiment/openpi/dataconfig/franka_dataconfig.py to your keys and action space.

Normalization statistics#

For a newly collected LeRobot dataset, compute norm_stats for state and actions before SFT (especially important for real-robot data):

# Local dataset directory (contains meta/info.json):
python toolkits/lerobot/calculate_norm_stats.py \
    --config-name pi0_realworld \
    --repo-id /path/to/realworld_franka_bin_relocation

# Or a Hugging Face repo id (cached under ~/.cache/huggingface/lerobot):
python toolkits/lerobot/calculate_norm_stats.py \
    --config-name pi0_realworld \
    --repo-id realworld_franka_bin_relocation

  • --repo-id is a local path or a LeRobot HF repo id; HF_LEROBOT_HOME changes the cache parent.

  • --config-name must match the training dataconfig.

  • calculate_norm_stats.py writes to ./assets/<config_name>/<repo_id>/norm_stats.json (TrainConfig.assets_dirs / repo_id). Training and eval do not look next to SFT full_weights.pt for that file. Set openpi_data.norm_stats_path to that norm_stats.json. If the key is omitted, OpenPI loads {assets_dir}/{asset_id}/norm_stats.json from the TrainConfig (after model_path is applied) and logs a warning. The RoboTwin recipe pins stats beside the base weights at physical-intelligence/robotwin/<task>/norm_stats.json.

If standard deviations are tiny or the q99–q01 range is very narrow, widening them often stabilizes training, especially when SFT is followed by online RL.

BEHAVIOR does not use that script. Point actor.model.openpi_data.norm_stats_path at an existing {assets}/{asset_id}/norm_stats.json.

Installation#

First, clone the RLinf repository:

# Mainland China users can use a mirror for faster cloning:
# git clone https://gh-proxy.com/github.com/RLinf/RLinf.git
git clone https://github.com/RLinf/RLinf.git
cd RLinf

Then set up the dependencies with one of the two methods below — a prebuilt Docker image (recommended) or a custom environment. The general setup (prerequisites, GPU drivers, the in-image switch_env helper, mirrors, and troubleshooting) is documented once in Installation; the commands in this recipe only differ in the Docker image tag and the --env value.

Option 1: Docker — image agentic-rlinf0.4-maniskill_libero:

docker run -it --rm --gpus all \
   --shm-size 20g \
   --network host \
   --name rlinf \
   -v .:/workspace/RLinf \
   rlinf/rlinf:agentic-rlinf0.4-maniskill_libero
   # Mainland China mirror: infinigence-ai-registry.cn-beijing.cr.aliyuncs.com/rlinf/rlinf:agentic-rlinf0.4-maniskill_libero

source switch_env openpi

Option 2: Custom environment:

# Add --use-mirror in mainland China.
bash requirements/install.sh embodied --model openpi --env maniskill_libero
source .venv/bin/activate

--model openpi installs the official OpenPI runtime (transforms and data pipeline). The trained network is still openpi.

LIBERO and real-world Franka#

Edit data paths and actor.model.model_path in the experiment YAML. Typical overrides:

data:
  train_data_paths: /path/to/libero-data

cluster:
  num_nodes: 1
  component_placement:
    actor,env,rollout: 0-0

actor:
  model:
    model_path: /path/to/pi0-model
    num_action_chunks: 4          # env-executed length, not network horizon
    is_lora: False                # set True and lora_rank for LoRA
    openpi:
      task: sft
      embed_dtype: bf16
      config_name: "pi0_libero"   # use pi0_realworld for Franka

Launch:

bash examples/sft/run_vla_sft.sh libero_sft_openpi
bash examples/sft/run_vla_sft.sh realworld_bin_relocation_sft_openpi

Logs and checkpoints go under runner.logger.log_path, saved every runner.save_interval steps at .../checkpoints/global_step_<N>/.

Pi0.5 Base Checkpoint Download and Conversion#

The π₀.₅ JAX base checkpoint is available at gs://openpi-assets/checkpoints/pi05_base. The downloaded directory contains params/ and must be converted to the RLinf OpenPI layout before use as actor.model.model_path. Replace the YAML’s /path/to/... placeholders with actual paths on your server.

Run from the RLinf repository root in the OpenPI environment with JAX / Orbax installed. Downloading requires the Google Cloud CLI:

mkdir -p /path/to/checkpoints/jax
gcloud storage cp --recursive \
  gs://openpi-assets/checkpoints/pi05_base /path/to/checkpoints/jax/

Prepare task-specific norm_stats.json first (see normalization statistics above). This example uses BEHAVIOR’s 32-step horizon. --input-model must point to the directory containing params/:

python -m rlinf.utils.ckpt_convertor.openpi.convert --mode jax_to_openpi \
  --input-model /path/to/checkpoints/jax/pi05_base \
  --input-norm-stats /path/to/task/norm_stats.json \
  --output-model /path/to/pi05_base_openpi \
  --output-norm-stats /path/to/pi05_base_openpi/task/norm_stats.json \
  --action-dim 32 \
  --action-horizon 32 \
  --max-token-len 200

The converter defaults to π₀.₅; do not add --no-pi05. Other π₀.₅ SFT recipes use the same workflow with --action-horizon matching their existing configs: 10 for ManiSkill RLT, and 20 for Franka RLT and dual-Franka. Conversion does not require changing action dimensions, image counts, or training settings in the YAML.

Verify that the output directory contains model.safetensors and config.json. Replace the actor.model.model_path placeholder with that output directory and point actor.model.openpi_data.norm_stats_path at the statistics for the same task, for example:

actor:
  model:
    model_path: /path/to/pi05_base_openpi
    openpi_data:
      norm_stats_path: /path/to/pi05_base_openpi/task/norm_stats.json

Checkpoints already in the RLinf OpenPI layout need no further conversion. Official OpenPI PyTorch checkpoints (with paligemma_with_expert.* keys) still require openpi_pytorch_to_openpi even when they contain model.safetensors. See rlinf/utils/ckpt_convertor/openpi/README.md for that command.

Pi0.5 + BEHAVIOR-1K#

Filesystem paths live in the experiment config; the model shape comes from model/pi0_5:

defaults:
  - model/pi0_5@actor.model
  - hybrid_engines/fsdp@actor.fsdp_config

The streaming loader reads only data: (no hidden defaults). Change at least:

data:
  train_data_paths: /path/to/2025-challenge-demos
  behavior_dataset_root: /path/to/2025-challenge-demos
  repo_id: "behavior-1k/2025-challenge-demos"
  modalities: ["rgb"]
  num_workers: 8
  tasks: ["turning_on_radio"]
  use_skill: false
  task_subtasks:
    turning_on_radio:
      - "move to radio"
      - "pick up radio from coffee table"
      - "press radio"
      - "place radio on coffee table"

actor:
  model:
    model_path: /path/to/pi05_base_openpi   # new-format fp32 base weights
    openpi:
      task: sft
      embed_dtype: bf16
    openpi_data:
      norm_stats_path: /path/to/assets/behavior-1k/2025-challenge-demos/norm_stats.json

  • use_skill: false trains on the main-task text; true uses per-frame skill text from task_subtasks (collapsed orchestrators in the dataset are not valid labels).

  • fine_grained_level / tolerance_s control streaming time alignment.

bash examples/sft/run_vla_sft.sh behavior_sft_openpi_pi05

Pi0 + RoboTwin#

Uses the official OpenPI / LeRobot map-style loader with config_name: pi0_aloha_robotwin. Example data: adjust_bottle; matching HF SFT weights: RLinf-Pi0-NEW-RoboTwin-SFT-adjust_bottle.

data:
  train_data_paths: /path/to/robotwin-data
  num_workers: 4

actor:
  model:
    model_path: /path/to/pi0_base_pytorch_new
    num_action_chunks: 50
    action_dim: 14
    openpi:
      task: sft
      embed_dtype: bf16
      config_name: "pi0_aloha_robotwin"
      num_images_in_input: 3
    openpi_data:
      norm_stats_path: ${actor.model.model_path}/physical-intelligence/robotwin/adjust_bottle/norm_stats.json

14-D ALOHA actions and three cameras; OpenPI pads actions to 32-D before the model. openpi_data.norm_stats_path pins SFT and eval to the same norm_stats.json. Change the task name in that path when switching tasks.

bash examples/sft/run_vla_sft.sh robotwin_adjust_bottle_sft_openpi

Evaluation#

SFT writes full_weights.pt under .../checkpoints/global_step_<N>/actor/model_state_dict/. openpi eval can load that .pt (or the global_step_<N> directory) directly: wrapper / FSDP prefixes are stripped in memory, and the eval YAML supplies the model shape. You do not need to convert to model.safetensors first.

Do not copy training openpi.task: sft into eval. Eval configs must use task: eval (that selects Pi0Eval and observation transforms). Keep config_name and num_action_chunks aligned with training.

Norm stats also do not travel with full_weights.pt. Stock eval YAMLs often interpolate openpi_data.norm_stats_path as ${rollout.model.model_path}/.../norm_stats.json, which assumes a converted directory that already bundles stats. If you only retarget model_path at an SFT checkpoint, also point stats back at the training copy (RoboTwin: next to the base weights at physical-intelligence/robotwin/<task>/norm_stats.json; BEHAVIOR: the training openpi_data.norm_stats_path).

# Example: evaluate the SFT .pt in place (keep training-time norm stats)
bash evaluations/run_eval.sh robotwin robotwin_adjust_bottle_openpi_eval \
  rollout.model.model_path=/path/to/logs/.../checkpoints/global_step_30000 \
  rollout.model.openpi_data.norm_stats_path=/path/to/pi0_base_pytorch_new/physical-intelligence/robotwin/adjust_bottle/norm_stats.json

Eval steps: BEHAVIOR-1K and RoboTwin. To export a shareable bare Pi0 directory (model.safetensors + config.json + copied norm_stats.json), use sft_to_openpi; commands are in rlinf/utils/ckpt_convertor/openpi/README.md.

Visualization#

Watch the training loss. Metric definitions: Training metrics.

tensorboard --logdir ./logs