Metadata-Version: 2.4
Name: JaxARC
Version: 1.0.1
Summary: MARL environment for ARC dataset in JAX
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 :: 4 - Beta
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: seaborn<0.14,>=0.13.2
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 platforms][pypi-platforms]][pypi-link]
[![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 experiment fast — with JIT
compilation giving you 100x+ speedups over Python loops.

If you're working on program synthesis, meta-learning, or hierarchical RL for
abstract reasoning, JaxARC gives you a solid foundation without the boilerplate.

## 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.

**Production-ready.** Type-safe configs, comprehensive tests, and functional
purity throughout. No hidden state, no surprises.

**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
- **100x+ Faster**: JIT compilation turns Python into XLA-optimized machine code
- **Configurable**: Multiple action spaces, reward functions, and observation
  formats
- **Four 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.

## Development

**Run tests:**

```bash
pixi run -e test test
```

**Lint code:**

```bash
pixi run lint
```

**Build docs:**

```bash
pixi run docs-serve
```

## 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:

```bibtex
@software{jaxarc2025,
  author = {Aadam},
  title = {JaxARC: JAX-based Reinforcement Learning for Abstract Reasoning},
  year = {2025},
  url = {https://github.com/aadimator/JaxARC}
}
```

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