概述

pi0 接收相机图像、任务文本和机器人 state，输出一段未来 action。PhyAI 的 ws1 路径把整条推理链留在一张 GPU 上：vision 与 language stack 先生成可复用的 prefix，action expert 再结合 prefix 和机器人 state，通过 flow matching 把噪声逐步还原成 action chunk。 PI0WS1Scheduler 只实现这条单卡路径，不承担 tensor parallel、continuous batching 或 preemption。

pi0 把 robot state 保留为 expert 侧的数值 token；pi0.5 则先把 state 离散成 bins，再写入语言 prompt。

架构

请求先通过进入 PhyAI。PI0Entry 负责构造模型和 scheduler，三个 runner 分别接管 vision、prefix 与 expert 计算。代码分布在这些文件中：

phyai/src/phyai/models/pi0

main_pi0.py

scheduler_ws1_pi0.py

model_runner_pi0.py

modeling_pi0.py

configuration_pi0.py

组件	职责
`PI0Entry`	注册 `"pi0"` engine plugin，构造 `PI0Model`，加载权重，并创建 scheduler
`PI0Config`	保存 vision、text、expert、action chunk、tokenizer 和相机数量相关的几何配置
`PI0Model`	持有 SigLIP/PaliGemma vision tower、PaliGemma text stack、Gemma expert stack、RoPE 以及 action/time heads
`PI0VisionRunner`	运行 vision tower，并支持 CUDA graph capture
`PI0LLMRunner`	运行 PaliGemma prefix pass，并把 prefix K/V 写入共享 cache
`PI0ExpertRunner`	在每个 flow-matching step 中运行 expert 的 state/action pass
`PI0WS1Scheduler`	在单张 GPU 上编排一次完整的 inference request

模型布局

代码分散在几个模块中，数据却始终沿着一条路径流动。SigLIP 把相机图像变成 tokens，PaliGemma 将它们与任务文本合并，较小的 Gemma expert 最后生成 action 轨迹。

Stack	默认形状	备注
Vision	SigLIP，27 层，224×224 图像，14×14 patch	生成图像 tokens，并投影到 PaliGemma text width
Text	PaliGemma/Gemma，18 层，hidden size 2048	处理 image + language prefix，并写入 prefix K/V
Expert	Gemma action expert，18 层，hidden size 1024	处理一个 state token 和完整 action chunk

三部分计算共用一组顶层几何配置：

字段	默认值	含义
`chunk_size`	`50`	每次 engine step 返回的 action token 数
`max_state_dim`	`32`	padding 后的 robot state 宽度
`max_action_dim`	`32`	padding 后的 action 宽度
`num_inference_steps`	`10`	Flow-matching Euler steps
`tokenizer_max_length`	`48`	右 padding 后的 PaliGemma task prompt 长度
`empty_cameras`	`0`	`num_images = 3 - empty_cameras`；pi0 支持 2 或 3 路相机

params_dtype 控制 language 与 expert stacks 的参数精度。Vision tower 使用独立的 vision_params_dtype，默认保留 fp32，以便和参考实现对齐。只有实验本身需要 bf16 vision 时，才传入 PI0Args(vision_params_dtype=torch.bfloat16)。

请求格式

运行时，PI0Request 承接 processor 的输出，并把模型所需的 tensor 交给 scheduler：

字段	Shape	备注
`pixel_values`	`(B, num_images, 3, image_size, image_size)`	已经 resize 和 normalize 的相机 tensor
`input_ids`	`(B, tokenizer_max_length)` int64	右 padding 后的 PaliGemma token ids
`lang_lens`	`(B,)` int64	每个样本真实的 task prompt 长度
`state`	`(B, max_state_dim)`	数值形式的 robot state，在进入 expert 前完成 padding
`noise`	`(B, chunk_size, max_action_dim)` 或 `None`	可选的初始 action noise；为 `None` 时 scheduler 内部采样 Gaussian noise

B 可以取 1 到 max_batch_size 之间的任意值。较小的 batch 会在 scheduler 内部补到 captured shape，返回前再切回 actual_B。

调度阶段

请求进入 engine.step 后，scheduler 按顺序执行这些阶段：

