JAxtar is a project with a JAX-native implementation of parallelizable a A* & Q* solver for neural heuristic search research. This project is inspired by mctx from Google DeepMind. If MCTS can be implemented entirely in pure JAX, why not A*?
This project provides the following algorithms:
- A*: Standard A* Search
- A* Deferred: A* Search with deferred node expansion (for heavy expansion costs)
- Bi-A*: Bidirectional A* Search
- Bi-A* Deferred: Bidirectional A* Search with deferred node expansion
- ID-A*: Iterative Deepening A* Search
- Beam Search: Standard Beam Search
- Q*: Q* Search (A* with learned Q-values)
- Bi-Q*: Bidirectional Q* Search
- ID-Q*: Iterative Deepening Q* Search (Experimental)
- Q-Beam: Beam Search with learned Q-values
MCTS, or tree search, is used in many RL algorithmic techniques, starting with AlphaGo, but graph search (not tree search) doesn't seem to have received much attention. Nevertheless, there are puzzle‐solving algorithms that use neural heuristics like DeepcubeA with A* or Q*(graph search).
However, the most frustrating aspect of my brief research (MSc) in this area is the time it takes to pass information back and forth between the GPU and CPU. When using a neural heuristic to evaluate a single node, the communication between the CPU and GPU, rather than the computation itself, can consume between 50% and 80% of the total processing time. Because of this communication overhead, DeepcubeA batches multiple nodes concurrently, which appears to work quite well.
However, these issues indicate that a more fundamental solution is needed. This led me to search for ways to perform A* directly on the GPU, but I discovered that most implementations suffer from the following problems.
- Many are written in pure C and CUDA, which is not well-suited for machine learning research.
- Some are written in JAX or PyTorch, but these are often limited to 2D grid environments or connectivity matrices, and cannot scale to an infinite number of different states that cannot all be held in memory.
- The implementation itself is often dependent on the specific definition of the state or problem.
To address these challenges, I decided to develop code based on the following principles:
- Pure JAX implementation
- Specifically for machine learning research.
- JAX-native priority queue
- The A* algorithm necessitates a priority queue to process nodes based on the lowest combined cost and heuristic estimate.
- However, standard Python heaps use lists, which are not JIT-compilable in JAX. Thus, a JAX-iterable heap is necessary.
- Hashable state representation and a hashtable for JAX operations.
- This is crucial for tracking node status (open/closed) in A* and efficiently retrieving parent state information.
- Hashing is optional for simple, indexable states. But for complex or infinite state spaces, hashing becomes essential for efficient indexing and retrieval of unique states.
- Fully batched and parallelized operations
- GPUs provide massive parallelism but have slower cores than CPUs. Therefore, algorithms for GPUs must be highly parallelized to leverage their architecture.
- Puzzle-agnostic implementation
- The implementation should be general enough to handle any puzzle with a defined state and action space.
- This generality enables wider research and allows for formalizing 'strict' behaviors in future implementations.
This project features specially written components, including:
Xtructure: A pip package providing JAX-compatible hash and priority queue implementations, originally developed as part of this project and later separated. This package includes:- a hashtable for parallel lookup and insertion operations
- a priority queue that supports batching, push, and pop operations
PuXle: All puzzle implementations have been moved to this separate high-performance library for parallelized puzzle environments built on JAX- PuXle world models registered for latent-space heuristic search
- Network heuristics and Q-functions designed for JIT-compilable integration with A* & Q* algorithm
- a fully JIT-compiled A* & Q* algorithm for puzzles
This project was quite challenging to develop, and it felt like performing acrobatics with JAX. However, I managed to create a fully functional version, and hopefully it will inspire you to discover something amazing as you delve into JAX.
JAxtar uses pyproject.toml for installation and command registration. PuXle and Xtructure are
installed from the pinned Git revisions in tool.uv.sources. From the JAxtar repository:
uv sync --extra gpu
uv run jaxtar --help
uv run jaxtar test --help
uv run astar_d --helpFor the complete GPU training pipeline, including logging backends and evaluation tools:
uv sync --extra trainThe console commands are checkout-only. External distributions intentionally contain only the
JAxtar core package; generic top-level modules such as cli, config, and helpers are not
packaged.
Omit --extra gpu for CPU-only use, or replace it with --extra tpu on TPU. Evaluation tools are
installed by default. The train extra adds the GPU runtime and all logging backends, so it
conflicts with tpu.
For detailed information on all available commands and their options, please refer to the official documentation.
We can find the optimal path using a jittable, batched A* search as shown below. This is not a super blazingly fast result, but it can be well integrated with heuristics using neural networks.
The following speed benchmarks were measured on an Nvidia RTX 5090 hosted at vast.ai.
You can easily test it yourself with the colab link below.
$ uv run astar
╭─────────────── Seed 0 ────────────────╮
│ Start State Target State │
│ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ │
│ ┃ 9 ┃ E ┃ 6 ┃ 8 ┃ ┃ 1 ┃ 2 ┃ 3 ┃ 4 ┃ │
│ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │
│ ┃ D ┃ 4 ┃ 7 ┃ ┃ ┃ 5 ┃ 6 ┃ 7 ┃ 8 ┃ │
│ ┣━━━╋━━━╋━━━╋━━━┫ → ┣━━━╋━━━╋━━━╋━━━┫ │
│ ┃ B ┃ 1 ┃ A ┃ C ┃ ┃ 9 ┃ A ┃ B ┃ C ┃ │
│ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │
│ ┃ 5 ┃ 3 ┃ F ┃ 2 ┃ ┃ D ┃ E ┃ F ┃ ┃ │
│ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ │
│ Dist: 34.00 │
╰───────────────────────────────────────╯
Search Result for Seed 0
┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
│ Status │ Solution Found │
│ Cost │ 50.0 │
│ Search Time │ 0.12 s │
│ Search States │ 738K │
│ States/s │ 6.15M │
└───────────────┴────────────────┘$ uv run astar --vmap_size 20
╭──────────────────────────────── Vmapped Search Setup ─────────────────────────────────╮
│ ╭────────────────────────────────── Batched State ──────────────────────────────────╮ │
│ │ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ │ │
│ │ ┃ 9 ┃ E ┃ 6 ┃ 8 ┃ ┃ 9 ┃ E ┃ 6 ┃ 8 ┃ ┃ 9 ┃ E ┃ 6 ┃ 8 ┃ ┃ 9 ┃ E ┃ 6 ┃ 8 ┃ │ │
│ │ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │ │
│ │ ┃ D ┃ 4 ┃ 7 ┃ ┃ ┃ D ┃ 4 ┃ 7 ┃ ┃ ┃ D ┃ 4 ┃ 7 ┃ ┃ ┃ D ┃ 4 ┃ 7 ┃ ┃ │ │
│ │ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ... ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │ │
│ │ ┃ B ┃ 1 ┃ A ┃ C ┃ ┃ B ┃ 1 ┃ A ┃ C ┃ ┃ B ┃ 1 ┃ A ┃ C ┃ ┃ B ┃ 1 ┃ A ┃ C ┃ │ │
│ │ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │ │
│ │ ┃ 5 ┃ 3 ┃ F ┃ 2 ┃ ┃ 5 ┃ 3 ┃ F ┃ 2 ┃ ┃ 5 ┃ 3 ┃ F ┃ 2 ┃ ┃ 5 ┃ 3 ┃ F ┃ 2 ┃ │ │
│ │ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ │ │
│ ╰────────────────────────────────── shape: (20,) ───────────────────────────────────╯ │
│ │
│ ↓ │
│ │
│ ╭─────────────────────────────── Batched SolveConfig ───────────────────────────────╮ │
│ │ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ ┏━━━┳━━━┳━━━┳━━━┓ │ │
│ │ ┃ 1 ┃ 2 ┃ 3 ┃ 4 ┃ ┃ 1 ┃ 2 ┃ 3 ┃ 4 ┃ ┃ 1 ┃ 2 ┃ 3 ┃ 4 ┃ ┃ 1 ┃ 2 ┃ 3 ┃ 4 ┃ │ │
│ │ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │ │
│ │ ┃ 5 ┃ 6 ┃ 7 ┃ 8 ┃ ┃ 5 ┃ 6 ┃ 7 ┃ 8 ┃ ┃ 5 ┃ 6 ┃ 7 ┃ 8 ┃ ┃ 5 ┃ 6 ┃ 7 ┃ 8 ┃ │ │
│ │ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ... ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │ │
│ │ ┃ 9 ┃ A ┃ B ┃ C ┃ ┃ 9 ┃ A ┃ B ┃ C ┃ ┃ 9 ┃ A ┃ B ┃ C ┃ ┃ 9 ┃ A ┃ B ┃ C ┃ │ │
│ │ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ ┣━━━╋━━━╋━━━╋━━━┫ │ │
│ │ ┃ D ┃ E ┃ F ┃ ┃ ┃ D ┃ E ┃ F ┃ ┃ ┃ D ┃ E ┃ F ┃ ┃ ┃ D ┃ E ┃ F ┃ ┃ │ │
│ │ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ ┗━━━┻━━━┻━━━┻━━━┛ │ │
│ ╰────────────────────────────────── shape: (20,) ───────────────────────────────────╯ │
│ │
╰───────────────────────────────────────────────────────────────────────────────────────╯
Vmapped Search Result
┏━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩
│ Search Time │ 3.18s (x26.5/20) │
│ Search States │ 14.8M (4.64M states/s) │
│ Speedup │ x0.8 │
│ Solutions Found │ 100.00% │
└─────────────────┴────────────────────────┘$ uv run astar -nn -h -p rubikscube -w 0.2
...
Search Result for Seed 0
┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
│ Status │ Solution Found │
│ Cost │ 22.0 │
│ Search Time │ 0.49 s │
│ Search States │ 1.85M │
│ States/s │ 3.78M │
└───────────────┴────────────────┘$ uv run qstar -nn -h -p rubikscube -w 0.2
...
Search Result for Seed 0
┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
│ Status │ Solution Found │
│ Cost │ 22.0 │
│ Search Time │ 0.16 s │
│ Search States │ 1.79M │
│ States/s │ 10.9M │
└───────────────┴────────────────┘$ uv run qstar -p rubikscube_world_model_optimized -nn -w 0.6
...
Search Result for Seed 0
┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━┩
│ Status │ Solution Found │
│ Cost │ 22.0 │
│ Search Time │ 0.17 s │
│ Search States │ 1.82M │
│ States/s │ 10.5M │
└───────────────┴────────────────┘| Rubikscube | Slidepuzzle | Lightsout | Sokoban |
|---|---|---|---|
 |
 |
 |
 |
 |
 |
 |
 |
| Maze | Pancake Sorting | Hanoi | TopSpin |
|---|---|---|---|
 |
 |
 |
 |
 |
 |
 |
 |
These types of puzzles are not strictly the kind that are typically solved with A*, but after some simple testing, it turns out that, depending on how the problem is defined, they can be solved. Furthermore, this approach can be extended to TSP and countless other COP problems, provided that with a good heuristic. The training method will need to be investigated further.
| Dotknot | TSP |
|---|---|
 |
 |
 |
 |
The world-model implementation and training APIs live in PuXle. JAxtar registers those models to perform A* & Q* search in the learned latent space, following "Learning Discrete World Models for Heuristic Search".
| Terminal View | Rubiks Cube | Rubiks Cube Reversed | Sokoban |
|---|---|---|---|
 |
 |
 |
 |
Please use this citation to reference this project.
@software{kyuseokjung2024jaxtar,
title = {JA^{xtar}: GPU-accelerated Batched parallel A* & Q* solver in pure JAX!},
author = {Kyuseok Jung},
url = {https://github.com/tinker495/JAxtar},
year = {2024},
}