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

TL;DR

JaxARC is a high-performance, scalable JAX-based environment for ARC that significantly accelerates AI research in human-like reasoning tasks by enabling large-scale reinforcement learning experiments.

Contribution

It introduces JaxARC, a novel, efficient, and flexible JAX-based environment for ARC, overcoming previous computational limitations and supporting large-scale RL research.

Findings

Achieves up to 5439x speedup over Gymnasium

Supports multiple ARC datasets and flexible action spaces

Enables large-scale RL experiments previously infeasible

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.