Phase	工作内容
`pi0.vision_loop`	把相机 tensor 移到 vision dtype，并对真实 batch 中每个样本运行一次 `PI0VisionRunner`
`pi0.lang_pack`	嵌入 language ids，然后把 image tokens 和 language tokens pack 到每个样本的 prefix buffer
`pi0.llm_prefix_plan`	重置 static caches，并准备 ragged prefix attention metadata
`pi0.llm_prefix_fwd`	运行 PaliGemma text stack，并把 prefix K/V 写入 `KVCachePool`
`pi0.expert_plan`	准备覆盖 prefix + suffix slots 的 state/action expert attention metadata
`pi0.expert_state_prefill`	对 robot state 做 embedding，并一次性缓存各层 K/V
`pi0.expert_loop`	初始化或复制 action noise，并运行 flow-matching integration
`pi0.expert_step`	`pi0.expert_loop` 内部的一次 expert velocity prediction 和 Euler update

Prefix 在 denoising 期间不会改变，因此每个 request 只需缓存一次。State query 与 action query 随后读取不同范围的 cache：

state query  -> prefix + state
action query -> prefix + state + action chunk

因此，pi0 的 suffix length 是 1 + chunk_size，由一个 state token 和随后的 action tokens 组成。

CUDA graphs

请求形状固定后，runner 才能复用 CUDA graph。启用 RuntimeConfig(use_cuda_graph=True) 时，每个 runner 都会在 scheduler.setup() 期间捕获自己的 graph：

Runner	捕获形状
`PI0VisionRunner`	`(num_images, 3, image_size, image_size)`
`PI0LLMRunner`	`(max_batch_size * n_per_sample, text_hidden_size)`
`PI0ExpertRunner`	固定 `max_batch_size` 下的 `state`、`x_t` 和 `time` buffers

每次处理请求时，runners 先刷新 static input buffers，再 replay 已捕获的 graph。Attention metadata 留在 captured region 之外，通过 backend 的 capture-aware buffers 写入。

需要在 Nsight Systems 中展开 trace 时，可以关闭 CUDA graph：

uv run python benchmark/bench_n_batch_ws1_pi0.py \
    --batch-sizes 4 \
    --no-cuda-graph

运行 pi0

准备权重

使用包含 config.json 和 model.safetensors 的 HF-style pi0 PyTorch checkpoint 目录。只有检查执行链路时，才省略 --checkpoint 并使用随机权重。

构造 Engine

选择 "pi0" plugin。构造 Engine 时会加载可选 checkpoint、准备 runners，并在启用时捕获 CUDA graph。

import torch
from pathlib import Path

from phyai.engine import Engine, EngineArgs
from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
from phyai.models.pi0.main_pi0 import PI0Args

engine = Engine(
    EngineArgs(
        plugin="pi0",
        plugin_args=PI0Args(
            checkpoint_dir=Path("/path/to/pi0_pytorch"),
            max_batch_size=4,
            vision_params_dtype=torch.float32,
        ),
        config=EngineConfig(
            device=DeviceConfig(target="cuda", params_dtype=torch.bfloat16),
            runtime=RuntimeConfig(use_cuda_graph=True),
        ),
    )
)

max_batch_size 会成为 captured shape 的一部分；修改它需要重建 Engine。

构造请求

phyai-utils-tools 中的 PI0Processor 负责把原始机器人观测转换成 PI0Request 所需的 tensor。

from phyai.models.pi0.scheduler_ws1_pi0 import PI0Request
from phyai_utils_tools.models.pi0 import PI0Processor

processor = PI0Processor(
    image_size=224,
    num_channels=3,
    num_images=3,
    tokenizer_max_length=48,
    max_state_dim=32,
    action_dim=7,
    device="cuda",
    params_dtype=torch.bfloat16,
)

processed = processor.preprocess(
    {
        "images": [cam0, cam1, cam2],
        "task": ["pick up the object"],
        "state": state,
    }
)

request = PI0Request(
    pixel_values=processed.pixel_values,
    input_ids=processed.input_ids,
    lang_lens=processed.lang_lens,
    state=processed.state,
)

运行一步

actions = engine.step(request)  # (B, chunk_size, max_action_dim)

若 processor 配置了 action_dim，processor.postprocess(actions) 会裁掉 padding 的 action 宽度；存在 dataset stats 时，还会把 action 还原到原始量纲。

关闭 Engine

engine.close()

端到端示例

examples/pi0/run_pi0.py 既能使用 processor 构造请求，也能直接走 raw smoke 路径。带 checkpoint 时这样运行：

