Package: juxai-s2
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 JAXJAX is a main dependency of JUX, and must be installed by user manually.pipinstall--upgrade"jax[cuda11_cudnn82]==0.4.7"-fhttps://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlYou 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:0Install JUXFinally, upgrade your pip and install JUX.pipinstall--upgradepip
pipinstalljuxai-s2UsageSeetutorial.ipynbfor a quick start. JUX is guaranteed to implement the same game logic asluxai_s2==3.0.0, if players' input actions are valid. When players' input actions are invalid, JUX and LuxAI-S2 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.LuxAI_S2 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 (luxai_s2==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.