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: The paper presents SEED RL, which is a scalable reinforcement learning agent. The approach restructure the interface / division of functionality between the actors (environments) and the learner as compared to the distributed approach in IMPALA (a state-of-the-art distributed RL framework). Most importantly, the model is only in the learner in SEED while it is distributed in IMPALA.