uv run python examples/pi0/run_pi0.py \
    --checkpoint /path/to/pi0_pytorch \
    --batch-size 1

只检查执行链路时，可以省略 checkpoint 并使用随机权重：

uv run python examples/pi0/run_pi0.py --raw --batch-size 1

Checkpoint 若留有一路 empty camera，应只传两路图像：

uv run python examples/pi0/run_pi0.py \
    --checkpoint /path/to/pi0_pytorch \
    --num-images 2

性能测试与分析

使用 benchmark/bench_n_batch_ws1_pi0.py，可以在同一套 scheduler 配置下比较不同 batch size：

uv run python benchmark/bench_n_batch_ws1_pi0.py \
    --batch-sizes 1 2 4 \
    --n-warmup 5 \
    --n-timed 30 \
    --result-file ./pi0_ws1_results.jsonl

同一脚本也能为 Nsight Systems 划出一段较短的 capture window：

nsys profile \
    --capture-range=cudaProfilerApi \
    --capture-range-end=stop \
    -o ./prof/pi0_ws1 \
    uv run python benchmark/bench_n_batch_ws1_pi0.py \
        --batch-sizes 4 \
        --profile-backend nsys \
        --profile-start-step 5 \
        --profile-num-steps 3

默认配置让 vision tower 保持 fp32。只有实验需要 bf16 vision latency 时，才添加 --vision-dtype bfloat16。

当前限制

CUDA graph 能够复用，前提是请求形状保持固定。PI0WS1Scheduler 构造单卡 Engine 时会固定 max_batch_size，graph capture 还会固定相机数量、图像尺寸和 tokenizer 长度。任何一项发生变化，都需要重建 Engine。Vision tower 则按真实 batch 中的样本逐个 replay，并不会把这些样本合成一次 vision batch。 Scheduler 的边界从预处理之后开始。图像 resize、tokenization、state padding 和 action unnormalize 都由 PI0Processor 完成；直接传给 Engine 的请求必须已经满足这一输入约定。

完整代码

from pathlib import Path

import torch

from phyai.engine import Engine, EngineArgs
from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
from phyai.models.pi0.configuration_pi0 import PI0Config
from phyai.models.pi0.main_pi0 import PI0Args
from phyai.models.pi0.scheduler_ws1_pi0 import PI0Request
from phyai.utils import load_config

CHECKPOINT_DIR = Path("/path/to/pi0_pytorch")
BATCH_SIZE = 1

cfg = load_config(CHECKPOINT_DIR, PI0Config)
device = torch.device("cuda")
dtype = torch.bfloat16

engine = Engine(
    EngineArgs(
        plugin="pi0",
        plugin_args=PI0Args(
            checkpoint_dir=CHECKPOINT_DIR,
            max_batch_size=BATCH_SIZE,
            vision_params_dtype=torch.float32,
        ),
        config=EngineConfig(
            device=DeviceConfig(target="cuda", params_dtype=dtype),
            runtime=RuntimeConfig(use_cuda_graph=True),
        ),
    )
)

try:
    input_ids = torch.zeros(
        BATCH_SIZE, cfg.tokenizer_max_length, dtype=torch.int64, device=device
    )
    input_ids[:, 0] = 2

    request = PI0Request(
        pixel_values=torch.rand(
            BATCH_SIZE,
            cfg.num_images,
            cfg.vision.num_channels,
            cfg.vision.image_size,
            cfg.vision.image_size,
            dtype=torch.float32,
            device=device,
        ),
        input_ids=input_ids,
        lang_lens=torch.ones(BATCH_SIZE, dtype=torch.int64, device=device),
        state=torch.rand(BATCH_SIZE, cfg.max_state_dim, dtype=dtype, device=device),
    )

    actions = engine.step(request)
    print(f"action chunk shape={tuple(actions.shape)}, dtype={actions.dtype}")
finally:
    engine.close()

快速开始

模型

量化

工具

开发者指南

单卡推理 PI0

概述

架构

模型布局

请求格式

调度阶段

CUDA graphs

运行 pi0

端到端示例

性能测试与分析

当前限制

完整代码

​概述

​架构

​模型布局

​请求格式

​调度阶段

​CUDA graphs

​运行 pi0

​端到端示例

​性能测试与分析

​当前限制

​完整代码

概述

架构

模型布局

请求格式

调度阶段

CUDA graphs

运行 pi0

端到端示例

性能测试与分析

当前限制

完整代码