| Name | juxai-s2 JSON |
| Version |
3.0.0
JSON |
| download |
| home_page | |
| Summary | JUX is a jax-accelerated engine for Lux-2022. |
| upload_time | 2023-09-27 03:35:26 |
| maintainer | |
| docs_url | None |
| author | |
| requires_python | <3.12,>=3.8 |
| 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. |
| keywords |
jax
luxai-s2
|
| VCS |
 |
| bugtrack_url |
|
| requirements |
No requirements were recorded.
|
| Travis-CI |
No Travis.
|
| coveralls test coverage |
No coveralls.
|
# 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 JAX
JAX is a main dependency of JUX, and must be installed by user manually.
```sh
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.
```sh
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.
```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==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.
| 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.
Raw data
{
"_id": null,
"home_page": "",
"name": "juxai-s2",
"maintainer": "",
"docs_url": null,
"requires_python": "<3.12,>=3.8",
"maintainer_email": "",
"keywords": "JAX,LuxAI-S2",
"author": "",
"author_email": "Qimai Li <qimaili@chaocanshu.ai>, Yuhao Jiang <yuhaojiang@chaocanshu.ai>",
"download_url": "https://files.pythonhosted.org/packages/44/95/965a89a2d5e6b6ecb7472812838c08aa6ae3c7373eaa69e77905274be3ea/juxai-s2-3.0.0.tar.gz",
"platform": null,
"description": "# JUX\nJUX 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.\n\n## Installation\n\n### Install JAX\nJAX is a main dependency of JUX, and must be installed by user manually.\n```sh\npip install --upgrade \"jax[cuda11_cudnn82]==0.4.7\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n```\n\nYou can test whether jax is installed successfully by running the following command.\n```sh\npython -c \"import jax.numpy as jnp; \\\n a = jnp.convolve(jnp.array([1, 2, 3]), jnp.array([0, 1, 0.5])); \\\n print(a); \\\n print(a.device());\"\n# You shall get something like this:\n# [0. 1. 2.5 4. 1.5]\n# gpu:0\n```\n\n### Install JUX\nFinally, upgrade your pip and install JUX.\n```sh\npip install --upgrade pip\npip install juxai-s2\n```\n\n## Usage\nSee [tutorial.ipynb](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.\n\n## Performance\nJUX 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.\n\nLuxAI_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.\n\n### Relative Throughput\nHere, 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.\n\n| Relative Throughput | | | | | | | | |\n|:-------------------- | ------------------:| ----------:| ---------:| ---------:| ---------:| ---------:| ---------:| ----------:|\n| GPU | GPU Mem. Bandwidth | Batch Size | UNITS=100 | UNITS=200 | UNITS=400 | UNITS=600 | UNITS=800 | UNITS=1000 |\n| A100-SXM4-40GB | 1555 GB/s | 20k | 1166x | 985x | 748x | 598x | 508x | 437x |\n| Tesla V100-SXM2-32GB | 900 GB/s | 20k | 783x | 647x | 480x | 375x | 317x | 269x |\n| Tesla T4 | 320 GB/s | 10k | 263x | 217x | 160x | 125x | 105x | 89x |\n| GTX 1660 Ti | 288 GB/s | 3k | 218x | 178x | 130x | 103x | 84x | 71x |\n\n\n| | Batch Size | Relative Throughput |\n|:----------------------------------------- | ----------:| -------------------:|\n| Intel\u00ae Core\u2122 i7-12700\t \t | 1 | 2.12x |\n| Intel\u00ae Xeon\u00ae Platinum 8255C CPU @ 2.50GHz\t| 1 | 1.00x |\n| Intel\u00ae Xeon\u00ae Gold 6133 @ 2.50GHz\t | 1 | 0.89x |\n\n\n\n### Absolute Throughput\nWe also report absolute throughput in steps per second here.\n| Throughput (steps/s) | | | | | | | |\n|:-------------------- | ----------:| ---------:| ---------:| ---------:| ---------:| ---------:| ----------:|\n| GPU | Batch Size | UNITS=100 | UNITS=200 | UNITS=400 | UNITS=600 | UNITS=800 | UNITS=1000 |\n| A100-SXM4-40GB | 20k | 363k | 307k | 233k | 186k | 158k | 136k |\n| Tesla V100-SXM2-32GB | 20k | 244k | 202k | 149k | 117k | 99k | 84k |\n| Tesla T4 | 10k | 82k | 68k | 50k | 39k | 33k | 28k |\n| GTX 1660 Ti | 3k | 68k | 56k | 40k | 32k | 26k | 22k |\n\n| | Batch Size | Throughput (steps/s) |\n|:----------------------------------------- | ----------:| --------------------:|\n| Intel\u00ae Core\u2122 i7-12700\t \t | 1 | 0.66k |\n| Intel\u00ae Xeon\u00ae Platinum 8255C CPU @ 2.50GHz\t| 1 | 0.31k |\n| Intel\u00ae Xeon\u00ae Gold 6133 @ 2.50GHz\t | 1 | 0.28k |\n\n## Contributing\nIf 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.\n",
"bugtrack_url": null,
"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. ",
"summary": "JUX is a jax-accelerated engine for Lux-2022.",
"version": "3.0.0",
"project_urls": {
"Bug Tracker": "https://github.com/RoboEden/jux/issues",
"repository": "https://github.com/RoboEden/jux"
},
"split_keywords": [
"jax",
"luxai-s2"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "8875da18b42894c9fb983f29bc9278c5d6129cd5c5acda029ece2efedbf54c6d",
"md5": "acd7620e5c91880056a784d585730be4",
"sha256": "4acd0a1cfb76e9016ee67a571e2bd50f338eed8e34d0ead6be2bafc7194a994b"
},
"downloads": -1,
"filename": "juxai_s2-3.0.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "acd7620e5c91880056a784d585730be4",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": "<3.12,>=3.8",
"size": 57172,
"upload_time": "2023-09-27T03:35:23",
"upload_time_iso_8601": "2023-09-27T03:35:23.427232Z",
"url": "https://files.pythonhosted.org/packages/88/75/da18b42894c9fb983f29bc9278c5d6129cd5c5acda029ece2efedbf54c6d/juxai_s2-3.0.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "4495965a89a2d5e6b6ecb7472812838c08aa6ae3c7373eaa69e77905274be3ea",
"md5": "8a78df1feaec9d866537684ce38caf28",
"sha256": "32a6a8985e1b7a0ee14c62efd7f348969d6d06402fa4369c108e0ff3b4120f53"
},
"downloads": -1,
"filename": "juxai-s2-3.0.0.tar.gz",
"has_sig": false,
"md5_digest": "8a78df1feaec9d866537684ce38caf28",
"packagetype": "sdist",
"python_version": "source",
"requires_python": "<3.12,>=3.8",
"size": 1137266,
"upload_time": "2023-09-27T03:35:26",
"upload_time_iso_8601": "2023-09-27T03:35:26.584672Z",
"url": "https://files.pythonhosted.org/packages/44/95/965a89a2d5e6b6ecb7472812838c08aa6ae3c7373eaa69e77905274be3ea/juxai-s2-3.0.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2023-09-27 03:35:26",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "RoboEden",
"github_project": "jux",
"travis_ci": false,
"coveralls": false,
"github_actions": false,
"tox": true,
"lcname": "juxai-s2"
}
