π0 源码深度解析
前言
π0 (OpenPI) 是一套用于通用机器人控制的视觉 - 语言 - 动作(VLA)模型框架。本文旨在深入剖析其源码实现,从模型架构设计、训练配置到 C/S 架构下的实际部署流程,帮助开发者理解如何基于 PaLI-Gemma 和扩散策略去噪生成动作。
第一部分 π0 模型架构的实现
1.1 models/model.py:核心基础模型定义
这是模型框架的核心文件,定义了基础的抽象类和数据结构。
BaseModelConfig: 所有模型配置的抽象基类BaseModel: 所有模型实现的抽象基类Observation: 保存模型输入的数据类Actions: 定义动作数据格式
1.1.1 基础组件和关键常量
模型类型枚举定义了两种支持的模型类型:
class ModelType(enum.Enum):
"""Supported model types."""
PI0 = "pi0"
PI0_FAST = "pi0_fast"
图像输入配置定义了模型期望接收三个视角的图像:
IMAGE_KEYS = (
"base_0_rgb",
"left_wrist_0_rgb",
"right_wrist_0_rgb",
)
1.1.2 Observation 类详解
Observation 类是 OpenPI 框架中的核心数据结构,用于存储和管理模型的输入数据。它包含了机器人感知系统收集的所有必要信息,包括低维度的机器人状态、图像数据及掩码等。
class Observation(Generic[ArrayT]):
"""Holds observations, i.e., inputs to the model."""
images: dict[str, at.Float[ArrayT, "*b h w c"]]
image_masks: dict[str, at.Bool[ArrayT, "*b"]]
state: at.Float[ArrayT, "*b s"]
# Tokenized prompt fields for FAST autoregressive model...
from_dict 方法负责从非结构化的字典数据创建 Observation 对象,并处理图像格式的转换(如将 uint8 转换为 float32)。
1.2 models/pi0.py 的实现
Pi0 是一个多模态扩散模型,继承自 BaseModel,使用 SigLIP 处理视觉输入、使用 Gemma 处理语言输入,实现了基于扩散的动作生成系统。


