libero_policy.py — LIBERO 平台的 PI0 策略适配器

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

文件概述

LIBERO 平台的 PI0 策略适配器，将 LIBERO 环境的观测格式转换为 PI0 模型需要的输入格式，以及将 PI0 的输出转换为 LIBERO 的动作格式。

关键代码

LiberoPi0Input

class LiberoPi0Input(Pi0Input):
    """将 LIBERO 观测转换为 PI0 输入"""

    @classmethod
    def from_env_obs(cls, obs, task_description):
        """从环境观测创建 PI0 输入

        LIBERO 观测包含：
        - full_image: 全局相机图像
        - wrist_image: 腕部相机图像
        - state: 机器人关节状态
        """
        images = [obs["full_image"]]  # PI0 目前只用全局图像
        img_masks = [torch.ones(batch_size, dtype=torch.bool)]
        state = obs["state"]
        task = task_description

        return cls(images=images, img_masks=img_masks,
                   task=task, state=state)

LiberoPi0Output

class LiberoPi0Output(Pi0Output):
    """将 PI0 输出转换为 LIBERO 动作"""

    @classmethod
    def from_model_output(cls, raw_actions, action_dim=7):
        """从模型输出提取 LIBERO 动作

        PI0 输出: (B, n_action_steps, max_action_dim)
        LIBERO 需要: (B, n_action_steps, 7)

        需要截取前 action_dim 维，因为 PI0 输出被填充到了
        max_action_dim（如 32 维）。
        """
        actions = raw_actions[:, :, :action_dim]
        return cls(action=actions)

核心类列表

名称	说明
LiberoPi0Input	LIBERO -> PI0 输入适配
LiberoPi0Output	PI0 输出 -> LIBERO 动作适配

与其他模块的关系

• 实现 Pi0Input/Pi0Output（base.py）
• 被 PI0 推理管道使用
• 依赖 LIBERO 环境的观测格式

小结

一个简洁的适配器层，处理 LIBERO 和 PI0 之间的数据格式差异。主要工作是图像选择（哪些相机的图像传给模型）和动作维度裁剪（PI0 输出的 max_action_dim -> LIBERO 的 7 维动作）。