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

2026-01-24

AI Summary

This work addresses the severe computational bottlenecks of existing Gymnasium-based reinforcement learning environments on tasks from the Abstraction and Reasoning Corpus (ARC), which hinder large-scale experimentation. We propose a high-performance, stateless, functional reinforcement learning environment built on JAX, leveraging its pure functional design to fully exploit vectorization, automatic parallelization, and just-in-time compilation for extreme parallelism. The environment supports multiple datasets, flexible action spaces, and composable wrappers. Compared to Gymnasium implementations under identical batch sizes, our approach achieves speedups ranging from 38× to 5,439×, with a peak throughput of 790 million steps per second—dramatically alleviating throughput constraints and enabling large-scale ARC reinforcement learning research.

Technology Category

• Machine Learning: Reinforcement Learning
• Multiagent Systems: Adversarial Agents
• Natural Language Processing: (Large) Language Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Practical large-scale studies of user experience

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.

Problem

• Abstraction and Reasoning Corpus
• reinforcement learning
• computational bottleneck
• experimental scale
• high-performance environment

Innovation

• JAX
• Abstraction and Reasoning Corpus
• high-performance reinforcement learning
• massive parallelism
• stateless architecture