JaxARC

JaxARC is a JAX-based reinforcement learning environment for the Abstraction and Reasoning Corpus (ARC) challenge. It's built for researchers who want to use extremely fast vectorized environments to explore reinforcement learning, and meta-learning techniques for abstract reasoning.

Why JaxARC?

Speed. Environments compile with jax.jit and vectorize with jax.vmap. Run thousands of episodes in parallel on GPU/TPU.

Flexible. Multiple action spaces (point-based, selection masks, bounding boxes). Multiple datasets (ARC-AGI, ConceptARC, MiniARC). Observation wrappers for different input formats. Configure everything via typed dataclasses or YAML.

Extensible. Clean parser interface for custom datasets. Wrapper system for custom observations and actions. Built with future HRL and Meta-RL experiments in mind.

Key Features

• JAX-Native: Pure functional API — every function is jax.jit-compatible
• Lightning Fast: JIT compilation turns Python into XLA-optimized machine code
• Configurable: Multiple action spaces, reward functions, and observation formats
• Multiple Datasets: ARC-AGI-1, ARC-AGI-2, ConceptARC, and MiniARC included
• Type-Safe: Full type hints with runtime validation
• Visual Debug: Terminal and SVG rendering for development

Installation

pip install jaxarc

Want to contribute?

git clone https://github.com/aadimator/JaxARC.git
cd JaxARC
pixi shell  # Sets up the environment
pixi run -e dev pre-commit install  # Hooks for code quality

See the tutorials for training loops, custom wrappers, and dataset management.

Stoix Integration

JaxARC uses the Stoa API, allowing seamless integration with Stoix, which is a JAX-based reinforcement learning codebase supporting various RL algorithms.

JaxARC provides a ready-to-use factory function for Stoix:

from jaxarc.stoix_adapter import make_jaxarc_env, jaxarc_custom_metrics

# Create Stoix-compatible environments
train_env, eval_env = make_jaxarc_env(config)

# Process episode metrics for logging
metrics = jaxarc_custom_metrics(raw_metrics)
# → success_rate, avg_steps_to_solve, truncation_rate, etc.

Domain-specific metrics (puzzle similarity, solve rate) are tracked via the ExtendedMetrics wrapper and flow through Stoix's standard logging pipeline.

You can explore jaxarc-baselines repository for example implementations of training agents on JaxARC environments using Stoix.

Contributing

Found a bug? Want a feature? Open an issue or submit a PR.

Related Work

JaxARC builds on great work from the community:

• ARC Challenge by François Chollet — The original dataset and challenge
• ARCLE — Python-based ARC environment (inspiration for our design)
• Stoix by Edan Toledo — Single-agent RL in JAX (we use their Stoa API)

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},
}

License

MIT License — see LICENSE for details.

Questions?

• Bugs/Features: GitHub Issues
• Discussions: GitHub Discussions
• Docs: jaxarc.readthedocs.io