用JAX加速经典游戏环境，让强化学习训练快十倍以上。
Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX
- 基于CHIP-8模拟器构建JAX原生游戏环境，支持GPU并行。
- 相比传统CPU模拟器，训练速度提升数个数量级。
- 适合大规模强化学习实验，支持自定义新游戏生成。
强化学习研究需要多样、挑战性强且可扩展的环境。尽管现代视频游戏动态丰富，但其依赖CPU执行，计算成本高，不适合大规模实验。本文提出Octax，一个基于芯片8（CHIP-8）模拟的高性能经典街机游戏环境套件，采用JAX实现，为JAX社区提供替代Atari游戏的端到端GPU方案。该环境覆盖益智、动作、策略等类型，所有游戏均可在现代GPU上大规模运行。基于JAX的实现相较传统CPU模拟器实现数量级加速。我们通过在多个游戏中训练强化学习智能体，验证了其在训练速度和可扩展性上的显著优势。模块化设计支持研究人员轻松扩展新游戏或使用大语言模型生成新环境，是大规模强化学习实验的理想平台。开源代码已发布于https://github.com/riiswa/octax/。
原文摘要 · Abstract (English)
Reinforcement learning (RL) research requires diverse, challenging environments that are both tractable and scalable. While modern video games may offer rich dynamics, they are computationally expensive and poorly suited for large-scale experimentation due to their CPU-bound execution. We introduce Octax, a high-performance suite of classic arcade game environments implemented in JAX, based on CHIP-8 emulation, a predecessor to Atari, which is widely adopted as a benchmark in RL research. Octax provides the JAX community with a long-awaited end-to-end GPU alternative to Atari games, offering image-based environments, spanning puzzle, action, and strategy genres, all executable at massive scale on modern GPUs. Our JAX-based implementation achieves orders-of-magnitude speedups over traditional CPU emulators. We demonstrate Octax's capabilities by training RL agents across multiple games, showing significant improvements in training speed and scalability compared to existing solutions. The environment's modular design enables researchers to easily extend the suite with new games or generate novel environments using large language models, making it an ideal platform for large-scale RL experimentation. Our open-source framework is available at https://github.com/riiswa/octax/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。