SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference

Github : http : //github.com/google-research/seed_rl . 1 INTRODUCTION . The field of reinforcement learning ( RL ) has recently seen impressive results across a variety of tasks . This has in part been fueled by the introduction of deep learning in RL and the introduction of accelerators such as GPUs . In the very recent history , focus on massive scale has been key to solve a number of complicated games such as AlphaGo ( Silver et al. , 2016 ) , Dota ( OpenAI , 2018 ) and StarCraft 2 ( Vinyals et al. , 2017 ) . The sheer amount of environment data needed to solve tasks trivial to humans , makes distributed machine learning unavoidable for fast experiment turnaround time . RL is inherently comprised of heterogeneous tasks : running environments , model inference , model training , replay buffer , etc . and current state-of-the-art distributed algorithms do not efficiently use compute resources for the tasks . The amount of data and inefficient use of resources makes experiments unreasonably expensive . The two main challenges addressed in this paper are scaling of reinforcement learning and optimizing the use of modern accelerators , CPUs and other resources . We introduce SEED ( Scalable , Efficient , Deep-RL ) , a modern RL agent that scales well , is flexible and efficiently utilizes available resources . It is a distributed agent where model inference is done centrally combined with fast streaming RPCs to reduce the overhead of inference calls . We show that with simple methods , one can achieve state-of-the-art results faster on a number of tasks . For optimal performance , we use TPUs ( cloud.google.com/tpu/ ) and TensorFlow 2 ( Abadi et al. , 2015 ) to simplify the implementation . The cost of running SEED is analyzed against IMPALA ( Espeholt et al. , 2018 ) which is a commonly used state-of-the-art distributed RL algorithm ( Veeriah et al . ( 2019 ) ; Li et al . ( 2019 ) ; Deverett et al . ( 2019 ) ; Omidshafiei et al . ( 2019 ) ; Vezhnevets et al . ( 2019 ) ; Hansen et al . ( 2019 ) ; Schaarschmidt et al . ; Tirumala et al . ( 2019 ) , ... ) . We show cost reductions of up to 80 % while being significantly faster . When scaling SEED to many accelerators , it can train on millions of frames per second . Finally , the implementation is open-sourced together with examples of running it at scale on Google Cloud ( see Appendix A.4 for details ) making it easy to reproduce results and try novel ideas . ∗Equal contribution 2 RELATED WORK . For value-based methods , an early attempt for scaling DQN was Nair et al . ( 2015 ) that used asynchronous SGD ( Dean et al. , 2012 ) together with a distributed setup consisting of actors , replay buffers , parameter servers and learners . Since then , it has been shown that asynchronous SGD leads to poor sample complexity while not being significantly faster ( Chen et al. , 2016 ; Espeholt et al. , 2018 ) . Along with advances for Q-learning such as prioritized replay ( Schaul et al. , 2015 ) , dueling networks ( Wang et al. , 2016 ) , and double-Q learning ( van Hasselt , 2010 ; Van Hasselt et al. , 2016 ) the state-of-the-art distributed Q-learning was improved with Ape-X ( Horgan et al. , 2018 ) . Recently , R2D2 ( Kapturowski et al. , 2018 ) achieved impressive results across all the Arcade Learning Environment ( ALE ) ( Bellemare et al. , 2013 ) games by incorporating value-function rescaling ( Pohlen et al. , 2018 ) and LSTMs ( Hochreiter & Schmidhuber , 1997 ) on top of the advancements of Ape-X . There have also been many approaches for scaling policy gradients methods . A3C ( Mnih et al. , 2016 ) introduced asynchronous single-machine training using asynchronous SGD and relied exclusively on CPUs . GPUs were later introduced in GA3C ( Mahmood , 2017 ) with improved speed but poor convergence results due to an inherently on-policy method being used in an off-policy setting . This was corrected by V-trace ( Espeholt et al. , 2018 ) in the IMPALA agent both for singlemachine training and also scaled using a simple actor-learner architecture to more than a thousand machines . PPO ( Schulman et al. , 2017 ) serves a similar purpose to V-trace and was used in OpenAI Rapid ( Petrov et al. , 2018 ) with the actor-learner architecture extended with Redis ( redis.io ) , an in-memory data store , and was scaled to 128,000 CPUs . For inexpensive environments like ALE , a single machine with multiple accelerators can achieve results quickly ( Stooke & Abbeel , 2018 ) . This approach was taken a step further by converting ALE to run on a GPU ( Dalton et al. , 2019 ) . A third class of algorithms is evolutionary algorithms . With simplicity and massive scale , they have achieved impressive results on a number of tasks ( Salimans et al. , 2017 ; Such et al. , 2017 ) . Besides algorithms , there exist a number of useful libraries and frameworks for reinforcement learning . ELF ( Tian et al. , 2017 ) is a framework for efficiently interacting with environments , avoiding Python global-interpreter-lock contention . Dopamine ( Castro et al. , 2018 ) is a flexible research focused RL framework with a strong emphasis on reproducibility . It has state of the art agent implementations such as Rainbow ( Hessel et al. , 2017 ) but is single-threaded . TF-Agents ( Guadarrama et al. , 2018 ) and rlpyt ( Stooke & Abbeel , 2019 ) both have a broader focus with implementations for several classes of algorithms but as of writing , they do not have distributed capability for largescale RL . RLLib ( Liang et al. , 2017 ) provides a number of composable distributed components and a communication abstraction with a number of algorithm implementations such as IMPALA and Ape-X . Concurrent with this work , TorchBeast ( Küttler et al. , 2019 ) was released which is an implementation of single-machine IMPALA with remote environments . SEED is closest related to IMPALA , but has a number of key differences that combine the benefits of single-machine training with a scalable architecture . Inference is moved to the learner but environments run remotely . This is combined with a fast communication layer to mitigate latency issues from the increased number of remote calls . The result is significantly faster training at reduced costs by as much as 80 % for the scenarios we consider . Along with a policy gradients ( V-trace ) implementation we also provide an implementation of state of the art Q-learning ( R2D2 ) . In the work we use TPUs but in principle , any modern accelerator could be used in their place . TPUs are particularly well-suited given they high throughput for machine learning applications and the scalability . Up to 2048 cores are connected with a fast interconnect providing 100+ petaflops of compute . 3 ARCHITECTURE . Before introducing the architecture of SEED , we first analyze the generic actor-learner architecture used by IMPALA , which is also used in various forms in Ape-X , OpenAI Rapid and others . An overview of the architecture is shown in Figure 1a . A large number of actors repeatedly read model parameters from the learner ( or parameter servers ) . Each actor then proceeds the local model to sample actions and generate a full trajectory of observations , actions , policy logits/Q-values . Finally , this trajectory along with recurrent state is transferred to a shared queue or replay buffer . Asynchronously , the learner reads batches of trajectories from the queue/replay buffer and optimizes the model . There are a number of reasons for why this architecture falls short : 1 . Using CPUs for neural network inference : The actor machines are usually CPU-based ( occasionally GPU-based for expensive environments ) . CPUs are known to be computationally inefficient for neural networks ( Raina et al. , 2009 ) . When the computational needs of a model increase , the time spent on inference starts to outweigh the environment step computation . The solution is to increase the number of actors which increases the cost and affects convergence ( Espeholt et al. , 2018 ) . 2 . Inefficient resource utilization : Actors alternate between two tasks : environment steps and inference steps . The compute requirements for the two tasks are often not similar which leads to poor utilization or slow actors . E.g . some environments are inherently single-threading while neural networks are easily parallelizable . 3 . Bandwidth requirements : Model parameters , recurrent state and observations are transferred between actors and learners . Relatively to model parameters , the size of the observation trajectory often only accounts for a few percents.1 Furthermore , memory-based models send large states , increase bandwidth requirements . While single-machine approaches such as GA3C ( Mahmood , 2017 ) and single-machine IMPALA avoid using CPU for inference ( 1 ) and do not have network bandwidth requirements ( 3 ) , they are restricted by resource usage ( 2 ) and the scale required for many types of environments . The architecture used in SEED ( Figure 1b ) solves the problems mentioned above . Inference and trajectory accumulation is moved to the learner which makes it conceptually a single-machine setup with remote environments ( besides handling failures ) . Moving the logic effectively makes the actors a small loop around the environments . For every single environment step , the observations are sent to the learner , which runs the inference and sends actions back to the actors . This introduces a new problem : 4 . Latency . To minimize latency , we created a simple framework that uses gRPC ( grpc.io ) - a high performance RPC library . Specifically , we employ streaming RPCs where the connection from actor to learner is kept open and metadata sent only once . Furthermore , the framework includes a batching module that efficiently batches multiple actor inference calls together . In cases where actors can fit on the same machine as learners , gRPC uses unix domain sockets and thus reduces latency , CPU and syscall overhead . Overall , the end-to-end latency , including network and inference , is faster for a number of the models we consider ( see Appendix A.7 ) . 1With 100,000 observations send per second ( 96 x 72 x 3 bytes each ) , a trajectory length of 20 and a 30MB model , the total bandwidth requirement is 148 GB/s . Transferring observations uses only 2 GB/s . The IMPALA and SEED architectures differ in that for SEED , at any point in time , only one copy of the model exists whereas for distributed IMPALA each actor has its own copy . This changes the way the trajectories are off-policy . In IMPALA ( Figure 2a ) , an actor uses the same policy πθt for an entire trajectory . For SEED ( Figure 2b ) , the policy during an unroll of a trajectory may change multiple times with later steps using more recent policies closer to the one used at optimization time . A detailed view of the learner in the SEED architecture is shown on Figure 3 . Three types of threads are running : 1 . Inference 2 . Data prefetching and 3 . Training . Inference threads receive a batch of observations , rewards and episode termination flags . They load the recurrent states and send the data to the inference TPU core . The sampled actions and new recurrent states are received , and the actions are sent back to the actors while the latest recurrent states are stored . When a trajectory is fully unrolled it is added to a FIFO queue or replay buffer and later sampled by data prefetching threads . Finally , the trajectories are pushed to a device buffer for each of the TPU cores taking part in training . The training thread ( the main Python thread ) takes the prefetched trajectories , computes gradients using the training TPU cores and applies the gradients on the models of all TPU cores ( inference and training ) synchronously . The ratio of inference and training cores can be adjusted for maximum throughput and utilization . The architecture scales to a TPU pod ( 2048 cores ) by roundrobin assigning actors to TPU host machines , and having separate inference threads for each TPU host . When actors wait for a response from the learner , they are idle so in order to fully utilize the machines , we run multiple environments on a single actor . To summarize , we solve the issues listed previously by : 1 . Moving inference to the learner and thus eliminating any neural network related computations from the actors . Increasing the model size in this architecture will not increase the need for more actors ( in fact the opposite is true ) . 2 . Batching inference on the learner and having multiple environments on the actor . This fully utilize both the accelerators on the learner and CPUs on the actors . The number of TPU cores for inference and training is finely tuned to match the inference and training workloads . All factors help reducing the cost of experiments . 3 . Everything involving the model stays on the learner and only observations and actions are sent between the actors and the learner . This reduces bandwidth requirements by as much as 99 % . 4 . Using streaming gRPC that has minimal latency and minimal overhead and integrating batching into the server module . We provide the following two algorithms implemented in the SEED framework : V-trace and Qlearning .

Summary: This paper presents a scalable reinforcement learning training architecture which combines a number of modern engineering advances to address the inefficiencies of prior methods. The proposed architecture shows good performance on a wide variety of benchmarks from ALE to DeepMind Lab and Google Research Football. Important to the community, authors also open source their code and provide an estimate which shows that the proposed framework is cheaper to run on cloud platforms.