# 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 分类 |

**仿真数据格式**：
```json
{
  "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 提取**：
```python
# video_process.py
ds = tfds.load("droid_100", data_dir="/playpen-ssd/dataset/", split="train")
# 提取 wrist_image_left + exterior_image_1_left
# 保存为 MP4 视频 + JSON 元数据
```

**从原始视频抽帧**：
```python
# extract_frames_to_images.py
# 根据标注的 frame_index 提取对应帧
# 支持 mirror/flat_hash 两种保存模式
```

### 2. GPT 标注

**改进版特性**：
- 动态任务描述（从 metadata JSON 提取）
- 滑动窗口分析（WINDOW_SIZE=5, STRIDE=3）
- 上下文传递（前一窗口结果作为上下文）
- 对齐仿真奖励指标

**输出格式**：
```json
{
  "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 数据格式

**训练格式**：
```json
{
  "conversation": [
    {
      "from": "human",
      "value": "<Image: path/to/frame.jpg>\nPlease output the execution stage and reward..."
    },
    {
      "from": "assistant",
      "value": "{\"stage\": \"reach\", \"reward\": 0.5, ...}"
    }
  ]
}
```

---

## 已完成工作

### 数据处理 ✓
- [x] DROID 数据集统计和分析
- [x] 视频帧提取脚本
- [x] GPT-4 批量标注脚本（基础版 + 改进版）
- [x] 滑动窗口分析模式
- [x] SFT 数据格式转换

### 标注系统 ✓
- [x] 与仿真奖励指标对齐
- [x] 6 阶段定义（reach, grasp, lift, move, place, retract）
- [x] 5 奖励组件（reachout, grasp, collision, fall, smooth）
- [x] 动态任务描述提取

### 数据管理 ✓
- [x] HuggingFace 数据集准备脚本
- [x] 采样策略支持（balanced, random, proportional）
- [x] 上传指南文档

### 模型训练 ✓
- [x] LoRA SFT 训练配置
- [x] 特殊 token 版本实验
- [x] 多个训练检查点

---

## 阶段与奖励定义

### 操作阶段 (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 标注配置

```python
MODEL_NAME = "gpt-4o"  # 推荐模型
FPS_SAMPLE = 2         # 帧采样率
USE_SLIDING_WINDOW = True
WINDOW_SIZE = 5        # 窗口大小
WINDOW_STRIDE = 3      # 滑动步长
```

### 数据筛选规则

```python
PROCESS_TYPE = "failure"     # failure/success/both
VIDEO_NAME_PATTERN = "2*.mp4"  # 以 2 开头
EXCLUDE_STEREO = True        # 排除 stereo 视频
```

---

## 快速开始

### 1. 环境配置

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

### 2. 设置 API Key

```bash
export OPENAI_API_KEY="your-key-here"
# 或创建 .env 文件
```

### 3. 运行标注

```bash
cd /home/jqliu/projects/RewardModel

# GPT 标注
python api_batch_improved.py

# 或 Qwen VL 标注
cd caption && python infer_caption_batch.py
```

### 4. 转换训练数据

```bash
python extract_frames_to_images.py
python convert_to_sft.py
```

---

## 参考文档

- [README_improved_api.md](./README_improved_api.md) - 改进版 API 详细说明
- [example_comparison.md](./example_comparison.md) - 新旧版本对比示例
- [caption/README_batch.md](./caption/README_batch.md) - 批量标注说明
- [data_sta/UPLOAD_GUIDE.md](./data_sta/UPLOAD_GUIDE.md) - HuggingFace 上传指南