Reward Model for Robotics - 项目概述
项目目标
训练一个**奖励模型 (Reward Model)**，用于为后续 Robotics VLA (Vision-Language-Action) 模型的 Policy 训练提供奖励信号。
核心理念
通过视觉观察机器人操作视频，评估每一帧/每一步的：
- **阶段识别 (Stage)**：reach → grasp → lift → move → place → retract
- 奖励组件：reachout, grasp, collision, fall, smooth
- 综合指标：reward, delta, success_prob, failure
数据来源
1. 真机数据 (Real Robot Data)
| 数据源 | 位置 | 说明 |
|---|---|---|
| DROID 数据集 | /playpen-ssd/dataset/droid_raw/1.0.1/ |
主要数据源，包含多个机构的机器人操作数据 |
| PennPAL | 同上 | 包含 failure/success 标注 |
| AUTOLab | 同上 | 包含多种任务类型 |
数据统计：
- 总轨迹数：15,157 条 (failure 数据)
- 涵盖任务类型：300+ 种
- 主要数据源：AUTOLab (3618), ILIAD (1303), IPRL (1298), 等
2. 仿真数据 (Simulation Data)
| 数据源 | 位置 | 说明 |
|---|---|---|
| fangyu 仿真数据 | /playpen-ssd/yufang/projects/reward/cokecan-50 |
~800-900 条轨迹 |
| yangyue 失败数据 | /playpen-ssd/dataset/libero_failures/dec_29_2025_v1 |
Libero 环境失败案例 |
| 本地仿真数据 | ./data/cokecan-50/ |
包含 fall, collision, grasp, smooth, success 分类 |
仿真数据格式：
{
"joint_poss": [[...], [...], ...], // 关节位置序列
"stages": [0, 0, 1, 1, 2, ...], // 阶段标签
"rewards": {
"reachout": [...],
"grasp": [...],
"collision": [...],
"fall": [...],
"smooth": [...]
}
}
项目架构
RewardModel/
├── api_batch_improved.py # GPT-4 视频标注（滑动窗口模式）
├── api_batch_hf_dataset.py # 处理 HuggingFace 数据集
├── video_process.py # 从 DROID TFDS 提取视频
├── extract_frames_to_images.py # 从视频提取帧图片
├── convert_to_sft.py # 转换为 SFT 训练格式
│
├── caption/ # 视频描述生成
│ ├── infer_caption.py # 单视频推理
│ └── infer_caption_batch.py # 批量推理（Qwen VL）
│
├── data_sta/ # 数据统计与上传
│ ├── count_failure_trajectories.py
│ ├── prepare_hf_dataset.py # 准备 HuggingFace 数据集
│ └── upload_to_huggingface.py
│
├── data/ # 本地数据
│ ├── cokecan-50/ # 仿真数据（按类型分类）
│ ├── frames/ # 提取的帧图片
│ └── sft_*.json # SFT 训练数据集
│
├── droid_processed/ # 处理后的 DROID 数据
│ └── episode_*.json # 每个 episode 的元数据
│
├── weights_reward_sft*/ # 训练权重
│ └── checkpoint-*/ # LoRA/Full 检查点
│
└── output/ # 标注输出
└── labels_*.jsonl # GPT/模型标注结果
数据处理流程
流程图
原始视频数据
↓
[video_process.py] 从 TFDS 提取视频和元数据
↓
[extract_frames_to_images.py] 抽帧保存为图片
↓
[api_batch_improved.py] 发送给 GPT-4 打分
↓
标注 JSONL 文件
↓
[convert_to_sft.py] 转换为 SFT 格式
↓
SFT 训练数据集 (.json)
↓
[训练脚本] LoRA/Full SFT 微调
↓
Reward Model 权重
1. 视频预处理
从 TFDS 提取：
# video_process.py
ds = tfds.load("droid_100", data_dir="/playpen-ssd/dataset/", split="train")
# 提取 wrist_image_left + exterior_image_1_left
# 保存为 MP4 视频 + JSON 元数据
从原始视频抽帧：
# extract_frames_to_images.py
# 根据标注的 frame_index 提取对应帧
# 支持 mirror/flat_hash 两种保存模式
2. GPT 标注
改进版特性：
- 动态任务描述（从 metadata JSON 提取）
- 滑动窗口分析（WINDOW_SIZE=5, STRIDE=3）
- 上下文传递（前一窗口结果作为上下文）
- 对齐仿真奖励指标
输出格式：
{
"video_id": "22008760",
"task": "Move object into container",
"t": 0,
"frame_index": 0,
"stage": 0,
"stage_name": "reach",
"reachout": 0.3,
"grasp": 0.0,
"collision": 0.0,
"fall": 0.0,
"smooth": 0.8,
"reward": 0.55,
"delta": 0.0,
"success_prob": 0.5,
"failure": 0,
"explanation": "Robot arm is reaching..."
}
3. SFT 数据格式
训练格式：
{
"conversation": [
{
"from": "human",
"value": "<Image: path/to/frame.jpg>\nPlease output the execution stage and reward..."
},
{
"from": "assistant",
"value": "{\"stage\": \"reach\", \"reward\": 0.5, ...}"
}
]
}
已完成工作
数据处理 ✓
- DROID 数据集统计和分析
- 视频帧提取脚本
- GPT-4 批量标注脚本（基础版 + 改进版）
- 滑动窗口分析模式
- SFT 数据格式转换
标注系统 ✓
- 与仿真奖励指标对齐
- 6 阶段定义（reach, grasp, lift, move, place, retract）
- 5 奖励组件（reachout, grasp, collision, fall, smooth）
- 动态任务描述提取
数据管理 ✓
- HuggingFace 数据集准备脚本
- 采样策略支持（balanced, random, proportional）
- 上传指南文档
模型训练 ✓
- LoRA SFT 训练配置
- 特殊 token 版本实验
- 多个训练检查点
阶段与奖励定义
操作阶段 (Stage)
| 索引 | 名称 | 描述 |
|---|---|---|
| 0 | reach | 机器人向目标物体伸手 |
| 1 | grasp | 机器人尝试抓取物体 |
| 2 | lift | 机器人抬起物体 |
| 3 | move | 机器人携带物体移动 |
| 4 | place | 机器人放置物体 |
| 5 | retract | 机器人收回到中性位置 |
奖励组件 (Reward Components)
| 组件 | 范围 | 描述 |
|---|---|---|
| reachout | 0.0-1.0 | 接近目标的进度 |
| grasp | 0.0-1.0 | 抓取质量（0=无接触, 1=稳固） |
| collision | 0.0-1.0 | 碰撞惩罚（0=无碰撞, 1=严重） |
| fall | 0.0-1.0 | 掉落惩罚（0=稳定, 1=掉落） |
| smooth | 0.0-1.0 | 运动平滑度（0=抖动, 1=平滑） |
综合指标
- reward: 综合奖励值 ≈ (reachout + grasp + smooth - collision - fall) / 3
- delta: 相对上一步的奖励变化（-1.0 到 1.0）
- success_prob: 最终成功概率估计（0.0 到 1.0）
- failure: 是否已失败（0 或 1）
模型训练权重
已保存检查点
| 版本 | 路径 | 说明 |
|---|---|---|
| v0 | weights_reward_sft/v0-20251105-171242/ |
初始版本 |
| v1 | weights_reward_sft/v1-20251110-230134/ |
改进版本 |
| v0 (special token) | weights_reward_sft_special_token/v0-20251110-232149/ |
特殊 token 版本 |
| v1 (special token) | weights_reward_sft_special_token/v1-20251110-234131/ |
特殊 token v1 |
| v2 (special token) | weights_reward_sft_special_token/v2-20251110-235126/ |
完整 merge 版本 |
关键配置
GPT 标注配置
MODEL_NAME = "gpt-4o" # 推荐模型
FPS_SAMPLE = 2 # 帧采样率
USE_SLIDING_WINDOW = True
WINDOW_SIZE = 5 # 窗口大小
WINDOW_STRIDE = 3 # 滑动步长
数据筛选规则
PROCESS_TYPE = "failure" # failure/success/both
VIDEO_NAME_PATTERN = "2*.mp4" # 以 2 开头
EXCLUDE_STEREO = True # 排除 stereo 视频
快速开始
1. 环境配置
pip install openai decord pillow python-dotenv tqdm tensorflow_datasets
2. 设置 API Key
export OPENAI_API_KEY="your-key-here"
# 或创建 .env 文件
3. 运行标注
cd /home/jqliu/projects/RewardModel
# GPT 标注
python api_batch_improved.py
# 或 Qwen VL 标注
cd caption && python infer_caption_batch.py
4. 转换训练数据
python extract_frames_to_images.py
python convert_to_sft.py
参考文档
- README_improved_api.md - 改进版 API 详细说明
- example_comparison.md - 新旧版本对比示例
- caption/README_batch.md - 批量标注说明
- data_sta/UPLOAD_GUIDE.md - HuggingFace 上传指南