{"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/jaxmarl-multi-agent-rl-environments-in-jax","title":"JaxMARL: Multi-Agent RL Environments and Algorithms in JAX","arxiv_id":"2311.10090","date":"2023-11-16","proceeding":null,"authors":["Alexander Rutherford","Benjamin Ellis","Matteo Gallici","Jonathan Cook","Andrei Lupu","Gardar Ingvarsson","Timon Willi","Ravi Hammond","Akbir Khan","Christian Schroeder de Witt","Alexandra Souly","Saptarashmi Bandyopadhyay","Mikayel Samvelyan","Minqi Jiang","Robert Tjarko Lange","Shimon Whiteson","Bruno Lacerda","Nick Hawes","Tim Rocktaschel","Chris Lu","Jakob Nicolaus Foerster"],"abstract":"Benchmarks are crucial in the development of machine learning algorithms, with available environments significantly influencing reinforcement learning (RL) research. Traditionally, RL environments run on the CPU, which limits their scalability with typical academic compute. However, recent advancements in JAX have enabled the wider use of hardware acceleration, enabling massively parallel RL training pipelines and environments. While this has been successfully applied to single-agent RL, it has not yet been widely adopted for multi-agent scenarios. In this paper, we present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments and popular baseline algorithms. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches, and up to 12500x when multiple training runs are vectorized. This enables efficient and thorough evaluations, potentially alleviating the evaluation crisis in the field. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge, which removes the need to run the StarCraft II game engine. This not only enables GPU acceleration, but also provides a more flexible MARL environment, unlocking the potential for self-play, meta-learning, and other future applications in MARL. The code is available at https://github.com/flairox/jaxmarl.","url_abs":"https://arxiv.org/abs/2311.10090v5","url_pdf":"https://arxiv.org/pdf/2311.10090v5.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":"jaxmarl-multi-agent-rl-environments-in-jax","repo_url":"https://github.com/flairox/jaxmarl","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"jaxmarl-multi-agent-rl-environments-in-jax","repo_url":"https://github.com/kfu02/jaxmarl","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"jaxmarl-multi-agent-rl-environments-in-jax","repo_url":"https://github.com/ucl-dark/pax","is_official":0,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[{"task_slug":null,"task_name":"CPU"},{"task_slug":null,"task_name":"GPU"},{"task_slug":"meta-learning","task_name":"Meta-Learning"},{"task_slug":"multi-agent-reinforcement-learning","task_name":"Multi-agent Reinforcement Learning"},{"task_slug":"reinforcement-learning-1","task_name":"Reinforcement Learning (RL)"},{"task_slug":"starcraft","task_name":"Starcraft"},{"task_slug":"starcraft-ii","task_name":"Starcraft II"}],"methods":[{"method_slug":"base","method_name":"BASE"},{"method_slug":null,"method_name":"Library"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2311.10090","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2311.10090"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+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/ucl-dark/pax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/flairox/jaxmarl","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/kfu02/jaxmarl","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"unverified":23},"by_repo_kind":{"listed":{"samples":23,"ran":0,"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":"6b4576fe27ed0a8e","entry":"add_batch_dim","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/utils.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/utils.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":"6b4576fe27ed0a8e"}},{"code_sha256_prefix":"8d1fa95afa90b12b","entry":"get_advantages","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/utils.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/utils.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":"8d1fa95afa90b12b"}},{"code_sha256_prefix":"7998dbf93494334f","entry":"global_setup","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/experiment.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/experiment.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":"7998dbf93494334f"}},{"code_sha256_prefix":"bdb8d47954719031","entry":"make_GRU","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/hyper/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/hyper/networks.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":"bdb8d47954719031"}},{"code_sha256_prefix":"61fc2fad0c403979","entry":"make_GRU_hypernetwork","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/hyper/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/hyper/networks.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":"61fc2fad0c403979"}},{"code_sha256_prefix":"0a49d3ae85eecc1f","entry":"make_GRU_ipd_att_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/shaper_att/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/shaper_att/networks.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":"0a49d3ae85eecc1f"}},{"code_sha256_prefix":"d64c512a602ec179","entry":"make_GRU_ipd_avg_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/shaper_att/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/shaper_att/networks.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":"d64c512a602ec179"}},{"code_sha256_prefix":"365cfd6131c2ac62","entry":"make_GRU_ipd_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/shaper_att/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/shaper_att/networks.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":"365cfd6131c2ac62"}},{"code_sha256_prefix":"607630c430f35d26","entry":"make_coingame_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/naive/network.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/naive/network.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":"607630c430f35d26"}},{"code_sha256_prefix":"ba7cdb74da41663e","entry":"make_cournot_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/ppo/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/ppo/networks.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":"ba7cdb74da41663e"}},{"code_sha256_prefix":"0c5f9da603bc1634","entry":"make_fishery_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/ppo/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/ppo/networks.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":"0c5f9da603bc1634"}},{"code_sha256_prefix":"db890d82d0ebeef0","entry":"make_ipd_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/ppo/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/ppo/networks.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":"db890d82d0ebeef0"}},{"code_sha256_prefix":"d48299dc9766e894","entry":"make_mfos_avg_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/mfos_ppo/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/mfos_ppo/networks.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":"d48299dc9766e894"}},{"code_sha256_prefix":"cb1865b00332fa19","entry":"make_mfos_continuous_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/mfos_ppo/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/mfos_ppo/networks.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":"cb1865b00332fa19"}},{"code_sha256_prefix":"2eb4d1417b0a6d2f","entry":"make_mfos_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/mfos_ppo/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/mfos_ppo/networks.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":"2eb4d1417b0a6d2f"}},{"code_sha256_prefix":"019e1da94cf6b12f","entry":"make_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/hyper/networks.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/hyper/networks.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":"019e1da94cf6b12f"}},{"code_sha256_prefix":"2cd8e496d080743d","entry":"make_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/lola/network.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/lola/network.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":"2cd8e496d080743d"}},{"code_sha256_prefix":"bd3555b465d1974f","entry":"make_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/naive/network.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/naive/network.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":"bd3555b465d1974f"}},{"code_sha256_prefix":"d2a9587eacd4ae54","entry":"make_policy_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/lola/network.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/lola/network.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":"d2a9587eacd4ae54"}},{"code_sha256_prefix":"8dfd96d7efaac91f","entry":"make_rice_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/naive/network.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/naive/network.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":"8dfd96d7efaac91f"}},{"code_sha256_prefix":"7dfc7920e7654bda","entry":"make_value_network","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/agents/lola/network.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/agents/lola/network.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":"7dfc7920e7654bda"}},{"code_sha256_prefix":"0c175a3d42cc09ab","entry":"reduce_outer_traj","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/runners/runner_marl_nplayer.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/runners/runner_marl_nplayer.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":"0c175a3d42cc09ab"}},{"code_sha256_prefix":"0c9e5aa92235c1ab","entry":"to_numpy","repo":"ucl-dark/pax","repo_kind":"listed","path":"pax/utils.py","file_url":"https://github.com/ucl-dark/pax/blob/HEAD/pax/utils.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":"0c9e5aa92235c1ab"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}