JUX

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

Installation

Install dependencies

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.

Install JUX

First, you need to clone the repository.

git clone https://github.com/RoboEden/jux.git
cd jux

Then, upgrade your pip and install JUX.

$ pip install --upgrade pip
$ pip install ./jux -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

For PyTorch users, you can install JUX with optional dependencies for PyTorch.

$ pip install ./jux[torch] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Usage

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

Performance

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.

LuxAI2022 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.

Relative Throughput

Here, we report the relative throughput over the original Python implementation (luxai2022==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.

Absolute Throughput

We also report absolute throughput in steps per second here.

Contributing

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.