面向 JAX 中的强化学习的加速 CHIP-8 arcade 环境 - Octax
机器学习
2025-10-06 v2
摘要
强化学习研究需要多样化、具备挑战性且既可实现又可扩展的环境。虽然现代视频游戏可能提供丰富的动态特性，但计算成本高昂且不适合大规模实验，因为其 CPU 绑定的执行方式。我们引入 Octax，一个在 JAX 中实现的经典街机游戏环境套件，基于 CHIP-8 模拟，这是 Atari 之前的流行基准，在 RL 研究中被广泛采用。Octax 为 JAX 社区提供了一个长期期待的端到端 GPU 替代方案，提供基于图像的环境，涵盖拼图、动作和策略等类型，全部可在现代 GPU 上大规模执行。我们的 JAX 实现相对于传统 CPU 模拟器实现了数量级的加速，同时保持对原始游戏机制的完美忠实度。我们通过在多个游戏上训练 RL 智能体来展示 Octax 的能力，显示出相对于现有解决方案在训练速度和可扩展性方面的显著优势。该环境的模块化设计使研究人员能够轻松扩展套件中的新游戏或使用大型语言模型生成新环境，使其成为大规模 RL 实验的理想平台。
引用
@article{arxiv.2510.01764,
title = {Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX},
author = {Waris Radji and Thomas Michel and Hector Piteau},
journal= {arXiv preprint arXiv:2510.01764},
year = {2025}
}