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 JAX

JAX is a main dependency of JUX, and must be installed by user manually.

pip install --upgrade "jax[cuda11_cudnn82]==0.4.7" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

You can test whether jax is installed successfully by running the following command.

python -c "import jax.numpy as jnp; \
    a = jnp.convolve(jnp.array([1, 2, 3]), jnp.array([0, 1, 0.5])); \
    print(a); \
    print(a.device());"
# You shall get something like this:
# [0.  1.  2.5 4.  1.5]
# gpu:0

Install JUX

Finally, upgrade your pip and install JUX.

pip install --upgrade pip
pip install juxai-s2

Usage

See tutorial.ipynb for a quick start. JUX is guaranteed to implement the same game logic as luxai_s2==3.0.0, if players' input actions are valid. When players' input actions are invalid, JUX and LuxAI-S2 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.

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.

Relative Throughput

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.

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.