TensorFlow Agents: Efficient Batched Reinforcement Learning in TensorFlow

TL;DR

TensorFlow Agents provides an efficient framework for parallel reinforcement learning by batching environment simulations and neural network computations, enabling scalable and synchronized training in TensorFlow.

Contribution

The paper introduces TensorFlow Agents, a new infrastructure for parallel RL that leverages batching and process separation, including the novel BatchPPO implementation.

Findings

Enables parallel environment simulation without interference

Improves efficiency of RL training in TensorFlow

Open sourced for community use and development

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.