概述
本系列为我在vla上的同步实践笔记，之前的实践中我大多都是在调用指令、修改参数，没有去深入了解相关的底层原理，因此在本系列中我将尽可能地探究底层原理。本文将展示我对PI05的基本原理和代码实践。
1. 原理探究
在进行代码实践之前，我又去回顾了一下pi05的相关基础原理，下面仅做出我浅薄的学习理解分享。
1.1. pi0
为什么需要先探究pi0的原理？就我的理解而言，pi0.5借鉴了pi0和下一小节会讲到的pi fast的思想，因此有必要在讲pi0.5之前先简单讲下pi0和pi fast。
pi0具体的原理我就不细讲了，很多文章和论文都有详解：
π0 细读 | 具身星图 (zhuyun97.github.io)
2410.24164 π_0: A Vision-Language-Action Flow Model for General Robot Control (arxiv.org)
就我的理解而言，关于pi0需要重点知道如下要点：
我们可作出如下总结：
pi0跟之前的VLA模型如RT2、openvla采用自回归逐个输出动作token不同，其一次性输出一长串动作，支持其输出一长串动作的核心机制是flow matching（流匹配）。其模型核心由负责语义和视觉理解的VLM和负责动作生成的动作专家组成。
下面，需要讲一下流匹配和扩散模型的对比。
1.2. 流匹配与扩散模型
流匹配和扩散模型都是生成领域的重要方法，比如图像生成、音频生成、机器人动作生成，二者都是有效的应用方法。
它们的起点和终点类似，都是从带噪声的初始出发，一点一点还原为期望的目标，只不过中间还原的方法不同。以当前的机器人动作生成举例子，其对比如下：
通俗点说，扩散模型预测的是噪声，通过消除噪声从而得到正确动作；而流匹配是在噪声场景下预测当前应有的速度变化，从而实现从噪声中走出"迷雾"，得到正确答案。
pi0的动作生成方法便是基于流匹配，初始拿到的是带噪声的动作，通过推理一点一点迭代成正确的动作序列。
1.3. pi fast
与pi0采用流匹配生成连续动作不同，pi fast反而是走的RT2、openvla的老路：采用自回归逐个生成动作token，但它却不如前者那般慢和无法做到精细控制，核心如下：
简单来讲，pi fast不再向之前的自回归模型那样老老实实地端到端输出单个token，而是采用一系列方法对长串动作token做压缩，在输出时给出压缩后的token，再解码为正常的完整动作token，如此便做到了使用自回归也能实现连续控制和精细控制。
编解码流程可以看看下面：
高频动作
→ 大量相邻帧非常相似
→ 离散后产生大量重复 token
→ 序列很长、有效信息密度很低
→ 模型甚至只复制上一个 token，也能得到较低 loss
→ 训练难、推理慢、动作细节学不好
一整段连续动作
→ 用 DCT 提取整体运动趋势和细节
→ 量化
→ 用 BPE 压缩
→ 变成更短但信息密度更高的 token 序列
模型生成较短的 FAST token
→ 反向 BPE 解码
→ 反量化
→ 逆 DCT
→ 还原成连续动作轨迹1.4. pi0.5与pi0和pi fast的关联
关于pi0.5的文章和论文说明也有很多，我这里也不做详细说明了：
【VLA 系列】万字详解 PI0.5 - 知乎 (zhihu.com)
我只说明pi0.5与pi0和pi fast的关联：
-
其模型结构与pi0一致，核心都是VLM+动作专家
-
其在预训练阶段，不包含动作专家，而是使用pi fast的压缩机制通过离散token吃下海量的预训练数据。
-
后训练阶段再加入动作专家，使用流匹配生成连续动作
2. 代码实践
大致了解完了原理，我便开始了代码实践，大致分成了以下6步：
-
数据预处理：数据在初始阶段经过了怎样的处理？
-
图像编码与前缀拼接：数据进入VLM之前，图像数据和文本数据分别经过了什么？
-
VLM前向与KV缓存：数据在VLM内部经过了什么？
-
动作专家单轮预测：数据在VLM经过处理后，对动作专家的帮助是什么？
-
完整动作生成与后处理：动作专家完整的工作流程是怎么样的？
-
训练前向与loss计算：训练时是如何指定目标和loss的？
2.1. 数据预处理
我进行了如下的代码实践：
import os
os.environ["HF_HOME"] = "/root/autodl-tmp/huggingface"
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["HF_DATASETS_OFFLINE"] = "1"
import json
from pathlib import Path
import torch
from torch.utils.data import default_collate
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.processor import PolicyProcessorPipeline
from lerobot.processor.converters import batch_to_transition, transition_to_batch
# 导入时注册π0.5专用的state→文本处理步骤
from lerobot.policies.pi05.processor_pi05 import (
Pi05PrepareStateTokenizerProcessorStep,
)
def main():
# 控制pytorch张量打印格式，小数点后5位，关闭科学计数法
torch.set_printoptions(precision=5, sci_mode=False)
checkpoint = Path(
"/root/autodl-tmp/workspace/outputs/week05/"
"b8_vlmlora_expertfull_20000steps_20261018/"
"checkpoints/020000/pretrained_model"
)
# 全局配置
config = json.loads((checkpoint / "train_config.json").read_text())
# 预处理配置
processor_config = json.loads(
(checkpoint / "policy_preprocessor.json").read_text()
)
root = Path(config["dataset"]["root"])
info = json.loads((root / "meta/info.json").read_text())
chunk_size = config["policy"]["chunk_size"]
dataset = LeRobotDataset(
repo_id=config["dataset"]["repo_id"],
root=root,
# 这个参数保证每一帧的数据还包括后面的50步动作
# [0/20, 1/20, 2/20, ..., 49/20]，lerobot会把对应时间偏移处的动作取出来组成序列
delta_timestamps={
"action": [i / info["fps"] for i in range(chunk_size)]
},
video_backend="pyav",
return_uint8=False,
)
# 取第0帧做观测对象
sample = dataset[0]
# 加batch维度，如此操作之后sample将带有batch并传给batch变量
batch = default_collate([sample])
print(f"数据目录: {root}")
print(f"episode_index: {batch["frame_index"]}")
print(f"原始任务: {batch["task"]}")
# 读取保存的处理流程和归一化统计量。
preprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=checkpoint,
config_filename="policy_preprocessor.json",
local_files_only=True,
# 数据入口转换
to_transition=batch_to_transition,
# 数据出口转换
to_output=transition_to_batch,
)
# step_through是现有接口：依次返回处理前和每一步处理后的数据。
for i, transition in enumerate(preprocessor.step_through(batch)):
current = transition_to_batch(transition)
if i == 0:
name = "原始batch"
else:
name = processor_config["steps"][i - 1]["registry_name"]
print(f"\n========== {i}. {name} ==========")
if i == 0 or name == "normalizer_processor":
print("state形状：", tuple(current["observation.state"].shape))
print("state数值：", current["observation.state"][0])
print("action形状：", tuple(current["action"].shape))
print("首步action：", current["action"][0, 0])
if name == "pi05_prepare_state_tokenizer_processor_step":
print("拼接后的完整文本：")
print(current["task"][0])
if name == "tokenizer_processor":
ids = current["observation.language.tokens"]
mask = current["observation.language.attention_mask"]
print("token形状：", tuple(ids.shape))
print("有效token数：", mask[0].sum().item())
print("前40个token ID：", ids[0, :40])
tokenizer = preprocessor.steps[i - 1].input_tokenizer
print("有效token解码后的文本：")
print(tokenizer.decode(ids[0][mask[0]].tolist()))
print("\n========== 预处理最终输出 ==========")
for key, value in current.items():
if torch.is_tensor(value):
print(f"{key}: shape={tuple(value.shape)}, dtype={value.dtype}")
assert current["observation.state"].shape == (1, 8)
assert current["action"].shape == (1, chunk_size, 7)
assert current["observation.language.tokens"].shape == (1, 200)
print("\nPIPER_PREPROCESS_CHECK_PASS")
if __name__ == "__main__":
main()代码选择了自己训练的checkpoint的预训练器，以及真实数据集，通过PolicyProcessorPipeline来观察每个阶段预处理后的数据状态。
输出如下：
(lerobot) root@autodl-container-33214eb7ac-384a3139:~/autodl-tmp/workspace# python -u "/root/autodl-tmp/workspace/my_test/week6/A2_1_piper_preprocess.py"
数据目录: /root/autodl-tmp/workspace/datasets/week5/piper_black_bowl_b1_v2_lerobot
episode_index: tensor([0])
原始任务: ['pick up the black bowl between the plate and the ramekin and place it on the plate']
========== 0. 原始batch ==========
state形状： (1, 8)
state数值： tensor([-0.00985, 0.00141, 1.21055, -0.00001, -4.22814, 0.00001, 0.03500,
-0.03500])
action形状： (1, 50, 7)
首步action： tensor([ 0.07500, -0.00000, 0.00000, 0.00000, -0.00000, -0.00000, -1.00000])
========== 1. rename_observations_processor ==========
========== 2. to_batch_processor ==========
========== 3. normalizer_processor ==========
state形状： (1, 8)
state数值： tensor([ 0.57175, -3.76730, 2.25397, 0.87199, -1.19538, 1.21714, 0.83498,
-0.83718])
action形状： (1, 50, 7)
首步action： tensor([-0.48262, -0.60115, -0.42069, 1.12819, 0.85930, -0.74884, -0.82515])
========== 4. pi05_prepare_state_tokenizer_processor_step ==========
拼接后的完整文本：
Task: pick up the black bowl between the plate and the ramekin and place it on the plate, State: 201 -1 255 239 -1 255 234 20;
Action:
========== 5. tokenizer_processor ==========
token形状： (1, 200)
有效token数： 56
前40个token ID： tensor([ 2, 7071, 235292, 4788, 908, 573, 2656, 14581, 1865,
573, 8811, 578, 573, 117064, 4074, 578, 2040, 665,
611, 573, 8811, 235269, 3040, 235292, 235248, 235284, 235276,
235274, 728, 235274, 235248, 235284, 235308, 235308, 235248, 235284,
235304, 235315, 728, 235274])
有效token解码后的文本：
<bos>Task: pick up the black bowl between the plate and the ramekin and place it on the plate, State: 201 -1 255 239 -1 255 234 20;
Action:
========== 6. device_processor ==========
========== 预处理最终输出 ==========
action: shape=(1, 50, 7), dtype=torch.float32
action_is_pad: shape=(1, 50), dtype=torch.bool
index: shape=(1,), dtype=torch.int64
task_index: shape=(1,), dtype=torch.int64
episode_index: shape=(1,), dtype=torch.int64
timestamp: shape=(1,), dtype=torch.float32
observation.images.image: shape=(1, 3, 256, 256), dtype=torch.float32
observation.images.image2: shape=(1, 3, 256, 256), dtype=torch.float32
observation.state: shape=(1, 8), dtype=torch.float32
observation.language.tokens: shape=(1, 200), dtype=torch.int64
observation.language.attention_mask: shape=(1, 200), dtype=torch.bool
PIPER_PREPROCESS_CHECK_PASS从输出可以看出，这一步主要做了以下几件事：
- 对输入的状态、动作数值做归一化操作，归一化方法为MEAN_STD，公式如下：
这一步是为了让不同维度的数值尺度接近，方便模型学习。因为原始数据中，位移、轴角、夹爪数据的大小和变化幅度可能不同，在训练时会造成不同的loss和梯度分布，即模型更可能去照顾变化幅度大的维度，进行归一化后能让数值各自的分布幅度接近，让模型"雨露均沾"。
-
8维状态数值被离散成了整数，拼在了任务文本后面。这一步实现了任务文本中夹带着状态信息，让模型学习时不仅知道当前任务是啥，也能知道当前状态怎么样。
-
拼好的任务文本+状态整数被统一转换为token id，这一步是为了在下一步通过查词表构建词嵌入向量，方便给VLM处理。
张量维度变化：
-
初始state为(1,8)，即1批次，8维，动作为(1,50,7)，即1批次，50个动作，每个动作有7维
-
处理之后，任务文本+state为(1,200)，即1批次，200维，其中200维并不全为有效信息，由同为(1,200)的attention_mask标记哪些位置的token才是有效的，哪些是填充的。
2.2. 图像编码与前缀拼接
代码如下：
import torch
from A2_1_piper_preprocess import main as prepare_sample
from lerobot.configs import PreTrainedConfig
from lerobot.policies.factory import make_policy
# 关闭梯度追踪，专门用于推理阶段，等价于 with torch.inference_mode():
@torch.inference_mode()
def main():
checkpoint, dataset, batch = prepare_sample()
# 1. 加载Piper第五轮模型：基座 + VLM LoRA + 全量动作专家。
# 取架构
config = PreTrainedConfig.from_pretrained(checkpoint)
# 取权重
config.pretrained_path = str(checkpoint)
config.device = "cuda"
print("\n========== 加载Piper模型 ==========")
assert config.use_peft
policy = make_policy(cfg=config, ds_meta=dataset.meta)
policy.eval()
policy.requires_grad_(False)
# 去掉PEFT的外层包装，访问PI05Policy及其内部模型。
# LoRA仍然保留在内部网络中，并没有被移除。
base_policy = policy.get_base_model()
model = base_policy.model
# 2. 调用模型实际使用的图像预处理函数。
# 两路真实图像缩放到224×224，范围从[0,1]变为[-1,1]；
# 按checkpoint配置补一条空相机，并将其mask设为False。
images, image_masks = base_policy._preprocess_images(batch)
expected_keys = list(base_policy.config.image_features)
image_keys = (
[key for key in expected_keys if key in batch]
+ [key for key in expected_keys if key not in batch]
)
print("\n========== 1. 模型内部图像预处理 ==========")
for key, image, mask in zip(image_keys, images, image_masks, strict=True):
print(f"\n相机：{key}")
print("形状：", tuple(image.shape))
print("数值范围：", image.min().item(), image.max().item())
print("相机是否有效：", mask.tolist())
tokens = batch["observation.language.tokens"]
text_mask = batch["observation.language.attention_mask"]
print("\n文本token形状：", tuple(tokens.shape))
print("有效文本token数：", text_mask[0].sum().item())
# 3. 调用真实的prefix构建函数。
# 内部执行：
# 图像 → SigLIP → 视觉投影
# 文本token → Embedding查表
# 最后沿序列维度拼接。
prefix, prefix_mask, prefix_att_mask = model.embed_prefix(
images, image_masks, tokens, text_mask
)
# 4. 从拼好的prefix中，分别取出每一路特征查看形状。
# 所有相机都已统一分辨率，因此每路视觉token数量相同。
vision_config = model.paligemma_with_expert.paligemma.config.vision_config
# 每个小方块的尺寸
patch_size = vision_config.patch_size
height, width = images[0].shape[-2:]
tokens_per_image = (height // patch_size) * (width // patch_size)
print("\n========== 2. 各路编码后的特征 ==========")
offset = 0
# 遍历每个相机
for key in image_keys:
# 第一次循环（offset=0）：切下 0 到 256 个 Token，这是第一张图的。
# 第二次循环（offset=256）：切下 256 到 512 个 Token，这是第二张图的。
image_features = prefix[:, offset:offset + tokens_per_image]
print(f"{key}：{tuple(image_features.shape)}")
offset += tokens_per_image
text_features = prefix[:, offset:]
print("文本Embedding：", tuple(text_features.shape))
print("\n========== 3. 拼接后的prefix(前缀) ==========")
print("prefix形状：", tuple(prefix.shape))
print("prefix有效位置mask形状：", tuple(prefix_mask.shape))
print("有效位置总数：", prefix_mask[0].sum().item())
expected_length = len(images) * tokens_per_image + tokens.shape[1]
expected_valid = (
sum(mask.sum().item() for mask in image_masks) * tokens_per_image
+ text_mask.sum().item()
)
assert prefix.shape[1] == expected_length
assert prefix_mask.shape == prefix.shape[:2]
assert prefix_mask[0].sum().item() == expected_valid
assert torch.equal(prefix_mask[:, offset:], text_mask)
print("\nPIPER_PREFIX_CHECK_PASS")
return model, batch, prefix, prefix_mask, prefix_att_mask
if __name__ == "__main__":
main()核心是调用了源码中的_preprocess_images、embed_prefix函数，让数据被处理。
核心输出：
========== 加载Piper模型 ==========
The PI05 model is a direct port of the OpenPI implementation.
This implementation follows the original OpenPI structure for compatibility.
Original implementation: https://github.com/Physical-Intelligence/openpi
Loading model from: /root/autodl-tmp/huggingface/hub/models--lerobot--pi05_libero_base/snapshots/a217bfd3b14673cf2ce597e69997ab21866438dd
✓ Loaded state dict from model.safetensors
WARNING:root:Vision embedding key might need handling: paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.patch_embedding.bias
WARNING:root:Vision embedding key might need handling: paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.patch_embedding.weight
Remapped 812 state dict keys
All keys loaded successfully!
========== 1. 模型内部图像预处理 ==========
相机：observation.images.image
形状： (1, 3, 224, 224)
数值范围： -0.9963985681533813 0.843137264251709
相机是否有效： [True]
相机：observation.images.image2
形状： (1, 3, 224, 224)
数值范围： -0.9831932783126831 0.7266906499862671
相机是否有效： [True]
相机：observation.images.empty_camera_0
形状： (1, 3, 224, 224)
数值范围： -1.0 -1.0
相机是否有效： [False]
文本token形状： (1, 200)
有效文本token数： 56
========== 2. 各路编码后的特征 ==========
observation.images.image：(1, 256, 2048)
observation.images.image2：(1, 256, 2048)
observation.images.empty_camera_0：(1, 256, 2048)
文本Embedding： (1, 200, 2048)
========== 3. 拼接后的prefix(前缀) ==========
prefix形状： (1, 968, 2048)
prefix有效位置mask形状： (1, 968)
有效位置总数： 568
PIPER_PREFIX_CHECK_PASS从输出可以看出，这一步主要做了以下几件事：
-
先对图像做归一化处理，缩放到224×224，让数值范围到-1到1
-
调用源码里的前缀构建函数embed_prefix，得到视觉信息、文本信息拼接好的token前缀。在embed_prefix内，具体经过了embed_image和embed_language_tokens两个函数，即视觉数据和语言数据经过不同的处理：
1. 图像经过SigLIP视觉网络＋视觉投影
img_emb = self.paligemma_with_expert.embed_image(img)
2. 文本token查Embedding表
lang_emb = self.paligemma_with_expert.embed_language_tokens(tokens)
3. 沿序列维度拼接
embs = torch.cat(embs, dim=1)
也就是说，这一步已经在用视觉模型处理图像，得到了提取出来的特征，拼在文本数据的后面，现在的数据已经包含"任务是什么"+"当前状态"+"当前视觉有哪些信息"。
- 从拼好的长前缀中，挨个取出图像、文本，核对一下
张量维度变化：
-
初始的图像维度为(1,3,256,256)，即1批，3通道，256x256，经过归一化处理和缩放后，变成了(1,3,224,224)，即数值进行了归一化处理，图像大小也变小了。
-
(1,3,224,224)的图像经过视觉编码器后，输出为(1, 256, 2048)，三路图像都经过了这样的处理
-
文本数据最开始为(1,200)，经过查表后变成了(1,200,2048)，维度与图像数据对齐，方便拼接
-
拼接后的维度变成(1, 256 * 3 + 200, 2048)=(1, 968, 2048)，这就是prefix(前缀)的最终维度，其有效的位置总数为968 - 之前文本padding的144 - 256的空图像 = 568。
2.3. VLM前向与KV缓存
代码如下：
import torch
from A2_2_piper_prefix import main as prepare_prefix
# 注意力掩码
from lerobot.policies.common.vla_utils import (
make_att_2d_masks,
prepare_attention_masks_4d,
)
# 关闭梯度追踪，专门用于推理阶段，等价于 with torch.inference_mode():
@torch.inference_mode()
def main():
# 接着A2_2继续，模型只加载一次。
model, batch, prefix, prefix_mask, prefix_att_mask = prepare_prefix()
vlm = model.paligemma_with_expert
# 1. 构建注意力矩阵：规定各位置可以读取哪些位置的信息。
# prefix_mask区分有效位置与padding；
# prefix_att_mask规定序列的注意力分块方式。
# prefix_att_mask 通常是一个一维的标记数组（形状 (B, L)），它给每一个 Token 打上了一个"标签"，用来标记每个token是属于图像还是文字
# make_att_2d_masks拿到这个prefix_att_mask后，能自己生成对应的二维注意力掩码矩阵
attention_2d = make_att_2d_masks(prefix_mask, prefix_att_mask)
# 转为Transformer接口需要的四维格式。
# 可以关注的位置填0，被屏蔽的位置填很大的负数。
attention_4d = prepare_attention_masks_4d(attention_2d)
# 位置编号按有效位置累计，padding不推进编号。
position_ids = torch.cumsum(prefix_mask, dim=1) - 1
print("\n========== 1. Gemma输入 ==========")
print("prefix形状：", tuple(prefix.shape))
print("prefix数据类型：", prefix.dtype)
print("注意力矩阵形状：", tuple(attention_2d.shape))
print("四维注意力矩阵形状：", tuple(attention_4d.shape))
print("位置编号形状：", tuple(position_ids.shape))
print("第0个位置可以关注的位置数：", attention_2d[0, 0].sum().item())
# 2. 沿用源码sample_actions中处理prefix的调用方式
# 找到底层的语言模型
language_model = vlm.paligemma.model.language_model
# 修改配置，保证 KV Cache 能提取、输入格式兼容
language_model.config._attn_implementation = "eager"
# [prefix, None]表示本次只处理prefix分支，先不处理动作专家分支。
# use_cache=True：保留每层注意力的Key/Value，供后续动作生成复用。
outputs, past_key_values = vlm.forward(
attention_mask=attention_4d,
# 传入位置编码，让模型知道token的先后顺序
position_ids=position_ids,
past_key_values=None,
inputs_embeds=[prefix, None],
# 保存KV
use_cache=True,
)
prefix_output = outputs[0]
print("\n========== 2. Gemma处理后的特征 ==========")
print("输入形状：", tuple(prefix.shape))
print("输出形状：", tuple(prefix_output.shape))
print("输出数据类型：", prefix_output.dtype)
print("Gemma层数：", len(language_model.layers))
print("输出数值片段：")
print(prefix_output[0, 0, :8])
# 3. 查看缓存结构，而不是打印全部缓存内容。
first_layer = past_key_values.layers[0]
keys = first_layer.keys
values = first_layer.values
print("\n========== 3. KV缓存 ==========")
print("缓存类型：", type(past_key_values).__name__)
print("缓存层数：", len(past_key_values.layers))
print("第0层Key形状：", tuple(keys.shape))
print("第0层Value形状：", tuple(values.shape))
print("缓存序列长度：", past_key_values.get_seq_length())
assert prefix_output.shape == prefix.shape
assert outputs[1] is None
assert len(past_key_values.layers) == len(language_model.layers)
assert past_key_values.get_seq_length() == prefix.shape[1]
assert keys.shape[-2] == prefix.shape[1]
assert values.shape == keys.shape
print("\nPIPER_GEMMA_PREFIX_CHECK_PASS")
return model, batch, prefix_mask, past_key_values
if __name__ == "__main__":
main()这里通过调用pi0.5源码中的PaliGemmaWithExpertModel.forward、PiGemmaModel.forward得到VLM模型处理数据后的结果，并打印分析。
核心输出如下：
========== 1. Gemma输入 ==========
prefix形状： (1, 968, 2048)
prefix数据类型： torch.float32
注意力矩阵形状： (1, 968, 968)
四维注意力矩阵形状： (1, 1, 968, 968)
位置编号形状： (1, 968)
第0个位置可以关注的位置数： 568
========== 2. Gemma处理后的特征 ==========
输入形状： (1, 968, 2048)
输出形状： (1, 968, 2048)
输出数据类型： torch.bfloat16
Gemma层数： 18
输出数值片段：
tensor([ 0.22168, 3.45312, -0.58203, -0.06592, -0.45312, -0.06641, 0.10400,
-0.48828], device='cuda:0', dtype=torch.bfloat16)
========== 3. KV缓存 ==========
缓存类型： DynamicCache
缓存层数： 18
第0层Key形状： (1, 1, 968, 256)
第0层Value形状： (1, 1, 968, 256)
缓存序列长度： 968
PIPER_GEMMA_PREFIX_CHECK_PASS从输出可以看出，这一步主要做了以下几件事：
- 通过上一步计算出的prefix_mask和prefix_att_mask生成注意力矩阵，也就是决定注意力机制中"谁能看谁"。prefix_mask用来区分前缀中哪些位置有效，prefix_att_mask类似于给有效的信息分块，比如说明568个位置中哪些部分属于image1，哪些属于image2，哪些属于文本。具体的注意力矩阵如下：
简单来说，就是有效的视觉和文本信息全都互相能看见，无效的信息都被排除在外。
- 经过模型主干的18层transformer，缓存前向过程中各层的KV数据(KV cache)，这个缓存非常重要，VLM对动作专家的贡献可以说就是这个KV缓存，因为K代表资料特征，V代表资料内容，KV联合一起就是在告诉动作专家当前场景的特征和内容。后面动作专家自己做推理的时候，需要联合这个缓存的KV跟自己的KV一起做注意力计算，最后得到一轮输出，每次输出50步更新方向，经过10轮这样的更新和积分，得到最终的50步动作。
也就是说，输入的数据到这里就不再进行下一步的流动了，核心的信息已经被保留到了KV缓存里，后续动作专家需要用到的只是这个缓存。
张量维度变化：
- 这一步中张量在输入输出维度没有发生变化，因为这一步并没有改变数据，而是提取了其信息，收获到的内容都被KV缓存保存了。
2.4. 动作专家单轮预测
代码如下：
import torch
from A2_3_piper_gemma import main as prepare_conditions
# 关闭梯度追踪，专门用于推理阶段，等价于 with torch.inference_mode():
@torch.inference_mode()
def main():
# 复用A2_1～A2_3，模型只加载一次。
model, _, prefix_mask, past_key_values = prepare_conditions()
device = prefix_mask.device
batch_size = prefix_mask.shape[0]
# 1. 从随机噪声开始，形状与模型内部动作块一致。
# 模型内部补齐为32维，最终输出时才裁回7维。pi0.5 内部统一把动作维度补齐到 32 维
torch.manual_seed(1000)
# 生成带噪动作块
x_t = model.sample_noise(
(
batch_size,
model.config.chunk_size,
model.config.max_action_dim,
),
device,
)
# t=1表示生成起点：当前动作块还是纯噪声。
# 这是生成过程的时间参数，不是仿真时间戳
# 在扩散模型/流匹配（Flow Matching）里，时间 t 的取值通常是 0 到 1：1代表纯噪声，100%的真实噪声，0%的真实动作
timestep = torch.ones(
batch_size, dtype=torch.float32, device=device
)
print("\n========== 1. 动作专家的原始输入 ==========")
print("带噪动作块形状：", tuple(x_t.shape))
print("时间参数：", timestep)
print("prefix缓存长度：", past_key_values.get_seq_length())
expert = model.paligemma_with_expert.gemma_expert.model
print("动作专家Transformer层数：", len(expert.layers))
observed = {}
# hook只观察数据，不修改网络输入输出。
# 当第 0 层准备执行 forward 时，先暂停一下，把传进来的参数看一眼
def inspect_expert_input(module, args, kwargs):
# 后缀suffix 前缀prefix
suffix = args[0]
time_condition = kwargs["adarms_cond"]
attention_mask = kwargs["attention_mask"]
observed["suffix_shape"] = tuple(suffix.shape)
print("\n========== 2. 专家第0层收到的数据 ==========")
print("动作suffix形状：", tuple(suffix.shape))
print("时间条件形状：", tuple(time_condition.shape))
print("注意力矩阵形状：", tuple(attention_mask.shape))
print(
"传入的prefix KV缓存长度：",
kwargs["past_key_values"].get_seq_length(),
)
print(
"第0个动作位置可以关注的位置数：",
(attention_mask[0, 0, 0] == 0).sum().item(),
)
# 在最终动作输出投影之前，查看专家处理后的特征。
def inspect_expert_output(module, args):
features = args[0]
observed["expert_shape"] = tuple(features.shape)
print("\n========== 3. 动作输出投影前 ==========")
print("专家处理后的特征形状：", tuple(features.shape))
# 注册狗子函数
handles = [
expert.layers[0].register_forward_pre_hook(
inspect_expert_input, with_kwargs=True
),
model.action_out_proj.register_forward_pre_hook(
inspect_expert_output
),
]
original_cache_length = past_key_values.get_seq_length()
try:
# 2. 调用源码中的完整单轮计算。
# 内部完成：
# 动作/时间编码 → 读取prefix KV → 专家Transformer → 输出投影。
velocity = model.denoise_step(
prefix_pad_masks=prefix_mask,
past_key_values=past_key_values,
x_t=x_t,
timestep=timestep,
)
finally:
for handle in handles:
handle.remove()
print("\n========== 4. 本轮输出 ==========")
print("更新方向v_t形状：", tuple(velocity.shape))
print("更新方向数据类型：", velocity.dtype)
print("首步v_t的前7维：", velocity[0, 0, :7])
print("原prefix KV缓存长度：", past_key_values.get_seq_length())
assert len(observed) == 2
assert velocity.shape == x_t.shape
assert torch.isfinite(velocity).all().item()
assert past_key_values.get_seq_length() == original_cache_length
print("\nPIPER_EXPERT_STEP_CHECK_PASS")
if __name__ == "__main__":
main()这一步使用钩子函数查看了动作专家对带噪动作块的一次前向处理结果，输出如下：
========== 1. 动作专家的原始输入 ==========
带噪动作块形状： (1, 50, 32)
时间参数： tensor([1.], device='cuda:0')
prefix KV缓存长度： 968
动作专家Transformer层数： 18
========== 2. 专家第0层收到的数据 ==========
动作suffix形状： (1, 50, 1024)
时间条件形状： (1, 1024)
注意力矩阵形状： (1, 1, 50, 1018)
传入的prefix KV缓存长度： 968
第0个动作位置可以关注的位置数： 618
========== 3. 动作输出投影前 ==========
专家处理后的特征形状： (1, 50, 1024)
========== 4. 本轮输出 ==========
更新方向v_t形状： (1, 50, 32)
更新方向数据类型： torch.float32
首步v_t的前7维： tensor([ 1.77773, 0.24216, -0.34321, -1.73192, -0.26545, 0.75156, 0.11495],
device='cuda:0')
原prefix KV缓存长度： 968
PIPER_EXPERT_STEP_CHECK_PASS从输出可以看出，这一步主要做了以下几件事：
-
构建原始输入：带噪声的动作块，动作专家正是要以这个动作块为起始，不断预测当前应有的速度变化，通过10次积分得到最终的动作块，本小节只展示一次前向过程，10次积分的实践放在2.5。
-
**进行注意力计算，**这一步相当重要，也比较复杂。Q(查询)是当前的50个动作块，K(特征)、V(值)的是当前的50个动作+前缀缓存的968个特征。也就是说，动作专家里，单层的注意力计算，就是要查询50个动作该有的速度，其对应可使用的特征和值是之前缓存好的视觉信息+文字信息+动作信息。想这样的计算进行了18层，最后得到单次前向的输出结果。
张量维度变化：
-
初始时的动作维度是(1,50,32)，即1个批次，50个动作，每个动作32维，进入动作专家后，32维被投影成了1024维，方便与1024维的时间条件做拼接。
-
注意力矩阵的维度为(1,1,50,1018)，即50个查询对应1018个候选位置，其中可关注位置数为618，刚好就是KV缓存中有效的568个位置加自身的50个动作位置。
-
模型输出前动作由1024维投影回32维，再剪裁为lerobot下的7维。
由下图再次说明一下注意力计算部分：
对应图上，Q为50x256，对应可查询的特征K和值V都为1018x256，算出A为50x1018，再算AV得到50x256。
2.5. 完整动作生成与后处理
代码如下：
import torch
from A2_3_piper_gemma import main as prepare_conditions
from lerobot.policies.common.flow_matching import euler_integrate
from lerobot.processor import PolicyProcessorPipeline
from lerobot.processor.converters import (
policy_action_to_transition,
transition_to_policy_action,
)
@torch.inference_mode()
def main():
# 复用前面的真实观测、模型和prefix KV。
model, _, prefix_mask, past_key_values = prepare_conditions()
device = prefix_mask.device
batch_size = prefix_mask.shape[0]
chunk_size = model.config.chunk_size
num_steps = model.config.num_inference_steps
action_dim = model.config.output_features["action"].shape[0]
original_cache_length = past_key_values.get_seq_length()
# 1. 与A2_4相同，从随机噪声动作块开始。
torch.manual_seed(1000)
noise = model.sample_noise(
(batch_size, chunk_size, model.config.max_action_dim),
device,
)
# 生成时间从1走向0，因此积分步长为负数。
dt = -1.0 / num_steps
step_times = []
print("\n========== 1. 完整动作生成 ==========")
print("初始噪声形状：", tuple(noise.shape))
print("生成迭代次数：", num_steps)
print("积分步长dt：", dt)
# 用回调函数打印信息，更新方向仍由源码denoise_step计算。
def denoise_and_inspect(x_t, timestep):
velocity = model.denoise_step(
prefix_pad_masks=prefix_mask,
past_key_values=past_key_values,
x_t=x_t,
timestep=timestep,
)
step_times.append(timestep[0].item())
# 模拟一次积分更新，看看结果。
# 整个动作块的积分由下面的euler_integrate负责。
preview = x_t[0, 0, :3] + dt * velocity[0, 0, :3]
values = ", ".join(f"{x:.5f}" for x in preview.tolist())
print(
f"第{len(step_times):02d}轮 "
f"t={step_times[-1]:.1f} "
f"更新后首步前3维=[{values}]"
)
return velocity
# 正式积分
actions_32 = euler_integrate(
denoise_fn=denoise_and_inspect,
noise=noise,
num_steps=num_steps,
)
print("\n========== 2. 积分完成 ==========")
print("内部动作块形状：", tuple(actions_32.shape))
print("首步前7维：", actions_32[0, 0, :action_dim])
# 3. 沿用predict_action_chunk中的裁维方式：32维裁回7维。
normalized_actions = actions_32[:, :, :action_dim].cpu()
print("\n========== 3. 裁回实际动作维度 ==========")
print("归一化动作块形状：", tuple(normalized_actions.shape))
print("归一化动作前3步：")
print(normalized_actions[0, :3])
# 4. 加载Piper第五轮保存的后处理器及归一化统计量。
checkpoint = (
"/root/autodl-tmp/workspace/outputs/week05/"
"b8_vlmlora_expertfull_20000steps_20261018/"
"checkpoints/020000/pretrained_model"
)
postprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=checkpoint,
config_filename="policy_postprocessor.json",
local_files_only=True,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
)
# 数据已经采用相对动作约定，不再进行额外表示转换。
assert model.config.use_relative_actions is False
relative_actions = postprocessor(normalized_actions)
print("\n========== 4. 反归一化后的预测动作 ==========")
print("最终动作块形状：", tuple(relative_actions.shape))
print("设备：", relative_actions.device)
print("相对动作前3步：")
print(relative_actions[0, :3])
# 5. 对应select_action中选取前n_action_steps步的逻辑。
# 本脚本仅展示这部分，不向环境发送动作。
n_action_steps = model.config.n_action_steps
actions_to_execute = relative_actions[:, :n_action_steps]
print("\n========== 5. 本次计划执行的部分 ==========")
print("执行步数配置：", n_action_steps)
print("计划执行动作形状：", tuple(actions_to_execute.shape))
print("prefix缓存长度：", past_key_values.get_seq_length())
assert len(step_times) == num_steps
assert actions_32.shape == noise.shape
assert relative_actions.shape == (batch_size, chunk_size, action_dim)
assert torch.isfinite(relative_actions).all().item()
assert 1 <= n_action_steps <= chunk_size
assert past_key_values.get_seq_length() == original_cache_length
print("\nPIPER_ACTION_CHUNK_CHECK_PASS")
if __name__ == "__main__":
main()这一步其实与2.4类似，只不过进行了10轮，每一轮算出当前这一时间步该做的速度变化后，进行积分，10次积分后得到最终的动作结果。输出如下：
========== 1. 完整动作生成 ==========
初始噪声形状： (1, 50, 32)
生成迭代次数： 10
积分步长dt： -0.1
第01轮 t=1.0 更新后首步前3维=[0.78876, 0.02687, -0.63014]
第02轮 t=0.9 更新后首步前3维=[0.61051, -0.00589, -0.59792]
第03轮 t=0.8 更新后首步前3维=[0.43412, -0.03900, -0.56977]
第04轮 t=0.7 更新后首步前3维=[0.27950, -0.06592, -0.54414]
第05轮 t=0.6 更新后首步前3维=[0.13950, -0.08696, -0.52619]
第06轮 t=0.5 更新后首步前3维=[-0.00140, -0.10882, -0.50954]
第07轮 t=0.4 更新后首步前3维=[-0.15156, -0.13194, -0.49515]
第08轮 t=0.3 更新后首步前3维=[-0.33468, -0.15508, -0.47558]
第09轮 t=0.2 更新后首步前3维=[-0.54430, -0.18642, -0.44891]
第10轮 t=0.1 更新后首步前3维=[-0.78131, -0.22164, -0.40382]
========== 2. 积分完成 ==========
内部动作块形状： (1, 50, 32)
首步前7维： tensor([-0.78131, -0.22164, -0.40382, 1.07502, 0.80763, -0.78789, -0.81078],
device='cuda:0')
========== 3. 裁回实际动作维度 ==========
归一化动作块形状： (1, 50, 7)
归一化动作前3步：
tensor([[-0.78131, -0.22164, -0.40382, 1.07502, 0.80763, -0.78789, -0.81078],
[-0.78978, -0.17264, -0.40275, 1.08662, 0.82191, -0.81087, -0.81602],
[-0.76964, -0.17851, -0.38636, 1.07582, 0.83681, -0.80586, -0.81198]])
========== 4. 反归一化后的预测动作 ==========
最终动作块形状： (1, 50, 7)
设备： cpu
相对动作前3步：
tensor([[ 0.02109, 0.04620, 0.00326, -0.00994, -0.01124, -0.00919, -0.98589],
[ 0.01956, 0.05217, 0.00347, -0.00777, -0.00814, -0.01459, -0.99104],
[ 0.02320, 0.05146, 0.00664, -0.00979, -0.00489, -0.01341, -0.98707]])
========== 5. 本次计划执行的部分 ==========
执行步数配置： 10
计划执行动作形状： (1, 10, 7)
prefix缓存长度： 968
PIPER_ACTION_CHUNK_CHECK_PASS2.6. 训练前向与loss计算
代码如下：
from functools import partial
import torch
from A2_1_piper_preprocess import main as prepare_sample
from lerobot.configs import PreTrainedConfig
from lerobot.policies.factory import make_policy
@torch.inference_mode()
def main():
checkpoint, dataset, batch = prepare_sample()
# 1. 加载同一个Piper第五轮模型。
config = PreTrainedConfig.from_pretrained(checkpoint)
config.pretrained_path = str(checkpoint)
config.device = "cuda"
policy = make_policy(cfg=config, ds_meta=dataset.meta)
policy.eval()
policy.requires_grad_(False)
base_policy = policy.get_base_model()
model = base_policy.model
# 2. 沿PI05Policy.forward源码准备输入。
images, image_masks = base_policy._preprocess_images(batch)
tokens = batch["observation.language.tokens"]
text_mask = batch["observation.language.attention_mask"]
# 真实动作已经由A2_1归一化，这里只补齐维度。
# action是数据集中的示范动作
actions = base_policy.prepare_action(batch)
action_dim = config.output_features["action"].shape[0]
torch.manual_seed(1000)
noise = model.sample_noise(actions.shape, actions.device)
time = model.sample_time(actions.shape[0], actions.device)
# 源码中的监督目标：噪声减去真实动作。
target_velocity = noise - actions
print("\n========== 1. 训练输入 ==========")
print("真实动作形状：", tuple(batch["action"].shape))
print("补齐后动作形状：", tuple(actions.shape))
print(
"补齐维度是否全0：",
(actions[:, :, action_dim:] == 0).all().item(),
)
print("随机噪声形状：", tuple(noise.shape))
print("本次采样的生成时间t：", time)
print("监督目标u_t形状：", tuple(target_velocity.shape))
print(
"本样本补齐的动作时间步数：",
batch["action_is_pad"].sum().item(),
)
observed = {}
# 只观察输入形状，不修改计算。
def inspect_input(name, module, args):
observed[name] = tuple(args[0].shape)
print(f"{name}：{observed[name]}")
def inspect_velocity(module, args, output):
observed["预测速度v_t"] = tuple(output.shape)
print("预测速度v_t形状：", tuple(output.shape))
vlm_layers = (
model.paligemma_with_expert.paligemma.model.language_model.layers
)
expert_layers = (
model.paligemma_with_expert.gemma_expert.model.layers
)
handles = [
model.action_in_proj.register_forward_pre_hook(
partial(inspect_input, "实际带噪动作x_t")
),
vlm_layers[0].self_attn.q_proj.register_forward_pre_hook(
partial(inspect_input, "VLM第0层的注意力输入")
),
expert_layers[0].self_attn.q_proj.register_forward_pre_hook(
partial(inspect_input, "专家第0层的注意力输入")
),
model.action_out_proj.register_forward_hook(inspect_velocity),
]
print("\n========== 2. 一次联合前向 ==========")
try:
# 直接调用源码的训练核心forward，只运行一次。
# 内部构造：
# x_t = t * noise + (1 - t) * actions
# u_t = noise - actions
#
# prefix和suffix联合经过网络，输出v_t，
# 再返回逐元素MSE：(u_t - v_t)^2。
losses_32 = model.forward(
images,
image_masks,
tokens,
text_mask,
actions,
noise,
time,
)
finally:
for handle in handles:
handle.remove()
# 3. 以下裁维和求均值，与PI05Policy.forward源码一致。
losses_7 = losses_32[:, :, :action_dim]
loss_per_dim = losses_7.mean(dim=(0, 1))
loss = losses_7.mean()
print("\n========== 3. loss计算 ==========")
print("逐元素MSE形状：", tuple(losses_32.shape))
print("裁到实际动作维度后：", tuple(losses_7.shape))
print("每个动作维度的平均loss：", loss_per_dim.cpu())
print("最终标量loss：", loss.item())
assert len(observed) == 4
assert observed["预测速度v_t"] == tuple(actions.shape)
assert losses_32.shape == actions.shape
assert losses_7.shape == batch["action"].shape
assert loss.ndim == 0
assert torch.isfinite(loss).item()
print("\nPIPER_TRAINING_FORWARD_CHECK_PASS")
if __name__ == "__main__":
main()这一次走了一个完整的训练过程：构造输入--前向--计算loss，输出如下：
All keys loaded successfully!
========== 1. 训练输入 ==========
真实动作形状： (1, 50, 7)
补齐后动作形状： (1, 50, 32)
补齐维度是否全0： True
随机噪声形状： (1, 50, 32)
本次采样的生成时间t： tensor([0.43525], device='cuda:0')
监督目标u_t形状： (1, 50, 32)
本样本补齐的动作时间步数： 0
========== 2. 一次联合前向 ==========
实际带噪动作x_t：(1, 50, 32)
VLM第0层的注意力输入：(1, 968, 2048)
专家第0层的注意力输入：(1, 50, 1024)
预测速度v_t形状： (1, 50, 32)
========== 3. loss计算 ==========
逐元素MSE形状： (1, 50, 32)
裁到实际动作维度后： (1, 50, 7)
每个动作维度的平均loss： tensor([0.66684, 0.21333, 0.03186, 0.01768, 0.15759, 0.04316, 0.00203])
最终标量loss： 0.16178372502326965
PIPER_TRAINING_FORWARD_CHECK_PASS
真实动作：(1,50,7)
↓ 补零
真实动作：(1,50,32)
↓ 按随机时间t=0.43525与噪声混合
带噪动作x_t：(1,50,32)
↓ 与图像、任务、state组成的prefix联合前向
预测更新方向v_t：(1,50,32)
↓ 与监督目标比较
逐元素MSE：(1,50,32)
↓ 只保留前7维，再求平均
loss = 0.16178372502326965这一步的重点在于知道训练时的loss是怎么定义的：
t=1时，表示xt为完全的噪声，t=0时，表示xt为干净的准确动作。我们这里选择t=0.43525，表示动作块初始是一个半对半错的含噪动作。
ut是当前xt对于时间t的求导，出来的结果便表示xt在当前应有的变化速度及方向，即"参考答案"。
最后算loss便就是用模型推理出来的速度与标准答案的速度求均方损失误差(MSE)。
3. 完整流程图
让ai生成了6组实践完整的流程图：
4. 总结
本文了解了pi05的基础原理，并通过代码探究了模型生成及训练过程中的数据处理流程。
本文仅做学习分享，如有谬误欢迎指正。