Metadata-Version: 2.4
Name: JaxARC
Version: 1.0.3
Summary: High-performance JAX-based ARC environment for RL research
Project-URL: Homepage, https://github.com/aadimator/JaxARC
Project-URL: Bug Tracker, https://github.com/aadimator/JaxARC/issues
Project-URL: Discussions, https://github.com/aadimator/JaxARC/discussions
Project-URL: Changelog, https://github.com/aadimator/JaxARC/releases
Author-email: Aadam <aadimator@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Typing :: Typed
Requires-Python: >=3.9
Requires-Dist: cairosvg<3,>=2.8.2
Requires-Dist: drawsvg<3,>=2.4.0
Requires-Dist: equinox<0.14,>=0.13.0
Requires-Dist: hydra-core<2,>=1.3.2
Requires-Dist: jaxtyping<0.4,>=0.3.2
Requires-Dist: loguru<0.8,>=0.7.3
Requires-Dist: pyprojroot<0.4,>=0.3.0
Requires-Dist: rich<15,>=14.0.0
Requires-Dist: stoa-env<0.2,>=0.1.2
Requires-Dist: tqdm<5,>=4.67.1
Requires-Dist: typer<0.17,>=0.16.0
Description-Content-Type: text/markdown

# JaxARC

<!-- [![Actions Status][actions-badge]][actions-link] -->

[![Documentation Status][rtd-badge]][rtd-link]
[![PyPI version][pypi-version]][pypi-link]
[![PyPI Downloads](https://static.pepy.tech/personalized-badge/jaxarc?period=total&units=INTERNATIONAL_SYSTEM&left_color=BLACK&right_color=GREEN&left_text=downloads)](https://pepy.tech/projects/jaxarc)
[![GitHub Discussion][github-discussions-badge]][github-discussions-link]

JaxARC is a JAX-based reinforcement learning environment for the
[Abstraction and Reasoning Corpus](https://github.com/fchollet/ARC) (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.

![JaxARC throughput as compared with ARCLE](docs/_static/images/sps_vs_envs_linear_h100.png)

**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

![JaxARC System Architecture](docs/_static/images/jaxarc_system_architecture.svg)

## Installation

```bash
pip install jaxarc
```

### Want to contribute?

```bash
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](https://jaxarc.readthedocs.io/en/latest/tutorials/)** for
training loops, custom wrappers, and dataset management.

## Stoix Integration

JaxARC uses the [Stoa API](https://github.com/EdanToledo/Stoa), allowing
seamless integration with [Stoix](https://github.com/EdanToledo/Stoix), which is
a JAX-based reinforcement learning codebase supporting various RL algorithms.

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

```python
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](https://github.com/aadimator/jaxarc-baselines) repository for
example implementations of training agents on JaxARC environments using Stoix.

## Contributing

Found a bug? Want a feature?
**[Open an issue](https://github.com/aadimator/JaxARC/issues)** or submit a PR.

## Related Work

JaxARC builds on great work from the community:

- **[ARC Challenge](https://github.com/fchollet/ARC)** by François Chollet — The
  original dataset and challenge
- **[ARCLE](https://github.com/ConfeitoHS/arcle)** — Python-based ARC
  environment (inspiration for our design)
- **[Stoix](https://github.com/EdanToledo/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](https://arxiv.org/abs/2601.17564):

```bibtex
@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](LICENSE) for details.

## Questions?

- **Bugs/Features**: [GitHub Issues](https://github.com/aadimator/JaxARC/issues)
- **Discussions**:
  [GitHub Discussions](https://github.com/aadimator/JaxARC/discussions)
- **Docs**: [jaxarc.readthedocs.io](https://jaxarc.readthedocs.io)

---

<!-- Links -->

[actions-badge]: https://github.com/aadimator/JaxARC/workflows/CI/badge.svg
[actions-link]: https://github.com/aadimator/JaxARC/actions
[conda-badge]: https://img.shields.io/conda/vn/conda-forge/JaxARC
[conda-link]: https://github.com/conda-forge/JaxARC-feedstock
[github-discussions-badge]:
  https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github
[github-discussions-link]: https://github.com/aadimator/JaxARC/discussions
[pypi-link]: https://pypi.org/project/JaxARC/
[pypi-platforms]: https://img.shields.io/pypi/pyversions/JaxARC
[pypi-version]: https://img.shields.io/pypi/v/JaxARC
[rtd-badge]: https://readthedocs.org/projects/JaxARC/badge/?version=latest
[rtd-link]: https://JaxARC.readthedocs.io/en/latest/?badge=latest
