Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Aadam, Monu Verma, Mohamed Abdel-Mottaleb

TL;DR

JaxARC presents a high-performance, JAX-based RL environment for ARC tasks, addressing severe scalability bottlenecks in prior ARCRL tooling. By enforcing a purely functional, stateless core and leveraging fixed-size grids, task buffers, and JAX vectorization, it achieves dramatic throughput gains over Python-based implementations and enables large-scale experiments for abstract reasoning. The work demonstrates substantial speedups across CPU and accelerator hardware, discusses integration with Stoix and BrainGridGame-inspired actions, and provides a pathway toward population-based training, architecture search, and meta-learning on ARC. While promising, the findings also indicate that current architectures struggle to leverage rich observation wrappers, motivating further advances in network design and learning paradigms for effective inductive reasoning.

Abstract

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.

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

TL;DR

JaxARC presents a high-performance, JAX-based RL environment for ARC tasks, addressing severe scalability bottlenecks in prior ARCRL tooling. By enforcing a purely functional, stateless core and leveraging fixed-size grids, task buffers, and JAX vectorization, it achieves dramatic throughput gains over Python-based implementations and enables large-scale experiments for abstract reasoning. The work demonstrates substantial speedups across CPU and accelerator hardware, discusses integration with Stoix and BrainGridGame-inspired actions, and provides a pathway toward population-based training, architecture search, and meta-learning on ARC. While promising, the findings also indicate that current architectures struggle to leverage rich observation wrappers, motivating further advances in network design and learning paradigms for effective inductive reasoning.

Abstract

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.
Paper Structure (25 sections, 5 figures, 4 tables)

This paper contains 25 sections, 5 figures, 4 tables.

Table of Contents

  1. Introduction
  2. JaxARC Architecture
  3. Performance Evaluation
  4. Configurations
  5. Results
  6. Impact
  7. Related Work
  8. Conclusion
  9. System Architecture
  10. Visualization and Tasks
  11. Usage Example
  12. Benchmark Details
  13. Log-Scale Throughput Analysis
  14. Detailed Speedup Analysis
  15. Action System Design
  16. ...and 10 more sections

Figures (5)

  • Figure 1: Throughput (steps/second) vs number of parallel environments on CPU (left), RTX 3090 (center), and H100 (right). JaxARC scales near-linearly while ARCLE saturates due to Python/multiprocessing overhead.
  • Figure 2: JaxARC System Architecture. The functional core enables seamless JAX transformations, while the task buffer ensures efficient data loading.
  • Figure 3: JaxARC Visualization Modes. (a) Input/Output pair view showing transformation. (b) Single grid view. (c) RL step visualization showing agent interaction. (d) Complete task view showing all training pairs and test input.
  • Figure 4: Throughput (steps/second) vs number of parallel environments on Log Scale for CPU (left), RTX 3090 (center), and H100 (right).
  • Figure 5: Success rate (left) and throughput (right) for PPO training with different observation wrapper configurations on MiniARC.