libero_env.py — LiberoEnv 是 LIBERO 仿真基准的 Gymnasium 环境封装

文件路径: verl/experimental/vla/envs/libero_env/libero_env.py 模块路径: verl.experimental.vla.envs.libero_env.libero_env

文件概述

LiberoEnv 是 LIBERO 仿真基准的 Gymnasium 环境封装。它管理多个仿真环境实例（向量化环境），支持任务/试验的动态切换、分步执行（chunk_step）和视频录制。

核心类：LiberoEnv

初始化

class LiberoEnv:
    def __init__(self, config, rank, world_size):
        self.config = config
        self.rank = rank
        self.num_envs = config.num_envs  # 每个 worker 管理的环境数

        # 加载 LIBERO 基准
        self.task_suite = get_benchmark("libero_10")()

        # 创建向量化子进程环境
        self.vec_env = ReconfigureSubprocEnv(
            env_fns=[make_env_fn(task_id) for _ in range(self.num_envs)]
        )

环境重置

def reset_envs_to_state_ids(self, state_ids, task_ids):
    """将环境重置到指定的任务和初始状态

    这是 LIBERO 的特殊需求：需要控制每个环境实例
    运行哪个任务、从哪个初始状态开始。
    """
    for i, (state_id, task_id) in enumerate(zip(state_ids, task_ids)):
        # 获取任务的初始状态
        init_state = self.task_suite.get_task_init_states(task_id)[trial_id]
        # 重置环境到该状态
        self.vec_env.reset_to_state(i, init_state, task_id)

    # 收集并返回初始观测
    obs = self._extract_observations()
    return obs, {}

分步执行（Chunk Step）

VLA 模型一次预测多步动作（action chunk），环境需要依次执行这些动作：

def chunk_step(self, chunk_actions):
    """执行一个动作 chunk（多步动作）

    Args:
        chunk_actions: (num_envs, num_chunks, action_dim)

    Returns:
        obs, rewards, terminations, truncations, infos
    """
    total_rewards = torch.zeros(self.num_envs)

    for step in range(self.num_action_chunks):
        action = chunk_actions[:, step, :]

        # 执行单步
        obs, reward, terminated, truncated, info = self.vec_env.step(action)

        # 累积奖励
        total_rewards += reward

        # 如果环境终止，停止该环境的执行
        ...

    return extracted_obs, total_rewards, terminations, truncations, infos

观测提取

def _extract_observations(self):
    """从原始观测中提取模型需要的信息

    返回:
        images_and_states: {
            full_image: (num_envs, H, W, C),  # 全局相机图像
            wrist_image: (num_envs, H, W, C),  # 腕部相机图像
            state: (num_envs, state_dim),       # 机器人状态
        }
        task_descriptions: list[str]  # 任务描述
    """
    ...

视频录制

def flush_video(self, video_sub_dir=""):
    """将录制的帧保存为视频文件"""
    save_rollout_video(self.recorded_frames, save_path, fps=10)

核心类/函数列表

名称	类型	说明
LiberoEnv	类	LIBERO 环境封装
reset_envs_to_state_ids	方法	按指定状态重置环境
chunk_step	方法	执行多步动作
_extract_observations	方法	提取模型所需观测
flush_video	方法	保存录制视频

与其他模块的关系

• 使用 ReconfigureSubprocEnv（venv.py）管理子进程环境
• 使用 utils.py 的图像提取和动作处理函数
• 被 EnvManager（workers/env/env_manager.py）创建和管理

小结

LiberoEnv 是连接 VLA 模型和 LIBERO 仿真器的桥梁。它的关键功能是 chunk_step（支持 action chunk 的分步执行）和 reset_envs_to_state_ids（支持精确控制环境初始条件），这两个功能是 VLA 强化学习训练的核心需求。