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

Read full text
Authors: Aadam, Monu Verma, Mohamed Saeed Abdel-Mottaleb
Cited by: 0

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.

Fields of Study & Subjects

Reinforcement Learning in Robotics Artificial Intelligence in Games AI-based Problem Solving and Planning

Publication details

Journal / source
arXiv (Cornell University)
Publication date
24 January 2026
Document type
Preprint