Repository: sail-sg/envpool (language: python)

Issue:
[Feature Request] Speed test for mujoco



Discussion / hints:

Fix patch:
diff --git a/Makefile b/Makefile
index a3b0dccd..cc81370d 100644
--- a/Makefile
+++ b/Makefile
@@ -1,11 +1,12 @@
-SHELL        = /bin/bash
-PROJECT_NAME = envpool
-PYTHON_FILES = $(shell find . -type f -name "*.py")
-CPP_FILES    = $(shell find $(PROJECT_NAME) -type f -name "*.h" -o -name "*.cc")
-COMMIT_HASH  = $(shell git log -1 --format=%h)
-COPYRIGHT    = "Garena Online Private Limited"
-BAZELOPT     =
-PATH         := $(HOME)/go/bin:$(PATH)
+SHELL          = /bin/bash
+PROJECT_NAME   = envpool
+PROJECT_FOLDER = $(PROJECT_NAME) third_party examples benchmark
+PYTHON_FILES   = $(shell find . -type f -name "*.py")
+CPP_FILES      = $(shell find $(PROJECT_NAME) -type f -name "*.h" -o -name "*.cc")
+COMMIT_HASH    = $(shell git log -1 --format=%h)
+COPYRIGHT      = "Garena Online Private Limited"
+BAZELOPT       =
+PATH           := $(HOME)/go/bin:$(PATH)
 
 # installation
 
@@ -111,7 +112,7 @@ bazel-clean: bazel-install
 # documentation
 
 addlicense: addlicense-install
-	addlicense -c $(COPYRIGHT) -l apache -y 2022 -check $(PROJECT_NAME) third_party examples
+	addlicense -c $(COPYRIGHT) -l apache -y 2022 -check $(PROJECT_FOLDER)
 
 docstyle: doc-install
 	pydocstyle $(PROJECT_NAME) && doc8 docs && cd docs && make html SPHINXOPTS="-W"
@@ -125,6 +126,10 @@ spelling: doc-install
 doc-clean:
 	cd docs && make clean
 
+doc-benchmark:
+	pandoc benchmark/README.md --from markdown --to rst -s -o docs/pages/benchmark.rst --columns 1000
+	cd benchmark && ./plot.py --suffix png && mv *.png ../docs/_static/images/throughput
+
 lint: buildifier flake8 py-format clang-format cpplint clang-tidy mypy docstyle spelling
 
 format: py-format-install clang-format-install buildifier-install addlicense-install
@@ -132,7 +137,7 @@ format: py-format-install clang-format-install buildifier-install addlicense-ins
 	yapf -ir $(PYTHON_FILES)
 	clang-format-11 -style=file -i $(CPP_FILES)
 	buildifier -r -lint=fix .
-	addlicense -c $(COPYRIGHT) -l apache -y 2022 $(PROJECT_NAME) third_party examples
+	addlicense -c $(COPYRIGHT) -l apache -y 2022 $(PROJECT_FOLDER)
 
 # Build docker images
 
@@ -163,4 +168,3 @@ release-test2:
 	cd examples && python3 make_env.py && python3 env_step.py
 
 release-test: release-test1 release-test2
