{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/podracer-architectures-for-scalable","title":"Podracer architectures for scalable Reinforcement Learning","arxiv_id":"2104.06272","date":"2021-04-13","proceeding":null,"authors":["Matteo Hessel","Manuel Kroiss","Aidan Clark","Iurii Kemaev","John Quan","Thomas Keck","Fabio Viola","Hado van Hasselt"],"abstract":"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).","url_abs":"https://arxiv.org/abs/2104.06272v1","url_pdf":"https://arxiv.org/pdf/2104.06272v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"podracer-architectures-for-scalable","repo_url":"https://github.com/instadeepai/flashbax","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"podracer-architectures-for-scalable","repo_url":"https://github.com/instadeepai/jumanji","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"podracer-architectures-for-scalable","repo_url":"https://github.com/sotetsuk/pgx","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[{"task_slug":"reinforcement-learning","task_name":"Reinforcement Learning"},{"task_slug":"reinforcement-learning-1","task_name":"Reinforcement Learning (RL)"},{"task_slug":"reinforcement-learning-2","task_name":"reinforcement-learning"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":"https://syntology.ai/paper/2104.06272","atlas_url":"https://app.syntology.ai/?focus=2104.06272","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2104.06272"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-25T09:33:49+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/instadeepai/jumanji","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/sotetsuk/pgx","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/instadeepai/flashbax","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"ran":7,"unverified":1},"by_repo_kind":{"listed":{"samples":8,"ran":7,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"ee392cd4e9cab16d","entry":"add_dim_to_args","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/utils.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/utils.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ee392cd4e9cab16d"}},{"code_sha256_prefix":"830f36b3edd5a3d4","entry":"get_max_divisible_length","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/buffers/prioritised_trajectory_buffer.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/buffers/prioritised_trajectory_buffer.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"830f36b3edd5a3d4"}},{"code_sha256_prefix":"76498514129c5242","entry":"get_num_items_in_row","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/buffers/prioritised_trajectory_buffer.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/buffers/prioritised_trajectory_buffer.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"76498514129c5242"}},{"code_sha256_prefix":"0483f1d005c133d0","entry":"get_sum_tree_capacity","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/buffers/prioritised_trajectory_buffer.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/buffers/prioritised_trajectory_buffer.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"0483f1d005c133d0"}},{"code_sha256_prefix":"90c5c94867432ea5","entry":"get_timestep_count","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/utils.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/utils.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"90c5c94867432ea5"}},{"code_sha256_prefix":"ba9774a1eb1e6f03","entry":"get_tree_depth","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/buffers/sum_tree.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/buffers/sum_tree.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ba9774a1eb1e6f03"}},{"code_sha256_prefix":"1ccd29a8a165f5a4","entry":"get_tree_shape_prefix","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/utils.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/utils.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"1ccd29a8a165f5a4"}},{"code_sha256_prefix":"d54c2ba6a3fac77d","entry":"get_tree_index","repo":"instadeepai/flashbax","repo_kind":"listed","path":"flashbax/buffers/sum_tree.py","file_url":"https://github.com/instadeepai/flashbax/blob/HEAD/flashbax/buffers/sum_tree.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"d54c2ba6a3fac77d"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}