OpenPI π0 源码深度剖析：从模型架构、扩散策略到 C/S 部署实战

from_dict 方法负责将非结构化的字典数据转换为结构化的 Observation 对象，并处理图像格式的转换（如从 uint8 转为 float32）。

Actions 类型定义为浮点数数组，表示批量维度、动作时间步长和动作维度。

1.2 models/pi0.py 的实现

Pi0 是一个多模态扩散模型，继承自 BaseModel，使用 SigLIP 处理视觉输入、使用 Gemma 处理语言输入，实现了基于扩散的动作生成系统。

1.2.1 make_attn_mask：注意力掩码生成函数

该函数生成 Transformer 中使用的注意力掩码，控制 token 之间的注意力流动方式，支持纯因果注意力、前缀语言模型注意力等多种模式。

1.2.2 posemb_sincos：位置编码函数

使用正弦余弦函数实现位置编码，确保模型能够理解序列中的相对位置信息。

1.2.3 class Pi0Config

定义了模型的配置参数，包括 PaLI-Gemma 变体（如 gemma_2b）、动作专家底层结构（gemma_300m）、动作维度（默认 32）和动作序列长度（默认 50）。

1.2.3.1 inputs_spec：定义了 π0 模型本身接收的输入数据格式

通过 inputs_spec 函数定义了模型接收的输入规格，包括视觉输入（三视角 RGB 图像）、语言输入（分词后的文本 prompt）和状态输入。输出则是一个时序动作序列。

1.2.3.2 get_freeze_filter：参数冻结器

该函数用于决定在微调时对 VLM 和 Action Expert 的哪部分进行冻结。支持仅对动作专家使用 LoRA、仅对 PaLI-Gemma 使用 LoRA 或两者都使用 LoRA 等场景。

1.2.4 class Pi0：含特征嵌入、损失函数、推理

核心模型类，实现了多模态输入处理、扩散过程（训练去噪、推理采样）以及注意力机制。

1.2.4.1 初始化方法 init

组合了多个核心组件：PaLI-Gemma 模型（结合 LLM 和视觉模型）、线性投影层（用于动作 - 时间混合等）。

1.2.4.2 特征嵌入方法：embed_prefix/embed_suffix

• embed_prefix：处理图像和文本输入，创建前缀 token，皆为双向注意力。
• embed_suffix：处理机器人状态信息和噪声化动作信息，创建后缀 token。其中状态为单个 token，第一个动作 token 设置为单向注意力，其余动作 tokens 之间设置为双向注意力。

1.2.4.3 损失函数 compute_loss

训练时，对原始动作数据加噪，让模型学习预测所添加的真实噪声，计算预测噪声与实际噪声之间的均方误差。这种方式比直接预测原始动作更稳定。

1.2.4.4 推理函数 sample_actions

基于扩散模型的逆向采样过程。从纯噪声开始，通过多步骤逐渐'去噪'，最终生成符合条件分布的机器人动作序列。过程中使用了 KV 缓存优化推理速度。

1.3 语言模型实现：models/gemma.py

实现了 Gemma 语言模型的核心组件，定义了 RMSNorm、Embedder、Attention、FeedForward 等模块。

1.4 视觉模型实现：models/siglip.py

实现了视觉编码器，基于 Vision Transformer (ViT)，支持不同大小的模型变体。

1.5 tokenizer.py

提供了文本 tokenization 功能，包括 PaligemmaTokenizer 和 FASTTokenizer。

1.5.1 PaligemmaTokenizer 类

专门处理文本 prompt，下载预训练的 PaliGemma 分词模型，支持 token 序列的最大长度限制及填充截断逻辑。

1.5.2 FASTTokenizer 类

可同时处理文本和动作数据，支持将动作 token 映射到 PaliGemma 词汇表中。

1.6 lora.py

实现了 LoRA (Low-Rank Adaptation) 微调方法，支持 Einsum 和 FeedForward 模块的 LoRA 适配。

1.7 vit.py

Vision Transformer 实现。

