JaxARC:一个用于抽象与推理研究的高性能 JAX 环境
人工智能
2026-01-27 v1 机器学习
摘要
抽象与推理语料库 (ARC) 测试 AI 系统从少量演示对中执行类人归纳推理的能力。现有的基于 Gymnasium 的强化学习环境由于计算瓶颈严重限制了实验规模。我们提出了 JaxARC,一个用 JAX 实现的开源、高性能 ARC 强化学习环境。其函数式、无状态的架构实现了大规模并行性,在匹配的批量大小下,比 Gymnasium 实现了 38-5439 倍的加速,峰值吞吐量达到每秒 7.9 亿步。JaxARC 支持多个 ARC 数据集、灵活的动作空间、可组合的包装器和配置驱动的可重复性,使得以前计算上不可行的大规模强化学习研究成为可能。JaxARC 可在 https://github.com/aadimator/JaxARC 获取。
引用
@article{arxiv.2601.17564,
title = {JaxARC: A High-Performance JAX-based Environment for Abstraction and Reasoning Research},
author = {Aadam and Monu Verma and Mohamed Abdel-Mottaleb},
journal= {arXiv preprint arXiv:2601.17564},
year = {2026}
}