# PuzzleJAX: A Benchmark for Reasoning and Learning

arXiv: https://arxiv.org/abs/2508.16821 · HTML: https://mlsanity.com/paper/2508.16821/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning benchmarks and evaluation, one of 17 directions in LLM reasoning.
- Builds on 27 of its 52 references that are in the map.
- First posted on arXiv 22 Aug 2025.

## Placed in

- LLM reasoning / [Reasoning benchmarks and evaluation](https://mlsanity.com/lr/reasoning-benchmarks-and-evaluation/)

## Builds on (most cited first, 27 of 27)

- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [Thinking Fast and Slow with Deep Learning and Tree Search](https://mlsanity.com/paper/1705.08439/) (2017)
- [Benchmarking the Spectrum of Agent Capabilities](https://mlsanity.com/paper/2109.06780/) (2021)
- [The NetHack Learning Environment](https://mlsanity.com/paper/2006.13760/) (2020)
- [Open-Ended Learning Leads to Generally Capable Agents](https://mlsanity.com/paper/2107.12808/) (2021)
- [BALROG: Benchmarking Agentic LLM and VLM Reasoning On Games](https://mlsanity.com/paper/2411.13543/) (2024)
- [SmartPlay: A Benchmark for LLMs as Intelligent Agents](https://mlsanity.com/paper/2310.01557/) (2023)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning](https://mlsanity.com/paper/2402.16801/) (2024)
- [Illuminating Generalization in Deep Reinforcement Learning through Procedural Level Generation](https://mlsanity.com/paper/1806.10729/) (2018)
- [Control of Memory, Active Perception, and Action in Minecraft](https://mlsanity.com/paper/1605.09128/) (2016)
- [Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2306.09884/) (2023)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learning](https://mlsanity.com/paper/2303.17503/) (2023)
- [General Video Game AI: a Multi-Track Framework for Evaluating Agents, Games and Content Generation Algorithms](https://mlsanity.com/paper/1802.10363/) (2018)
- [Werewolf Arena: A Case Study in LLM Evaluation via Social Deduction](https://mlsanity.com/paper/2407.13943/) (2024)
- [Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks](https://mlsanity.com/paper/2410.23208/) (2024)
- [Deep Reinforcement Learning for General Video Game AI](https://mlsanity.com/paper/1806.02448/) (2018)
- [GAVEL: Generating Games Via Evolution and Language Models](https://mlsanity.com/paper/2407.09388/) (2024)
- [DSGBench: A Diverse Strategic Game Benchmark for Evaluating LLM-based Agents in Complex Decision-Making Environments](https://mlsanity.com/paper/2503.06047/) (2025)
- [Deep Interactive Evolution](https://mlsanity.com/paper/1801.08230/) (2018)
- [TALES: Text Adventure Learning Environment Suite](https://mlsanity.com/paper/2504.14128/) (2025)
- [DreamGarden: A Designer Assistant for Growing Games from a Single Prompt](https://mlsanity.com/paper/2410.01791/) (2024)
- [GameTraversalBenchmark: Evaluating Planning Abilities Of Large Language Models Through Traversing 2D Game Maps](https://mlsanity.com/paper/2410.07765/) (2024)
- [Beating the World's Best at Super Smash Bros. with Deep Reinforcement Learning](https://mlsanity.com/paper/1702.06230/) (2017)

## Built on by (most cited first, 4 of 4)

- [A Survey of Reinforcement Learning for Large Reasoning Models](https://mlsanity.com/paper/2509.08827/) (2025)
- [Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX](https://mlsanity.com/paper/2510.01764/) (2025)
- [AI for Games in the Foundation Model Era](https://mlsanity.com/paper/2609.16679/) (2026)
- [PlayTrain: An Efficient Reinforcement Learning Framework for LLM-Generated Adaptable JavaScript Games](https://mlsanity.com/paper/2609.09059/) (2026)

## Abstract (arXiv)

We introduce PuzzleJAX, a GPU-accelerated puzzle game engine and description language designed to support rapid benchmarking of tree search, reinforcement learning, and LLM reasoning abilities. Unlike existing GPU-accelerated learning environments that provide hard-coded implementations of fixed sets of games, PuzzleJAX allows dynamic compilation of any game expressible in its domain-specific language (DSL). This DSL follows PuzzleScript, which is a popular and accessible online game engine for designing puzzle games. In this paper, we validate in PuzzleJAX several hundred of the thousands of games designed in PuzzleScript by both professional designers and casual creators since its release in 2013, thereby demonstrating PuzzleJAX's coverage of an expansive, expressive, and human-relevant space of tasks. By analyzing the performance of search, learning, and language models on these games, we show that PuzzleJAX can naturally express tasks that are both simple and intuitive to understand, yet often deeply challenging to master, requiring a combination of control, planning, and high-level insight.