执行摘要

• 一句话：修复 Pi0.5 语言嵌入缩放缺失，对齐 OpenPI 输出
• 推荐动作：值得精读。重点看 build_openpi_model_inputs 与 _preprocess_image 两个设计点：前者演示了如何用参考实现自身的输入变换构造一致性对比，后者用显式标记避免图像归一化歧义。该 PR 为「与官方实现的数值一致性」提供了可执行、可回归的测试范本。

功能与动机

官方 OpenPI PyTorch Pi0.5 实现（src/openpi/models_pytorch/pi0_pytorch.py 的 embed_prefix）在拼接图像嵌入前将语言 token 嵌入乘以 sqrt(hidden_dim)，SGLang 遗漏该缩放，导致 action 输出与官方存在系统性偏差。PR body 指出原有 CI 未发现该问题：pi05_action_http 使用的旧 GT 由旧版 SGLang 生成，Python 一致性测试只检查可重复性与缓存复用，手动 benchmark 又对比了经过机器人后处理的 action，均掩盖了模型空间的差异。

实现拆解

1. 核心修复：语言嵌入缩放

• 在 python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py 的 Pi05CoreModel.embed_prefix 中，对 embed_language_tokens 的输出补乘 math.sqrt(lang_emb.shape[-1])，对齐 OpenPI 官方语义。改动位于前缀构建路径，不触达 denoise 循环。

2. 图像预处理对齐

• pi05_preprocess.py 的 _tensor_from_image 返回值从 torch.Tensor 改为 (torch.Tensor, bool, bool)，携带「是否字节缩放、是否已归一化」两个标记。
• _resize_with_pad_image_tensor 新增 round_to_uint8 与 pad_value 参数，字节图 resize 后先 round 再映射到 [-1, 1]。
• 新增 _preprocess_image 统一入口：已归一化的 float32 图像只做 clamp 并保留 -1 letterbox 填充，避免二次归一化。

3. 基准升级到归一化模型空间

• bench_pi05_openpi.py 新增 build_openpi_model_inputs，调用 OpenPI policy 的 _input_transform 得到与官方完全相同的模型输入（含 token、token mask、image mask、camera order）。
• 拆出 _openpi_model_batch / _openpi_model_actions，直接获取模型空间 action 与 SGLang 对比，绕过 OpenPI 的机器人 unnormalization 与 horizon slicing。
• --deterministic-noise 模式下用 OpenPI 变换后的输入喂给 SGLang；新增 --action-max-abs-diff / --action-mean-abs-diff 阈值参数；aloha 的 state 不再 pad 到 32 维，直接使用 14 维真实状态。

4. 测试与 CI 配套

• test_pi05_runtime_helpers.py 新增 3 个单测：语言嵌入缩放、uint8 resize 先取整、归一化浮点图不二次归一化。
• test_pi05_e2e.py 将 test_pi05_python_action_consistency 重命名为 test_pi05_python_action_repeatability_and_cache，明确可重复性语义，state 默认维度回归 14。
• testcase_configs.py 的 CI 用例 state_dim 从 32 改为 14。
• test_utils.py 将 SGL_TEST_FILES_CI_DATA_REVISION 固定到 d05810e3ea3eff1d137dec723f6e66d9c11b470f，使一致性 CI 基于官方 provonance-backed GT。
• docs/cookbook/vla/OpenPI/Pi0.5.mdx 更新官方 parity 数据与对比方法说明。

关键文件：

