JaxARC Documentation

Welcome to JaxARC - A JAX-native Reinforcement Learning environment for the Abstraction and Reasoning Corpus (ARC).

What is JaxARC?

JaxARC provides a high-performance, functionally-pure environment for training AI agents on abstract reasoning puzzles from the ARC challenge. Built entirely in JAX, it enables researchers and developers to leverage modern hardware acceleration (GPUs/TPUs) while maintaining clean, composable code.

The ARC challenge tests an AI system’s ability to solve novel reasoning tasks by observing a few examples and inferring the underlying transformation rule. JaxARC makes it easy to experiment with different agent architectures and training strategies on this challenging benchmark.

Why JaxARC?

  • JAX-Native Performance: Full support for JIT compilation, vectorization (vmap), and parallelization (pmap) means you can scale from single environments to thousands of parallel rollouts with minimal code changes.

  • Fast & Efficient: Leverage GPU/TPU acceleration for environment simulation. Run hundreds of thousands of steps per second on modern hardware.

  • Flexible & Composable: Clean functional API makes it easy to combine with other JAX libraries like Optax, Flax, and Haiku for agent training.

  • Reproducible: Explicit PRNG key management and immutable state ensure your experiments are perfectly reproducible across runs and platforms.

Key Features

  • Multiple Environments: Support for ARC-AGI-1, ARC-AGI-2, ConceptARC, and MiniARC datasets

  • Parser Registry: Extensible system for loading and parsing different ARC dataset formats

  • Wrappers: Transform observations, rewards, and states with composable wrappers

  • Visualization: Rich terminal and SVG rendering for debugging and analysis

  • Type Safety: Precise array shape annotations with jaxtyping for better error messages

  • Comprehensive Testing: Extensive test suite ensuring JAX compatibility and correctness

Architecture Overview

JaxARC follows a modular, functional design that keeps the core environment fast (JIT-compatible) while allowing flexible customization:

JaxARC System Architecture

Key Components:

  • Datasets: ARC-AGI-1/2, ConceptARC, MiniARC with JSON task definitions

  • Parsers: Convert JSON → JAX arrays with validation and normalization

  • Task Buffer: JIT-compatible stacked arrays with fixed shapes

  • Registry: Manages dataset variants and task subsets

  • Environment: Pure functional reset/step API following the Stoa standard

  • State: Immutable dataclass with working grid, input, target, and clipboard

  • Actions: Multiple representations (mask, point, bbox) with transformations

  • Wrappers: Modular observation and action space transforms

  • Visualization: Terminal and SVG rendering for debugging

Quick Navigation

🚀 Getting Started

Install JaxARC and run your first environment in under 10 minutes.

Getting Started with JaxARC
📚 Tutorials

Step-by-step guides for common tasks like downloading datasets, creating agents, and using wrappers.

Tutorials
📖 API Reference

Complete API documentation for all modules, classes, and functions.

API Reference

Installation

# Using Pixi (recommended)
pixi add jaxarc

# Or using pip
pip install jaxarc

Quick Example

import jax
import jaxarc

# Create an environment
env, env_params = jaxarc.make("Mini-Most_Common_color_l6ab0lf3xztbyxsu3p")

# Reset to get initial state
key = jax.random.PRNGKey(0)
state, timestep = env.reset(key, env_params)

# Take a step
action_space = env.action_space(env_params)
action = action_space.sample(key)
next_state, next_timestep = env.step(state, action)

print(f"Observation shape: {next_timestep.observation.shape}")
print(f"Reward: {next_timestep.reward}")

Next Steps

Contents

Documentation

Citation

If you use JaxARC in your research, please cite our paper:

@misc{aadam2026jaxarchighperformancejaxbasedenvironment,
      title={JaxARC: A High-Performance JAX-based Environment for Abstraction and Reasoning Research},
      author={Aadam and Monu Verma and Mohamed Abdel-Mottaleb},
      year={2026},
      eprint={2601.17564},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2601.17564},
}

Community & Support