用于世界建模的扩散模型：Atari 游戏中的视觉细节至关重要（NeurIPS 2024 Spotlight）

[简而言之] 💎 DIAMOND（DIffusion As a Model Of eNvironment Dreams）是一种完全在扩散世界模型中训练的强化学习智能体。

🌍 项目页面 • 🤓 论文 • 𝕏 Atari 相关推文 • 𝕏 CSGO 相关推文 • 💬 Discord

使用 miniconda 快速安装以试用我们的 预训练世界模型：

git clone https://github.com/eloialonso/diamond.git
cd diamond
conda create -n diamond python=3.10
conda activate diamond
pip install -r requirements.txt

对于 Atari（世界模型 + 强化学习智能体）：

python src/play.py --pretrained

对于 CSGO（仅世界模型）：

git checkout csgo
python src/play.py

然后按下 m 键即可接管控制权（默认情况下由策略控制）！

警告：Atari ROM 将随依赖项一起下载，这意味着您确认自己拥有使用这些 ROM 的许可。

CSGO

编辑：请查看 csgo 分支，试用我们基于 反恐精英：全球攻势 训练的 DIAMOND 世界模型！

git checkout csgo
python src/play.py

注意：在 Apple Silicon 上，您必须启用 MPS 后端的 CPU 回退功能，即设置 PYTORCH_ENABLE_MPS_FALLBACK=1 后再运行 python src/play.py。

快速链接

• 试用我们的可玩扩散世界模型
• 启动训练流程
• 配置
• 可视化
  • 游玩模式（默认）
  • 数据集模式（添加 -d）
  • 其他选项，适用于游玩和数据集模式
• 运行文件夹结构
• 结果
• 引用
• 致谢

⬆️ 试用我们的可玩扩散世界模型

python src/play.py --pretrained

然后选择一款游戏，系统将从我们在 Hugging Face Hub上的仓库 🤗 下载并缓存在您的机器上，预训练好的 Atari 100k 数据集上的世界模型和策略。

您可以尝试以下操作：

• 按下 m 键在智能体和人类之间切换控制方式（默认为智能体控制）。
• 按下 ↑/↓ 键调整想象视野（默认值为 50）。

要调整已训练扩散世界模型的采样参数（去噪步数、随机性、顺序等），例如在采样速度和质量之间进行权衡，请编辑 config/trainer.yaml 文件中的 world_model_env.diffusion_sampler 部分。

更多关于可用命令和选项的详细信息，请参阅 可视化 部分。

⬆️ 启动训练流程

要在 cuda:0 上使用论文中使用的超参数进行训练，请运行：

python src/main.py env.train.id=BreakoutNoFrameskip-v4 common.devices=0

这将在 outputs/YYYY-MM-DD/hh-mm-ss/ 目录下创建一个新的运行文件夹。

若要恢复意外中断的训练，请进入该运行文件夹并执行：

./scripts/resume.sh

⬆️ 配置

我们使用 Hydra 进行配置管理。

所有配置文件都位于 config 文件夹中：

• config/trainer.yaml：主配置文件。
• config/agent/default.yaml：架构超参数。
• config/env/atari.yaml：环境超参数。

您可以在 config/trainer.yaml 文件的 wandb 部分开启 weights & biases 日志记录功能。

将 config/trainer.yaml 文件中的 training.model_free=true 设置为 true，即可“断开”世界模型的连接，从而执行标准的无模型强化学习。

⬆️ 可视化

⬆️ 游玩模式（默认）

要可视化您最近的检查点，请从 运行文件夹 中运行：

python src/play.py

默认情况下，您将看到策略在世界模型中运行。如果您想亲自操控或切换到真实环境，可以使用下面介绍的控制键。

控制键（游玩模式）

（游戏特定的命令将在启动时显示）

⏎   ：重置环境

m   ：切换控制器（策略/人类）
↑/↓ ：想象视野（+1/-1）
←/→ ：下一个环境 [世界模型 ←→ 测试真实环境 ←→ 训练真实环境]

.   ：暂停/继续
e   ：单步执行（暂停时）

添加 -r 可以切换“录制模式”（仅在游玩模式下有效）。每个完成的回合都会保存在 dataset/rec_<env_name>_<controller> 中。例如：

• dataset/rec_wm_π：策略在世界模型中运行。
• dataset/rec_wm_H：人类在世界模型中运行。
• dataset/rec_test_H：人类在测试真实环境中运行。

随后，您可以使用下一节中描述的“数据集模式”来回放存储的回合。

⬆️ 数据集模式（添加 -d）

在运行文件夹中，要可视化 dataset 子文件夹中的数据集，请添加 -d 切换到“数据集模式”：

python src/play.py -d

您可以使用下面介绍的控制键来浏览数据集和回合。

