Podracer architectures for scalable Reinforcement Learning
Open accessAbstract
Supporting state-of-the-art AI research requires balancing rapid prototyping, ease of use, and quick iteration, with the ability to deploy experiments at a scale traditionally associated with production systems.Deep learning frameworks such as TensorFlow, PyTorch and JAX allow users to transparently make use of accelerators, such as TPUs and GPUs, to offload the more computationally intensive parts of training and inference in modern deep learning systems. Popular training pipelines that use these frameworks for deep learning typically focus on (un-)supervised learning. How to best train reinforcement learning (RL) agents at scale is still an active research area. In this report we argue that TPUs are particularly well suited for training RL agents in a scalable, efficient and reproducible way. Specifically we describe two architectures designed to make the best use of the resources available on a TPU Pod (a special configuration in a Google data center that features multiple TPU devices connected to each other by extremely low latency communication channels).
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
Claims other papers make about this one
Each is checked against this paper’s own text: does it say what the citing paper says it does?
Due to full environment compatibility with JAX, the Anakin architecture can be used, the entire training loop can be jit-compiled, and parallelization across multiple devices using jax.pmap is easy.
“Due to the full environment compatibility with JAX, we can use the Anakin architecture (Hessel et al., 2021), jit-compile the entire training loop, and easily parallelize across multiple devices using jax.pmap transformation.”
From XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX (2024) · cites this paper as Hessel et al., 2021
Cites
ImageNet classification with deep convolutional neural networks
Krizhevsky, A., Sutskever, I., & Hinton, G. E. (2017) · Communications of the ACM · doi:10.1145/3065386 · Cited by 112,155 · see citing works
Human-level control through deep reinforcement learning
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015) · Nature · doi:10.1038/nature14236 · Cited by 30,357 · see citing works
Reinforcement Learning: An Introduction
Sutton, R. S., & Barto, A. (2005) · IEEE Transactions on Neural Networks · doi:10.1109/tnn.2004.842673 · Cited by 25,263 · see citing works
Harris, C., Millman, K., van der Walt, S., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N., et al. (2020) · RePEc: Research Papers in Economics · Cited by 17,844 · see citing works
Mastering the game of Go with deep neural networks and tree search
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., van den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. (2016) · Nature · doi:10.1038/nature16961 · Cited by 15,612 · see citing works
Continuous control with deep reinforcement learning
Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., & Wierstra, D. (2015) · arXiv · doi:10.48550/arxiv.1509.02971 · Cited by 10,782 · see citing works
(2012) · Cited by 9,062 · see citing works
(2006) · Springer eBooks · Cited by 8,123 · see citing works
Asynchronous Methods for Deep Reinforcement Learning
Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T. P., Harley, T., Silver, D., & Kavukcuoglu, K. (2016) · arXiv · doi:10.48550/arxiv.1602.01783 · Cited by 7,132 · see citing works
In-Datacenter Performance Analysis of a Tensor Processing Unit
Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., et al. (2017) · Proceedings of the 44th Annual International Symposium on Computer Architecture · doi:10.1145/3079856.3080246 · Cited by 4,537 · see citing works
Large Scale Distributed Deep Networks
Dean, J. B., Corrado, G. S., Monga, R., Chen, K., Devin, M., Mao, M., Ranzato, M., Senior, A., Tucker, P. A., Yang, K., et al. (2012) · Cited by 3,189 · see citing works
Rainbow: Combining Improvements in Deep Reinforcement Learning
Hessel, M., Modayil, J., Van Hasselt, H., Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., & Silver, D. (2018) · Proceedings of the AAAI Conference on Artificial Intelligence · doi:10.1609/aaai.v32i1.11796 · Cited by 1,846 · see citing works
Self-citationSome Computer Organizations and Their Effectiveness
Flynn, M. J. (1972) · IEEE Transactions on Computers · doi:10.1109/tc.1972.5009071 · Cited by 1,720 · see citing works
Steps toward Artificial Intelligence
Minsky, M. (1961) · Proceedings of the IRE · doi:10.1109/jrproc.1961.287775 · Cited by 1,706 · see citing works
Recursive functions of symbolic expressions and their computation by machine, Part I
McCarthy, J. (1960) · Communications of the ACM · doi:10.1145/367177.367199 · Cited by 1,635 · see citing works
Mastering Atari, Go, chess and shogi by planning with a learned model
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. (2020) · RePEc: Research Papers in Economics · Cited by 1,310 · see citing works
Distributed Prioritized Experience ReplayPreprint
Horgan, D., Quan, J., Budden, D. M., Barth-Maron, G., Hessel, M., van Hasselt, H. P., & Silver, D. (2018) · arXiv (Cornell University) · doi:10.48550/arxiv.1803.00933 · Cited by 557 · see citing works
Self-citationMassively Parallel Methods for Deep Reinforcement LearningPreprint
Nair, A., Srinivasan, P., Blackwell, S., Alcicek, C., Fearon, R., De Maria, A., Panneershelvam, V., Suleyman, M., Beattie, C., Petersen, S., et al. (2015) · arXiv (Cornell University) · doi:10.48550/arxiv.1507.04296 · Cited by 390 · see citing works
A simple automatic derivative evaluation program
Wengert, R. E. (1964) · Communications of the ACM · doi:10.1145/355586.364791 · Cited by 389 · see citing works
Reinforcement Learning in Continuous State and Action Spaces
van Hasselt, H. P. (2012) · Adaptation, learning, and optimization · doi:10.1007/978-3-642-27645-3_7 · Cited by 193 · see citing works
RLlib: Abstractions for Distributed Reinforcement Learning
Liang, E., Liaw, R., Nishihara, R., Moritz, P., Fox, R., Goldberg, K. Y., Gonzalez, J. E., Jordan, M. I., & Stoica, I. G. (2018) · International Conference on Machine Learning · Cited by 175 · see citing works
Proving program refinements and transformations
Ward, M. P. (1986) · Cited by 83 · see citing works
Discovering Reinforcement Learning AlgorithmsPreprint
Oh, J., Hessel, M., Czarnecki, W. M., Xu, Z., van Hasselt, H. P., Singh, S. P., & Silver, D. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2007.08794 · Cited by 71 · see citing works
Self-citationSEED RL: Scalable and Efficient Deep-RL with Accelerated Central InferencePreprint
Espeholt, L., Marinier, R., Stańczyk, P., Wang, K., & Michalski, M. (2019) · arXiv (Cornell University) · doi:10.48550/arxiv.1910.06591 · Cited by 71 · see citing works
A Self-Tuning Actor-Critic AlgorithmPreprint
Zahavy, T., Xu, Z., Veeriah, V., Hessel, M., Oh, J., van Hasselt, H. P., Silver, D., & Singh, S. P. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2002.12928 · Cited by 46 · see citing works
Self-citationMeta-Gradient Reinforcement Learning with an Objective Discovered OnlinePreprint
Xu, Z., van Hasselt, H. P., Hessel, M., Oh, J., Singh, S. P., & Silver, D. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2007.08433 · Cited by 42 · see citing works
Self-citationA micro controlled peripheral processor
Barron, E., & Glorioso, R. M. (1973) · IEEE/ACM International Symposium on Microarchitecture (MICRO) · doi:10.1145/800203.806247 · Cited by 4 · see citing works
Flagged first, then by citations
1 of the 28 references couldn’t be matched to a known paper.
Cited by
Sorted by citations
FinRL-podracer: high performance and scalable deep reinforcement learning for quantitative finance
Li, Z., Liu, X.-Y., Zheng, J., Wang, Z., Walid, A., & Guo, J. (2021) · Proceedings of the Second ACM International Conference on AI in Finance · doi:10.1145/3490354.3494413 · Cited by 28 · see citing works
“In recent years, distributed DRL frameworks and massively parallel simulations have been recognized as the critical software development for the RLOps paradigm [12, 19, 20, 27].”
MethodMuesli: Combining Improvements in Policy OptimizationPreprint
Hessel, M., Danihelka, I., Viola, F., Guez, A., Schmitt, S., Sifre, L., Weber, T., Silver, D., & van Hasselt, H. P. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2104.06159 · Cited by 18 · see citing works
Self-citationReinforcement Learning Controllers for Soft Robots Using Learned Environments
Berdica, U., Jackson, M., Veronese, N. E., Foerster, J., & Maiolino, P. (2024) · IEEE International Conference on Soft Robotics (RoboSoft) · doi:10.1109/robosoft60065.2024.10522003 · Cited by 7 · see citing works
Learning Expected Emphatic Traces for Deep RL
Jiang, R., Zhang, S., Chelu, V., White, A., & van Hasselt, H. P. (2022) · Proceedings of the AAAI Conference on Artificial Intelligence · doi:10.1609/aaai.v36i6.20660 · Cited by 6 · see citing works
Proper Value EquivalencePreprint
Grimm, C., Barreto, A. D. M. S., Farquhar, G., Silver, D., & Singh, S. P. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2106.10316 · Cited by 6 · see citing works
Meta-X design: from tools to partners in human–machine co–creation for industry 5.0
Yu, C., Yan, T., & Chen, Y. (2025) · Journal of Engineering Design · doi:10.1080/09544828.2025.2571005 · Cited by 4 · see citing works
XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX
Nikulin, A., Kurenkov, V., Zisman, I., Agarkov, A., Sinii, V., & Kolesnikov, S. (2024) · Advances in Neural Information Processing Systems 37 · doi:10.52202/079017-1390 · Cited by 3 · see citing works
“This makes it possible to use the Anakin architecture (Hessel et al., 2021) and easily scale to multiple devices using the jax.pmap transformation.”
MethodMassively High-Throughput Reinforcement Learning for Classic Control on GPUs
Sha, X., & Lan, T. (2024) · IEEE Access · doi:10.1109/access.2024.3441242 · Cited by 3 · see citing works
Self-Consistent Models and ValuesPreprint
Farquhar, G., Baumli, K., Marinho, Z., Filos, A., Hessel, M., van Hasselt, H. P., & Silver, D. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2110.12840 · Cited by 2 · see citing works
Self-citationKinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks
Matthews, M., Beukman, M., Lu, C., & Foerster, J. (2024) · arXiv · doi:10.48550/arxiv.2410.23208 · Cited by 1 · see citing works
“…of GPU-based environments (Lange, 2022; Koyamada et al., 2023; Rutherford et al., 2023; Nikulin et al., 2023; Matthews et al., 2024; Kazemkhani et al., 2024; Bonnet et al., 2024; Pignatelli et al., 2024), allowing the entire RL pipeline to run on a hardware accelerator (Hessel et al., 2021).”
- Authors
- Hessel, M., Kroiss, M., Clark, A., Kemaev, I., Quan, J., Keck, T., Viola, F., & van Hasselt, H.
- Venue
- arXiv
- Publisher
- arXiv
- Published
- 2021
- ISSN
- 2331-8422 (ISSN-L)
- arXiv
- 2104.06272
- OpenAlex
- W3153278392
- Cited by
- 8
- References
- 28
- Citations by year
- 2025: 1 · 2024: 2 · 2022: 1 · 2021: 4
- Type
- posted content
- Language
- en
- Open access
- green
- arxiv.org
- Locations
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University)
- Topics
- Physical Sciences › Computer Science › Artificial Intelligence › Reinforcement Learning in Robotics
- Physical Sciences › Computer Science › Hardware and Architecture › Parallel Computing and Optimization Techniques
- Physical Sciences › Computer Science › Information Systems › Cloud Computing and Resource Management
- Keywords
- reinforcement learning, scalable training, TPU, reproducibility, deep learning framework
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
