prepare_libero_dataset.py — 这是一个数据预处理脚本

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

文件概述

这是一个数据预处理脚本，将 LIBERO 机器人操作基准数据集转换为 Parquet 格式。它实现了训练/测试集划分（包括分布内测试和分布外测试），为 VLA 训练提供标准化的数据输入。

LIBERO 基准简介

LIBERO 是一个桌面机器人操作的仿真基准，包含 10 个任务（libero_10），每个任务有 50 个初始状态（共 500 个 state）。每个任务是一个不同的操作技能（如"打开抽屉"、"放置物品"等）。

关键代码

数据集划分策略

# 10 个任务中随机选 9 个做训练，1 个做分布外测试
all_task_ids = list(range(task_suite.get_num_tasks()))  # [0, 1, ..., 9]
train_task_ids = sorted(random.sample(all_task_ids, 9))
ood_test_task_id = list(set(all_task_ids) - set(train_task_ids))[0]

三种数据划分

# 1. 训练集：9 个任务 × 40 个初始状态 = 360 条
for task_id in train_task_ids:
    selected_train_trials = all_trials[:40]
    for state_id in selected_train_trials:
        train_metadata.append({
            "task_id": task_id,
            "state_id": state_id,
            "data_source": "train"
        })

# 2. 分布内测试集：9 个任务 × 10 个初始状态 = 90 条
for task_id in train_task_ids:
    selected_id_test_trials = all_trials[40:][:10]
    for state_id in selected_id_test_trials:
        test_metadata.append({
            "task_id": task_id,
            "state_id": state_id,
            "data_source": "test_in_distribution"
        })

# 3. 分布外测试集：1 个任务 × 20 个初始状态 = 20 条
ood_all_trials = get_state_ids_for_task(ood_test_task_id)
for state_id in ood_all_trials[:20]:
    test_metadata.append({
        "task_id": ood_test_task_id,
        "state_id": state_id,
        "data_source": "test_out_of_distribution"
    })

数据格式

每条数据包含：

data = {
    "data_source": "train",           # 数据来源标识
    "prompt": task.language,           # 任务的自然语言描述
    "state_ids": state_id,            # 环境初始状态 ID
    "task_ids": task_id,              # 任务 ID
    "ability": "robot",              # 能力类型标识
    "extra_info": { ... }            # 额外元信息
}

核心类/函数列表

名称	类型	说明
patched_get_task_init_states	函数	修补 LIBERO 加载初始状态的方法
compute_total_num_group_envs	函数	计算总环境数和累积分箱
get_state_ids_for_task	函数	获取某任务的所有 state ID 范围
map_and_process	函数	将元数据映射为完整数据记录

与其他模块的关系

• 输出的 Parquet 文件被 RobRayPPOTrainer 读取作为训练数据
• state_ids 和 task_ids 被 EnvWorker 用于环境重置

小结

这个脚本将 LIBERO 基准组织为强化学习训练所需的格式。关键设计是三层数据划分：训练集（同任务同状态）、分布内测试（同任务不同状态）、分布外测试（不同任务），这使得可以评估模型的泛化能力。