控制键（数据集模式）

m   ：下一个数据集（如果有多个数据集，如录制内容等）
↑/↓ ：下一个/上一个回合
←/→ ：回合中的下一个/上一个时间步
PgUp：前进 10 个时间步
PgDn：后退 10 个时间步
⏎   ：返回第一个时间步

⬆️ 其他选项，适用于游玩和数据集模式

--fps FPS             目标帧率（默认 15）。
--size SIZE           窗口大小（默认 800）。
--no-header           移除标题栏。

⬆️ 运行文件夹结构

每次新的运行都会保存在 outputs/YYYY-MM-DD/hh-mm-ss/ 目录下。该文件夹的结构如下：

outputs/YYYY-MM-DD/hh-mm-ss/
│
└─── checkpoints
│   │   state.pt  # 完整的训练状态
│   │
│   └─── agent_versions
│       │   ...
│       │   agent_epoch_00999.pt
│       │   agent_epoch_01000.pt  # 仅包含智能体权重
│
└─── config
│   |   trainer.yaml
|
└─── dataset
│   │
│   └─── train
│   |   │   info.pt
│   |   │   ...
|   |
│   └─── test
│       │   info.pt
│       │   ...
│
└─── scripts
│   │   resume.sh
|   |   ...
|
└─── src
|   |   main.py
|   |   ...
|
└─── wandb
    |   ...

⬆️ 结果

文件 results/data/DIAMOND.json 包含了论文中所使用的所有游戏和随机种子的实验结果。

用于论文第5.1节的DDPM代码可以在 ddpm 分支中找到。

⬆️ 引用

@inproceedings{alonso2024diffusionworldmodelingvisual,
      title={Diffusion for World Modeling: Visual Details Matter in Atari},
      author={Eloi Alonso and Adam Jelley and Vincent Micheli and Anssi Kanervisto and Amos Storkey and Tim Pearce and François Fleuret},
      booktitle={Thirty-eighth Conference on Neural Information Processing Systems}}
      year={2024},
      url={https://arxiv.org/abs/2405.12399},
}

⬆️ 致谢

• https://github.com/crowsonkb/k-diffusion/
• https://github.com/huggingface/huggingface_hub
• https://github.com/google-research/rliable
• https://github.com/pytorch/pytorch

DIAMOND 快速上手指南

DIAMOND (DIffusion As a Model Of eNvironment Dreams) 是一个完全在扩散世界模型中训练的强化学习智能体，支持 Atari 和 CS:GO 等环境。

环境准备

• 操作系统: Linux / macOS (Windows 需自行配置兼容环境)
• 硬件要求: 推荐配备 NVIDIA GPU (CUDA 支持)；Apple Silicon (M1/M2/M3) 用户需注意后端配置。
• 前置依赖:
  • Miniconda 或 Anaconda
  • Git
  • Python 3.10

注意: 运行 Atari 游戏时，脚本会自动下载对应的 ROM 文件。请确保您拥有使用这些 ROM 的合法授权。

安装步骤

1. 克隆项目代码

   git clone https://github.com/eloialonso/diamond.git
   cd diamond
   

2. 创建并激活 Conda 环境

   conda create -n diamond python=3.10
   conda activate diamond
   

3. 安装依赖包

   国内加速建议: 如果 pip 下载速度慢，建议使用清华或阿里镜像源。

   # 使用默认源
   pip install -r requirements.txt
   
   # 或使用清华镜像源 (推荐国内用户)
   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

1. 体验预训练模型 (Atari)

直接运行以下命令，系统将自动从 Hugging Face 下载预训练的世界模型和策略模型（首次运行需联网）。

python src/play.py --pretrained

• 操作说明:
  • 启动后选择游戏。
  • 默认由 AI 策略 自动游玩。
  • 按 m 键：切换控制权（在 AI 策略与人类玩家之间切换）。
  • 按 ↑ / ↓ 键：调整想象步长 (imagination horizon)。
  • 按 ← / → 键：切换环境模式 (世界模型 ↔ 真实测试环境 ↔ 真实训练环境)。

2. 体验 CS:GO 世界模型

CS:GO 模型位于独立分支，需切换分支后运行。

git checkout csgo
python src/play.py

• Apple Silicon 用户特别注意: 在 M 系列芯片 Mac 上运行 CS:GO 模型时，必须启用 MPS 后端回退机制：

  PYTORCH_ENABLE_MPS_FALLBACK=1 python src/play.py
  

3. 开始训练 (可选)

若需在本地重新训练模型（以 Breakout 游戏为例，使用 CUDA 0）：

python src/main.py env.train.id=BreakoutNoFrameskip-v4 common.devices=0

训练产出将保存在 outputs/YYYY-MM-DD/hh-mm-ss/ 目录下。