• python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py（模块 模型核心；类别 source；类型 core-logic；符号 embed_prefix）: 核心修复点：embed_prefix 中补上缺失的 sqrt(hidden_dim) 语言嵌入缩放，一行改动影响所有 Pi0.5 输出。
• python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/pi05_preprocess.py（模块 预处理器；类别 source；类型 data-contract；符号 _tensor_from_image, _preprocess_image, _resize_with_pad_image_tensor）: 图像预处理契约重构：区分字节图与已归一化浮点图，匹配 OpenPI 的 resize 取整与 letterbox 填充行为。
• python/sglang/multimodal_gen/benchmarks/bench_pi05_openpi.py（模块 基准脚本；类别 source；类型 dependency-wiring；符号 build_openpi_model_inputs, _openpi_direct_batch, _openpi_model_batch, _openpi_model_actions）: 手动对比基准大幅升级：新增 OpenPI 模型输入构造与归一化模型空间对比，是验证修复效果的基石工具。
• python/sglang/multimodal_gen/test/unit/test_pi05_runtime_helpers.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_prefix_language_embedding_matches_openpi_scale, test_uint8_resize_rounds_before_normalization, test_normalized_float_image_is_not_normalized_twice）: 新增 3 个针对性单测，锁定缩放、uint8 取整与归一化不做二次映射的行为。
• python/sglang/multimodal_gen/test/single_test_file/test_pi05_e2e.py（模块 端到端测试；类别 test；类型 test-coverage；符号 test_pi05_python_action_repeatability_and_cache）: 一致性测试改名落地 repeatability 语义，state 默认维度从 32 回归 14，与真实 Aloha 状态一致。
• python/sglang/multimodal_gen/test/server/testcase_configs.py（模块 测试配置；类别 test；类型 configuration）: CI 用例 state_dim 从 32 改为 14，使服务端一致性用例与真实 Aloha 观测一致。
• python/sglang/multimodal_gen/test/test_utils.py（模块 测试工具；类别 test；类型 configuration）: 固定官方 GT 的 ci-data-diffusion revision，使一致性 CI 依赖可复现的 provenance-backed 数据。
• docs/cookbook/vla/OpenPI/Pi0.5.mdx（模块 文档；类别 docs；类型 documentation）: 更新官方 parity 数据与对比方法说明，记录归一化模型空间的验证流程。

关键符号：Pi05CoreModel.embed_prefix, _preprocess_image, _resize_with_pad_image_tensor, _tensor_from_image, build_openpi_model_inputs, _openpi_model_batch, _openpi_model_actions

关键源码片段

python/sglang/multimodal_gen/runtime/models/vlas/pi05_core.py

核心修复点：embed_prefix 中补上缺失的 sqrt(hidden_dim) 语言嵌入缩放，一行改动影响所有 Pi0.5 输出。

