base.py — 定义了 PI0 策略的输入/输出抽象基类 Pi0Input 和 Pi0Output

文件路径: verl/experimental/vla/models/pi0_torch/policy/base.py 模块路径: verl.experimental.vla.models.pi0_torch.policy.base

文件概述

定义了 PI0 策略的输入/输出抽象基类 Pi0Input 和 Pi0Output，为不同机器人平台提供统一的数据接口。

关键代码

Pi0Input - 策略输入抽象类

class Pi0Input(ABC):
    """PI0 模型的输入数据抽象

    不同机器人平台（LIBERO、ALOHA 等）的观测格式不同，
    但都需要转换为统一的 PI0 输入格式。
    """
    images: list[torch.Tensor]     # 图像列表
    img_masks: list[torch.Tensor]  # 图像有效性 mask
    task: str                       # 任务描述文本
    state: torch.Tensor            # 机器人状态

Pi0Output - 策略输出抽象类

class Pi0Output(ABC):
    """PI0 模型的输出数据抽象"""
    action: torch.Tensor  # 预测的动作

设计思想

LIBERO 观测格式 ──> LiberoPi0Input ──> Pi0Input (统一接口) ──> PI0Model
ALOHA 观测格式  ──> AlohaPi0Input  ──>
其他平台        ──> ...            ──>

与其他模块的关系

• 被 libero_policy.py 的 LiberoPi0Input/Output 实现
• 被 PI0ForActionPrediction 使用

小结

策略适配器模式的抽象接口，使得 PI0 模型可以对接不同的机器人平台而不需要修改模型代码。