EfficientZero（NeurIPS 2021）

EfficientZero 的开源代码库，源自 NeurIPS 2021 的论文《用有限数据掌握 Atari 游戏》（https://arxiv.org/abs/2111.00210）。

环境要求

EfficientZero 需要 Python 3（>=3.6）和 PyTorch（>=1.8.0），并安装开发头文件。

我们建议使用 PyTorch AMP（--amp_type torch_amp）来加速训练。

前提条件

在开始训练之前，您需要构建 C++/Cython 样式的外部依赖包。（需 GCC 7.5+ 版本。）

cd core/ctree
bash make.sh

该代码库的分布式框架基于 Ray 构建。

安装

对于该代码库所需的其他依赖包，请运行 pip install -r requirements.txt。

使用方法

快速入门

• 训练：python main.py --env BreakoutNoFrameskip-v4 --case atari --opr train --amp_type torch_amp --num_gpus 1 --num_cpus 10 --cpu_actor 1 --gpu_actor 1 --force
• 测试：python main.py --env BreakoutNoFrameskip-v4 --case atari --opr test --amp_type torch_amp --num_gpus 1 --load_model --model_path model.p

Bash 脚本

我们提供了 train.sh 和 test.sh 用于训练和评估。

• 训练：
  • 使用 4 张显卡（3090）：bash train.sh
• 测试：bash test.sh

必选参数	描述
--env	环境名称
--case {atari}	用于切换不同领域（默认：atari）
--opr {train,test}	选择要执行的操作
--amp_type {torch_amp,none}	使用 PyTorch AMP 加速

其他参数	描述
--force	将重写结果目录
--num_gpus 4	可用的 GPU 数量
--num_cpus 96	可用的 CPU 数量
--cpu_actor 14	CPU 工作进程数
--gpu_actor 20	GPU 工作进程数
--seed 0	随机种子
--use_priority	在回放缓冲区采样中使用优先级
--use_max_priority	对新收集的数据使用最大优先级
--amp_type 'torch_amp'	使用 PyTorch AMP 加速
--info 'EZ-V0'	为您的实验添加一些标签
--p_mcts_num 8	设置自博弈中环境的并行数量
--revisit_policy_search_rate 0.99	设置重新分析策略的比率
--use_root_value	在价值目标中使用根值（需要更多 GPU 工作进程）
--render	在评估时渲染画面
--save_video	保存评估视频

架构设计

训练流水线的架构如下所示：

一些建议

• 如果想使用更小的模型，可以在配置中选择较小的投影层维度（例如：256/64）和 LSTM 隐藏层维度（例如：64）。
• 对于内存为 10G 而非 20G 的显卡，可以在 core/reanalyze_worker.py 中为每个 GPU 分配 0.25 的 GPU 资源（@ray.remote(num_gpus=0.25)）。

新环境注册

如果您想将 EfficientZero 应用于新的环境，比如 mujoco，请按以下步骤进行注册：

1. 复制 config/atari 目录，并在 config/mujoco 中创建对应环境的目录。
2. 实现您的 MujocoConfig(BaseConfig) 类，并实现模型以及环境包装器。
3. 在 main.py 中注册该案例。

结果

使用 32 个种子进行了 3 次不同运行的评估（不同的种子）。

引用

如果您觉得这个仓库有用，请引用我们的论文：

@inproceedings{ye2021mastering,
  title={用有限数据掌握 Atari 游戏},
  author={叶伟瑞、刘少怀、塔纳德·库鲁塔奇、皮特·阿比尔、高阳},
  booktitle={NeurIPS},
  year={2021}
}

联系方式

如果您有任何问题或希望使用该代码，请联系 ywr20@mails.tsinghua.edu.cn。

致谢

我们非常感谢以下 GitHub 仓库提供的宝贵代码实现：

https://github.com/koulanurag/muzero-pytorch

https://github.com/werner-duvaud/muzero-general

https://github.com/pytorch/ELF

EfficientZero 快速上手指南

环境准备

• 系统要求：Python 3.6+，GCC 7.5+
• 深度学习框架：PyTorch 1.8.0+
• 推荐加速：使用 --amp_type torch_amp 启用 Torch AMP 加速训练
• 分布式框架：基于 Ray

安装步骤

1. 克隆仓库并进入核心目录：

   cd core/ctree
   bash make.sh
   

2. 安装 Python 依赖（推荐使用清华源加速）：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

训练模型（以 Breakout 为例）

python main.py --env BreakoutNoFrameskip-v4 --case atari --opr train --amp_type torch_amp --num_gpus 1 --num_cpus 10 --cpu_actor 1 --gpu_actor 1 --force

测试模型

python main.py --env BreakoutNoFrameskip-v4 --case atari --opr test --amp_type torch_amp --num_gpus 1 --load_model --model_path model.p

使用预置脚本（推荐）

• 训练（4×3090 GPU）：

  bash train.sh
  

• 测试：

  bash test.sh