Metadata-Version: 2.1
Name: juxai-s2
Version: 2.1.1
Summary: JUX is a jax-accelerated engine for Lux-2022.
Author-email: Qimai Li <qimaili@chaocanshu.ai>, Yuhao Jiang <yuhaojiang@chaocanshu.ai>
License: MIT License
        
        Copyright (c) 2022 Qimai Li, Yuhao Jiang.
        
        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
        
Project-URL: repository, https://github.com/RoboEden/jux
Project-URL: Bug Tracker, https://github.com/RoboEden/jux/issues
Keywords: JAX,LuxAI-S2
Classifier: Programming Language :: Python :: 3
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: License :: OSI Approved :: MIT License
Classifier: Topic :: Games/Entertainment :: Real Time Strategy
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: <3.11,>=3.7
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: luxai-s2 (==2.1.9)
Requires-Dist: chex
Provides-Extra: dev
Requires-Dist: toml ; extra == 'dev'
Requires-Dist: pre-commit ; extra == 'dev'
Requires-Dist: pytype ; extra == 'dev'
Requires-Dist: tox ; extra == 'dev'
Requires-Dist: tox-conda ; extra == 'dev'
Requires-Dist: pytest ; extra == 'dev'
Requires-Dist: pytest-cov ; extra == 'dev'
Requires-Dist: isort ; extra == 'dev'
Requires-Dist: autoflake ; extra == 'dev'
Requires-Dist: yapf ; extra == 'dev'
Requires-Dist: bump2version ; extra == 'dev'
Requires-Dist: rich ; extra == 'dev'
Requires-Dist: pandas ; extra == 'dev'
Requires-Dist: ipywidgets (<8.0,>=7.0) ; extra == 'dev'
Requires-Dist: notebook ; extra == 'dev'
Requires-Dist: build ; extra == 'dev'
Requires-Dist: twine ; extra == 'dev'
Provides-Extra: torch
Requires-Dist: torch ; extra == 'torch'

# JUX
JUX is a <ins>J</ins>ax-accelerated game core for L<ins>ux</ins> 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.
```sh
conda install -c nvidia cuda-nvcc cuda-python
conda install cudnn
```
For docker users, you can use the [NVIDIA CUDA docker image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/cuda) or the [PyTorch docker image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch), which has all of them ready and compatible with each other.

### Install JAX
Please follow the [official installation guide](https://github.com/google/jax#installation) 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]"`.

### Install JUX
Finally, upgrade your pip and install JUX.
```sh
pip install --upgrade pip
pip install juxai-s2
```

## Usage
See [tutorial.ipynb](tutorial.ipynb) for a quick start. JUX is guaranteed to implement the same game logic as `luxai_s2==2.1.9`, 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.

| Relative Throughput  |                    |            |           |           |           |           |           |            |
|:-------------------- | ------------------:| ----------:| ---------:| ---------:| ---------:| ---------:| ---------:| ----------:|
| GPU                  | GPU Mem. Bandwidth | Batch Size | UNITS=100 | UNITS=200 | UNITS=400 | UNITS=600 | UNITS=800 | UNITS=1000 |
| A100-SXM4-40GB       |          1555 GB/s |        20k |     1166x |      985x |      748x |      598x |      508x |       437x |
| Tesla V100-SXM2-32GB |           900 GB/s |        20k |      783x |      647x |      480x |      375x |      317x |       269x |
| Tesla T4             |           320 GB/s |        10k |      263x |      217x |      160x |      125x |      105x |        89x |
| GTX 1660 Ti          |           288 GB/s |         3k |      218x |      178x |      130x |      103x |       84x |        71x |


|                                           | Batch Size | Relative Throughput |
|:----------------------------------------- | ----------:| -------------------:|
| Intel® Core™ i7-12700	 	                |          1 |               2.12x |
| Intel® Xeon® Platinum 8255C CPU @ 2.50GHz	|          1 |               1.00x |
| Intel® Xeon® Gold 6133 @ 2.50GHz	        |          1 |               0.89x |



### Absolute Throughput
We also report absolute throughput in steps per second here.
| Throughput (steps/s) |            |           |           |           |           |           |            |
|:-------------------- | ----------:| ---------:| ---------:| ---------:| ---------:| ---------:| ----------:|
| GPU                  | Batch Size | UNITS=100 | UNITS=200 | UNITS=400 | UNITS=600 | UNITS=800 | UNITS=1000 |
| A100-SXM4-40GB       |        20k |      363k |      307k |      233k |      186k |      158k |       136k |
| Tesla V100-SXM2-32GB |        20k |      244k |      202k |      149k |      117k |       99k |        84k |
| Tesla T4             |        10k |       82k |       68k |       50k |       39k |       33k |        28k |
| GTX 1660 Ti          |         3k |       68k |       56k |       40k |       32k |       26k |        22k |

|                                           | Batch Size | Throughput (steps/s) |
|:----------------------------------------- | ----------:| --------------------:|
| Intel® Core™ i7-12700	 	                |          1 |                0.66k |
| Intel® Xeon® Platinum 8255C CPU @ 2.50GHz	|          1 |                0.31k |
| Intel® Xeon® Gold 6133 @ 2.50GHz	        |          1 |                0.28k |

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