第二部分 策略适配接口：src 下 policy 的全面分析与解读

src/openpi/policies 目录包含 BasePolicy、AlohaPolicy、DroidPolicy 等文件，定义了特定于机器人的输入和输出转换函数。

2.1 policy.py

实现了 Policy 类和 PolicyRecorder 类。

2.1.1 Policy 类

继承自 openpi_client.base_policy.BasePolicy，负责输入转换、模型推理和输出转换。

def infer(self, obs: dict) -> dict:
    inputs = self._input_transform(inputs)
    outputs = { "state": inputs["state"], "actions": self._sample_actions(...) }
    return self._output_transform(outputs)

2.1.2 PolicyRecorder

装饰器类，包装基础策略并在执行的同时将输入输出保存到磁盘。

2.2 policy_config.py

定义了 PolicyConfig 数据和 create_trained_policy 函数，用于从检查点创建策略实例。

2.3 policies/aloha_policy.py

实现了 Aloha 策略的输入输出处理和数据转换，包括 make_aloha_example、AlohaInputs 和 AlohaOutputs。

第三部分 模型训练的配置：src 下 training 模块的全面分析与解读

training 模块负责训练相关功能，包含 checkpoints、config、data_loader、optimizer 等。

3.1 配置系统 (config.py)

定义了 TrainConfig、DataConfigFactory 等配置类型，支持 ALOHA、DROID、LIBERO 等环境的预定义配置。

3.2 数据加载系统 data_loader.py

实现了数据集和数据加载器的接口，支持真实数据集（LeRobot）和模拟数据（FakeDataset）。

3.2.2 create_dataset

根据配置参数创建适合模型训练的数据集，处理时间戳计算和数据源加载。

3.2.3 transform_dataset

对数据集应用转换，包括数据清洗、归一化和模型特定转换。

3.3 优化器系统 (optimizer.py)

定义了多种学习率调度策略和优化器配置，如 AdamW、SGD。

3.4 检查点系统 (checkpoints.py)

负责模型状态的保存和恢复，使用 Orbax 库实现高效的检查点存储。

3.5 模型分片系统 (sharding.py)

实现分布式训练时的模型参数分片，提供 FSDP 的实现。

3.6 权重加载系统 (weight_loaders.py)

定义了 WeightLoader 协议，支持从检查点加载完整权重或部分加载。

3.7 辅助工具 (utils.py)

定义了 TrainState 数据类，封装训练过程的状态。

第四部分 模型的训练与部署：基于客户端 - 服务器 C/S 架构

packages/openpi-client 和 scripts 模块分别负责客户端通信和服务器端功能。

4.1 packages/openpi-client

提供与策略服务器通信的接口，使用 WebSocketClientPolicy 连接服务器，处理观察数据和动作序列。

4.2 scripts(策略服务器)

包含数据处理、模型训练、模型推理的多个脚本。

4.2.3 serve_policy.py

启动策略服务，用于模型推理。支持定义特定任务的文本指令 prompt。

def main(args: Args) -> None:
    policy = create_policy(args)
    server = websocket_policy_server.WebsocketPolicyServer(...)
    server.serve_forever()

4.2.5 train.py

训练模型入口点，集成了模型初始化、训练循环、日志记录等功能。

第五部分 examples：各种机器人平台及策略客户端的示例实现

examples 模块提供了不同平台的示例，包括 aloha_real、aloha_sim、droid、libero 等。

5.1 aloha_real

用于控制真实 ALOHA 双臂机器人的完整实现。

5.1.1 核心架构

包含硬件常量定义、数据转换工具、环境接口（AlohaRealEnvironment）和主控制流程（main.py）。

5.1.2 系统工作流程与部署方式

1. 初始化阶段：启动 ROS 节点，初始化双臂机器人，建立 WebSocket 连接。
2. 运行时循环：获取观察数据，发送到策略服务器，接收动作序列，执行动作。
3. 部署方式：支持 Docker 部署和本地部署（需启动三个终端分别运行客户端、ROS 驱动和推理服务）。