Sample Factory: 高吞吐量的同步和异步强化学习框架

Sample Factory是一个专注于高效实现同步和异步策略梯度算法(如PPO)的强化学习框架。它由Aleksei Petrenko等人开发,旨在提供极高的训练吞吐量和样本效率。

主要特点

Sample Factory具有以下关键特性:

高度优化的算法架构,可实现最大学习吞吐量
支持同步和异步训练模式
单进程串行模式,便于调试
在基于CPU和GPU加速环境中都能实现最佳性能
支持单智能体和多智能体训练、自我对弈、多策略并行训练
基于人口的训练(PBT)
支持离散、连续及混合动作空间
支持向量、图像及字典观察空间
自动根据动作/观察空间创建模型架构,同时支持自定义模型
设计为可导入其他项目,支持自定义环境
详细的WandB和Tensorboard可视化,支持自定义指标
集成HuggingFace,可上传训练模型和指标
提供多个环境集成示例及调优参数

适用场景

Sample Factory特别适合以下应用场景:

需要在单机上进行大规模强化学习训练,追求最快训练速度。
进行多智能体或基于人口的训练实验。
在VizDoom、DMLab等支持良好的环境中训练智能体。
作为分布式强化学习系统的单节点原型或异步算法的参考实现。

安装与使用

Sample Factory可通过pip直接安装:

pip install sample-factory

基本使用方法如下:

from sample_factory.algorithms.appo.train_appo import train_appo
from sample_factory.algorithms.appo.enjoy_appo import enjoy_appo

# 训练
train_appo(env='doom_basic', train_for_env_steps=3000000, num_workers=20)

# 评估
enjoy_appo(env='doom_basic')

对于VizDoom、DMLab等特定环境,需要安装额外依赖,具体可参考官方文档。

性能展示

Sample Factory在多个强化学习基准测试中展现了卓越的性能。以下是在VizDoom、IsaacGym、DMLab-30、Megaverse、Mujoco和Atari等环境中训练的智能体展示:

VizDoom agents trained using Sample Factory 2.0 IsaacGym agents trained using Sample Factory 2.0

这些结果展示了Sample Factory在各种复杂环境中的强大学习能力。

配置与监控

Sample Factory通过命令行参数进行配置。可以使用以下命令查看完整的参数列表:

python -m sample_factory.algorithms.appo.train_appo --algo=APPO --env=doom_battle --help

训练过程可以通过Tensorboard进行监控:

tensorboard --logdir=./train_dir

此外,Sample Factory还提供了一个方便的脚本来监控多个实验:

python -m sample_factory.tb '*custom_experiment*' '*another*custom*experiment_name'

高级功能

Sample Factory支持多种高级功能,如多智能体训练、人口基础训练(PBT)、自定义模型架构等。例如,以下是一个在VizDoom环境中进行多智能体PBT训练的示例:

python -m sample_factory.algorithms.appo.train_appo --env=doom_duel --num_workers=72 --num_envs_per_worker=16 --num_policies=8 --with_pbt=True

结语

Sample Factory是一个强大而灵活的强化学习框架,特别适合需要高吞吐量训练的场景。它不仅提供了卓越的性能,还支持多种高级功能,使研究人员和实践者能够轻松进行复杂的强化学习实验。无论是在学术研究还是工业应用中,Sample Factory都是一个值得考虑的工具。

如果您在研究中使用了Sample Factory,请引用以下论文:

@inproceedings{petrenko2020sf,
  title={Sample Factory: Egocentric 3D Control from Pixels at 100000 FPS with Asynchronous Reinforcement Learning},
  author={Petrenko, Aleksei and Huang, Zhehui and Kumar, Tushar and Sukhatme, Gaurav and Koltun, Vladlen},
  booktitle={ICML},
  year={2020}
}

Sample Factory是开源项目,欢迎社区贡献。如有任何问题或建议,可以在GitHub仓库提出issue或pull request。让我们一起推动强化学习的发展!