去噪扩散策略优化

用于论文《使用强化学习训练扩散模型》（https://rl-diffusion.github.io/）的训练代码。该代码库已在 Google Cloud TPUs 上测试过（RWR 使用 v3，DDPO 使用 v4）；尚未在 GPU 上进行测试。

更新： 我们现在提供了一个支持 GPU 和 LoRA 以实现低内存训练的 PyTorch 实现，详情请见 这里！

prompt_fn	filter_field	权重与演示
imagenet_animals	jpeg	ddpo-compressibility
imagenet_animals	neg_jpeg	ddpo-incompressibility
from_file(assets/common_animals.txt)	aesthetic	ddpo-aesthetic
nouns_activities(assets/common_animals.txt, assets/activities_v0.txt)	llava_bertscore	ddpo-alignment

安装

conda env create -f environment_tpu.yml
conda activate ddpo-tpu
pip install -e .

运行 DDPO

python pipeline/policy_gradient.py --dataset compressed-animals

--dataset 标志可以替换为 config/base.py 中定义的任何配置。第一个配置字典 base 定义了通用参数，这些参数会在后续的具体配置中被覆盖。部分参数在不同方法间共享；DDPO 特有的超参数位于 pg 字段中。

最重要的两个参数是 prompt_fn 和 filter_field，它们分别定义了提示分布和奖励函数。提示函数可在 training/prompts.py 中找到，奖励函数则位于 training/callbacks.py。

运行 RWR

对于标准 RWR，其中权重是奖励的 softmax 归一化结果：

bash pipeline/run-rwr.sh

对于 RWR-sparse，仅保留并训练奖励分布中高于特定百分位数的样本：

bash pipeline/run-sparse.sh

这些方法使用 Bash 而不是 Python 来运行最外层的训练循环。它们会执行 pipeline/sample.py 脚本以收集样本和奖励数据集，然后运行 pipeline/finetune.py 来基于最新数据集训练模型，并重复若干次迭代。采样步骤和微调步骤使用不同的配置，在 config/base.py 中分别标记为 "sample" 和 "train"。

运行 LLaVA 推理

LLaVA 推理是通过向独立的 GPU 服务器发送 HTTP 请求来完成的。客户端代码可在 training/callbacks.py 中的 llava_bertscore 奖励函数中找到，而服务器端代码则位于 此仓库。

参考文献

@inproceedings{black2023ddpo,
      title={Training Diffusion Models with Reinforcement Learning},
      author={Kevin Black and Michael Janner and Yilun Du and Ilya Kostrikov and Sergey Levine},
      year={2023},
      eprint={2305.13301},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

DDPO 快速上手指南

DDPO (Denoising Diffusion Policy Optimization) 是一个利用强化学习训练扩散模型的开源项目。以下指南基于官方 README 整理，旨在帮助开发者快速搭建环境并运行基础示例。

环境准备

• 硬件要求：
  • 官方代码主要针对 Google Cloud TPUs 进行测试（RWR 使用 v3，DDPO 使用 v4）。
  • 注意：当前仓库代码未在 GPU 上进行测试。如果您需要使用 GPU 或进行低显存训练（支持 LoRA），请使用社区提供的 PyTorch 版本：ddpo-pytorch。
• 软件依赖：
  • Conda 包管理器
  • Python 环境（由 environment_tpu.yml 定义）

安装步骤

请按照以下命令创建虚拟环境并安装依赖：

conda env create -f environment_tpu.yml
conda activate ddpo-tpu
pip install -e .

提示：国内用户若下载 pip 依赖较慢，可在安装命令后追加国内镜像源，例如：pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple。

基本使用

运行 DDPO 训练

使用以下命令启动基础的策略梯度训练流程。该示例使用 compressed-animals 数据集配置：

python pipeline/policy_gradient.py --dataset compressed-animals

关键参数说明

• --dataset：可替换为 config/base.py 中定义的任何配置名称。
• 核心参数：
  • prompt_fn：定义提示词分布（见 training/prompts.py）。
  • filter_field：定义奖励函数（见 training/callbacks.py）。
• 超参数：DDPO 特有的超参数位于配置字典的 pg 字段中。

其他运行模式

如果需要运行标准的 RWR (Reward Weighted Regression) 或稀疏 RWR，可使用以下脚本（这些脚本在 Bash 层面控制采样与微调的循环）：

• 标准 RWR：

  bash pipeline/run-rwr.sh
  

• 稀疏 RWR (RWR-sparse)：

  bash pipeline/run-sparse.sh