IROS 2026 | Paper | Project Page
PAKR is a massively parallel kinodynamic planner that uses JAX and the XLA compiler to achieve GPU-accelerated sampling-based planning entirely through standard Python tooling. Combined with the AO-x meta-algorithm, it achieves asymptotic optimality via fast iterative replanning.
Build and run the Docker container with GPU access:
docker build -t pakr .
docker run --gpus all -it pakrThis is for the 9 kinopax problems. Run rrt.py. Inside it you can change batch size, branching factor, env, and dynamic model. Can also use the flags --env and --motion. Envs for tree, house, and narrow. For quadrotor, use quadtree, quadhouse, and quadnarrow envs since everything is scaled 100x.
To run the AO-x version of the kinopax problems, run ao/aorrt.py with the same parameters that can be modified. There is an additional parameter for the cost threshold. The reason for this is because JAX does not allow us to interrupt a program, so where normally if the best cost is too low and the resource allotment is exceeded without finding a solution and the program exits, here we cannot interrupt and the entire program hangs. So instead, we have to set a cost threshold that can find a solution in a reasonable time, so if we get "lucky" with initial runs and get a very low best cost, it exits earlier.
Can be found in ./benchmarks. Working experiments include acrobot.py, cartpole.py, di2d.py (2d double integrator), qc.py, unicycle_acc.py, and unicycle_vel.py. Their corresponding ao scripts are also in the same folder. There are visualization scripts for unicycle_vel and di2d to show the solution trajectories.
MJX models are found in ./models. There are only 2 mjx experiments: eeonly, which is the block push, and cartpole. Acrobot does not use MJX but a simpler custom propagator, but can use the mjx models to visualize. To visualize mjx solutions, first run parse_solutions_mjx.py. This reruns the mjx propagator with the solution trajectory and saves each state. This is necessary because the mjx simulator has significantly different outputs compared to the mujoco sim, so we can't just run the solution in mujoco, but we want to use the mujoco viewer (mjx does not have one since everything is batched)
PAKR is designed to be modular. To define your own system, you replace these core functions:
- Propagation — forward-simulates the system given a state and action over a time step
- Validation — determines whether a state is valid. This can mean within joint/vel limits, or collision free if needed
- Distance — defines the metric used for nearest-neighbor selection in the tree.
- Sample States — determines the range of values to sample from for each dim
- Sample Actions — same thing but for valid actions
- Goal Condition — can be as simple as distance < threshold, but sometimes we dont want to check all dimensions
- Params — things such as simulator step size, bounds etc.
- Model — For mjx experiments, create your own xml.
In the files in ./benchmark, all of these custom functions are in the same file. For the original kinopax experiments, these function were in helper.py, with value ranges defined in params.py. It's not clean but there are just too many different experiments to make it modular. Easier to just copy paste what's needed so each file is its independent module.
@inproceedings{gao2026fastasymptoticallyoptimalkinodynamic,
title={Fast Asymptotically Optimal Kinodynamic Planning via Vectorization},
author={Gao, Yitian and Lu, Andrew and Kingston, Zachary},
booktitle={Proceedings of the IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
year={2026},
note={To appear},
url={https://arxiv.org/abs/2607.03987}
}