SRL: Scaling Distributed Reinforcement Learning to Over Ten Thousand Cores

Summary

The authors introduce SRL, a system for distributed RL training. This system is built upon a new architecture where *workers* host *task handlers*. There are 3 types of task handlers: actor workers, which execute black box environment programs, policy workers, which handle policy inference, and trainer workers, which compute policy updates. A parameter service is used to store model parameters, and sample and inference streams are defined to send relevant data between the task handlers. The authors evaluate their system by training on a variety of environments and demonstrating the performance of their system versus baselines in both the single-machine and distributed setting. They also reproduce the hide & seek experiment from OpenAI, which requires a large amount of experience.

Strengths

- The authors describe their system well and the paper is clearly written - The experiments clearly demonstrate the effectiveness of their system, both in terms of its correctness and scalability - The design they propose is logical and seems more flexible than comparable approaches.

Weaknesses

- It's unfortunate that the authors were not able to compare with baselines other than RLLIB in the distributed setting, although this is completely understandable given the large-scale distributed training frameworks are closed-source.

Questions

- Can the actor workers run GPU-based environments? These have been shown to have very impressive speed-ups over CPU-based environments [1], and although frameworks like JAX allow for parallelisation across GPUs on the same machine, this is not possible across multiple machines. How does the training speed compare with 8 A100 GPUs (i.e. a single machine with GPU based environments)? [1] https://github.com/luchris429/purejaxrl/tree/main [2] https://github.com/RobertTLange/gymnax

Rating

8: accept, good paper

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

In this work the authors implement and evaluate a scalable and efficient framework for highly distributed RL training, known as SRL. Although highly distributed RL training has been studied in previous works they were dominated by industry-led closed-source implementations that yielded very few insights into the correct decompositions and dataflows required to fully utilize the computational hardware on a given system. The authors noted several limitations of existing open-source implementations, stemming from low resource efficiency, inadequate optimizations, or difficult adaptability due to highly coupled algorithm-system interfaces. SRL addresses these issues by presenting a high-level abstraction of the work required to perform RL training into a set of generic workers performing specific tasks connected by different streams of data and coordinated using several services. This abstraction is mapped onto a specific instance of system resources in a flexible manner that facilitates both a large degree of flexibility and resource-aware optimizations to fully exploit the available hardware. The RL-specific components are presented as 3 types of workers to interact with the environment and provide policy updates, 2 streams, and a set of services to perform parameter updating and synchronization. Some optimizations that are specific to RL include the usage of multiple environments to ensure CPU instances can switch between environments without waiting for the next actions to be performed by the policy workers. A wide array of experiments using multiple environments with a mix of CPU and GPU resources demonstrates the utility of SRL to both generate a large number of training frames per second and execute a large number of policy updates per second. Notably, the authors achieve performance that cannot be replicated using any existing open-source RL training frameworks and apply this method to extremely challenging environments, such as hide-and-seek, that require a great deal of data to achieve reasonable performance.

Strengths

- The most noteworthy and interesting aspect of SRL is the demonstrated scalability to support training on an extremely large number of distributed computational resources. Although this level of performance has been noted in other works it was always achieved using RL systems that were not fully released or open-sourced for usage by the general public. - The flexibility of SRL is evident from the evaluations using a variety of RL algorithms, such as PPO, DQN, and VDN on a number of different system architectures and achieving reasonable performance in all cases. The difference in the necessary computational primitives required by on-vs-off policy methods and single-vs-multi agent environments makes this level of generality challenging and/or overly complicated using many existing RL libraries. - Well-defined user-friendly extensibility interface appears to be quite attractive for future extensions to scenarios that may not fit neatly into the existing designs typically used common, or the most popular, algorithms used at any particular time. - The number of evaluations performed on multiple system configurations using a variety of RL environments and incorporating a number of different algorithms is impressive. Figures 3, 4, and 7 do a particularly good job illustrating the level of performance achieved by SRL, comparing that with existing work, and conveying the impact this new level of performance could have on the pace of RL research that may be conducted using open-source software moving forward. - Many interesting details regarding the implementation specifics weren't clear to me until I read through the appendices. Unfortunately, with a system this wide in scope it is hard to fully understand some aspects without reading through all the additional notes but I was glad the authors provided so many details.

Weaknesses

Major: - The lack of fair comparison with other large-scale RL training implementations is unfortunate. Though this is by no means the fault of the authors it leaves a level of unknown with regard to a proper comparison of the proposed system against other pre-existing closed-source systems. This is unavoidable and it was good to see the authors attempt to reimplement previous work using their current system to provide some semblance of fairness for comparison. Minor: - The text could use a few more passes to fix minor typos in several sections. Nothing major just more editing required. - Formatting obviously needs to be updated for the code presented in Code 1 and 2 in section 3.4.

Questions

I incorporated all my comments into the strengths and weaknesses sections I have no questions at this time.

Rating

8: accept, good paper

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

This paper proposes a scalable, efficient, and extensible distributed RL framework which can easily scale to ten thousand cores.

Strengths

This paper presents a dataflow abstraction for RL training which allows framework to allocate suitable computing resources in a cluster with heterogeneous hardware. It scales well to ten thousand CPU cores.

Weaknesses

1) MSRL also introduced a data flow abstraction and supported IMPALA and SeedRL distributed architectures. The worker abstraction in SRL is similar to the "fragment" in MSRL. Environment ring is also introduced in EnvPool. Please provide further information about the novelty of SRL. 2) In order to provide a comprehensive understanding of where any gains come from, it would be valuable to include a breakdown ablation study in the evaluation for each optimization and design. 3) It would be beneficial to include MSRL in the experiments to evaluate the gain compared to the state-of-the-art solutions. This could provide deeper insights into where SRL stands in relation to other established techniques.

Questions

1) Please address the above weaknesses.

Rating

6: marginally above the acceptance threshold

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Summary

The paper introduces SRL, a (really) scalable reinforcement learning framework, designed to parallelize DRL training to large clusters. SRL decomposes DRL training into heterogeneous workers of different types, interconnects them, and schedules them appropriately in order to optimize training throughput. This abstraction is general and enables one to easily implement many different DRL training algorithms. The evaluation demonstrates impressive scalability and training throughput without compromising learning quality.

Strengths

1. SRL is a practical system that seems like it would be of significant benefit to the community. Source code is already (anonymously) available. While the largest scales are likely of interest only to relatively few groups, this capability is nevertheless important, and SRL demonstrates improved performance even at smaller scales (e.g., 32-64 cores). 2. The paper is clear and well-written. 3. The benchmarks cover a number of baselines and consistently show SRL matching or outperforming other systems for DRL training. Results also demonstrate that the system does not compromise learning quality.

Weaknesses

1. While practical, this is primarily "engineering" work and there is not much novelty. (I nevertheless think the practical benefits outweigh this.) In particular, the overall design of the SRL seems very reminiscent of classic task-based programming models from parallel and scientific computing (e.g., Charm++). The paper would benefit from some discussion of this, and it may be a source of inspiration for additional optimizations. 2. The performance optimizations discussed in the main paper, "environment ring" and "trainer pre-fetching", are standard, widely-used optimizations in deep learning frameworks. The environment ring seems to be a case of double-buffering (or pipelining); and many frameworks support prefetching and staging data in advance to the GPU (e.g., while it is not specifically for DRL, the DALI library does this). 3. Performance results lack error bars.

Questions

1. Please clarify or contextualize the novelty of the work (see above). 2. Please add error bars to the performance results.

Rating

8: accept, good paper

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good