Reward Model for Robotics

训练视觉奖励模型，为机器人 VLA Policy 训练提供奖励信号

项目目标

训练一个视觉奖励模型，用于评估机器人操作视频：

输入：机器人操作视频帧
输出：操作阶段、奖励分数、失败检测
应用：为 VLA (Vision-Language-Action) 模型的强化学习提供奖励信号

项目结构

RewardModel_v2/
│
├── scripts/                    # 核心脚本
│   ├── annotate_gpt.py         # GPT 视频标注（主要使用）
│   ├── annotate_hf_dataset.py  # HuggingFace 数据集标注
│   ├── extract_frames.py       # 从视频提取帧图片
│   ├── convert_to_sft.py       # 转换为 SFT 训练格式
│   ├── train_reward_model.sh   # ms-swift LoRA 训练
│   └── caption/                # Qwen-VL 视频描述
│       ├── infer_caption.py
│       └── infer_caption_batch.py
│
├── tools/                      # 工具脚本
│   ├── count_failure_trajectories.py    # 统计失败轨迹
│   ├── check_video_framerate.py         # 检查视频帧率
│   ├── prepare_hf_dataset.py            # 准备 HuggingFace 数据集
│   ├── upload_to_huggingface.py         # 上传到 HuggingFace
│   └── quick_start.sh                   # 快速开始脚本
│
├── data/                       # 数据目录
│   ├── raw/                    # 原始数据
│   │   ├── cokecan-50/         # 仿真数据（按类型分类）
│   │   └── droid/              # DROID 元数据
│   ├── training/               # 训练数据
│   │   └── sft_*.json          # SFT 格式数据
│   ├── system_prompt.txt       # 推理 prompt
│   └── tokens.txt              # 特殊 token 定义
│
├── outputs/                    # 输出目录
│   ├── labels/                 # 标注结果 (.jsonl)
│   ├── frames/                 # 提取的帧图片
│   └── weights/                # 训练权重
│
├── models/                     # 模型权重
│   ├── weights_reward_sft/               # 标准版本
│   └── weights_reward_sft_special_token/ # Special Token 版本
│
├── docs/                       # 文档
│   ├── QUICK_START_FOR_AGENT.md  # AI Agent 快速开始
│   ├── PROJECT_STATUS.md         # 项目状态快照
│   ├── CODE_GUIDE.md             # 代码详细说明
│   ├── PROJECT_OVERVIEW.md       # 项目完整概述
│   ├── DEVELOPMENT_LOG.md        # 开发历程
│   ├── dataset_source.md         # 数据源说明
│   └── ...                       # 其他文档
│
├── archive/                    # 归档（旧版本/历史文件）
│   ├── old_scripts/            # 旧脚本
│   ├── old_outputs/            # 旧输出
│   ├── statistics/             # 统计结果
│   └── droid_processed/        # 处理过的 DROID 样本
│
└── README.md                   # 本文件

快速开始

1. 环境配置

pip install openai decord pillow python-dotenv tqdm tensorflow_datasets

2. 运行标注

cd /home/jqliu/projects/RewardModel_v2

# 设置 API Key
export OPENAI_API_KEY="your-key-here"

# GPT 标注（推荐）
python scripts/annotate_gpt.py

# 或 Qwen-VL 标注
python scripts/caption/infer_caption_batch.py

3. 数据处理流程

# 1. 标注视频 → outputs/labels/
python scripts/annotate_gpt.py

# 2. 提取帧图片 → outputs/frames/
python scripts/extract_frames.py

# 3. 转换为训练格式 → data/training/
python scripts/convert_to_sft.py

4. 模型训练

# 安装 ms-swift（首次使用）
pip install ms-swift

# 标准模式训练（JSON 输出格式）
bash scripts/train_reward_model.sh

# Special Token 模式训练（结构化输出格式）
bash scripts/train_reward_model.sh --special-token

# 仅打印命令，不执行（用于检查参数）
bash scripts/train_reward_model.sh --dry-run

参数配置

通过环境变量覆盖默认参数：

# 使用更大模型
MODEL=Qwen/Qwen2.5-VL-7B-Instruct bash scripts/train_reward_model.sh

# 调整 LoRA 和训练参数
EPOCHS=5 LORA_RANK=16 LR=2e-5 bash scripts/train_reward_model.sh

# 使用自定义数据集
DATASET=data/training/my_dataset.json bash scripts/train_reward_model.sh

# 指定 GPU
GPU=0,1 bash scripts/train_reward_model.sh

关键参数说明

参数	默认值	说明
`MODEL`	Qwen2.5-VL-3B-Instruct	基座模型
`DATASET`	自动选择	训练数据路径（相对路径）
`LORA_RANK`	8	LoRA 秩，越大容量越强
`LORA_ALPHA`	32	LoRA 缩放系数
`EPOCHS`	1	训练轮数
`GRAD_ACCUM`	16	梯度累积步数
`LR`	1e-5	学习率
`MAX_LENGTH`	1024	最大序列长度

数据量与参数建议

数据规模	建议 EPOCHS	建议 LORA_RANK
< 500 条	5 ~ 10	8
500 ~ 5000 条	1 ~ 3	16 ~ 32
> 5000 条	1	8 ~ 16

训练权重输出到 outputs/weights/ 目录，自动带时间戳版本号。

核心脚本说明

脚本	功能	输出
`scripts/annotate_gpt.py`	GPT 视频标注（滑动窗口）	`outputs/labels/*.jsonl`
`scripts/annotate_hf_dataset.py`	HuggingFace 数据集标注	`outputs/labels/*.jsonl`
`scripts/extract_frames.py`	提取视频帧为图片	`outputs/frames/*/.jpg`
`scripts/convert_to_sft.py`	转换为 SFT 训练格式	`data/training/*.json`
`scripts/train_reward_model.sh`	ms-swift LoRA 训练	`outputs/weights/`

项目状态

模块	状态
数据处理流水线	✅ 完成
GPT 标注系统	✅ 完成
小规模模型训练	✅ 完成
大规模数据标注	⏳ 进行中
模型评估系统	📋 待启动

详见：docs/PROJECT_STATUS.md

文档导航

文档	用途
QUICK_START_FOR_AGENT.md	AI Agent 快速上手
PROJECT_STATUS.md	项目状态快照
CODE_GUIDE.md	所有脚本详细说明
dataset_source.md	数据源位置

数据源

真机数据

DROID 数据集：/playpen-ssd/dataset/droid_raw/1.0.1/
failure 轨迹：15,157 条

仿真数据

本地 cokecan-50：data/raw/cokecan-50/
fangyu 仿真：/playpen-ssd/yufang/projects/reward/

详见：docs/dataset_source.md

归档说明

archive/ 目录包含旧版本文件，保留供参考：

目录	内容
`old_scripts/`	早期脚本版本（api.py, api_batch.py 等）
`old_outputs/`	历史输出文件
`statistics/`	数据统计结果
`droid_processed/`	处理过的 DROID 样本数据

更新日志

v2.0.0 (2026-01-21)

重构项目结构，更加清晰
重命名脚本，语义更明确
归档旧版本文件
统一输出目录

v1.0.0 (2025-11-10)

完成基础功能开发
完成小规模模型训练