Source-linked AI summary
RLlib: Abstractions for Distributed Reinforcement Learning
Eric Liang, Richard Liaw, Philipp Moritz, Robert Nishihara, Roy Fox, Ken Goldberg, Joseph E. Gonzalez, Michael I. Jordan, Ion Stoica
TL;DR
RL systems must support highly irregular, nested computation, but existing abstractions provide limited composability. The paper proposes hierarchical control with encapsulated parallelism and implements it in RLlib using scalable task-based primitives. RLlib supports broad algorithm coverage with high performance, scalability, and code reuse, while the authors note that composability must be achieved together with scalability.
Problem
RL workloads combine irregular computation, nested parallelism, and varied resource requirements, while existing systems provide comparatively limited abstractions for composing such components.
Method
The paper uses hierarchical, logically centralized control to encapsulate component parallelism and resource requirements within distributed task-based computations, implemented in RLlib.
Results
RLlib provides scalable abstractions and reference implementations across a broad range of RL workloads, including near-linear Ape-X scaling to 160k environment frames per second with 256 workers.
Takeaways & Limitations
Composable hierarchical abstractions can support scalable RL algorithm implementations while enabling substantial reuse of existing components.
Takeaways & Limitations
The paper emphasizes that composability without scalability is insufficient, so both objectives must be achieved simultaneously.
Abstract
from arXiv · showhide
Reinforcement learning (RL) algorithms involve the deep nesting of highly irregular computation patterns, each of which typically exhibits opportunities for distributed computation. We argue for distributing RL components in a composable way by adapting algorithms for top-down hierarchical control, thereby encapsulating parallelism and resource requirements within short-running compute tasks. We demonstrate the benefits of this principle through RLlib: a library that provides scalable software primitives for RL. These primitives enable a broad range of algorithms to be implemented with high performance, scalability, and substantial code reuse. RLlib is available at https://rllib.io/.
1. Introduction
RL workloads combine irregular, nested computation and diverse resource demands, while existing distributed frameworks offer limited composability. The paper proposes hierarchical, logically centralized control and presents RLlib as a scalable library built on these principles.
- 1. Introduction: Existing RL libraries often rely on long-running distributed replicas, making component composition difficult and requiring communication logic to be reimplemented.This limits scalable reuse across algorithms despite progress in RL libraries.
- 1.1. Irregularity of RL training workloads: RL training spans irregular task durations, communication patterns, resource requirements, and nested computations.These demands range from millisecond updates to larger rollout batches and from synchronous to asynchronous execution.
- 1.2. Logically centralized control for distributed RL: Hierarchical control extends centralized program control so worker processes can delegate simulations, gradient computation, and other work to sub-workers.This model supports nested computations while keeping control logic in a single driver structure.
- 1.2. Logically centralized control for distributed RL: Encapsulating component parallelism enables code reuse, separates resource requirements, and simplifies implementation of distributed RL algorithms.The model also supports high performance because bulk data transfer can occur outside the driver.
- 1. Introduction: The paper proposes a general composable hierarchical programming model and RLlib abstractions for a broad range of RL algorithms.The authors report state-of-the-art performance across varied RL workloads and emphasize rapid development through scalable abstractions.
2. Hierarchical Parallel Task Model
The paper builds hierarchical RL systems on flexible task-based programming models, allowing distributed components to remain reusable while being invoked as nested remote tasks. RLlib uses Ray to provide this delegation, scheduling, and communication infrastructure.
- 2. Hierarchical Parallel Task Model: Task-based programming models schedule fine-grained asynchronous subroutines on workers and pass results between processes.This supports composing distributed RL components without rewriting their internal implementations.
- 2.2. Ray implementation of hierarchical control: RLlib uses Ray, whose scheduler supports nested computation without a central task scheduling bottleneck.Ray distributes Python tasks across large clusters and fits the proposed hierarchical control model.
- 2. Hierarchical Parallel Task Model: RL libraries built with hierarchical and logically centralized control can compose distributed hyperparameter search with distributed training components.Short-running tasks allow control decisions between tasks without inserting coordination points into long-running programs.
- 2.2. Ray implementation of hierarchical control: Ray actors can launch additional actors and schedule tasks during method calls, implementing hierarchical delegation.Ray also provides aggregate and broadcast primitives plus zero-copy sharing through a shared-memory object store.
3. Abstractions for Reinforcement Learning
RLlib separates algorithm-specific policy definitions from distributed execution, using policy graphs, evaluators, and interchangeable optimizers to compose scalable RL algorithms. Its abstractions support parallel evaluation, replay, gradient optimization, multiple execution strategies, and implementations spanning derivative-free, hybrid, and high-throughput methods.
- Core abstractions: RLlib requires algorithms to declare a policy, experience postprocessor, and loss, which can be specified in TensorFlow or PyTorch.Policy evaluators and optimizers then provide distributed policy evaluation and training strategies.
- Policy graphs: A policy graph maps observations and optional recurrent states to actions and next recurrent states, while exposing user-defined outputs for training.These outputs can include quantities such as value predictions or TD error.
- Experience processing: Trajectory postprocessors transform experience batches for uses including advantage estimation and goal relabeling, with access to other agents’ batches in multi-agent environments.The postprocessor operates on trajectory tuples containing observations, actions, rewards, and subsequent observations.
- Policy evaluation: Policy evaluators wrap policy graphs and environments, can run as Ray remote actors across a cluster, and collect experience batches in parallel through sample.remote().A minimal TensorFlow policy-gradient graph can therefore be replicated across multiple evaluators for parallel sampling.
- Policy optimization: Policy optimizers operate on a local policy graph and remote evaluator replicas, incorporating RL sampling into optimization and supporting interchangeable execution strategies.The abstraction separates policy and loss definitions from choices such as synchronous versus asynchronous execution, allreduce versus parameter servers, and CPUs versus GPUs.
- Generality and performance: RLlib expresses diverse algorithms through reusable abstractions, including Ape-X, Evolution Strategies, and PPO-ES, while retaining scalable implementations.Ape-X reaches nearly 160k environment frames per second with 256 workers; PPO-ES required approximately 50 lines and improved convergence over PPO on Walker2d-v1.
4. Framework Performance
RLlib uses Ray primitives to support nested, resource-aware parallelism for reinforcement-learning workloads. These primitives address state sharing, data movement, task scheduling, fault tolerance, straggler mitigation, and compression.
- Ray actors let tasks share mutable state, supporting simulators and neural-network weights that must be updated across tasks.
- The shared-memory object store passes large rollout and weight objects without a central bottleneck or same-machine copies.
- Remote calls take approximately 200µs on the same machine and around 1ms across saturated nodes.
- RLlib supports nested parallel calls through lightweight remote tasks, while resource-aware scheduling preserves component performance.
- RLlib adds fault tolerance, straggler mitigation, vectorized evaluation, and LZ4 compression for scalable execution.LZ4 reduces image-observation network traffic and memory usage by more than an order of magnitude at approximately 1 GB/s/core.
5. Evaluation
RLlib’s evaluation shows scalable sampling and competitive performance across ES, PPO, and A3C. Its policy optimizers and evaluator architecture support high throughput across CPU, GPU, and multi-GPU settings.
- Pendulum-CPU exceeds 1.5 million actions/s, while Pong-GPU approaches 200k actions/s using four intermediate aggregation actors.
- RLlib’s ES reaches reward 6000 on Humanoid-v1 in a median 3.7 minutes using 8192 cores, over twice as fast as the best published result.
- A3C solves PongDeterministic-v4 in 12 minutes with an asynchronous optimizer and 9 minutes with a sharded parameter-server optimizer, matching a tuned baseline.
- A specialized multi-GPU PPO optimizer outperforms distributed allreduce when data fits entirely into GPU memory.
- Policy evaluation scales nearly linearly from 1 to 128 cores, reaching approximately 200k Pong actions/s and 1.5M Pendulum actions/s.
- RLlib implementations of ES and PPO outperform highly optimized reference implementations on Humanoid-v1.
6. Related work
RLlib differs from prior reinforcement-learning libraries by emphasizing hierarchical composition rather than long-running program replicas. This supports broader composition of distributed components, while building on abstractions from deep learning systems.
- Many RL libraries coordinate distributed computation through long-running program replicas, which do not generalize well to complex architectures.
- RLlib instead uses hierarchical control with short-running tasks so components can control their own distributed execution.
- ONNX, NNVM, and Gluon explore composition between deep-learning model specifications and hardware, while deep-learning libraries provide gradient-optimization components used in RL.
7. Conclusion
RLlib is an open-source reinforcement-learning library combining fine-grained nested parallelism with reference algorithms and abstractions for composing new algorithms.
- RLlib offers reference algorithms and scalable abstractions for composing new reinforcement-learning algorithms across broad workloads.