# 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. 环境配置

```bash
pip install openai decord pillow python-dotenv tqdm tensorflow_datasets
```

### 2. 运行标注

```bash
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. 数据处理流程

```bash
# 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. 模型训练

```bash
# 安装 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
```

#### 参数配置

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

```bash
# 使用更大模型
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](docs/PROJECT_STATUS.md)

---

## 文档导航

| 文档 | 用途 |
|------|------|
| [QUICK_START_FOR_AGENT.md](docs/QUICK_START_FOR_AGENT.md) | AI Agent 快速上手 |
| [PROJECT_STATUS.md](docs/PROJECT_STATUS.md) | 项目状态快照 |
| [CODE_GUIDE.md](docs/CODE_GUIDE.md) | 所有脚本详细说明 |
| [dataset_source.md](docs/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](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)
- 完成基础功能开发
- 完成小规模模型训练