Paper: Cleanba: A Reproducible and Efficient Distributed Reinforcement Learning Platform (ICLR-2024)

Reviewer comment (Experiments):
The performance of the new approach is tested only in Atari, which is not a real-world problem where reproducibility is a key challenge. The paper can greatly benefit from a real-world application.

Author response:
> Q: The paper can greatly benefit from a real-world application. Atari is a great way to benchmark but does not present all the variability such as in robotics applications for example to truly highlight issues related to the reproducibility.

The reproducibility issues related to working with physical robots seem to be in a completely different realm of reproducibility issues dealt with in this paper. Our work focuses on reproducibility issues related to scaling caused by distributed data structures and implementation.

Reviewer comment (Writing):
The language is a bit off in several sections, specifically in the first paragraph, which makes it difficult to read at times. There are several grammatical errors as well.

Author response:
> Q: The language is a bit off in several sections, specifically in the first paragraph, which makes it difficult to read at times. There are several grammatical errors as well.

We thank this constructive feedback and will get the paper proofread by a native speaker should the paper get accepted.

Reviewer comment (Reproducibility):
It is important to consider a real-world example when reproducibility is one of the main challenges. Atari is a great way to benchmark but does not present all the variability such as in robotics applications, which could truly highlight issues related to reproducibility.

Author response:
> Q: The paper can greatly benefit from a real-world application. Atari is a great way to benchmark but does not present all the variability such as in robotics applications for example to truly highlight issues related to the reproducibility.

The reproducibility issues related to working with physical robots seem to be in a completely different realm of reproducibility issues dealt with in this paper. Our work focuses on reproducibility issues related to scaling caused by distributed data structures and implementation.

Reviewer comment (Experiments):
In Section 5.1, the authors state that Cleanba’s IMPALA is 6.8x faster than monobeast’s IMPALA, mostly because Cleanba actors run on GPUs whereas monobeast’s actors run on CPUs. The performance metrics should be based on the same compute hardware — CPU only or GPU only.

Author response:
> Q: In Section 5.1, the authors state that Cleanba’s IMPALA is 6.8x faster than monobeast’s IMPALA, mostly because Cleanba actors run on GPUs, whereas monobeast’s actors run on CPUs. I think the performance metrics should be based on the same compute hardware CPU only or GPU only.

This is a great observation! GPU only is difficult because both monobeast’s IMPALA and the original IMPALA does not support GPU actors by default. Moolib however is a more fair comparison though, since it runs the actor on GPUs.

Reviewer comment (Writing):
The abstract’s first sentences feel a bit choppy.

Author response:
> Q: The abstract’s first sentences feel a bit choppy

We changed it to "Distributed Deep Reinforcement Learning (DRL) aims to train autonomous agents in less wall-clock time by leveraging more computational resources". This puts the parallel construction closer so less difficult to read.

Reviewer comment (Experiments):
Fig 2 could motivate the problem better, given that in the current figure both versions end up at very similar scores. If there was a different environment just highlighting that this 1 second lag has a meaningful detriment to final performance that could be really compelling.

Author response:
> Fig 2 could motivate the problem better, given that in the current figure both versions end up at very similar scores. If there was a different environment just highlighting that this 1 second lag has a meaningful detriment to the final performance could be really compelling.

Thank you for this suggestion. Our cluster is under heavy usage at the moment but happy to explore this later (unfortunately after the discussion period). We may just run these experiments on Breakout since we expect to see a much larger difference.

Reviewer comment (Presentation):
In figure 3, the optimality gap x tickers are overlapping.

Author response:
> In figure 3, the optimality gap x tickers are overlapping

Thank you. This is fixed.

Reviewer comment (Experiments):
Moolib is a good point of comparison, but I see that this repo is now archived. Are there other DDRL libraries worth comparing to that are more actively maintained?

Author response:
> Moolib is a good point of comparison, but I see that this repo is now archived. Are there other DDRL libraries worth comparing to that are more actively maintained?

There are some DDRL libraries that we thought about comparing but chose not to do it for the following reasons. 

1. Ray’s RLlib: we are not aware of its distributed IMPALA reproducing the level of performance in Atari. For example, the following table shows RLlib’s official benchmark at https://github.com/ray-project/rl-experiments#impala-and-a2c ran for 4 (out of 57) games and significantly underperformed IMPALA Shallow under an hour of training

| | RLlib IMPALA 32-workers | IMPALA Shallow (Table C.1, Espeholt et al 2018) |
| --- | --- | --- |
| BeamRider | 3,181 | 32,463.47 |
| Breakout | 538 | 640.43 |
| QBert | 10,850 | 351,200.12 |
| Spacelnvaders | 843 | 43,595.78 |

