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.
Figures & tables
Fig. 1 : The PuzzleJAX framework facilitates the use of AI agents, search, and optimization for playing PuzzleScript games. A dataset of over 26,000 game files is collected via public gists. Search is used to generate (partial) solutions to valid games in the original JavaScript engine, and a validation harness ensures functional equivalency of multi-threaded/GPU-accelerated reimplementations in C++ and JAX.
Fig. 2 : Example games from the framework that Select dataset, showcasing the diversity of PuzzleScript games.
Fig. 3 : Parallelizing PuzzleScript for Python . Frames-per-second of random rollouts on a 36-core i9 with 2 RTX 4090 GPUs. The original (headless) engine runs relatively fast ( NodeJS (native) ), but slows considerably when the outer loop is conducted over a Python-NodeJS bridge ( NodeJS ). Even parallelizing individual steps on the NodeJS side of the bridge ( NodeJS (batched) ) does not close the gap. Swapping in a multi-threaded C++ reimplementation ( C++ (batched) of the engine, however, makes for considerably faster rollouts in Python. Running a JIT-able JAX reimplementation of PuzzleScript within a JIT-ed outer loop, running on the GPU, also enables significant speedups. (For C++, we sweep over multiple num_threads for each batch size and take the best.)
Fig. 4 : PuzzleJAX validation results . We use search to find solutions (or best action sequences according to Win Conditions -derived heuristics) in the original PuzzleScript engine in NodeJS, and replay action sequences in JAX, logging any discrepancies in the terminal state, its successfully triggering the win flag, terminal or intermediary score/heuristic values, and compilation errors. The majority of games in our PuzzleScript corpus are successfully validated in this fashion. The majority of solution and state errors result from randomness, and the majority of compilation errors involve the rigid keyword, which we exclude for simplicity. (The C++ backend achieves a >99% overall success rate.)
Game
N. rules
Description
Blocks
0
The player simultaneously controls multiple avatars, navigating obstacles to guide them to target positions.
Sokoban
1
The player navigates a top-down grid of traversible and wall tiles, pushing crates onto targets. The challenge is to sequence moves such that crates do not wind up “deadlocked” in a position (e.g. a corner) from which they cannot be moved.
Sokoban Match 3
2
Sokoban, except that when 3 crates are arranged in a horizontal or vertical line, they disappear.
Not Snake
2
When the player moves onto a cell, they flip it between one of two colors. One must walk a hamiltonian path or strategically retrace one’s steps to produce a successful grid-coloring.
Zen Puzzle Garden
3
Tasks the player with “raking” each cell in a central square of sand without retracing its steps, while at the same time avoiding increasingly complex arrangements of obstacles within the sand patch. The player may freely navigate around the border of the sand patch.
Multi-Word Dictionary Game
4
The player arranges letters by either pushing and pulling them in different directions to correctly spell an English word (“Cat” or “Act”).
TABLE I : Description of the Select PuzzleScript games used for benchmarking.
Fig. 5 : Win rate of search, RL, and LLM-driven player agents over PuzzleJAX Select , averaged across levels. For RL (PPO) and LLMs, we additionally average across trial random seeds (5 for RL, 10 for LLMs). Games ordered by cumulative win-rate across methods. We vary the maximum nodes allotted to search and ExIt’s search phase, and the maximum number of episode steps in RL.
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Fig. 6 : Average length of solutions generated by BFS over our test suite of PuzzleJAX games.
Fig. 7 : Win rate of search, RL, and LLM-driven player agents over a subset of PuzzleJAX games, averaged across levels.
Game
# Levels
Total Plays
Avg Solution Length
Simple Block Pushing Game
2
14
32.36
Match 3 Block Push
2
19
88.47
Travelling salesman
6
12
14.33
Atlas Shrank
9
94
98.79
constellation z
11
53
35.02
Multi-word Dictionary Game
1
15
57.93
Appendix
TABLE II : Results from preliminary playtesting of PuzzleScript games with human players.
Figure 10
Fig. 9 : In Blocks , a PPO Reinforcement Learning agent quickly learns to increase score according to the heuristic, but falls into a sub-optimal strategy in which one of the Player blocks is trapped in a dead-end corridor adjacent to the last remaining target.
Fig. 10 : Comparison of RL against breadth-first search in Sokoban . Episode rollouts from RL are pictured on the right. Here, the agent greedily maximizes the heuristic (the sum of manhattan distances between targets and their nearest crates), preventing discovery of optimal solutions.
Fig. 11 : Comparison of RL against breadth-first search in Sokoban Match 3 .
Fig. 12 : Comparison of RL against breadth-first search in Notsnake . Episode rollouts from RL is pictured on the right.
Fig. 13 : Comparison of RL against breadth-first search in Limerick . Episode rollouts from RL are pictured on the right. Agents only master levels with a relatively straightforward path to the goal. They do not generally uncover strategies involving significant roundabouts away from the goal, and can fall prey to “obvious” traps along the more direct path.
Fig. 14 : Comparison of RL against breadth-first search in Slidings . Here, RL rapidly discovers the solutions found via search, though it often exhibits some instability later in training.
Fig. 15 : Comparison of RL against breadth-first search in Kettle . RL agents are able to find optimal solutions, which involve a short sequence of actions, though the time taken to learn this optimal strategy steadily increases as levels (and optimal action sequences) grow and complexify.
Fig. 16 : Comparison of RL against breadth-first search in Take Heart Lass . RL can handily find solutions to early levels which involve effectively evolve running away from encroaching despair and toward a goal, but it has difficulty in later levels that introduce the use of pushable hearts to strategically block the despair’s advance.
Fig. 17 : Comparison of RL against breadth-first search in Multi-Word Dictionary Game .
Vision-Language(-Action) Models (VLMs) are increasingly applied to interactive environments, yet existing benchmarks often overlook the complex physical reasoning required for point-and-click puzzle games. This paper introduces Vision-Language Against The Incredible Machine (VLATIM), a benchmark designed to evaluate human-like logical problem-solving capabilities within the classic physics puzzle game The Incredible Machine 2 (TIM). Unlike existing benchmarks, VLATIM specifically targets the critical gap between high-level logical reasoning and continuous action spaces requiring precise mouse interactions. This benchmark is structured into five progressive parts, assessing capabilities that range from basic visual grounding and domain understanding to multi-step manipulation and full puzzle solving. Our results reveal a significant disparity between reasoning and execution. While large proprietary models demonstrate superior planning abilities, they struggle with precise visual grounding. Consequently, they do not yet show human-like problem-solving capabilities.
Maximilian Triebel, Marco Menner, Dominik Helfenstein
Institute of Artificial Intelligence, University of Stuttgart, Stuttgart, Germany.
Compositional generalization, the ability to understand and produce novel combinations of known components, remains a fundamental challenge for modern artificial intelligence. While few benchmarks exist, many focus on linguistic tasks and lack complex, explicit compositional structures. We introduce ClassicLogic, a new benchmark suite designed to evaluate an agent's ability to learn and compose problem-solving strategies. The benchmark consists of four classic logic puzzles: Sudoku, KenKen, Kakuro, and Futoshiki. Its core innovation is a hierarchical, explicit knowledge base for each game, where complex solving strategies are formally defined as compositions of simpler, foundational strategies. This structure allows for fine-grained evaluation of an agent's reasoning capabilities, from learning basic rules to applying multi-step compositional strategies to solve puzzles of increasing, mathematically validated difficulty. The open-source benchmark provides a challenging new testbed for advancing neuro-symbolic and other advanced AI reasoning systems.
Discovery---formulating novel generalizations---is a central part of the scientific process. Despite its importance, there is a gap in the current AI benchmark landscape, with few benchmarks directly probing the capacity for discovering new knowledge with experimentation in controlled environments where the objective is unknown. To address this gap, we release a new benchmark: DiG-bench (Discovery in Games). DiG-bench consists of a set of 70 independent games. Each game is encoded as a short string and has unique transformation rules that must be discovered through interaction and experimentation. The levels of the game present a series of challenges to test whether the rules have been discovered, where the win conditions for each level are also unknown. We provide games at seven tiers of difficulty for AI agents. The lowest tier is routinely solvable by multiple models, while the highest tier challenges the best models in agentic harnesses. All 70 games were solved by at least one human on first attempt. A subset of 21 games is released publicly, and the remainder is held private for secure evaluation.
Ruairidh M. Battleday, Kai Sandbrink, Jimi Cullen-Drohan +13
1Thinking About Thinking · 2Independent · University of Oxford +5