Package: juxai2022
JUXJUX is a Jax-accelerated game core for Lux AI Challenge Season 2, aimed to maximize game environment throughput for reinforcement learning (RL) training.InstallationInstall dependenciesOne 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.$condainstall-cnvidiacuda-nvcccuda-python$condainstallcudnnFor docker users, you can use theNVIDIA CUDA docker imageor thePyTorch docker image, which has all of them ready and compatible with each other.Install JUXFirst, you need to clone the repository.git clone https://github.com/RoboEden/jux.gitcd juxThen, upgrade your pip and install JUX.$pipinstall--upgradepip$pipinstall./jux-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlFor PyTorch users, you can install JUX with optional dependencies for PyTorch.$pipinstall./jux[torch]-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlUsageSeetutorial.ipynbfor a quick start. JUX is guaranteed to implement the same game logic asluxai2022==1.1.4, if players' input actions are valid. When players' input actions are invalid, JUX and LuxAI2022 may process them differently.PerformanceJUX 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, becausejax.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, soMAX_N_UNITS=1500is a safe choice. However, we found that no player builds more than 200 units by watching game replays, soMAX_N_UNITS=200is a practical choice.Relative ThroughputHere, we report the relative throughput over the original Python implementation (luxai2022==1.1.3), on several different GPUs with differentMAX_N_UNITSsettings. 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.Relative ThroughputGPUGPU Mem. BandwidthBatch SizeUNITS=100UNITS=200UNITS=400UNITS=600UNITS=800UNITS=1000A100-SXM4-40GB1555 GB/s20k1166x985x748x598x508x437xTesla V100-SXM2-32GB900 GB/s20k783x647x480x375x317x269xTesla T4320 GB/s10k263x217x160x125x105x89xGTX 1660 Ti288 GB/s3k218x178x130x103x84x71xBatch SizeRelative ThroughputIntel® Core™ i7-1270012.12xIntel® Xeon® Platinum 8255C CPU @ 2.50GHz11.00xIntel® Xeon® Gold 6133 @ 2.50GHz10.89xAbsolute ThroughputWe also report absolute throughput in steps per second here.Throughput (steps/s)GPUBatch SizeUNITS=100UNITS=200UNITS=400UNITS=600UNITS=800UNITS=1000A100-SXM4-40GB20k363k307k233k186k158k136kTesla V100-SXM2-32GB20k244k202k149k117k99k84kTesla T410k82k68k50k39k33k28kGTX 1660 Ti3k68k56k40k32k26k22kBatch SizeThroughput (steps/s)Intel® Core™ i7-1270010.66kIntel® Xeon® Platinum 8255C CPU @ 2.50GHz10.31kIntel® Xeon® Gold 6133 @ 2.50GHz10.28kContributingIf you find any bugs or have any suggestions, please feel free to open an issue or a pull request. SeeCONTRIBUTING.mdfor setting up a developing environment.