RLlib Flow: Distributed Reinforcement Learning is a Dataflow Problem

1 Introduction . The past few years have seen the rise of deep reinforcement learning ( RL ) as a new , powerful optimization method for solving sequential decision making problems . As with deep supervised learning , researchers and practitioners frequently leverage parallel computation , which has led to the development of numerous distributed RL algorithms and systems as the field rapidly evolves . However , despite the high-level of abstraction that RL algorithms are defined in ( i.e. , as a couple dozen lines of update equations ) , their implementations have remained quite low level ( i.e. , at the level of message passing ) . This is particularly true for distributed RL algorithms , which are typically implemented directly on low-level message passing systems or actor frameworks [ 14 ] . Libraries such as Acme [ 15 ] , RLgraph [ 26 ] , RLlib [ 19 ] , and Coach [ 2 ] provide unified abstractions for defining single-agent RL algorithms , but their user-facing APIs only allow algorithms to execute within the bounds to predefined distributed execution patterns or “ templates ” . While the aforementioned libraries have been highly successful at replicating a large number of novel RL algorithms introduced over the years , showing the generality of their underlying actor or graph-based computation models , the needs of many researchers and practitioners are often not met by their abstractions . We have observed this firsthand from users of open source RL libraries : First , RL practitioners are typically not systems engineers . They are not well versed with code that mixes together the logical dataflow of the program and system concerns such as performance and bounding memory usage . This leads to a high barrier of entry for most RL users to experimenting with debugging existing distributed RL algorithms or authoring new distributed RL approaches . ∗ indicates equal contributions . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Second , even when an RL practitioner is happy with a particular algorithm , they may wish to customize it in various ways . This is especially important given the diversity of RL tasks ( e.g. , singleagent , multi-agent , meta-learning ) . While many customizations within common RL environments can be anticipated and made available as configuration options ( e.g. , degree of parallelism , batch size ) , it is difficult for a library author to provide enough options to cover less common tasks that necessarily alter the distributed pattern of the algorithm ( e.g. , interleaved training of different distributed algorithms , different replay strategies ) . Our experience is that when considering the needs of users for novel RL applications and approaches , RL development requires a significant degree of programming flexibility . Advanced users want to tweak or add various distributed components ( i.e. , they need to write programs ) . In contrast to supervised learning , it is more difficult to provide a fixed set of abstractions for scaling RL training . As a result , it is very common for RL researchers or practitioners to eschew existing infrastructure , either sticking to non-parallel approaches , which are inherently easier to understand and customize [ 3 , 6 ] , or writing their own distributed framework that fits their needs . The large number of RL frameworks in existence today is evidence of this , especially considering the number of these frameworks aiming to be “ simpler ” versions of other frameworks . In this paper , we re-examine the challenges posed by distributed RL in the light of these user requirements , drawing inspiration from prior work in the field of data processing and distributed dataflow . To meet these challenges , we propose RLlib Flow , a hybrid actor-dataflow model for distributed RL . Like streaming data systems , RLlib Flow provides a small set of operator-like primitives that can be composed to express distributed RL algorithms . Unlike data processing systems , RLlib Flow explicitly exposes references to actor processes participating in the dataflow , permitting limited message passing between them in order to more simply meet the requirements of RL algorithms . The interaction of dataflow and actor messages is managed via special sequencing and concurrency operators . The contributions of our paper are as follows : 1 . We examine the needs of distributed RL algorithms and RL practitioners from a dataflow perspective , identifying key challenges ( Section 2 and 3 ) . 2 . We propose RLlib Flow , a hybrid actor-dataflow programming model that can simply and efficiently express distributed RL algorithms , and enables composition of multi-agent algorithms not possible by end users before without writing low-level systems code . ( Section 4 and 5 ) . 3 . We port all the algorithms of a production RL library ( RLlib ) to RLlib Flow , providing 2-9× savings in distributed execution code , compare its performance with the original implementation , and show performance benefits over systems such as Spark Streaming ( Section 6 ) . 2 Distributed Reinforcement Learning . We first discuss the relevant computational characteristics of distributed RL algorithms , starting with the common single-agent training scenario , where the goal is to optimize a single agent ’ s performance in an environment , and then discuss the computational needs of emerging multi-agent , model-based , and meta-learning training patterns . 2.1 RL Algorithm Basics . The goal of an RL algorithm is typically to improve the performance of a policy with respect to an objective defined through an environment ( e.g. , simulator ) . The policy is usually defined as a deep neural network , which can range from several KB to several hundred MB in size . RL algorithms can be generally broken down into the following basic steps : Rollout : To generate experiences , the policy , which outputs actions to take given environment observations , is run against the environment to collect batches of data . The batch consists of observations , actions , rewards , and episode terminals and can vary in size ( 10s to 10000s of steps ) . Replay : On-policy algorithms ( e.g. , PPO [ 27 ] , A3C [ 22 ] ) collect new experiences from the current policy to learn . On the other hand , off-policy algorithms ( e.g. , DQN [ 23 ] , SAC [ 12 ] ) can leverage experiences from past versions of the policy as well . For these algorithms , a replay buffer of past experiences can be used . The size of these buffers ranges from a few hundred to millions of steps . Optimization : Experiences , either freshly collected or replayed , can be used to improve the policy . Typically this is done by computing and applying a gradient update to the policy and value neural networks . While in many applications a single GPU suffices to compute gradient updates , it is sometimes desirable to leverage multiple GPUs within a single node , asynchronous computation of gradients on multiple CPUs [ 22 ] , or many GPUs spread across a cluster [ 29 ] . 2.2 RL Algorithm Variants . Single-Agent Training . Training a single RL agent—the most basic and common scenario—consists of applying the steps of rollout , replay , and optimization repeatedly until the policy reaches the desired performance . Synchronous algorithms such as A2C [ 22 ] and PPO apply the steps strictly sequentially . Parallelism may be leveraged internally within each step . Asynchronous algorithm variations such as A3C [ 22 ] , Ape-X [ 16 ] , APPO [ 21 ] , and IMPALA [ 8 ] , pipeline and overlap the rollout and optimization steps asynchronously to hit higher data throughputs . Rate limiting [ 15 ] can be applied to control learning dynamics in the asynchronous setting . Multi-Agent Training . In multi-agent training , there are multiple acting entities in the environment ( e.g. , cooperating or competing agents ) . While there is a rich literature on multi-agent algorithms , we note that the dataflow structure of multi-agent training is similar to that of single-agent—as long as all entities are being trained with the same algorithm and compatible hyperparameters.However , problems arise should it be required to customize the training of any of the agents in the environment . For example , in a two-agent environment , one agent may desire to be optimized at a higher frequency ( i.e. , smaller batch size ) . This fundamentally alters the training dataflow—there are now two iterative loops executing at different frequencies . Furthermore , if these agents are trained with entirely different algorithms , there is a need to compose two different distributed dataflows . Model-Based and Meta-Learning Algorithms . Model-based algorithms seek to learn transition dynamics of the environment to improve the sample efficiency of training . This can be thought of as adding a supervised training step on top of standard distributed RL , where an ensemble of one or more dynamics models are trained from environment-generated data . Handling the data routing , replay , optimization , and stats collection for these models naturally adds complexity to the distributed dataflow graph , “ breaking the mold ” of standard model-free RL algorithms and hard to be implemented in low-level systems . Using RLlib Flow , we have implemented two state-of-the-art model-based algorithms : MB-MPO [ 5 ] and Dreamer [ 13 ] . 2.3 A Case for a Higher Level Programming Model . Given that existing distributed RL algorithms are already implementable using low level actor and RPC primitives , it is worth questioning the value of defining a higher level computation model . Our experience is that RL is more like data analytics than supervised learning . Advanced users want to tweak or add various distributed components ( i.e. , they need to program ) , and there is no way to have a “ one size fits all ” ( i.e. , Estimator interface from supervised learning ) . We believe that , beyond the ability to more concisely and cleanly capture single-agent RL algorithms , the computational needs of more advanced RL training patterns motivate higher level programming models like RLlib Flow . 3 Reinforcement Learning vs Data Streaming . The key observation behind RLlib Flow is that the dataflow graph of RL algorithms are quite similar to those of data streaming applications . Indeed , RL algorithms can be captured in general purpose dataflow programming models . However , due to several characteristics , they are not a perfect fit , even for dataflow programming models that support iterative computation . In this section we examine the dataflow of the A3C algorithm ( Figure 4 ) to compare and contrast RL with streaming dataflow . A3C starts with ( 1 ) parallel rollouts across many experiences . Policy gradients are computed in parallel based on rollouts in step ( 2 ) . In step ( 3 ) , the gradients are asynchronously gathered and applied on a central model , which is then used to update rollout worker weights . Importantly , each box or operator in this dataflow may be stateful ( e.g. , ParallelRollouts holds environment state as well as the current policy snapshot ) . Distributed Workflow API Similar to data processing topologies , A3C is applying a transformation to a data stream ( of rollouts ) in parallel ( to compute gradients ) . This is denoted by the black arrow between ( 1 ) and ( 2 ) . There is also a non-parallel transformation to produce metrics from the computation , denoted by the black arrow between ( 3 ) and ( 4 ) However , zooming out to look at the entire dataflow graph , a few differences emerge : Asynchronous Dependencies : RL algorithms often leverage asynchronous computation to reduce update latencies and eliminate stragglers [ 22 ] . In RLlib Flow , we represent these with a pink arrow between a parallel and sequential iterator . This means items will be fetched into the sequential iterator as soon as they are available , instead of in a deterministic ordering . The level of asynchrony can be configured to increase pipeline parallelism . Message Passing : RL algorithms , like all iterative algorithms , need to update upstream operator state during execution ( e.g. , update policy weights ) . Unlike iterative algorithms , these updates may be fine-grained and asynchronous ( i.e. , update the parameters of a particular worker ) , as well as coarse-grained ( i.e. , update all workers at once after a global barrier ) . RLlib Flow allows method calls ( messages ) to be sent to any actor in the dataflow . Ordering of messages in RLlib Flow with respect to dataflow steps is guaranteed if synchronous data dependencies ( black arrows ) fully connect the sender to the receiver , providing barrier semantics . Consistency and Durability : Unlike data streaming , which has strict requirements such as exactlyonce processing of data [ 30 ] , RL has less strict consistency and durability requirements . This is since on a fault , the entire computation can be restarted from the last checkpoint with minimal loss of work . Message or data loss can generally be tolerated without adverse affect on training . Individual operators can be restarted on failure , discarding any temporary state . This motivates a programming model that minimizes overhead ( e.g. , avoids state serialization and logging cost ) .

Summary: This manuscript proposes AnonFlow, a set of dataflow abstractions that allow for easy specification and implementation of distributed RL algorithms. AnonFlow consists of several composable dataflow operators/iterators implemented on top of RLlib. The manuscript discusses the implementation of several popular RL algorithms and quantifies the effectiveness of AnonFlow in terms of LoC and overall system performance.