def embed_prefix(
    self,
    images: list[torch.Tensor],
    image_masks: list[torch.Tensor],
    tokens: torch.Tensor,
    token_masks: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    embs = []
    pad_masks = []
    att_masks = []
    image_embs = self.paligemma_with_expert.embed_images(images)
    for image_emb, image_mask in zip(image_embs, image_masks, strict=True):
        batch_size, num_image_embs = image_emb.shape[:2]
        embs.append(image_emb)
        pad_masks.append(image_mask[:, None].expand(batch_size, num_image_embs))
        att_masks += [0] * num_image_embs
​
    lang_emb = self.paligemma_with_expert.embed_language_tokens(tokens)
    # 关键修复：OpenPI 官方在拼接图像嵌入前对语言嵌入乘以 sqrt(hidden_dim)。
    # 此前 SGLang 遗漏该缩放，导致语言分支与图像分支幅度失配，
    # action 输出相对官方实现出现系统性偏差。
    lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
    embs.append(lang_emb)
    pad_masks.append(token_masks)
    att_masks += [0] * lang_emb.shape[1]
​
    embs = torch.cat(embs, dim=1)
    pad_masks = torch.cat(pad_masks, dim=1)
    att_masks_t = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
    att_masks_t = att_masks_t[None, :].expand(pad_masks.shape[0], len(att_masks))
    return embs, pad_masks, att_masks_t

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/pi05_preprocess.py

图像预处理契约重构：区分字节图与已归一化浮点图，匹配 OpenPI 的 resize 取整与 letterbox 填充行为。

def _preprocess_image(value: Any, size: tuple[int, int]) -> torch.Tensor:
    # _tensor_from_image 返回 ( 张量 , 是否字节缩放 , 是否已归一化到 [-1, 1])
    tensor, is_byte_scaled, is_normalized = _tensor_from_image(value)
    tensor = _resize_with_pad_image_tensor(
        tensor,
        size,
        round_to_uint8=is_byte_scaled,
        pad_value=-1.0 if is_normalized else 0.0,
    )
    if is_normalized:
        # 已归一化的 float32 图像：只 clamp，避免二次 2*x-1 映射
        return tensor.clamp_(-1.0, 1.0)
    return tensor * 2.0 - 1.0
​
​
def _resize_with_pad_image_tensor(
    tensor: torch.Tensor,
    size: tuple[int, int],
    *,
    round_to_uint8: bool = False,
    pad_value: float = 0.0,
) -> torch.Tensor:
    height, width = size
    if tensor.shape[-2:] == (height, width):
        return tensor
    _, cur_height, cur_width = tensor.shape
    ratio = max(cur_width / width, cur_height / height)
    resized_height = int(cur_height / ratio)
    resized_width = int(cur_width / ratio)
    tensor = F.interpolate(
        tensor[None],
        size=(resized_height, resized_width),
        mode="bilinear",
        align_corners=False,
    )[0]
    if round_to_uint8:
        # 与 OpenPI 对齐：字节图 resize 后先 round 再映射到 [-1, 1]
        tensor = torch.round(tensor * 255.0).clamp_(0.0, 255.0) / 255.0
    pad_h0, rem_h = divmod(height - resized_height, 2)
    pad_w0, rem_w = divmod(width - resized_width, 2)
    # 归一化图用 -1 填充 letterbox，字节图用 0 填充后再做 2*x-1
    return F.pad(
        tensor,
        (pad_w0, pad_w0 + rem_w, pad_h0, pad_h0 + rem_h),
        mode="constant",
        value=pad_value,
    )

评论区精华

本 PR 没有实质性 review 讨论：mickqian 直接 approve，唯一的 bot 评论是 Gemini Code Assist 停服说明；另一条评论仅触发 tag-and-rerun-ci。关键的验收证据链在 PR body 中：官方 GT 由 OpenPI revision 15a9616... 生成两次、SHA256 完全一致（319c1a1e...）；rebased head 在 RTX 5090 直连 GPU 流式下（13.5 GiB / 1.81 s）通过同一一致性用例。

• 旧 CI 为何未拦截语言嵌入缩放缺失 (correctness): 引入 ci-data-diffusion#1 的官方 GT（OpenPI revision 15a9616、LeRobot checkpoint b211f3d），并改为在归一化模型空间、相同输入与确定性噪声下对比。

风险与影响

• 风险：
  • 输出行为变更：embed_prefix 的缩放会改变所有 Pi0.5 请求的 prefix 嵌入与最终 action，依赖旧输出的用户需重新验证。
  • 预处理契约变更：_tensor_from_image 返回类型从 torch.Tensor 变为三元组，虽有单测覆盖，但其他潜在调用方需同步适配。
  • GT 依赖固定 revision：test_utils.py 固定 ci-data-diffusion 到 d05810e，若上游 OpenPI 实现或 LeRobot checkpoint 再变化，需重新生成 GT，否则 CI 会开始失败。
  • 无性能回归：缩放是 prefix 构建期的一次逐元素乘法，不进入 denoise 循环。
• 影响：
  • 用户：所有基于 Pi0.5 的 VLA 输出 action 数值会整体变化（向官方实现对齐），属于预期内的正确性修复。
  • 系统：影响面集中在 multimodal_gen 的 Pi0.5 模型族，图像预处理路径的返回契约变化需要模块内同步。
  • 团队：确立了「官方模型空间对比 + 可复现 GT provenance」的验证范式，后续 VLA 模型接入可复用 build_openpi_model_inputs 模式。
  • 风险标记：核心推理路径变更, 图像预处理契约变更, CI GT 依赖固定 revision, 旧行为输出不兼容

关联脉络

• PR #33317 [diffusion] Fix component accuracy topology reuse: 同为 diffusion 测试正确性修复，本 PR 在该方向基础上补充了模型空间对齐验证。
• PR #33281 [CI] Add MiniMax-H3 2-GPU consistency coverage: 建立了 diffusion 一致性 CI 机制，本 PR 在 Pi0.5 上落地官方 GT 回归。
• PR #33453 [diffusion] Restrict request-level quality to two validated tiers: lossless (default) and high: 同属 multimodal_gen 配置与文档演进，扩大 diffusion 模型支持面。