# JaxARC: A High-Performance JAX-based Environment for Abstraction and Reasoning Research

arXiv: https://arxiv.org/abs/2601.17564 · HTML: https://mlsanity.com/paper/2601.17564/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 16 references that are in the map.
- First posted on arXiv 24 Jan 2026.

## Builds on (most cited first, 12 of 12)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [A Tutorial on Meta-Reinforcement Learning](https://mlsanity.com/paper/2301.08028/) (2023)
- [Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2306.09884/) (2023)
- [Podracer architectures for scalable Reinforcement Learning](https://mlsanity.com/paper/2104.06272/) (2021)
- [Reasoning Abilities of Large Language Models: In-Depth Analysis on the Abstraction and Reasoning Corpus](https://mlsanity.com/paper/2403.11793/) (2024)
- [Neural-guided, Bidirectional Program Search for Abstraction and Reasoning](https://mlsanity.com/paper/2110.11536/) (2021)
- [CodeIt: Self-Improving Language Models with Prioritized Hindsight Replay](https://mlsanity.com/paper/2402.04858/) (2024)
- [Abstract Visual Reasoning Enabled by Language](https://mlsanity.com/paper/2303.04091/) (2023)
- [A Neurodiversity-Inspired Solver for the Abstraction & Reasoning Corpus (ARC) Using Visual Imagery and Program Synthesis](https://mlsanity.com/paper/2302.09425/) (2023)
- [ARCLE: The Abstraction and Reasoning Corpus Learning Environment for Reinforcement Learning](https://mlsanity.com/paper/2407.20806/) (2024)

## Abstract (arXiv)

The Abstraction and Reasoning Corpus (ARC) tests AI systems' ability to perform human-like inductive reasoning from a few demonstration pairs. Existing Gymnasium-based RL environments severely limit experimental scale due to computational bottlenecks. We present JaxARC, an open-source, high-performance RL environment for ARC implemented in JAX. Its functional, stateless architecture enables massive parallelism, achieving 38-5,439x speedup over Gymnasium at matched batch sizes, with peak throughput of 790M steps/second. JaxARC supports multiple ARC datasets, flexible action spaces, composable wrappers, and configuration-driven reproducibility, enabling large-scale RL research previously computationally infeasible. JaxARC is available at https://github.com/aadimator/JaxARC.