WarpDrive：在GPU上实现极速端到端深度多智能体强化学习
机器学习
2021-10-12 v3 多智能体系统
摘要
深度强化学习（RL）是在复杂环境中训练决策模型的强大框架。然而，RL可能很慢，因为它需要与环境的模拟进行反复交互。特别是，在具有多个智能体以及高维状态、观测或动作空间的复杂环境中使用RL时，存在关键的系统工程瓶颈。我们提出WarpDrive，一个灵活、轻量且易用的开源RL框架，它在单个GPU（图形处理单元）上实现端到端深度多智能体RL，构建于PyCUDA和PyTorch之上。利用GPU的极致并行化能力，WarpDrive相比混合CPU模拟与GPU模型的常见实现实现了数量级更快的RL。我们的设计并行运行各模拟及其中的智能体，消除了CPU与GPU之间的数据拷贝，并采用GPU上单一模拟数据存储进行安全的原地更新。WarpDrive提供轻量Python接口和灵活的环境封装，易于使用和扩展。综上，这使用户能轻松运行数千个并发多智能体模拟并在极大经验批次上训练。通过大量实验，我们验证了WarpDrive提供高吞吐且几乎线性扩展到多智能体与并行环境。例如，在基准Tag模拟中，WarpDrive以2000个环境和1000个智能体产生每秒290万环境步数（至少比CPU实现高100倍吞吐）。因此，WarpDrive是一个快速且可扩展的多智能体RL平台，可显著加速研发。
引用
@article{arxiv.2108.13976,
title = {WarpDrive: Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning on a GPU},
author = {Tian Lan and Sunil Srinivasa and Huan Wang and Stephan Zheng},
journal= {arXiv preprint arXiv:2108.13976},
year = {2021}
}
备注
TL and SS contributed equally. Code is available at https://www.github.com/salesforce/warp-drive. 11 pages, 14 figures