naive_rollout_pi05.py — PI0RolloutRob 扩展了 NaiveRolloutRob

文件路径: verl/experimental/vla/sac/naive_rollout_pi05.py 模块路径: verl.experimental.vla.sac.naive_rollout_pi05

文件概述

PI0RolloutRob 扩展了 NaiveRolloutRob，为 PI0/PI0.5 模型提供推理支持。主要区别是输入处理方式不同（PI0 使用归一化状态而非离散 token）。

关键代码

class PI0RolloutRob(NaiveRolloutRob):
    """PI0/PI0.5 模型的推理封装"""

    def __init__(self, model_config, module=None):
        # 不调用父类 __init__（OpenVLA 的初始化不适用）
        self.model_config = model_config
        if module is not None:
            self.module = module
        else:
            self.module = PI0ForActionPrediction.from_pretrained(
                model_config["path"], trust_remote_code=True
            )
        self.module.eval()

    @torch.no_grad()
    def generate_sequences(self, prompts: DataProto) -> DataProto:
        """PI0 推理：从观测生成动作

        与 OpenVLA 不同，PI0 不需要文本 tokenization，
        而是直接使用归一化的状态向量和图像。
        """
        # 准备 PI0 输入
        images = prompts.batch["full_image"]
        states = prompts.batch["state"]
        task_descriptions = prompts.non_tensor_batch["task_descriptions"]

        # 调用 PI0 的 sample_actions
        actions = self.module.sample_actions(
            images=images,
            states=states,
            task_descriptions=task_descriptions
        )

        return DataProto.from_dict(tensors={"action": actions})

PI0RolloutRob vs NaiveRolloutRob 对比

特性	NaiveRolloutRob (OpenVLA)	PI0RolloutRob (PI0)
模型	OpenVLAForActionPrediction	PI0ForActionPrediction
输入预处理	文本 tokenize + 图像 crop	状态归一化 + 图像 resize
动作生成	离散 token -> 连续动作	Flow Matching 直接生成
Processor	PrismaticProcessor	内置变换

与其他模块的关系

• 继承/替代 NaiveRolloutRob（naive_rollout_rob.py）
• 使用 PI0ForActionPrediction（modeling_pi0_torch.py）
• 被 SAC 训练流程使用

小结

PI0 版本的推理封装。与 OpenVLA 的推理相比，PI0 不需要文本 tokenization 和离散化解码，输入输出都是连续值，流程更简洁。