JUX is a Jax-accelerated game core for Lux AI Challenge Season 2, aimed to maximize game environment throughput for reinforcement learning (RL) training.

One of the main dependencies is JAX, which in turn relies on NVCC, CUDA Toolkit and cuDNN. There are two ways to get them ready, either by conda or docker (recommended).

For conda users, you can install them with the following commands.

conda install -c nvidia cuda-nvcc cuda-python
conda install cudnn

For docker users, you can use the NVIDIA CUDA docker image or the PyTorch docker image, which has all of them ready and compatible with each other.

Please follow the official installation guide to install JAX. Note: JAX must be compatible with your cuDNN.

• If cudnn >= 8.2 and < 8.6, please install "jax[cuda11_cudnn82]".
• If cudnn >= 8.6, please install "jax[cuda11_cudnn86]".

Finally, upgrade your pip and install JUX.

pip install --upgrade pip
pip install juxai-s2

See tutorial.ipynb for a quick start. JUX is guaranteed to implement the same game logic as luxai_s2==2.1.0, if players' input actions are valid. When players' input actions are invalid, JUX and LuxAI-S2 may process them differently.

JUX maps all game logic to array operators in JAX so that we can harvest the computational power of modern GPUs and support tons of environments running in parallel. We benchmarked JUX on several different GPUs, and increased the throughput by hundreds to thousands of times, compared with the original single-thread Python implementation.

LuxAI_S2 is a game with a dynamic number of units, making it hard to be accelerated by JAX, because jax.jit() only supports arrays with static shapes. As a workaround, we allocate a large buffer with static length to store units. The buffer length (buf_cfg.MAX_N_UNITS) greatly affects the performance. Theoretically, no player can build more than 1500 units under current game configs, so MAX_N_UNITS=1500 is a safe choice. However, we found that no player builds more than 200 units by watching game replays, so MAX_N_UNITS=200 is a practical choice.

Here, we report the relative throughput over the original Python implementation (luxai_s2==1.1.3), on several different GPUs with different MAX_N_UNITS settings. The original single-thread Python implementation running on an 8255C CPU serves as the baseline. We can observe that the throughput is proportional to GPU memory bandwidth because the game logic is mostly memory-bound, not compute-bound. Byte-access operators take a large portion of the game logic in JUX implementation.

We also report absolute throughput in steps per second here.

If you find any bugs or have any suggestions, please feel free to open an issue or a pull request. See CONTRIBUTING.md for setting up a developing environment.