PuzzleJAX: A Benchmark for Reasoning and Learning
On the map: LLM reasoning → Reasoning benchmarks and evaluation
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.
Key facts. Cited by 4 papers in the map (data through 7 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.
Builds on
27 of 52 references.
References · 27
- 2013Playing Atari with Deep Reinforcement Learningcited by 3,211
- 2019Dota 2 with Large Scale Deep Reinforcement Learningcited by 1,084
- 2023Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Taskscited by 320
- 2017Thinking Fast and Slow with Deep Learning and Tree Searchcited by 245
- 2020The NetHack Learning Environmentcited by 196
- 2021Benchmarking the Spectrum of Agent Capabilitiescited by 193
- 2021Open-Ended Learning Leads to Generally Capable Agentscited by 177
- 2024BALROG: Benchmarking Agentic LLM and VLM Reasoning On Gamescited by 141
- 2023SmartPlay: A Benchmark for LLMs as Intelligent Agentscited by 123
- 2022Discovered Policy Optimisationcited by 103
- 2018Illuminating Generalization in Deep Reinforcement Learning through Procedural Level Generationcited by 80
- 2024Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learningcited by 80
15 more
- 2016Control of Memory, Active Perception, and Action in Minecraftcited by 61
- 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 54
- 2023Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAXcited by 52
- 2018General Video Game AI: a Multi-Track Framework for Evaluating Agents, Games and Content Generation Algorithmscited by 45
- 2023Pgx: Hardware-Accelerated Parallel Game Simulators for Reinforcement Learningcited by 45
- 2024Werewolf Arena: A Case Study in LLM Evaluation via Social Deductioncited by 42
- 2024Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Taskscited by 40
- 2018Deep Reinforcement Learning for General Video Game AIcited by 25
- 2024GAVEL: Generating Games Via Evolution and Language Modelscited by 24
- 2025DSGBench: A Diverse Strategic Game Benchmark for Evaluating LLM-based Agents in Complex Decision-Making Environmentscited by 22
- 2018Deep Interactive Evolutioncited by 13
- 2024DreamGarden: A Designer Assistant for Growing Games from a Single Promptcited by 10
- 2025TALES: Text Adventure Learning Environment Suitecited by 9
- 2024GameTraversalBenchmark: Evaluating Planning Abilities Of Large Language Models Through Traversing 2D Game Mapscited by 8
- 2017Beating the World's Best at Super Smash Bros. with Deep Reinforcement Learningcited by 2
Built on by
Citing papers per year: 2025: 2; 2026: 2.
4 papers.
- 2025A Survey of Reinforcement Learning for Large Reasoning Modelscited by 175
- 2025Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAXcited by 5
- 2026AI for Games in the Foundation Model Eracited by 2
- 2026PlayTrain: An Efficient Reinforcement Learning Framework for LLM-Generated Adaptable JavaScript Gamescited by 0
- 2026AI for Games in the Foundation Model Era15 Sep 2026
- 2026PlayTrain: An Efficient Reinforcement Learning Framework for LLM-Generated Adaptable JavaScript Games8 Sep 2026
- 2025Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX2 Oct 2025
- 2025A Survey of Reinforcement Learning for Large Reasoning Models10 Sep 2025
Similar, not cited
Close in wording, no citation either way
- 2025PuzzlePlex: Benchmarking Foundation Models on Reasoning and Planning with Puzzles
- 2025PuzzleWorld: A Benchmark for Multimodal, Open-Ended Reasoning in Puzzlehunts
- 2025HardcoreLogic: Challenging Large Reasoning Models with Long-tail Logic Puzzle Games
- 2024PUZZLES: A Benchmark for Neural Algorithmic Reasoning
- 2026The Token Games: Evaluating Language Model Reasoning with Puzzle Duels
- 2026ClassicLogic: A Knowledge-Driven Benchmark of Classic Puzzle Games for Evaluating Compositional Generalization
- 2026LongPuzzleBench: Evaluating GUI Agents on Long-Horizon Visual Puzzles
- 2025TextGames: Learning to Self-Play Text-Based Puzzle Games via Language Model Reasoning