单卡推理 PI0.5

PhyAI 如何在单卡上推理 pi0.5

​

概述

π0.5 是 Physical Intelligence 推出的视觉-语言-动作（VLA）模型，基于机器人演示数据和大规模多模态数据共同训练，能够在未见过的真实开放世界环境中执行长时程任务，并具备泛化能力。

本页文档专注于 ws1 即 world_size=1,单 rank、不走分布式。本页全部内容都围绕这套单卡配置。 entry 是 PI05Scheduler。

​

架构

PhyAI 的 引擎 + 插件契约把 pi0.5 推理拆成四块协作组件:

phyai/src/phyai/models/pi05

main_pi05.py

scheduler_pi05.py

model_runner_pi05.py

modeling_pi05.py

configuration_pi05.py

img_preprocess_pi05.py

tokenization_pi05.py

下面这张动图展示了 phyai 的 3 个 model runner 是怎么和 scheduler 协作的,以及 engine 初始化是怎么衔接到 scheduler.setup() 和 scheduler.step() 的。

​

运行 pi0.5

1. 拿到权重

   准备一份 pi05_base safetensors checkpoint, 可以从 huggingface 下载:

   https://huggingface.co/lerobot/pi05_base
   
2. 构造 Engine

   插件名是 "pi05"。引擎一次性完成 setup、权重加载和 graph 捕获。

   import torch
   from pathlib import Path
   from phyai.engine import Engine, EngineArgs
   from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
   from phyai.models.pi05.main_pi05 import PI05Args
   
   engine = Engine(
       EngineArgs(
           plugin="pi05",
           plugin_args=PI05Args(
               checkpoint_dir=Path("/path/to/pi05_base/"),
               max_batch_size=4,
           ),
           config=EngineConfig(
               device=DeviceConfig(target="cuda", params_dtype=torch.bfloat16),
               runtime=RuntimeConfig(use_cuda_graph=True),
           ),
       )
   )
   

   max_batch_size 固定捕获图的 batch 维度。按你即将提交的最大 batch 来选;比这小的 batch 在内部自动填充。

3. 构造请求

   PI05Request 携带每步推理的输入:

   字段	Shape	备注
   pixel_values	(B, 3, 3, H, W)	每个 robot 3 路摄像头 × 3 通道,H = W = image_size
   input_ids	(B, tokenizer_max_length) int64	右侧用 0 填充
   lang_lens	(B,) int64	每个样本未填充前的真实长度
   noise	(B, chunk_size, max_action_dim) 或 None	可选;为 None 时调度器内部采新的 Gaussian

   B 可以是 [1, max_batch_size] 区间内的任意值。张量构造在引擎所在的 device 上;调度器会校验 shape,不一致会立即抛错。

4. 运行一步

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

   填充在返回前已经切掉 —— 你拿到的张量首维就是真实的 batch。

5. 关闭引擎

   engine.close()
   

   释放调度器侧的缓冲,拆掉捕获的 cuda graph。

​

端到端示例

examples/pi05/run_pi05.py 用确定性 dummy 输入跑了 max_batch_size ∈ {1, 4} 的全路径,并包含多 batch 等价性检查。运行命令:

python examples/pi05/run_pi05.py --checkpoint /path/to/pi0.5

脚本会打印每阶段的 latency 统计(3 次预热 + 30 次计时的 mean / median / std / min / max)以及等价性检查的 PASS 行。把 --checkpoint 后的路径改成你本地的 checkpoint 路径即可。

​

训练 rollout 集成

外部 RL controller 可以读取 diffusion trajectory，也可以替换内存中的 Actor 权重，无需重建 engine。使用完整 checkpoint 初始化 engine，并在构造 PI05Args 时设置 capture_rollout=True；训练期间的更新可以只包含 controller 选择的参数。

Controller 负责管理更新事务：

engine.begin_weight_update()
try:
    for weight_batch in actor_weight_batches:
        engine.update_weights(weight_batch)
    engine.finish_weight_update(version=actor_version)
except Exception:
    engine.abort_weight_update()
    raise

rollout = engine.rollout_step(request)

每个 weight_batch 都是 HuggingFace 参数名到 tensor 的映射。每次完成更新后都会刷新依赖权重的 scheduler 状态，并重新捕获 rollout graph。rollout_step 接收 PI05RolloutRequest，返回 PI05RolloutOutput，其中包含 action、denoising chain、逐 transition log probability 和可选的 value。请求必须提供 sigmas；noise 和 step_noise 是可选的确定性输入，未提供时由 runtime 内部采样。

RLinf BF16 集成使用可选的兼容策略 examples/configs/kernel_policies/pi05/rlinf_bf16.yaml。应通过 controller 配置或 PHYAI_KERNEL_CONFIG 显式选择；PI0.5 默认不会自动加载该策略。

部分权重写入后如果更新失败，不要继续用这个 Engine 推理。此时参数可能混合了两个版本，Engine 会把实例标记为不可用。终止事务后重新创建 engine。

​

当前限制

• 仅支持单卡。Tensor parallel、continuous batching、preemption 都不在 PI05Scheduler 的范围内。
• max_batch_size 在引擎构造时就固定。要改尺寸,必须把引擎拆掉重建。
• Vision tower 是按真实 robot 数顺序 replay 的,没在摄像头维度上 batch。

​

完整代码

from pathlib import Path

import torch

from phyai.engine import Engine, EngineArgs
from phyai.engine_config import DeviceConfig, EngineConfig, RuntimeConfig
from phyai.models.pi05.configuration_pi05 import PI05Config
from phyai.models.pi05.main_pi05 import PI05Args
from phyai.models.pi05.scheduler_pi05 import PI05Request
from phyai.utils import load_config

CHECKPOINT_DIR = Path("/path/to/pi05_base/")  # 改成你本地的权重目录
BATCH_SIZE = 1

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

# 1. 构造 Engine —— 一次性完成 setup、权重加载、CUDA graph 捕获
engine = Engine(
    EngineArgs(
        plugin="pi05",
        plugin_args=PI05Args(
            checkpoint_dir=CHECKPOINT_DIR,
            max_batch_size=BATCH_SIZE,
        ),
        config=EngineConfig(
            device=DeviceConfig(target="cuda", params_dtype=dtype),
            runtime=RuntimeConfig(use_cuda_graph=True),
        ),
    )
)
try:
    # 2. 构造 dummy request: 随机像素 + 单 token prompt
    input_ids = torch.zeros(
        BATCH_SIZE, cfg.tokenizer_max_length, dtype=torch.int64, device=device
    )
    input_ids[:, 0] = 2  # 任意非 pad token id
    request = PI05Request(
        pixel_values=torch.rand(
            BATCH_SIZE,
            3,
            3,
            cfg.vision.image_size,
            cfg.vision.image_size,
            dtype=dtype,
            device=device,
        ),
        input_ids=input_ids,
        lang_lens=torch.ones(BATCH_SIZE, dtype=torch.int64, device=device),
    )

    # 3. 跑一步推理
    actions = engine.step(request)
    print(f"action chunk shape={tuple(actions.shape)}, dtype={actions.dtype}")
finally:
    # 4. 释放 scheduler 缓冲、拆掉捕获的 cuda graph
    engine.close()