2. DeepMind/acme: we noted that its technical report [V2](https://arxiv.org/pdf/2006.00979.pdf) contained only Atari evaluations on 5 (out of 57) games and revealed no runtime information. Its [v1](https://arxiv.org/pdf/2006.00979v1.pdf) revealed runtime information IMPALA but it seems quite different from the original IMPALA results
  * For example, ACME’s IMPALA takes 12 hours to get 100k in Asteroids, 24 hours to get 750 in Breakout, but original IMPALA (deep) reported 100k in Asteroids and 783 in Breakout in ~2 hours (the runtime is interpolated based on Figure 2 in https://openreview.net/pdf?id=r1lyTjAqYX)
  * We also tried running ACME’s IMPALA and got very low FPS (like 1200 FPS) both in single GPU mode and distributed node for some reason… We suspect a higher performance of ACME might be coupled with DeepMind’s internal infrastructure. 

We chose to evaluate against moolib and tourchbeast because they, to our knowledge, are the **only** open-source DDRL infrastructures that have evaluated their performance on 50+ Atari games using the 200M frames. 200M frames also make the computational resources reasonable for us (about 2 hours per game with a single GPU or 20-30 minutes per game with 8 GPUs), as opposed to Seed RL which is evaluated on 40B frames (about 40 hours using their settings — 8xTPUv3, 213 CPUs)

The lack of highly reproducible and easy-to-understand DDRL infrastructure is what partly motivated this work — we aim to make DDRL more accessible and transparent for the research community.

Reviewer comment (Reproducibility):
One concern I have is that this paper tries too hard to be an algorithm paper. It presents cleanba, which is a software package that the authors want people to use, but not a lot of time is dedicated to that. More time is dedicated to the analysis of results from cleanba. While these results are somewhat interesting, I think they could be consolidated more (e.g. move the individual atari games to the appendix, and just present HNS, just an idea). Then use that space to actually describe the package of cleanba more (e.g. how does a user interact with it? Extend it? UML diagram perhaps? Etc.). The paper is called “cleanba” for the package, but I can’t say I really know programmatically much about cleanba after reading the paper.

Author response:
> Q: I can’t evaluate the code, since there is no anonymized access to it, [...] The paper is called “cleanba” for the package, but I can’t say I really know programmatically much about cleanba after reading the paper.

We made the anonymized version of Cleanba available at https://anonymous.4open.science/r/cleanba-test-9868/README.md. It’s basically two files: cleanba_ppo.py and cleanba_impala.py, both of which are self-contained and within 800 lines of code. 

We just want to provide simple-to-understand DDRL infrastructure which is battle-tested and scalable in various scenarios, so we ran extensive evaluation on 57 Atari games and reproducibility under different hardware scenarios.

On the other end of the spectrum, we have libraries like RLlib / ACME, which are monolithic. Have tons of abstraction and modules (ACME has 500 python files, and Ray (including) has 2.8k Python files). A motivating (dummy) question for us is “Do we really need to read hundreds of interconnected files to understand how DDRL works from end to end?”

The main target audience is students and researchers, so they can learn DDRL with a principled approach and hack it for their respective research, similar to CleanRL’s approach to DRL infrastructure (https://www.jmlr.org/papers/v23/21-1342.html).

Reviewer comment (Reproducibility):
I can’t evaluate the code, since there is no anonymized access to it, but I feel given that this introduces a package, that code quality is relevant. Is the code documented well and complete with type hints and docstrings and the like?

Author response:
> Is the code documented well and complete with type hints and docstrings and the like?

Cleanba’s implementations are extremely lightweight: they have minimal lines of code and some necessary comments but come with extensive documentation, evaluation, and documentation in the repo/paper. We note that since each Cleanba variant is only 800 lines of code end-to-end without extra modules, adding extensive docstring may not be necessary. 

The users can easily trace the flow of the program all happening in a single file (e.g., cleanba_ppo.py), as opposed to DRL libraries like SB3, where it flows through dozens of files like below. In SB3, since there are hundreds of classes and methods involved, docstrings are necessary. In cleanba_ppo.py, there are only 7 classes and 16 functions, all of which are self-explanatory.

The list of involved files in SB3 (taken from Appendix B in https://www.jmlr.org/papers/volume23/21-1342/21-1342.pdf)
1. stable baselines3/ppo/ppo.py — 315 lines of code (LOC), 51 lines of docstring (LOD)
2. stable baselines3/common/on policy algorithm.py — 280 LOC, 49 LOD
3. stable baselines3/common/base class.py — 819 LOC, 231 LOD
4. stable baselines3/common/utils.py — 506 LOC, 195 LOD
5. stable baselines3/common/env util.py — 157 LOC, 43 LOD
6. stable baselines3/common/atari wrappers.py — 249 LOC, 84 LOD
7. stable baselines3/common/vec env/ init .py — 73 LOC, 24 LOD
8. stable baselines3/common/vec env/dummy vec env.py — 126 LOC, 25 LOD
9. stable baselines3/common/vec env/base vec env.py — 375 LOC, 112 LOD
10. stable baselines3/common/vec env/util.py — 77 LOC, 31 LOD
11. stable baselines3/common/vec env/vec frame stack.py — 65 LOC, 14 LOD
12. stable baselines3/common/vec env/stacked observations.py — 267 LOC, 74 LOD
13. stable baselines3/common/preprocessing.py — 217 LOC, 68 LOD
14. stable baselines3/common/buffers.py — 770 LOC, 183 LOD
15. stable baselines3/common/policies.py — 962 LOC, 336 LOD
16. stable baselines3/common/torch layers.py — 318 LOC, 97 LOD
17. stable baselines3/common/distributions.py — 700 LOC, 228 LOD
18. stable baselines3/common/monitor.py — 240 LOC, 76 LOD
19. stable baselines3/common/logger.py — 640 LOC, 201 LOD
20. stable baselines3/common/callbacks.py — 603 LOC, 150 LOD

Reviewer comment (Presentation):
The Cleanba shown in Figure 1 does not match the “learner always learns from the rollout data obtained from the second latest policy” (line 7).

Author response:
> Q: The Cleanba shown in Figure 1 does not match the “learner always learns from the rollout data obtained from the second latest policy” (line 7).

There is a pre-condition “**starting from iteration 3** the learner always learns from the rollout data obtained from the second latest policy”; before iteration 3, the learner does not necessarily learn from the second latest policy.

Reviewer comment (Theory):
“ensuring the learner performs gradient updates with rollout data of second latest policy” assumption is too strong. It may slow down the IMPALA training.

Author response:
> Q: “ensuring the learner performs gradient updates with rollout data of second latest policy” assumption is too strong. It may slow down the IMPALA training.

Though this may be a possibility, our experiments show otherwise: Cleanba IMPALA has obtained equivalent scores and faster training time than monobeast and moolib’s IMPALA. Cleanba’s IMPALA is faster than monobeast’s IMPALA because Cleanba IMPALA run actors on GPU while monobeast’s IMPALA run actors on CPU. Cleanba’s IMPALA is faster than moolib’s IMPALA likely because Cleanba’s JAX + EnvPool is faster than Moolib’s PyTorch + various C++ async operations (e.g., async accumulation of gradients, async environment stepping, async network evaluation).

Reviewer comment (Reproducibility):
It would be beneficial to provide more reproducibility results in relation to different RL algorithms such as Ape-X, SeedRL, etc.

Author response:
> Q: It would be beneficial to provide more reproducibility results in relation to different RL algorithms such as Ape-X, SeedRL, etc.

To the best of our knowledge, there is no open-source Ape-X replication that has public results showing that they could reproduce 350% median HNS in 20 hours. Without reliable Ape-X infrastructure, it’s not viable to provide reproducibility results. 

Reproducibility results with Seed RL is challenging for several reasons:
* Seed RL is already unmaintained and running its code can be a challenge. 
  * We tried to follow its setup instructions “./run_local.sh atari r2d2 4 4" and it immediately errored out with some GPG errors `W: GPG error: https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64  InRelease: The following  signatures couldn't be verified because the public key is not available: NO_PUBKEY A4B469963BF863CC`
  * We then tried fixing it with https://github.com/NVIDIA/nvidia-docker/issues/1632#issuecomment-1112667716 and got 
      ```ERROR: Could not find a version that satisfies the requirement ale-py~=0.8.0; extra == "atari" (from gym[atari]) (from versions: 0.6.0.dev20200207, 0.7rc0, 0.7rc1, 0.7rc2, 0.7rc3, 0.7rc4, 0.7, 0.7.1)```
  * We then kept trying to fix it for 30 minutes to no avail and it was just an error every step of the way (e.g., cmake errors)

Its repository contains many reproducibility-related issues: 
* unable to reproduce FPS (https://github.com/google-research/seed_rl/issues/80) 
* unable to reproduce Pong results (https://github.com/google-research/seed_rl/issues/75, https://github.com/google-research/seed_rl/issues/51) 
* does not work as expected with TPU / never tested end-to-end results with GPUs (https://github.com/google-research/seed_rl/issues/78#issuecomment-1258515133)
* error with multiple GPUs (https://github.com/google-research/seed_rl/issues/16)

We chose to evaluate against moolib and tourchbeast because they, to our knowledge, are the only open-source DDRL infrastructures which have evaluated their performance on 50+ Atari games using the 200M frames. 200M frames also make the computational resources reasonable for us (about 2 hours per game with a single GPU or 20-30 minutes per game with 8 GPUs), as opposed to Seed RL which is evaluated on 40B frames (about 40 hours using their settings — 8xTPUv3, 213 CPUs)

Reviewer comment (Experiments):
The authors claim to have built a distributed learning framework, but have only evaluated their method on a single machine. This is a major flaw in their evaluations, and the reason I recommend rejection for this paper. If this is included I will increase my score.

Author response:
We thank the reviewer for this constructive feedback. Cleanba **can** scale to multi-node and we have included evaluations of multi-node training in Appendix F. Multi-node training is computationally expensive, so we focus evaluation specifically on a representative Atari game Breakout instead of the full Atari suite. We have also mentioned the multi-node scaling results briefly in the main text.

Reviewer comment (Writing):
At the end of section 4.2 you write `jax.distibuted` -- I think this should be `jax.distributed` (extra r).

Author response:
Thanks for the correction on our typo.