-
diff --git a/README.md b/README.md
index 2c4bf664..73d22d15 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@
 [![GitHub forks](https://img.shields.io/github/forks/sail-sg/envpool)](https://github.com/sail-sg/envpool/network)
 [![GitHub license](https://img.shields.io/github/license/sail-sg/envpool)](https://github.com/sail-sg/envpool/blob/master/LICENSE)
 
-**EnvPool** is a C++-based batched environment pool with pybind11 and thread pool. It has high performance (\~1M raw FPS in DGX on Atari games) and compatible APIs (supports both gym and dm\_env, both sync and async, both single and multi player environment). Currently it supports:
+**EnvPool** is a C++-based batched environment pool with pybind11 and thread pool. It has high performance (\~1M raw FPS with Atari games, \~3M raw FPS with Mujoco simulator on DGX-A100) and compatible APIs (supports both gym and dm\_env, both sync and async, both single and multi player environment). Currently it supports:
 
 - [x] Atari games
 - [x] Mujoco
@@ -32,13 +32,13 @@ Here are EnvPool's several highlights:
 - Support both synchronous execution and asynchronous execution;
 - Support both single player and multi-player environment;
 - Easy C++ developer API to add new envs;
-- **1 Million** Atari frames per second simulation with 256 CPU cores, **~13x** throughput of Python subprocess-based vector env;
+- **1 Million** Atari frames / **3 Miillion** Mujoco steps per second simulation with 256 CPU cores, **~20x** throughput of Python subprocess-based vector env;
 - **~3x** throughput of Python subprocess-based vector env on low resource setup like 12 CPU cores;
 - Comparing with existing GPU-based solution ([Brax](https://github.com/google/brax) / [Isaac-gym](https://developer.nvidia.com/isaac-gym)), EnvPool is a **general** solution for various kinds of speeding-up RL environment parallelization;
 - Compatible with some existing RL libraries, e.g., [Stable-Baselines3](https://github.com/DLR-RM/stable-baselines3), [Tianshou](https://github.com/thu-ml/tianshou), or [CleanRL](https://github.com/vwxyzjn/cleanrl).
   - Stable-Baselines3 [`Pendulum-v0` example](https://github.com/sail-sg/envpool/blob/master/examples/sb3_examples/ppo.py);
   - Tianshou [`CartPole` example](https://github.com/sail-sg/envpool/blob/master/examples/tianshou_examples/cartpole_ppo.py) and [`Pendulum-v0` example](https://github.com/sail-sg/envpool/blob/master/examples/tianshou_examples/pendulum_ppo.py);
-  - CleanRL [`Pong-v5` example](https://github.com/sail-sg/envpool/blob/master/examples/cleanrl_examples/ppo_atari_envpool.py) (solving Pong in 5 mins ([tracked experiment](https://wandb.ai/costa-huang/cleanRL/runs/opk2dmta)));
+  - CleanRL [`Pong-v5` example](https://github.com/sail-sg/envpool/blob/master/examples/cleanrl_examples/ppo_atari_envpool.py) ([Solving Pong in 5 mins](https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/#solving-pong-in-5-minutes-with-ppo--envpool) ([tracked experiment](https://wandb.ai/costa-huang/cleanRL/runs/opk2dmta)));
 - Support [customized C++ environment integration](https://envpool.readthedocs.io/en/latest/pages/env.html).
 
 ## Installation
@@ -70,35 +70,38 @@ Please refer to the [guideline](https://envpool.readthedocs.io/en/latest/pages/b
 
 The tutorials and API documentation are hosted on [envpool.readthedocs.io](https://envpool.readthedocs.io).
 
-The example scripts are under [examples/](https://github.com/sail-sg/envpool/tree/master/examples) folder.
+The example scripts are under [examples/](https://github.com/sail-sg/envpool/tree/master/examples) folder; benchmark scripts are under [benchmark/](https://github.com/sail-sg/envpool/tree/master/benchmark) folder.
 
 ## Benchmark Results
 
-We perform our benchmarks with ALE Atari environment (with environment wrappers) on different hardware setups, including a TPUv3-8 virtual machine (VM) of 96 CPU cores and 2 NUMA nodes, and an NVIDIA DGX-A100 of 256 CPU cores with 8 NUMA nodes. Baselines include 1) naive Python for-loop; 2) the most popular RL environment parallelization execution by Python subprocess, e.g., [gym.vector_env](https://github.com/openai/gym/blob/master/gym/vector/vector_env.py); 3) to our knowledge, the fastest RL environment executor [Sample Factory](https://github.com/alex-petrenko/sample-factory) before EnvPool. 
+We perform our benchmarks with ALE Atari environment (with environment wrappers) and Mujoco environment on different hardware setups, including a TPUv3-8 virtual machine (VM) of 96 CPU cores and 2 NUMA nodes, and an NVIDIA DGX-A100 of 256 CPU cores with 8 NUMA nodes. Baselines include 1) naive Python for-loop; 2) the most popular RL environment parallelization execution by Python subprocess, e.g., [gym.vector_env](https://github.com/openai/gym/blob/master/gym/vector/vector_env.py); 3) to our knowledge, the fastest RL environment executor [Sample Factory](https://github.com/alex-petrenko/sample-factory) before EnvPool. 
 
-We report EnvPool performance with sync mode, async mode, and NUMA + async mode, compared with the baselines on different number of workers (i.e., number of CPU cores). As we can see from the results, EnvPool achieves significant improvements over the baselines on all settings. On the high-end setup, EnvPool achieves 1 Million frames per second on 256 CPU cores, which is 13.3x of the `gym.vector_env` baseline. On a typical PC setup with 12 CPU cores, EnvPool's throughput is 2.8x of `gym.vector_env`.
+We report EnvPool performance with sync mode, async mode, and NUMA + async mode, compared with the baselines on different number of workers (i.e., number of CPU cores). As we can see from the results, EnvPool achieves significant improvements over the baselines on all settings. On the high-end setup, EnvPool achieves 1 Million frames per second with Atari and 3 Million frames per second with Mujoco on 256 CPU cores, which is 14.9x / 19.6x of the `gym.vector_env` baseline. On a typical PC setup with 12 CPU cores, EnvPool's throughput is 3.1x / 2.9x of `gym.vector_env`.
 
-Our benchmark script is in [examples/benchmark.py](https://github.com/sail-sg/envpool/blob/master/examples/benchmark.py). The detail configurations of 4 types of system are:
+|  Atari Highest FPS   | Laptop (12) | Workstation (32) | TPU-VM (96) | DGX-A100 (256) |
+| :------------------: | :---------: | :--------------: | :---------: | :------------: |
+|       For-loop       |    4,893    |      7,914       |    3,993    |     4,640      |
+|      Subprocess      |   15,863    |      47,699      |   46,910    |     71,943     |
+|    Sample-Factory    |   28,216    |     138,847      |   222,327   |    707,494     |
+|    EnvPool (sync)    |   37,396    |     133,824      |   170,380   |    427,851     |
+|   EnvPool (async)    | **49,439**  |   **200,428**    |   359,559   |    891,286     |
+| EnvPool (numa+async) |      /      |        /         | **373,169** | **1,069,922**  |
 
-- Personal laptop: 12 core `Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz` 
-- TPU-VM: 96 core `Intel(R) Xeon(R) CPU @ 2.00GHz`
-- Apollo: 96 core `AMD EPYC 7352 24-Core Processor`
-- DGX-A100: 256 core `AMD EPYC 7742 64-Core Processor`
-
-|     Highest FPS      | Laptop (12) | TPU-VM (96) | Apollo (96) | DGX-A100 (256) |
-| :------------------: | :---------: | :---------: | :---------: | :------------: |
-|       For-loop       |    4,876    |    3,817    |    4,053    |     4,336      |
-|      Subprocess      |   18,249    |   42,885    |   19,560    |     79,509     |
-|    Sample Factory    |   27,035    |   192,074   |   262,963   |    639,389     |
-|    EnvPool (sync)    |   40,791    |   175,938   |   159,191   |    470,170     |
-|   EnvPool (async)    | **50,513**  |   352,243   |   410,941   |    845,537     |
-| EnvPool (numa+async) |      /      | **367,799** | **458,414** | **1,060,371**  |
+|  Mujoco Highest FPS  | Laptop (12) | Workstation (32) | TPU-VM (96) | DGX-A100 (256) |
+| :------------------: | :---------: | :--------------: | :---------: | :------------: |
+|       For-loop       |   12,861    |      20,298      |   10,474    |     11,569     |
+|      Subprocess      |   36,586    |     105,432      |   87,403    |    163,656     |
+|    Sample-Factory    |   62,510    |     309,264      |   461,515   |   1,573,262    |
+|    EnvPool (sync)    |   66,622    |     380,950      |   296,681   |    949,787     |
+|   EnvPool (async)    | **105,126** |   **582,446**    |   887,540   |   2,363,864    |
+| EnvPool (numa+async) |      /      |        /         | **896,830** | **3,134,287**  |
 
 <p float="center">
-<img width="49%" height="auto" src="https://i.imgur.com/wHu7m4C.png">
-<img width="48%" height="auto" src="https://i.imgur.com/JP5RApq.png">
+<img width="48%" height="auto" src="https://envpool.readthedocs.io/en/latest/_images/throughput/Atari_DGX-A100.png">
+<img width="48%" height="auto" src="https://envpool.readthedocs.io/en/latest/_images/throughput/Mujoco_DGX-A100.png">
 </p>
 
+Please refer to the [benchmark](https://envpool.readthedocs.io/en/latest/pages/benchmark.html) page for more details.
 
 ## API Usage
 
@@ -168,7 +171,7 @@ Besides `num_envs`, there is one more argument `batch_size`. While `num_envs` de
 envpool.make("Pong-v5", env_type="gym", num_envs=64, batch_size=16)
 ```
 
-There are other configurable arguments with `envpool.make`; please check out [envpool interface introduction](https://envpool.readthedocs.io/en/latest/pages/interface.html).
+There are other configurable arguments with `envpool.make`; please check out [EnvPool Python interface introduction](https://envpool.readthedocs.io/en/latest/pages/interface.html).
 
 ## Contributing
 
diff --git a/benchmark/.gitignore b/benchmark/.gitignore
new file mode 100644
index 00000000..ecf5e870
--- /dev/null
+++ b/benchmark/.gitignore
@@ -0,0 +1,2 @@
+*.pdf
+*.png
diff --git a/benchmark/README.md b/benchmark/README.md
new file mode 100644
index 00000000..d14b90b4
--- /dev/null
+++ b/benchmark/README.md
@@ -0,0 +1,285 @@
+# Benchmark
+
+The following results are generated from four types of machine:
+
+1. Personal laptop: 12 core ``Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz``, GTX1060
+2. Personal workstation: 32 core ``AMD Ryzen 9 5950X 16-Core Processor``, 2x RTX3090
+3. TPU-VM: 96 core ``Intel(R) Xeon(R) CPU @ 2.00GHz``, 2 NUMA core, TPU v3-8
+4. DGX-A100: 256 core ``AMD EPYC 7742 64-Core Processor``, 8 NUMA core, 8x A100
+
+We use `PongNoFrameskip-v4` (with environment wrappers from [OpenAI baselines](https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py)) and `Ant-v3` for Atari/Mujoco environment benchmark test with `envpool==0.5.3.post1`. Other packages' versions are all in `requirements.txt`:
+
+```bash
+$ pip install -r requirements.txt
+```
+
+To align with other baseline results, FPS is multiplied with `frame_skip` (4 for `PongNoFrameskip-v4` and 5 for `Ant-v3`).
+
+## Highest FPS Overview
+
+
+| Atari Highest FPS    | Laptop (12) | Workstation (32) | TPU-VM (96) | DGX-A100 (256) |
+| -------------------- | ----------- | ---------------- | ----------- | -------------- |
+| For-loop             | 4,893       | 7,914            | 3,993       | 4,640          |
+| Subprocess           | 15,863      | 47,699           | 46,910      | 71,943         |
+| Sample-Factory       | 28,216      | 138,847          | 222,327     | 707,494        |
+| EnvPool (sync)       | 37,396      | 133,824          | 170,380     | 427,851        |
+| EnvPool (async)      | **49,439**  | **200,428**      | 359,559     | 891,286        |
+| EnvPool (numa+async) | /           | /                | **373,169** | **1,069,922**  |
+
+| Mujoco Highest FPS   | Laptop (12) | Workstation (32) | TPU-VM (96) | DGX-A100 (256) |
+| -------------------- | ----------- | ---------------- | ----------- | -------------- |
+| For-loop             | 12,861      | 20,298           | 10,474      | 11,569         |
+| Subprocess           | 36,586      | 105,432          | 87,403      | 163,656        |
+| Sample-Factory       | 62,510      | 309,264          | 461,515     | 1,573,262      |
+| EnvPool (sync)       | 66,622      | 380,950          | 296,681     | 949,787        |
+| EnvPool (async)      | **105,126** | **582,446**      | 887,540     | 2,363,864      |
+| EnvPool (numa+async) | /           | /                | **896,830** | **3,134,287**  |
+
+![](../_static/images/throughput/throughput.png)
+
+## Testing Method and Command
+
+All of the scripts are under [benchmark/](https://github.com/sail-sg/envpool/tree/master/benchmark) folder. When increasing the number of envs, we also adjust the total number of steps to make each test run for about one minute.
+
+### For-loop
+
+Command to run:
+
+```bash
+# atari
+python3 test_gym.py --env atari --num-envs 12 --total-step 6000
+# mujoco
+python3 test_gym.py --env mujoco --num-envs 12 --total-step 12000
+```
+
+### Subprocess (gym.vector_env)
+
+Command to run:
+
+```bash
+# atari
+python3 test_gym.py --env atari --async_ --num-envs 10 --total-step 20000
+# mujoco
+python3 test_gym.py --env mujoco --async_ --num-envs 10 --total-step 50000
+```
+
+### Sample Factory
+
+To run with Ant-v3 in Sample Factory, add one line in `sample_factory/envs/mujoco/mujoco_utils.py`:
+
+```diff
+ MUJOCO_ENVS = [
++    MujocoSpec('mujoco_ant', 'Ant-v3'),
+     MujocoSpec('mujoco_hopper', 'Hopper-v2'),
+     MujocoSpec('mujoco_halfcheetah', 'HalfCheetah-v2'),
+     MujocoSpec('mujoco_humanoid', 'Humanoid-v2'),
+ ]
+```
+
+and finally use FPS \* 5 as the result.
+
+Command to run:
+
+```bash
+# atari
+python3 -m sample_factory.run_algorithm --algo=DUMMY_SAMPLER --env=atari_pong --env_frameskip=4 --num_workers=12 --num_envs_per_worker=1 --sample_env_frames=1600000 --experiment=test
+# mujoco
+python3 -m sample_factory.run_algorithm --algo=DUMMY_SAMPLER --env=mujoco_ant --env_frameskip=1 --num_workers=12 --num_envs_per_worker=1 --sample_env_frames=1000000 --experiment=test
+```
+
+We found that `num_envs_per_worker == 1` is best for all scenarios.
+
+<!--
+
+```python
+def run_sf(w, fac=312500, frame_skip=1, task="atari_pong"):
+    cmd = f"python3 -m sample_factory.run_algorithm --algo=DUMMY_SAMPLER --env={task} --env_frameskip={frame_skip} --num_workers={w} --num_envs_per_worker=1 --sample_env_frames={fac * w} --experiment=test"
+    p = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT)
+    return float([i for i in p.decode().splitlines() if "avg FPS" in i][0].split("FPS: ")[-1].split("\x1b")[0])
+
+for i in num_workers:
+    print(i, run_sf(i, frame_skip=4, task="atari_pong", fac=fac))
+for i in num_workers:
+    print(i, run_sf(i, frame_skip=1, task="mujoco_ant", fac=fac) * 5)
+```
+
+-->
+
+### EnvPool
+
+<!--
+
+```bash
+for i in num_workers:
+    for j in [1, 2.5, 2.6, 3, 4]:
+        print(i, j)
+        os.system(f"python3 test_envpool.py --env mujoco --num-envs {int(i * j)} --batch-size {int(i)} 2>/dev/null > tmp")
+        os.system("grep FPS tmp")
+
+numa_cnt = 8
+for i in num_workers:
+    x = i // numa_cnt
+    if x == 0:
+        continue
+    for j in [2.5, 3, 4]:
+        os.system(f"./numa_test.sh {numa_cnt} python3 test_envpool.py --env mujoco --num-envs {int(x * j)} --batch-size {x} --thread-affinity-offset -1")
+        print(i, x, int(x * j), f'{sum([float([i for i in open(f"log{i}").read().splitlines() if "EnvPool FPS" in i][0].split("=")[-1]) for i in range(numa_cnt)]):.2f}')
+```
+
+-->
+
+#### sync
+
+```bash
+# atari
+python3 test_envpool.py --env atari --num-envs 12 --batch-size 12
+# mujoco
+python3 test_envpool.py --env mujoco --num-envs 12 --batch-size 12
+```
+
+#### async
+
+```bash
+# atari
+python3 test_envpool.py --env atari --num-envs 36 --batch-size 12
+# mujoco
+python3 test_envpool.py --env mujoco --num-envs 36 --batch-size 12
+```
+
+#### numa+async
+
+Use `numactl -s` to determine the number of NUMA cores.
+
+```bash
+# atari
+./numa_test.sh 8 python3 test_envpool.py --env atari --num-envs 100 --batch-size 32 --thread-affinity-offset -1
+# mujoco
+./numa_test.sh 8 python3 test_envpool.py --env mujoco --num-envs 100 --batch-size 32 --thread-affinity-offset -1
+```
+
+### Brax and Isaac-gym (Mujoco only)
+
+TODO
+
+## Result
+
+### Atari
+
+<!-- Atari - Laptop -->
+
+| Atari - Laptop  | 1        | 2        | 3        | 4        | 6        | 8        | 10       | 12       |
+| --------------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- |
+| For-loop        | 4745.54  | 4796.03  | 4694.94  | 4776.76  | 4811.98  | 4892.70  | 4795.49  | 4830.31  |
+| Subprocess      | 4006.04  | 7274.79  | 10028.28 | 11251.66 | 12235.83 | 13280.10 | 15863.42 | 15658.02 |
+| Sample-Factory  | 5844.7   | 11148.0  | 15567.5  | 18236.7  | 25879.3  | 26695.2  | 28216.4  | 28034.7  |
+| EnvPool (sync)  | 7887.51  | 14605.92 | 20288.29 | 26427.86 | 33587.28 | 28602.50 | 34311.75 | 37395.68 |
+| EnvPool (async) | 10213.75 | 18880.65 | 26599.45 | 36375.89 | 48390.40 | 46921.23 | 47184.54 | 49438.56 |
+
+<!-- Atari - Laptop -->
+
+![](../_static/images/throughput/Atari_Laptop.png)
+
+<!-- Atari - Workstation -->
+
+| Atari - Workstation | 1        | 2        | 4        | 8         | 12        | 16        | 20        | 24        | 28        | 32        |
+| ------------------- | -------- | -------- | -------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- |
+| For-loop            | 7739.15  | 7900.56  | 7853.82  | 7865.10   | 7914.04   | 7855.68   | 7587.67   | 7857.92   | 7635.10   | 7868.14   |
+| Subprocess          | 7126.57  | 13086.18 | 23402.05 | 33733.84  | 39766.60  | 42567.05  | 30384.52  | 37224.14  | 46132.40  | 47699.40  |
+| Sample-Factory      | 9259.5   | 18429.2  | 36776.8  | 71435.0   | 101555.5  | 106382.5  | 127522.5  | 131653.0  | 136605.7  | 138847.2  |
+| EnvPool (sync)      | 12623.93 | 23416.68 | 44527.99 | 78612.10  | 105459.54 | 126382.48 | 106088.13 | 117524.07 | 127986.00 | 133824.37 |
+| EnvPool (async)     | 14577.17 | 28383.39 | 55106.44 | 106992.10 | 153258.47 | 188554.16 | 192034.45 | 196540.73 | 200427.90 | 199684.50 |
+
+<!-- Atari - Workstation -->
+
+![](../_static/images/throughput/Atari_Workstation.png)
+
+<!-- Atari - TPU-VM -->
+
+| Atari - TPU-VM       | 1       | 2        | 4        | 8        | 16        | 24        | 32        | 48        | 64        | 80        | 96        |
+| -------------------- | ------- | -------- | -------- | -------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- |
+| For-loop             | 3830.19 | 3942.33  | 3993.01  | 3987.62  | 3967.83   | 3990.12   | 3976.47   | 3986.15   | 3946.44   | 3964.18   | 3973.26   |
+| Subprocess           | 3361.86 | 6586.32  | 12341.66 | 21547.19 | 34152.83  | 34864.23  | 38675.01  | 45471.75  | 41927.33  | 45893.35  | 46910.45  |
+| Sample-Factory       | 4906.3  | 9751.2   | 19450.3  | 38828.2  | 76206.7   | 108471.7  | 137571.6  | 203113.6  | 210596.9  | 217512.9  | 222327.4  |
+| EnvPool (sync)       | 7213.41 | 13827.95 | 27057.69 | 47143.35 | 71660.49  | 98892.99  | 123136.03 | 148110.55 | 141873.23 | 159635.70 | 170380.26 |
+| EnvPool (async)      | 8836.44 | 17815.91 | 35524.72 | 69888.53 | 127106.74 | 184798.27 | 246497.85 | 352195.40 | 354203.40 | 356793.59 | 359558.61 |
+| EnvPool (numa+async) | /       | 17976.26 | 35761.01 | 71967.27 | 136663.09 | 196424.25 | 253789.56 | 368680.81 | 371798.47 | 373169.33 | 362744.14 |
+
+<!-- Atari - TPU-VM -->
+
+![](../_static/images/throughput/Atari_TPU-VM.png)
+
+<!-- Atari - DGX-A100 -->
+
+| Atari - DGX-A100     | 1       | 2        | 4        | 8        | 16        | 32        | 64        | 96        | 128       | 160       | 192       | 224        | 256        |
+| -------------------- | ------- | -------- | -------- | -------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- | ---------- | ---------- |
+| For-loop             | 4449.38 | 4587.37  | 4620.44  | 4635.26  | 4617.21   | 4639.16   | 4618.30   | 4594.96   | 4629.90   | 4616.15   | 4640.20   | 4596.57    | 4620.50    |
+| Subprocess           | 4052.06 | 7832.98  | 12460.71 | 18306.28 | 24754.34  | 33336.38  | 43208.56  | 52435.64  | 42449.85  | 32958.90  | 45312.39  | 45767.11   | 71942.74   |
+| Sample-Factory       | 5563.2  | 11003.0  | 21976.3  | 43891.1  | 87702.0   | 175408.8  | 350855.5  | 476048.4  | 505494.8  | 616958.7  | 651428.8  | 679186.5   | 707494.3   |
+| EnvPool (sync)       | 7723.96 | 14865.81 | 28499.79 | 52681.02 | 91970.45  | 155386.07 | 243231.45 | 304423.24 | 358549.95 | 367559.69 | 388419.70 | 427851.27  | 427395.89  |
+| EnvPool (async)      | 8790.69 | 17866.75 | 36089.43 | 70749.63 | 139540.29 | 278186.45 | 451858.26 | 677504.68 | 817738.45 | 838174.97 | 881210.42 | 891286.00  | 874802.04  |
+| EnvPool (numa+async) | /       | /        | /        | 70629.88 | 140528.93 | 279113.15 | 555426.41 | 762417.99 | 936443.47 | 955620.20 | 998668.02 | 1032953.80 | 1069921.98 |
+
+<!-- Atari - DGX-A100 -->
+
+![](../_static/images/throughput/Atari_DGX-A100.png)
+
+### Mujoco
+
+<!-- Mujoco - Laptop -->
+
+| Mujoco - Laptop | 1        | 2        | 3        | 4        | 6        | 8         | 10        | 12        |
+| --------------- | -------- | -------- | -------- | -------- | -------- | --------- | --------- | --------- |
+| For-loop        | 12325.95 | 12453.54 | 12861.30 | 12517.09 | 12467.92 | 12447.57  | 12631.33  | 12576.39  |
+| Subprocess      | 8377.65  | 14851.20 | 18479.33 | 23137.12 | 26667.67 | 29260.77  | 36586.01  | 31952.74  |
+| Sample-Factory  | 13270.0  | 25452.0  | 34882.0  | 41666.5  | 58892.0  | 60657.5   | 62509.5   | 60847.0   |
+| EnvPool (sync)  | 15641.44 | 30409.65 | 40063.78 | 43126.54 | 58395.28 | 53269.71  | 63424.83  | 66622.24  |
+| EnvPool (async) | 20922.70 | 41279.93 | 57362.56 | 73119.43 | 95542.45 | 105126.36 | 100771.24 | 101603.31 |
+
+<!-- Mujoco - Laptop -->
+
+![](../_static/images/throughput/Mujoco_Laptop.png)
+
+<!-- Mujoco - Workstation -->
+
+| Mujoco - Workstation | 1        | 2        | 4         | 8         | 12        | 16        | 20        | 24        | 28        | 32        |
+| -------------------- | -------- | -------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- |
+| For-loop             | 19472.04 | 19251.41 | 19902.03  | 20076.99  | 19959.82  | 19513.40  | 19460.23  | 19724.42  | 20297.76  | 19797.03  |
+| Subprocess           | 14428.85 | 26943.13 | 48700.27  | 71303.02  | 89901.77  | 102833.40 | 93676.48  | 97473.05  | 105432.15 | 102533.10 |
+| Sample-Factory       | 20854.0  | 40113.5  | 78408.5   | 156563.0  | 225075.0  | 268005.5  | 284237.5  | 296082.5  | 305235.0  | 309264.5  |
+| EnvPool (sync)       | 25725.25 | 50531.72 | 90808.85  | 180372.40 | 212389.98 | 309341.24 | 282954.27 | 326454.83 | 357376.48 | 380950.25 |
+| EnvPool (async)      | 34500.65 | 68382.03 | 133496.84 | 265710.65 | 383015.28 | 478845.88 | 511142.63 | 538558.16 | 566014.54 | 582445.50 |
+
+<!-- Mujoco - Workstation -->
+
+![](../_static/images/throughput/Mujoco_Workstation.png)
+
+<!-- Mujoco - TPU-VM -->
+
+| Mujoco - TPU-VM      | 1        | 2        | 4        | 8         | 16        | 24        | 32        | 48        | 64        | 80        | 96        |
+| -------------------- | -------- | -------- | -------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- | --------- |
+| For-loop             | 9960.98  | 10239.58 | 10186.08 | 10473.73  | 10201.70  | 10370.85  | 10454.78  | 10460.48  | 10455.71  | 10360.71  | 10386.68  |
+| Subprocess           | 7236.32  | 13788.93 | 25054.73 | 40668.40  | 64148.06  | 60409.58  | 70747.21  | 78947.79  | 87403.16  | 79734.62  | 81964.35  |
+| Sample-Factory       | 11008.0  | 21368.0  | 42730.0  | 83475.5   | 153976.0  | 222311.5  | 280664.5  | 406916.5  | 432212.0  | 449143.0  | 461515.0  |
+| EnvPool (sync)       | 13706.61 | 26587.92 | 49074.86 | 92444.28  | 155288.26 | 181397.00 | 231293.39 | 283748.86 | 250586.54 | 268296.99 | 296680.68 |
+| EnvPool (async)      | 18195.81 | 37359.25 | 78337.13 | 148284.57 | 259915.75 | 386448.09 | 512987.78 | 745083.58 | 801768.88 | 857586.18 | 887539.80 |
+| EnvPool (numa+async) | /        | 35804.57 | 75467.72 | 147281.29 | 284323.79 | 412165.16 | 516120.17 | 755509.66 | 816405.50 | 868455.12 | 896830.21 |
+
+<!-- Mujoco - TPU-VM -->
+
+![](../_static/images/throughput/Mujoco_TPU-VM.png)
+
+<!-- Mujoco - DGX-A100 -->
+
+| Mujoco - DGX-A100    | 1        | 2        | 4        | 8         | 16        | 32        | 64         | 96         | 128        | 160        | 192        | 224        | 256        |
+| -------------------- | -------- | -------- | -------- | --------- | --------- | --------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- | ---------- |
+| For-loop             | 11018.57 | 11269.45 | 11059.39 | 11250.06  | 11505.15  | 11328.79  | 11568.72   | 11485.74   | 11245.55   | 11478.49   | 11430.16   | 11151.71   | 11199.28   |
+| Subprocess           | 8814.10  | 17201.64 | 27106.27 | 44383.63  | 62785.60  | 83054.19  | 151352.88  | 158797.86  | 148815.92  | 116200.41  | 163656.36  | 147653.41  | 161599.97  |
+| Sample-Factory       | 11870.0  | 24602.0  | 48577.0  | 96826.5   | 193800.5  | 381208.5  | 761752.0   | 985909.0   | 1249369.5  | 1332128.5  | 1397427.5  | 1318249.0  | 1573262.0  |
+| EnvPool (sync)       | 16024.43 | 31899.44 | 61605.04 | 114488.28 | 228492.88 | 388624.94 | 656277.80  | 832101.96  | 949787.15  | 858298.85  | 945808.57  | 813799.36  | 849410.96  |
+| EnvPool (async)      | 21177.71 | 44025.65 | 92312.35 | 176135.82 | 354006.02 | 700052.08 | 1167838.03 | 1678787.71 | 1730102.62 | 2052844.58 | 2185146.77 | 2355604.96 | 2363863.67 |
+| EnvPool (numa+async) | /        | /        | /        | 170348.47 | 340269.34 | 693793.45 | 1388410.00 | 1920762.84 | 2341562.20 | 2569997.03 | 2776143.15 | 2964886.91 | 3134286.77 |
+
+<!-- Mujoco - DGX-A100 -->
+
+![](../_static/images/throughput/Mujoco_DGX-A100.png)
diff --git a/benchmark/atari_wrappers.py b/benchmark/atari_wrappers.py
new file mode 100644
index 00000000..d2624a3c
--- /dev/null
+++ b/benchmark/atari_wrappers.py
@@ -0,0 +1,269 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Atari wrappers copy from tianshou's example.
+
+https://github.com/thu-ml/tianshou/blob/master/examples/atari/atari_wrapper.py
+"""
+
+from collections import deque
+
+import cv2
+import gym
+import numpy as np
+
+
+class NoopResetEnv(gym.Wrapper):
+  """Sample initial states by taking random number of no-ops on reset.
+
+  No-op is assumed to be action 0.
+
+  :param gym.Env env: the environment to wrap.
+  :param int noop_max: the maximum value of no-ops to run.
+  """
+
+  def __init__(self, env, noop_max=30):
+    super().__init__(env)
+    self.noop_max = noop_max
+    self.noop_action = 0
+    assert env.unwrapped.get_action_meanings()[0] == "NOOP"
+
+  def reset(self):
+    self.env.reset()
+    noops = self.unwrapped.np_random.integers(1, self.noop_max + 1)
+    for _ in range(noops):
+      obs, _, done, _ = self.env.step(self.noop_action)
+      if done:
+        obs = self.env.reset()
+    return obs
+
+
+class MaxAndSkipEnv(gym.Wrapper):
+  """Return only every `skip`-th frame (frameskipping) using most recent raw
+  observations (for max pooling across time steps)
+
+  :param gym.Env env: the environment to wrap.
+  :param int skip: number of `skip`-th frame.
+  """
+
+  def __init__(self, env, skip=4):
+    super().__init__(env)
+    self._skip = skip
+
+  def step(self, action):
+    """Step the environment with the given action.
+
+    Repeat action, sum reward, and max over last observations.
+    """
+    obs_list, total_reward, done = [], 0., False
+    for _ in range(self._skip):
+      obs, reward, done, info = self.env.step(action)
+      obs_list.append(obs)
+      total_reward += reward
+      if done:
+        break
+    max_frame = np.max(obs_list[-2:], axis=0)
+    return max_frame, total_reward, done, info
+
+
+class EpisodicLifeEnv(gym.Wrapper):
+  """Make end-of-life == end-of-episode, but only reset on true game over.
+
+  It helps the value estimation.
+
+  :param gym.Env env: the environment to wrap.
+  """
+
+  def __init__(self, env):
+    super().__init__(env)
+    self.lives = 0
+    self.was_real_done = True
+
+  def step(self, action):
+    obs, reward, done, info = self.env.step(action)
+    self.was_real_done = done
+    # check current lives, make loss of life terminal, then update lives to
+    # handle bonus lives
+    lives = self.env.unwrapped.ale.lives()
+    if 0 < lives < self.lives:
+      # for Qbert sometimes we stay in lives == 0 condition for a few
+      # frames, so its important to keep lives > 0, so that we only reset
+      # once the environment is actually done.
+      done = True
+    self.lives = lives
+    return obs, reward, done, info
+
+  def reset(self):
+    """Calls the Gym environment reset, only when lives are exhausted.
+
+    This way all states are still reachable even though lives are episodic,
+    and the learner need not know about any of this behind-the-scenes.
+    """
+    if self.was_real_done:
+      obs = self.env.reset()
+    else:
+      # no-op step to advance from terminal/lost life state
+      obs = self.env.step(0)[0]
+    self.lives = self.env.unwrapped.ale.lives()
+    return obs
+
+
+class FireResetEnv(gym.Wrapper):
+  """Take action on reset for environments that are fixed until firing.
+
+  Related discussion: https://github.com/openai/baselines/issues/240
+
+  :param gym.Env env: the environment to wrap.
+  """
+
+  def __init__(self, env):
+    super().__init__(env)
+    assert env.unwrapped.get_action_meanings()[1] == "FIRE"
+    assert len(env.unwrapped.get_action_meanings()) >= 3
+
+  def reset(self):
+    self.env.reset()
+    return self.env.step(1)[0]
+
+
+class WarpFrame(gym.ObservationWrapper):
+  """Warp frames to 84x84 as done in the Nature paper and later work.
+
+  :param gym.Env env: the environment to wrap.
+  """
+
+  def __init__(self, env):
+    super().__init__(env)
+    self.size = 84
+    self.observation_space = gym.spaces.Box(
+      low=np.min(env.observation_space.low),
+      high=np.max(env.observation_space.high),
+      shape=(self.size, self.size),
+      dtype=env.observation_space.dtype
+    )
+
+  def observation(self, frame):
+    """Returns the current observation from a frame"""
+    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
+    return cv2.resize(
+      frame, (self.size, self.size), interpolation=cv2.INTER_AREA
+    )
+
+
+class ScaledFloatFrame(gym.ObservationWrapper):
+  """Normalize observations to 0~1.
+
+  :param gym.Env env: the environment to wrap.
+  """
+
+  def __init__(self, env):
+    super().__init__(env)
+    low = np.min(env.observation_space.low)
+    high = np.max(env.observation_space.high)
+    self.bias = low
+    self.scale = high - low
+    self.observation_space = gym.spaces.Box(
+      low=0., high=1., shape=env.observation_space.shape, dtype=np.float32
+    )
+
+  def observation(self, observation):
+    return (observation - self.bias) / self.scale
+
+
+class ClipRewardEnv(gym.RewardWrapper):
+  """clips the reward to {+1, 0, -1} by its sign.
+
+  :param gym.Env env: the environment to wrap.
+  """
+
+  def __init__(self, env):
+    super().__init__(env)
+    self.reward_range = (-1, 1)
+
+  def reward(self, reward):
+    """Bin reward to {+1, 0, -1} by its sign. Note: np.sign(0) == 0."""
+    return np.sign(reward)
+
+
+class FrameStack(gym.Wrapper):
+  """Stack n_frames last frames.
+
+  :param gym.Env env: the environment to wrap.
+  :param int n_frames: the number of frames to stack.
+  """
+
+  def __init__(self, env, n_frames):
+    super().__init__(env)
+    self.n_frames = n_frames
+    self.frames = deque([], maxlen=n_frames)
+    shape = (n_frames,) + env.observation_space.shape
+    self.observation_space = gym.spaces.Box(
+      low=np.min(env.observation_space.low),
+      high=np.max(env.observation_space.high),
+      shape=shape,
+      dtype=env.observation_space.dtype
+    )
+
+  def reset(self):
+    obs = self.env.reset()
+    for _ in range(self.n_frames):
+      self.frames.append(obs)
+    return self._get_ob()
+
+  def step(self, action):
+    obs, reward, done, info = self.env.step(action)
+    self.frames.append(obs)
+    return self._get_ob(), reward, done, info
+
+  def _get_ob(self):
+    # the original wrapper use `LazyFrames` but since we use np buffer,
+    # it has no effect
+    return np.stack(self.frames, axis=0)
+
+
+def wrap_deepmind(
+  env,
+  episode_life=True,
+  clip_rewards=True,
+  frame_stack=4,
+  scale=False,
+  warp_frame=True,
+):
+  """Configure environment for DeepMind-style Atari.
+
+  The observation is channel-first: (c, h, w) instead of (h, w, c).
+
+  :param str env: the raw atari environment.
+  :param bool episode_life: wrap the episode life wrapper.
+  :param bool clip_rewards: wrap the reward clipping wrapper.
+  :param int frame_stack: wrap the frame stacking wrapper.
+  :param bool scale: wrap the scaling observation wrapper.
+  :param bool warp_frame: wrap the grayscale + resize observation wrapper.
+  :return: the wrapped atari environment.
+  """
+  assert "NoFrameskip" in env.spec.id
+  env = NoopResetEnv(env, noop_max=30)
+  env = MaxAndSkipEnv(env, skip=4)
+  if episode_life:
+    env = EpisodicLifeEnv(env)
+  if "FIRE" in env.unwrapped.get_action_meanings():
+    env = FireResetEnv(env)
+  if warp_frame:
+    env = WarpFrame(env)
+  if scale:
+    env = ScaledFloatFrame(env)
+  if clip_rewards:
+    env = ClipRewardEnv(env)
+  if frame_stack:
+    env = FrameStack(env, frame_stack)
+  return env
diff --git a/benchmark/plot.py b/benchmark/plot.py
new file mode 100755
index 00000000..f85c1fae
--- /dev/null
+++ b/benchmark/plot.py
@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+
+import matplotlib.ticker as ticker
+import pandas as pd
+import seaborn as sns
+
+data = {}
+
+
+def reset_data() -> None:
+  global data
+  data = {
+    "Num. Workers": [],
+    "FPS": [],
+    "Env": [],
+    "System": [],
+    "Method": [],
+  }
+
+
+def parse_table(env: str, system: str, suffix: str) -> None:
+  private_copy = {
+    "Num. Workers": [],
+    "FPS": [],
+    "Env": [],
+    "System": [],
+    "Method": [],
+  }
+  sep = f"<!-- {env} - {system} -->"
+  raw = open("README.md").read().split(sep)[1].strip().splitlines()
+  worker_num = list(map(int, raw[0].split("|")[2:-1]))
+  for line in raw[2:]:
+    line = line.split("|")[1:-1]
+    method = line.pop(0).strip()
+    for w, f in zip(worker_num, line):
+      for d in [data, private_copy]:
+        d["Num. Workers"].append(w)
+        d["FPS"].append(None if f.strip() == "/" else float(f))
+        d["Env"].append(env)
+        d["System"].append(system)
+        d["Method"].append(method)
+  d = pd.DataFrame(private_copy)
+  plot = sns.lineplot(
+    x="Num. Workers", y="FPS", hue="Method", data=d, marker="o"
+  )
+  plot.xaxis.set_major_formatter(ticker.EngFormatter())
+  plot.yaxis.set_major_formatter(ticker.EngFormatter())
+  plot.legend(fontsize=9)
+  plot.set_title(f"{env} throughput, {system}")
+  plot.set_xlabel("Num. Workers")
+  frame_skip = {"Atari": 4, "Mujoco": 5}[env]
+  plot.set_ylabel(f"FPS, frameskip = {frame_skip}")
+  plot.get_figure().savefig(f"{env}_{system}.{suffix}")
+  plot.get_figure().clf()
+
+
+def benchmark(suffix: str) -> None:
+  global data
+  reset_data()
+  for env in ["Atari", "Mujoco"]:
+    for system in ["Laptop", "Workstation", "TPU-VM", "DGX-A100"]:
+      parse_table(env, system, suffix)
+  data = pd.DataFrame(data)
+  print(data.groupby(["Env", "Method", "System"]).max())
+
+  def mapping(x, y, **kwargs):
+    plot = sns.lineplot(x=x, y=y, **kwargs)
+    plot.xaxis.set_major_formatter(ticker.EngFormatter())
+    plot.yaxis.set_major_formatter(ticker.EngFormatter())
+
+  g = sns.FacetGrid(
+    data,
+    row="Env",
+    col="System",
+    hue="Method",
+    height=3,
+    aspect=1.6,
+    sharex=False,
+    sharey=False,
+  )
+  g.map(mapping, "Num. Workers", "FPS", marker="o")
+  g.add_legend(bbox_to_anchor=(0.52, 1.02), ncol=6)
+  axes = g.axes.flatten()
+  alphabet = "abcdefgh"
+  for ax, i in zip(axes, alphabet):
+    env, system = ax.get_title().split("|")
+    env = env.split("=")[-1].strip()
+    system = system.split("=")[-1].strip()
+    ax.set_title(f"({i}) {env} throughput, {system}")
+    if ax.get_ylabel():
+      frame_skip = {"Atari": 4, "Mujoco": 5}[env]
+      ax.set_ylabel(f"FPS, frameskip = {frame_skip}")
+  g.savefig(f"throughput.{suffix}")
+
+
+if __name__ == "__main__":
+  pd.options.display.float_format = '{:,.0f}'.format
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--suffix", type=str, default="png")
+  args = parser.parse_args()
+  benchmark(args.suffix)
diff --git a/benchmark/requirements.txt b/benchmark/requirements.txt
new file mode 100644
index 00000000..88f63618
--- /dev/null
+++ b/benchmark/requirements.txt
@@ -0,0 +1,7 @@
+gym[accept-rom-license]==0.23.1
+ale-py==0.7.5
+mujoco==2.1.5
+envpool==0.5.3.post1
+sample-factory==1.123.0
+mujoco_py==2.1.2.14
+tqdm
diff --git a/docs/_static/images/atari_throughput_apollo.png b/docs/_static/images/atari_throughput_apollo.png
deleted file mode 100644
index 659344cd..00000000
Binary files a/docs/_static/images/atari_throughput_apollo.png and /dev/null differ
diff --git a/docs/_static/images/atari_throughput_dgx.png b/docs/_static/images/atari_throughput_dgx.png
deleted file mode 100644
index 4c510482..00000000
Binary files a/docs/_static/images/atari_throughput_dgx.png and /dev/null differ
diff --git a/docs/_static/images/atari_throughput_tpu.png b/docs/_static/images/atari_throughput_tpu.png
deleted file mode 100644
index 30b81321..00000000
Binary files a/docs/_static/images/atari_throughput_tpu.png and /dev/null differ
diff --git a/docs/_static/images/throughput/Atari_DGX-A100.png b/docs/_static/images/throughput/Atari_DGX-A100.png
new file mode 100644
index 00000000..bba3dae3
Binary files /dev/null and b/docs/_static/images/throughput/Atari_DGX-A100.png differ
diff --git a/docs/_static/images/throughput/Atari_Laptop.png b/docs/_static/images/throughput/Atari_Laptop.png
new file mode 100644
index 00000000..3b3d184a
Binary files /dev/null and b/docs/_static/images/throughput/Atari_Laptop.png differ
diff --git a/docs/_static/images/throughput/Atari_TPU-VM.png b/docs/_static/images/throughput/Atari_TPU-VM.png
new file mode 100644
index 00000000..f740ac01
Binary files /dev/null and b/docs/_static/images/throughput/Atari_TPU-VM.png differ
diff --git a/docs/_static/images/throughput/Atari_Workstation.png b/docs/_static/images/throughput/Atari_Workstation.png
new file mode 100644
index 00000000..5b650d23
Binary files /dev/null and b/docs/_static/images/throughput/Atari_Workstation.png differ
diff --git a/docs/_static/images/throughput/Mujoco_DGX-A100.png b/docs/_static/images/throughput/Mujoco_DGX-A100.png
new file mode 100644
index 00000000..6d9b3e61
Binary files /dev/null and b/docs/_static/images/throughput/Mujoco_DGX-A100.png differ
diff --git a/docs/_static/images/throughput/Mujoco_Laptop.png b/docs/_static/images/throughput/Mujoco_Laptop.png
new file mode 100644
index 00000000..1eb1df01
Binary files /dev/null and b/docs/_static/images/throughput/Mujoco_Laptop.png differ
diff --git a/docs/_static/images/throughput/Mujoco_TPU-VM.png b/docs/_static/images/throughput/Mujoco_TPU-VM.png
new file mode 100644
index 00000000..4eeb02c0
Binary files /dev/null and b/docs/_static/images/throughput/Mujoco_TPU-VM.png differ
diff --git a/docs/_static/images/throughput/Mujoco_Workstation.png b/docs/_static/images/throughput/Mujoco_Workstation.png
new file mode 100644
index 00000000..298c8594
Binary files /dev/null and b/docs/_static/images/throughput/Mujoco_Workstation.png differ
diff --git a/docs/_static/images/throughput/throughput.png b/docs/_static/images/throughput/throughput.png
new file mode 100644
index 00000000..7102a1b9
Binary files /dev/null and b/docs/_static/images/throughput/throughput.png differ
diff --git a/docs/api/mujoco.rst b/docs/api/mujoco.rst
index a96f7814..15d77dbf 100644
--- a/docs/api/mujoco.rst
+++ b/docs/api/mujoco.rst
@@ -24,7 +24,10 @@ Ant-v3/v4
 - Observation space: ``(111)``, first 13 elements for ``qpos[2:]``, next 14
   elements for ``qvel``, other elements for clipped ``cfrc_ext`` (com-based
   external force on body, a.k.a. contact force);
-- Action space: ``(8)``, with range ``[-1, 1]``.
+- Action space: ``(8)``, with range ``[-1, 1]``;
+- ``frame_skip``: 5;
+- ``max_episode_steps``: 1000;
+- ``reward_threshold``: 6000.0;
 
 
 HalfCheetah-v3/v4
@@ -38,7 +41,10 @@ HalfCheetah-v3/v4
 
 - Observation space: ``(17)``, first 8 elements for ``qpos[1:]``, next 9
   elements for ``qvel``;
-- Action space: ``(6)``, with range ``[-1, 1]``.
+- Action space: ``(6)``, with range ``[-1, 1]``;
+- ``frame_skip``: 5;
+- ``max_episode_steps``: 1000;
+- ``reward_threshold``: 4800.0;
 
 
 Hopper-v3/v4
@@ -52,7 +58,10 @@ Hopper-v3/v4
 
 - Observation space: ``(11)``, first 5 elements for ``qpos[1:]``, next 6
   elements for ``qvel``;
-- Action space: ``(3)``, with range ``[-1, 1]``.
+- Action space: ``(3)``, with range ``[-1, 1]``;
+- ``frame_skip``: 4;
+- ``max_episode_steps``: 1000;
+- ``reward_threshold``: 6000.0;
 
 
 Humanoid-v3/v4, HumanoidStandup-v2/v4
@@ -75,7 +84,9 @@ Humanoid-v3/v4, HumanoidStandup-v2/v4
   inertia and mass), next 84 elements for ``cvel`` (com-based velocity [3D
   rot; 3D tran]), next 23 elements for ``qfrc_actuator`` (actuator force),
   next 84 elements for ``cfrc_ext`` (com-based external force on body);
-- Action space: ``(17)``, with range ``[-0.4, 0.4]``.
+- Action space: ``(17)``, with range ``[-0.4, 0.4]``;
+- ``frame_skip``: 5;
+- ``max_episode_steps``: 1000;
 
 
 InvertedDoublePendulum-v2/v4
@@ -90,7 +101,10 @@ InvertedDoublePendulum-v2/v4
 - Observation space: ``(11)``, first 1 element for ``qpos[0]``, next 2
   elements for ``sin(qpos[1:])``, next 2 elements for ``cos(qpos[1:])``,
   next 3 elements for ``qvel``, next 3 elements for ``qfrc_constraint``;
-- Action space: ``(1)``, with range ``[-1, 1]``.
+- Action space: ``(1)``, with range ``[-1, 1]``;
+- ``frame_skip``: 5;
+- ``max_episode_steps``: 1000;
+- ``reward_threshold``: 9100.0;
 
 
 InvertedPendulum-v2/v4
@@ -104,7 +118,10 @@ InvertedPendulum-v2/v4
 
 - Observation space: ``(4)``, first 2 elements for ``qpos``, next 2 elements
   for ``qvel``;
-- Action space: ``(1)``, with range ``[-3, 3]``.
+- Action space: ``(1)``, with range ``[-3, 3]``;
+- ``frame_skip``: 2;
+- ``max_episode_steps``: 1000;
+- ``reward_threshold``: 950.0;
 
 
 Pusher-v2/v4
@@ -119,7 +136,10 @@ Pusher-v2/v4
 - Observation space: ``(23)``, first 7 elements for ``qpos[:7]``, next 7
   elements for ``qvel[:7]``, next 3 elements for ``tips_arm``, next 3
   elements for ``object``, next 3 elements for ``goal``;
-- Action space: ``(7)``, with range ``[-2, 2]``.
+- Action space: ``(7)``, with range ``[-2, 2]``;
+- ``frame_skip``: 5;
+- ``max_episode_steps``: 100;
+- ``reward_threshold``: 0.0;
 
 
 Reacher-v2/v4
@@ -135,7 +155,10 @@ Reacher-v2/v4
   elements for ``sin(qpos[:2])``, next 2 elements for ``qpos[2:]``, next 2
   elements for ``qvel[:2]``, next 3 elements for ``dist``, a.k.a.
   ``fingertip - target``;
-- Action space: ``(2)``, with range ``[-1, 1]``.
+- Action space: ``(2)``, with range ``[-1, 1]``;
+- ``frame_skip``: 2;
+- ``max_episode_steps``: 50;
+- ``reward_threshold``: -3.75;
 
 
 Swimmer-v3/v4
@@ -149,7 +172,10 @@ Swimmer-v3/v4
 
 - Observation space: ``(8)``, first 3 elements for ``qpos[2:]``, next 5
   elements for ``qvel``;
-- Action space: ``(2)``, with range ``[-1, 1]``.
+- Action space: ``(2)``, with range ``[-1, 1]``;
+- ``frame_skip``: 4;
+- ``max_episode_steps``: 1000;
+- ``reward_threshold``: 360.0;
 
 
 Walker2d-v3/v4
@@ -163,4 +189,6 @@ Walker2d-v3/v4
 
 - Observation space: ``(17)``, first 8 elements for ``qpos[1:]``, next 9
   elements for ``qvel``;
-- Action space: ``(6)``, with range ``[-1, 1]``.
+- Action space: ``(6)``, with range ``[-1, 1]``;
+- ``frame_skip``: 4;
+- ``max_episode_steps``: 1000;
diff --git a/docs/pages/benchmark.rst b/docs/pages/benchmark.rst
index 141f1083..4a019a70 100644
--- a/docs/pages/benchmark.rst
+++ b/docs/pages/benchmark.rst
@@ -1,36 +1,364 @@
-Benchmark Results on Atari
-==========================
+Benchmark
+=========
 
 The following results are generated from four types of machine:
 
-1. Personal laptop: 12 core ``Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz``
-2. TPU-VM: 96 core ``Intel(R) Xeon(R) CPU @ 2.00GHz``
-3. Apollo: 96 core ``AMD EPYC 7352 24-Core Processor``
-4. DGX-A100: 256 core ``AMD EPYC 7742 64-Core Processor``
-
-FPS is based on ALE frame (Pong). Using ``envpool==0.4.0``
-
-+----------------------+----------------------+-------------+-------------+----------------+
-| Highest FPS          | i7-8750H laptop (12) | TPU-VM (96) | Apollo (96) | DGX-A100 (256) |
-+======================+======================+=============+=============+================+
-| For-loop             | 4,876                | 3,817       | 4,053       | 4,336          |
-+----------------------+----------------------+-------------+-------------+----------------+
-| Subprocess           | 18,249               | 42,885      | 19,560      | 79,509         |
-+----------------------+----------------------+-------------+-------------+----------------+
-| Sample Factory       | 27,035               | 192,074     | 262,963     | 639,389        |
-+----------------------+----------------------+-------------+-------------+----------------+
-| EnvPool (sync)       | 40,791               | 175,938     | 159,191     | 470,170        |
-+----------------------+----------------------+-------------+-------------+----------------+
-| EnvPool (async)      | 50,513               | 352,243     | 410,941     | 845,537        |
-+----------------------+----------------------+-------------+-------------+----------------+
-| EnvPool (numa+async) | /                    | 367,799     | 458,414     | 1,060,371      |
-+----------------------+----------------------+-------------+-------------+----------------+
-
-.. image:: ../_static/images/atari_throughput_tpu.png
-    :align: center
-
-.. image:: ../_static/images/atari_throughput_apollo.png
-    :align: center
-
-.. image:: ../_static/images/atari_throughput_dgx.png
-    :align: center
+1. Personal laptop: 12 core ``Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz``, GTX1060
+2. Personal workstation: 32 core ``AMD Ryzen 9 5950X 16-Core Processor``, 2x RTX3090
+3. TPU-VM: 96 core ``Intel(R) Xeon(R) CPU @ 2.00GHz``, 2 NUMA core, TPU v3-8
+4. DGX-A100: 256 core ``AMD EPYC 7742 64-Core Processor``, 8 NUMA core, 8x A100
+
+We use ``PongNoFrameskip-v4`` (with environment wrappers from `OpenAI baselines <https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py>`__) and ``Ant-v3`` for Atari/Mujoco environment benchmark test with ``envpool==0.5.3.post1``. Other packages’ versions are all in ``requirements.txt``:
+
+.. code:: bash
+
+   $ pip install -r requirements.txt
+
+To align with other baseline results, FPS is multiplied with ``frame_skip`` (4 for ``PongNoFrameskip-v4`` and 5 for ``Ant-v3``).
+
+Highest FPS Overview
+--------------------
+
+==================== =========== ================ =========== ==============
+Atari Highest FPS    Laptop (12) Workstation (32) TPU-VM (96) DGX-A100 (256)
+==================== =========== ================ =========== ==============
+For-loop             4,893       7,914            3,993       4,640
+Subprocess           15,863      47,699           46,910      71,943
+Sample-Factory       28,216      138,847          222,327     707,494
+EnvPool (sync)       37,396      133,824          170,380     427,851
+EnvPool (async)      **49,439**  **200,428**      359,559     891,286
+EnvPool (numa+async) /           /                **373,169** **1,069,922**
+==================== =========== ================ =========== ==============
+
+==================== =========== ================ =========== ==============
+Mujoco Highest FPS   Laptop (12) Workstation (32) TPU-VM (96) DGX-A100 (256)
+==================== =========== ================ =========== ==============
+For-loop             12,861      20,298           10,474      11,569
+Subprocess           36,586      105,432          87,403      163,656
+Sample-Factory       62,510      309,264          461,515     1,573,262
+EnvPool (sync)       66,622      380,950          296,681     949,787
+EnvPool (async)      **105,126** **582,446**      887,540     2,363,864
+EnvPool (numa+async) /           /                **896,830** **3,134,287**
+==================== =========== ================ =========== ==============
+
+|image0|
+
+Testing Method and Command
+--------------------------
+
+All of the scripts are under ``benchmark/`` folder. When increasing the number of envs, we also adjust the total number of steps to make each test run for about one minute.
+
+For-loop
+~~~~~~~~
+
+Command to run:
+
+.. code:: bash
+
+   # atari
+   python3 test_gym.py --env atari --num-envs 12 --total-step 6000
+   # mujoco
+   python3 test_gym.py --env mujoco --num-envs 12 --total-step 12000
+
+Subprocess (gym.vector_env)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Command to run:
+
+.. code:: bash
+
+   # atari
+   python3 test_gym.py --env atari --async_ --num-envs 10 --total-step 20000
+   # mujoco
+   python3 test_gym.py --env mujoco --async_ --num-envs 10 --total-step 50000
+
+Sample Factory
+~~~~~~~~~~~~~~
+
+To run with Ant-v3 in Sample Factory, add one line in ``sample_factory/envs/mujoco/mujoco_utils.py``:
+
+.. code:: diff
+
+    MUJOCO_ENVS = [
+   +    MujocoSpec('mujoco_ant', 'Ant-v3'),
+        MujocoSpec('mujoco_hopper', 'Hopper-v2'),
+        MujocoSpec('mujoco_halfcheetah', 'HalfCheetah-v2'),
+        MujocoSpec('mujoco_humanoid', 'Humanoid-v2'),
+    ]
+
+and finally use FPS \* 5 as the result.
+
+Command to run:
+
+.. code:: bash
+
+   # atari
+   python3 -m sample_factory.run_algorithm --algo=DUMMY_SAMPLER --env=atari_pong --env_frameskip=4 --num_workers=12 --num_envs_per_worker=1 --sample_env_frames=1600000 --experiment=test
+   # mujoco
+   python3 -m sample_factory.run_algorithm --algo=DUMMY_SAMPLER --env=mujoco_ant --env_frameskip=1 --num_workers=12 --num_envs_per_worker=1 --sample_env_frames=1000000 --experiment=test
+
+We found that ``num_envs_per_worker == 1`` is best for all scenarios.
+
+.. raw:: html
+
+   <!--
+
+   ```python
+   def run_sf(w, fac=312500, frame_skip=1, task="atari_pong"):
+       cmd = f"python3 -m sample_factory.run_algorithm --algo=DUMMY_SAMPLER --env={task} --env_frameskip={frame_skip} --num_workers={w} --num_envs_per_worker=1 --sample_env_frames={fac * w} --experiment=test"
+       p = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT)
+       return float([i for i in p.decode().splitlines() if "avg FPS" in i][0].split("FPS: ")[-1].split("\x1b")[0])
+
+   for i in num_workers:
+       print(i, run_sf(i, frame_skip=4, task="atari_pong", fac=fac))
+   for i in num_workers:
+       print(i, run_sf(i, frame_skip=1, task="mujoco_ant", fac=fac) * 5)
+   ```
+
+   -->
+
+EnvPool
+~~~~~~~
+
+.. raw:: html
+
+   <!--
+
+   ```bash
+   for i in num_workers:
+       for j in [1, 2.5, 2.6, 3, 4]:
+           print(i, j)
+           os.system(f"python3 test_envpool.py --env mujoco --num-envs {int(i * j)} --batch-size {int(i)} 2>/dev/null > tmp")
+           os.system("grep FPS tmp")
+
+   numa_cnt = 8
+   for i in num_workers:
+       x = i // numa_cnt
+       if x == 0:
+           continue
+       for j in [2.5, 3, 4]:
+           os.system(f"./numa_test.sh {numa_cnt} python3 test_envpool.py --env mujoco --num-envs {int(x * j)} --batch-size {x} --thread-affinity-offset -1")
+           print(i, x, int(x * j), f'{sum([float([i for i in open(f"log{i}").read().splitlines() if "EnvPool FPS" in i][0].split("=")[-1]) for i in range(numa_cnt)]):.2f}')
+   ```
+
+   -->
+
+sync
+^^^^
+
+.. code:: bash
+
+   # atari
+   python3 test_envpool.py --env atari --num-envs 12 --batch-size 12
+   # mujoco
+   python3 test_envpool.py --env mujoco --num-envs 12 --batch-size 12
+
+async
+^^^^^
+
+.. code:: bash
+
+   # atari
+   python3 test_envpool.py --env atari --num-envs 36 --batch-size 12
+   # mujoco
+   python3 test_envpool.py --env mujoco --num-envs 36 --batch-size 12
+
+numa+async
+^^^^^^^^^^
+
+Use ``numactl -s`` to determine the number of NUMA cores.
+
+.. code:: bash
+
+   # atari
+   ./numa_test.sh 8 python3 test_envpool.py --env atari --num-envs 100 --batch-size 32 --thread-affinity-offset -1
+   # mujoco
+   ./numa_test.sh 8 python3 test_envpool.py --env mujoco --num-envs 100 --batch-size 32 --thread-affinity-offset -1
+
+Brax and Isaac-gym (Mujoco only)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+TODO
+
+Result
+------
+
+Atari
+~~~~~
+
+.. raw:: html
+
+   <!-- Atari - Laptop -->
+
+=============== ======== ======== ======== ======== ======== ======== ======== ========
+Atari - Laptop  1        2        3        4        6        8        10       12
+=============== ======== ======== ======== ======== ======== ======== ======== ========
+For-loop        4745.54  4796.03  4694.94  4776.76  4811.98  4892.70  4795.49  4830.31
+Subprocess      4006.04  7274.79  10028.28 11251.66 12235.83 13280.10 15863.42 15658.02
+Sample-Factory  5844.7   11148.0  15567.5  18236.7  25879.3  26695.2  28216.4  28034.7
+EnvPool (sync)  7887.51  14605.92 20288.29 26427.86 33587.28 28602.50 34311.75 37395.68
+EnvPool (async) 10213.75 18880.65 26599.45 36375.89 48390.40 46921.23 47184.54 49438.56
+=============== ======== ======== ======== ======== ======== ======== ======== ========
+
+.. raw:: html
+
+   <!-- Atari - Laptop -->
+
+|image1|
+
+.. raw:: html
+
+   <!-- Atari - Workstation -->
+
+=================== ======== ======== ======== ========= ========= ========= ========= ========= ========= =========
+Atari - Workstation 1        2        4        8         12        16        20        24        28        32
+=================== ======== ======== ======== ========= ========= ========= ========= ========= ========= =========
+For-loop            7739.15  7900.56  7853.82  7865.10   7914.04   7855.68   7587.67   7857.92   7635.10   7868.14
+Subprocess          7126.57  13086.18 23402.05 33733.84  39766.60  42567.05  30384.52  37224.14  46132.40  47699.40
+Sample-Factory      9259.5   18429.2  36776.8  71435.0   101555.5  106382.5  127522.5  131653.0  136605.7  138847.2
+EnvPool (sync)      12623.93 23416.68 44527.99 78612.10  105459.54 126382.48 106088.13 117524.07 127986.00 133824.37
+EnvPool (async)     14577.17 28383.39 55106.44 106992.10 153258.47 188554.16 192034.45 196540.73 200427.90 199684.50
+=================== ======== ======== ======== ========= ========= ========= ========= ========= ========= =========
+
+.. raw:: html
+
+   <!-- Atari - Workstation -->
+
+|image2|
+
+.. raw:: html
+
+   <!-- Atari - TPU-VM -->
+
+==================== ======= ======== ======== ======== ========= ========= ========= ========= ========= ========= =========
+Atari - TPU-VM       1       2        4        8        16        24        32        48        64        80        96
+==================== ======= ======== ======== ======== ========= ========= ========= ========= ========= ========= =========
+For-loop             3830.19 3942.33  3993.01  3987.62  3967.83   3990.12   3976.47   3986.15   3946.44   3964.18   3973.26
+Subprocess           3361.86 6586.32  12341.66 21547.19 34152.83  34864.23  38675.01  45471.75  41927.33  45893.35  46910.45
+Sample-Factory       4906.3  9751.2   19450.3  38828.2  76206.7   108471.7  137571.6  203113.6  210596.9  217512.9  222327.4
+EnvPool (sync)       7213.41 13827.95 27057.69 47143.35 71660.49  98892.99  123136.03 148110.55 141873.23 159635.70 170380.26
+EnvPool (async)      8836.44 17815.91 35524.72 69888.53 127106.74 184798.27 246497.85 352195.40 354203.40 356793.59 359558.61
+EnvPool (numa+async) /       17976.26 35761.01 71967.27 136663.09 196424.25 253789.56 368680.81 371798.47 373169.33 362744.14
+==================== ======= ======== ======== ======== ========= ========= ========= ========= ========= ========= =========
+
+.. raw:: html
+
+   <!-- Atari - TPU-VM -->
+
+|image3|
+
+.. raw:: html
+
+   <!-- Atari - DGX-A100 -->
+
+==================== ======= ======== ======== ======== ========= ========= ========= ========= ========= ========= ========= ========== ==========
+Atari - DGX-A100     1       2        4        8        16        32        64        96        128       160       192       224        256
+==================== ======= ======== ======== ======== ========= ========= ========= ========= ========= ========= ========= ========== ==========
+For-loop             4449.38 4587.37  4620.44  4635.26  4617.21   4639.16   4618.30   4594.96   4629.90   4616.15   4640.20   4596.57    4620.50
+Subprocess           4052.06 7832.98  12460.71 18306.28 24754.34  33336.38  43208.56  52435.64  42449.85  32958.90  45312.39  45767.11   71942.74
+Sample-Factory       5563.2  11003.0  21976.3  43891.1  87702.0   175408.8  350855.5  476048.4  505494.8  616958.7  651428.8  679186.5   707494.3
+EnvPool (sync)       7723.96 14865.81 28499.79 52681.02 91970.45  155386.07 243231.45 304423.24 358549.95 367559.69 388419.70 427851.27  427395.89
+EnvPool (async)      8790.69 17866.75 36089.43 70749.63 139540.29 278186.45 451858.26 677504.68 817738.45 838174.97 881210.42 891286.00  874802.04
+EnvPool (numa+async) /       /        /        70629.88 140528.93 279113.15 555426.41 762417.99 936443.47 955620.20 998668.02 1032953.80 1069921.98
+==================== ======= ======== ======== ======== ========= ========= ========= ========= ========= ========= ========= ========== ==========
+
+.. raw:: html
+
+   <!-- Atari - DGX-A100 -->
+
+|image4|
+
+Mujoco
+~~~~~~
+
+.. raw:: html
+
+   <!-- Mujoco - Laptop -->
+
+=============== ======== ======== ======== ======== ======== ========= ========= =========
+Mujoco - Laptop 1        2        3        4        6        8         10        12
+=============== ======== ======== ======== ======== ======== ========= ========= =========
+For-loop        12325.95 12453.54 12861.30 12517.09 12467.92 12447.57  12631.33  12576.39
+Subprocess      8377.65  14851.20 18479.33 23137.12 26667.67 29260.77  36586.01  31952.74
+Sample-Factory  13270.0  25452.0  34882.0  41666.5  58892.0  60657.5   62509.5   60847.0
+EnvPool (sync)  15641.44 30409.65 40063.78 43126.54 58395.28 53269.71  63424.83  66622.24
+EnvPool (async) 20922.70 41279.93 57362.56 73119.43 95542.45 105126.36 100771.24 101603.31
+=============== ======== ======== ======== ======== ======== ========= ========= =========
+
+.. raw:: html
+
+   <!-- Mujoco - Laptop -->
+
+|image5|
+
+.. raw:: html
+
+   <!-- Mujoco - Workstation -->
+
+==================== ======== ======== ========= ========= ========= ========= ========= ========= ========= =========
+Mujoco - Workstation 1        2        4         8         12        16        20        24        28        32
+==================== ======== ======== ========= ========= ========= ========= ========= ========= ========= =========
+For-loop             19472.04 19251.41 19902.03  20076.99  19959.82  19513.40  19460.23  19724.42  20297.76  19797.03
+Subprocess           14428.85 26943.13 48700.27  71303.02  89901.77  102833.40 93676.48  97473.05  105432.15 102533.10
+Sample-Factory       20854.0  40113.5  78408.5   156563.0  225075.0  268005.5  284237.5  296082.5  305235.0  309264.5
+EnvPool (sync)       25725.25 50531.72 90808.85  180372.40 212389.98 309341.24 282954.27 326454.83 357376.48 380950.25
+EnvPool (async)      34500.65 68382.03 133496.84 265710.65 383015.28 478845.88 511142.63 538558.16 566014.54 582445.50
+==================== ======== ======== ========= ========= ========= ========= ========= ========= ========= =========
+
+.. raw:: html
+
+   <!-- Mujoco - Workstation -->
+
+|image6|
+
+.. raw:: html
+
+   <!-- Mujoco - TPU-VM -->
+
+==================== ======== ======== ======== ========= ========= ========= ========= ========= ========= ========= =========
+Mujoco - TPU-VM      1        2        4        8         16        24        32        48        64        80        96
+==================== ======== ======== ======== ========= ========= ========= ========= ========= ========= ========= =========
+For-loop             9960.98  10239.58 10186.08 10473.73  10201.70  10370.85  10454.78  10460.48  10455.71  10360.71  10386.68
+Subprocess           7236.32  13788.93 25054.73 40668.40  64148.06  60409.58  70747.21  78947.79  87403.16  79734.62  81964.35
+Sample-Factory       11008.0  21368.0  42730.0  83475.5   153976.0  222311.5  280664.5  406916.5  432212.0  449143.0  461515.0
+EnvPool (sync)       13706.61 26587.92 49074.86 92444.28  155288.26 181397.00 231293.39 283748.86 250586.54 268296.99 296680.68
+EnvPool (async)      18195.81 37359.25 78337.13 148284.57 259915.75 386448.09 512987.78 745083.58 801768.88 857586.18 887539.80
+EnvPool (numa+async) /        35804.57 75467.72 147281.29 284323.79 412165.16 516120.17 755509.66 816405.50 868455.12 896830.21
+==================== ======== ======== ======== ========= ========= ========= ========= ========= ========= ========= =========
+
+.. raw:: html
+
+   <!-- Mujoco - TPU-VM -->
+
+|image7|
+
+.. raw:: html
+
+   <!-- Mujoco - DGX-A100 -->
+
+==================== ======== ======== ======== ========= ========= ========= ========== ========== ========== ========== ========== ========== ==========
+Mujoco - DGX-A100    1        2        4        8         16        32        64         96         128        160        192        224        256
+==================== ======== ======== ======== ========= ========= ========= ========== ========== ========== ========== ========== ========== ==========
+For-loop             11018.57 11269.45 11059.39 11250.06  11505.15  11328.79  11568.72   11485.74   11245.55   11478.49   11430.16   11151.71   11199.28
+Subprocess           8814.10  17201.64 27106.27 44383.63  62785.60  83054.19  151352.88  158797.86  148815.92  116200.41  163656.36  147653.41  161599.97
+Sample-Factory       11870.0  24602.0  48577.0  96826.5   193800.5  381208.5  761752.0   985909.0   1249369.5  1332128.5  1397427.5  1318249.0  1573262.0
+EnvPool (sync)       16024.43 31899.44 61605.04 114488.28 228492.88 388624.94 656277.80  832101.96  949787.15  858298.85  945808.57  813799.36  849410.96
+EnvPool (async)      21177.71 44025.65 92312.35 176135.82 354006.02 700052.08 1167838.03 1678787.71 1730102.62 2052844.58 2185146.77 2355604.96 2363863.67
+EnvPool (numa+async) /        /        /        170348.47 340269.34 693793.45 1388410.00 1920762.84 2341562.20 2569997.03 2776143.15 2964886.91 3134286.77
+==================== ======== ======== ======== ========= ========= ========= ========== ========== ========== ========== ========== ========== ==========
+
+.. raw:: html
+
+   <!-- Mujoco - DGX-A100 -->
+
+|image8|
+
+.. |image0| image:: ../_static/images/throughput/throughput.png
+.. |image1| image:: ../_static/images/throughput/Atari_Laptop.png
+.. |image2| image:: ../_static/images/throughput/Atari_Workstation.png
+.. |image3| image:: ../_static/images/throughput/Atari_TPU-VM.png
+.. |image4| image:: ../_static/images/throughput/Atari_DGX-A100.png
+.. |image5| image:: ../_static/images/throughput/Mujoco_Laptop.png
+.. |image6| image:: ../_static/images/throughput/Mujoco_Workstation.png
+.. |image7| image:: ../_static/images/throughput/Mujoco_TPU-VM.png
+.. |image8| image:: ../_static/images/throughput/Mujoco_DGX-A100.png
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index f9ce2514..692ca856 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -51,3 +51,4 @@ tran
 Reacher
 golang
 uncomment
+Brax


Test patch:
diff --git a/benchmark/numa_test.sh b/benchmark/numa_test.sh
index f1ee3bf8..dfcdb5a0 100755
--- a/benchmark/numa_test.sh
+++ b/benchmark/numa_test.sh
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-total=`seq 0 $1`
+total=`seq 0 $(expr $1 - 1)`
 shift
 for i in $total
 do
diff --git a/benchmark/benchmark.py b/benchmark/test_envpool.py
similarity index 69%
rename from benchmark/benchmark.py
rename to benchmark/test_envpool.py
index 5f0365eb..fcc1ff69 100644
--- a/benchmark/benchmark.py
+++ b/benchmark/test_envpool.py
@@ -26,13 +26,13 @@
 ====
 
 To test with NUMA, first use ``numactl -s`` to see the number of NUMA core;
-then use ``./numa_test.sh `YOUR_CORE_NUM-1` python3 benchmark.py [args]``.
+then use ``./numa_test.sh `YOUR_CORE_NUM` python3 test_envpool.py [args]``.
 
 In DGX-A100 (8 NUMA core), we use the following script to achieve overall
 1M+ FPS:
 ::
 
-  ./numa_test.sh 7 python3 benchmark.py --num-envs 100 --batch-size 32 \
+  ./numa_test.sh 8 python3 test_envpool.py --num-envs 100 --batch-size 32 \
     --thread-affinity-offset -1
 
 Note: When using NUMA, it's better to disable thread affinity by setting
@@ -43,35 +43,50 @@
 import time
 
 import numpy as np
+import tqdm
 
 import envpool
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--task", type=str, default="Pong-v5")
+  parser.add_argument(
+    "--env", type=str, default="atari", choices=["atari", "mujoco", "vizdoom"]
+  )
   parser.add_argument("--num-envs", type=int, default=645)
   parser.add_argument("--batch-size", type=int, default=248)
   # num_threads == 0 means to let envpool itself determine
   parser.add_argument("--num-threads", type=int, default=0)
   # thread_affinity_offset == -1 means no thread affinity
   parser.add_argument("--thread-affinity-offset", type=int, default=0)
-  parser.add_argument("--total-iter", type=int, default=50000)
+  parser.add_argument("--total-step", type=int, default=50000)
+  parser.add_argument("--seed", type=int, default=0)
   args = parser.parse_args()
-  env = envpool.make_gym(
-    args.task,
+  print(args)
+  task_id = {
+    "atari": "Pong-v5",
+    "mujoco": "Ant-v3",
+    "vizdoom": "HealthGathering-v1",
+  }[args.env]
+  kwargs = dict(
     num_envs=args.num_envs,
     batch_size=args.batch_size,
     num_threads=args.num_threads,
     thread_affinity_offset=args.thread_affinity_offset,
-    use_inter_area_resize=False,
   )
+  if args.env in ["atari", "vizdoom"]:
+    kwargs.update(use_inter_area_resize=False)
+  env = envpool.make_gym(task_id, **kwargs)
   env.async_reset()
-  action = np.ones(args.batch_size, dtype=np.int32)
+  env.action_space.seed(args.seed)
+  action = np.array(
+    [env.action_space.sample() for _ in range(args.batch_size)]
+  )
   t = time.time()
-  for _ in range(args.total_iter):
+  for _ in tqdm.trange(args.total_step):
     info = env.recv()[-1]
     env.send(action, info["env_id"])
   duration = time.time() - t
-  fps = args.total_iter * args.batch_size / duration * 4
+  frame_skip = getattr(env.spec.config, "frame_skip", 1)
+  fps = args.total_step * args.batch_size / duration * frame_skip
   print(f"Duration = {duration:.2f}s")
   print(f"EnvPool FPS = {fps:.2f}")
diff --git a/benchmark/test_gym.py b/benchmark/test_gym.py
new file mode 100644
index 00000000..4c26316d
--- /dev/null
+++ b/benchmark/test_gym.py
@@ -0,0 +1,56 @@
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import time
+
+import gym
+import tqdm
+from atari_wrappers import wrap_deepmind
+
+
+def run(env, num_envs, total_step, async_):
+  if env == "atari":
+    task_id = "PongNoFrameskip-v4"
+    frame_skip = 4
+    env = gym.vector.make(
+      task_id, num_envs, async_, lambda e:
+      wrap_deepmind(e, episode_life=False, clip_rewards=False, frame_stack=4)
+    )
+  elif env == "mujoco":
+    task_id = "Ant-v3"
+    frame_skip = 5
+    env = gym.vector.make(task_id, num_envs, async_)
+  else:
+    raise NotImplementedError(f"Unknown env {env}")
+  env.seed(0)
+  env.reset()
+  action = env.action_space.sample()
+  t = time.time()
+  for _ in tqdm.trange(total_step):
+    env.step(action)
+  print(f"FPS = {frame_skip * total_step * num_envs / (time.time() - t):.2f}")
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+    "--env", type=str, default="atari", choices=["atari", "mujoco"]
+  )
+  parser.add_argument("--async_", action="store_true")
+  parser.add_argument("--num-envs", type=int, default=10)
+  parser.add_argument("--total-step", type=int, default=5000)
+  args = parser.parse_args()
+  print(args)
+  run(args.env, args.num_envs, args.total_step, args.async_)