
说明本文是用于在debug pi0代码库时的笔记更新中pi0 代码库 pi0官方代码库整体结构这个项目可以按三层看环境与客户端main.py启动 LIBERO 仿真、整理图像和机器人状态并通过 WebSocket 把 observation 发给策略服务。它负责评测流程不直接运行模型。策略服务与数据变换serve_policy.py创建策略和 WebSocket 服务policy_config.py加载配置、checkpoint 和输入输出变换policy.py的Policy.infer()把输入变成模型所需格式再调用模型采样动作。模型与训练models放模型实现和配置config.py定义诸如pi0_libero、pi0_libero_low_mem_finetune的训练配置。数据变换、训练流程等分别在transforms.py和training。实际推理大致是LIBERO main.py→client.infer(element)→Policy.infer()→Pi0.sample_actions()→ 生成一段动作 → 返回给仿真环境执行。看 PaliGemma 的推荐顺序从pi0.py的Pi0.__init__()、embed_prefix()、embed_suffix()和sample_actions()开始。这里能看到图像、指令、机器人状态如何变成模型输入以及模型如何生成动作。看siglip.py的 ViTModule了解图像如何编码成视觉 tokens。看gemma.py的Config/get_config()、Embedder和Module。这里实现 Gemma 文本 embedding 和 TransformerModule支持多个 expert。看pi0_config.py的Pi0Config确认模型 variant、动作维度和输入规格。标准配置默认是paligemma_variantgemma_2b、action_expert_variantgemma_300m。再看config.py中 LIBERO 对应的训练配置确认你正在跑的是pi0_libero、LoRA 变体还是pi05_libero。这里有个容易混淆的点JAX 主路径里没有一个单独叫PaliGemma的类。Pi0 把SigLIP 图像编码器和Gemma 语言 Transformer组合起来并在同一套模型结构里加入 action expert。文本通过gemma.Module.embed()转 embedding图像通过 SigLIP 编成视觉 tokens随后在Pi0.embed_prefix()中组合。如果服务端加载的是 PyTorch checkpoint再看gemma_pytorch.py和pi0_pytorch.py这一分支使用 Hugging Face 的PaliGemmaForConditionalGeneration。policy_config.py会根据 checkpoint 文件选择 JAX 或 PyTorch 路径。数据流转路径按推理时的数据流从头读最直观。可以先把它想成仿真环境产生 observation策略服务把它转换成模型输入Pi0 生成一段动作再把动作交还给环境。仿真环境整理输入看main.py。重点是eval_libero()里构造element的部分主相机图像、腕部图像、机器人状态和任务指令都在这里打包client.infer(element)是送去模型服务的调用点。模型返回的actions会放进action_plan按步执行。WebSocket 收发从serve_policy.py的main()看服务启动再看websocket_policy_server.py如何接收请求并调用策略。这样能把 LIBERO 客户端和模型所在进程连起来。输入变成模型可用的数据看policy_config.py的create_trained_policy()确认策略由配置、checkpoint 和 transforms 组装而成接着看policy.py的Policy.infer()这是输入变换、加 batch、创建Observation、调用模型和输出变换的总入口。LIBERO 专属适配先读libero_policy.py的LiberoInputs和LiberoOutputs。它把 LIBERO 的键映射成模型统一使用的图像、状态和 prompt并把模型动作裁成 LIBERO 所需的 7 维。再看config.py里的LeRobotLiberoDataConfig.create()了解这些 transforms 是怎么接入策略的。这里有个区别dataset 的RepackTransform用于训练数据推理时客户端已经按约定发送键所以通常直接应用 LIBERO 输入适配和模型 transforms。通用输入预处理与文本 token 化看transforms.py中相关 transforms以及config.py的ModelTransformFactory。然后看model.py的Observation和preprocess_observation()理解模型真正收到的结构以及图像、mask 等如何进入模型。模型如何用图像和语言生成动作回到pi0.py按Pi0.__init__()→embed_prefix()→embed_suffix()→sample_actions()阅读。重点跟两路输入图像经self.PaliGemma.img(...)编成视觉 tokensprompt 经self.PaliGemma.llm(..., methodembed)编成文本 tokens。它们组成 prefix和状态、带噪动作组成的 suffix 一起经过 Transformer最后由动作头和 Flow Matching 采样得到动作序列。跟动作回到环境回看Policy.infer()的输出变换以及main.py里取出并执行actions的部分。这样数据流就闭环了。建议第一遍先略过 Transformer 内部数学细节只追踪每一站的键名、张量形状和变换前后含义。最短阅读路线是main.py→serve_policy.py/ WebSocket server →policy.py→libero_policy.py→ModelTransformFactory/model.py→pi0.py→ 回到policy.py和main.py。如果目标转向训练数据流再补看LeRobotLiberoDataConfig.create()和training/config.py中对应的pi0_libero配置它们决定训练集如何转换、使用哪个模型配置和 checkpoint。注意【一般代码仓库不要整体看要先看数据流转的代码在看一些工程实现比如是怎么实现通讯的等这些细枝末节】