[论文解读] TensorFlow Agents: Efficient Batched Reinforcement Learning in TensorFlow

Danijar Hafner, James Davidson+1 more|arXiv (Cornell University)|Sep 8, 2017

Adversarial Robustness in Machine Learning参考文献 9被引用 29

一句话总结

本文介绍了 TensorFlow Agents，一个用于在 TensorFlow 中高效进行批量强化学习的框架，通过并行化环境滚动和批量处理神经网络推理来加速训练。它提出了 BatchPPO，一种对近端策略优化（Proximal Policy Optimization）的优化实现，在仅使用 CPU 资源的情况下，其性能与或优于现有文献在 MuJoCo 运动控制任务上的结果。

ABSTRACT

We introduce TensorFlow Agents, an efficient infrastructure paradigm for building parallel reinforcement learning algorithms in TensorFlow. We simulate multiple environments in parallel, and group them to perform the neural network computation on a batch rather than individual observations. This allows the TensorFlow execution engine to parallelize computation, without the need for manual synchronization. Environments are stepped in separate Python processes to progress them in parallel without interference of the global interpreter lock. As part of this project, we introduce BatchPPO, an efficient implementation of the proximal policy optimization algorithm. By open sourcing TensorFlow Agents, we hope to provide a flexible starting point for future projects that accelerates future research in the field.

研究动机与目标

• 通过为 TensorFlow 中的强化学习算法训练提供统一且高效的基础设施，加速强化学习研究。
• 通过使用独立的 Python 进程并行模拟多个环境，解决强化学习中环境交互的瓶颈问题。
• 通过批量处理观测并针对整个批次执行神经网络推理，而非单个转移样本，提升计算效率。
• 提供一个可重用、可扩展的框架，将算法实现与环境交互解耦，以实现更快的原型开发。
• 发布 BatchPPO，作为未来强化学习研究的高性能基线。

提出的方法

• 扩展 OpenAI Gym 接口以支持批量环境，多个环境在独立的 Python 进程中并行推进，避免全局解释器锁（GIL）争用。
• 将批量环境集成到 TensorFlow 计算图中，支持通过单次会话调用实现端到端可微分训练。
• 使用重要性采样校正分布偏移问题，当使用先前策略收集的离策略数据更新策略时。
• 采用双流神经网络架构：一个用于动作均值预测，一个用于价值函数估计，两者均使用 ReLU 激活函数，输出层使用 tanh 以限制动作范围。
• 使用流式统计量对观测和奖励进行归一化，并引入 KL 散度惩罚以稳定训练过程。
• 使用固定学习率的 Adam 优化器，每批次 25 场景的 25 次梯度更新，同时更新策略网络和价值网络。

实验结果

研究问题

• RQ1 在 TensorFlow 中构建统一的批量基础设施，能否显著加速强化学习算法的训练？
• RQ2 与顺序执行或受 GIL 限制的执行方式相比，并行化多个独立进程中的环境模拟，如何提升训练吞吐量？
• RQ3 批量实现的 PPO 在标准 MuJoCo 控制基准测试中，其性能在多大程度上能与或超越已有文献结果？
• RQ4 通过适当的批量处理和并行化，是否仅使用 CPU 资源即可实现高效且可扩展的强化学习训练？
• RQ5 将环境模拟与神经网络推理整合到单一 TensorFlow 计算图中，如何简化训练循环？

主要发现

• BatchPPO 在 HalfCheetah-v1、Walker2d-v1 和 Humanoid-v1 任务上，性能与原始 PPO 实现相比达到或超越已有结果。
• 该框架可在仅使用 6 个 CPU 核心的单台机器上实现高效训练，且通过批量环境交互实现训练性能的有效扩展。
• 使用平均动作进行评估可获得稳定且高回报的结果，而训练期间采样则在三个随机种子下均表现出一致的学习进展。
• 结果图中 25% 至 75% 分位数的置信区间（阴影区域）显示方差较低，表明性能可靠且可复现。
• 使用独立进程进行环境步进消除了 GIL 瓶颈，实现了环境交互的真正并行化。
• 该框架的设计支持通过控制张量实现训练与评估阶段的无缝切换，包括重置和日志记录。