PuzzleJAX: A Benchmark for Reasoning and Learning
This provides a new benchmark for evaluating AI reasoning and learning abilities, though it is incremental as it builds on existing puzzle game engines.
The authors introduced PuzzleJAX, a GPU-accelerated puzzle game engine and description language for benchmarking tree search, reinforcement learning, and LLM reasoning, validating it on hundreds of existing games to cover an expressive and human-relevant task space.
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.