Package: envpool
EnvPoolis 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:Atari gamesMujoco (gym)Classic control RL envs: CartPole, MountainCar, Pendulum, AcrobotToy text RL envs: Catch, FrozenLake, Taxi, NChain, CliffWalking, BlackjackViZDoom single playerDeepMind Control SuiteBox2DProcgenMinigridHere are EnvPool's several highlights:Compatible with OpenAIgymAPIs, DeepMinddm_envAPIs, andgymnasiumAPIs;Manage a pool of envs, interact with the envs in batched APIs by default;Support both synchronous execution and asynchronous execution;Support both single player and multi-player environment;Easy C++ developer API to add new envs:Customized C++ environment integration;Free~2xspeedup with only single environment;1 MillionAtari frames /3 MillionMujoco steps per second simulation with 256 CPU cores,~20xthroughput of Python subprocess-based vector env;~3xthroughput of Python subprocess-based vector env on low resource setup like 12 CPU cores;Comparing with existing GPU-based solution (Brax/Isaac-gym), EnvPool is ageneralsolution for various kinds of speeding-up RL environment parallelization;XLA support with JAX jit function:XLA Interface;Compatible with some existing RL libraries, e.g.,Stable-Baselines3,Tianshou,ACME,CleanRL, orrl_games.Stable-Baselines3Pendulum-v1example;TianshouCartPoleexample,Pendulum-v1example,Atari example,Mujoco example, andintegration guideline;ACMEHalfCheetahexample;CleanRLPong-v5example(Solving Pong in 5 mins(tracked experiment));rl_gamesAtari example(2 mins Pongand15 mins Breakout) andMujoco example(5 minsAntandHalfCheetah).Check out ourarXiv paperfor more details!InstallationPyPIEnvPool is currently hosted onPyPI. It requires Python >= 3.7.You can simply install EnvPool with the following command:$pipinstallenvpoolAfter installation, open a Python console and typeimportenvpoolprint(envpool.__version__)If no error occurs, you have successfully installed EnvPool.From SourcePlease refer to theguideline.DocumentationThe tutorials and API documentation are hosted onenvpool.readthedocs.io.The example scripts are underexamples/folder; benchmark scripts are underbenchmark/folder.Benchmark ResultsWe perform our benchmarks with ALE Atari environmentPongNoFrameskip-v4(with environment wrappers fromOpenAI Baselines) and Mujoco environmentAnt-v3on 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; 3) to our knowledge, the fastest RL environment executorSample Factorybefore 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 with Atari and 3 Million frames per second with Mujoco on 256 CPU cores, which is 14.9x / 19.6x of thegym.vector_envbaseline. On a typical PC setup with 12 CPU cores, EnvPool's throughput is 3.1x / 2.9x ofgym.vector_env.Atari Highest FPSLaptop (12)Workstation (32)TPU-VM (96)DGX-A100 (256)For-loop4,8937,9143,9934,640Subprocess15,86347,69946,91071,943Sample-Factory28,216138,847222,327707,494EnvPool (sync)37,396133,824170,380427,851EnvPool (async)49,439200,428359,559891,286EnvPool (numa+async)//373,1691,069,922Mujoco Highest FPSLaptop (12)Workstation (32)TPU-VM (96)DGX-A100 (256)For-loop12,86120,29810,47411,569Subprocess36,586105,43287,403163,656Sample-Factory62,510309,264461,5151,573,262EnvPool (sync)66,622380,950296,681949,787EnvPool (async)105,126582,446887,5402,363,864EnvPool (numa+async)//896,8303,134,287Please refer to thebenchmarkpage for more details.API UsageThe following content shows both synchronous and asynchronous API usage of EnvPool. You can also run the full script atexamples/env_step.pySynchronous APIimportenvpoolimportnumpyasnp# make gym envenv=envpool.make("Pong-v5",env_type="gym",num_envs=100)# or use envpool.make_gym(...)obs=env.reset()# should be (100, 4, 84, 84)act=np.zeros(100,dtype=int)obs,rew,term,trunc,info=env.step(act)Under the synchronous mode,envpoolclosely resemblesopenai-gym/dm-env. It has theresetandstepfunctions with the same meaning. However, there is one exception inenvpool: batch interaction is the default. Therefore, during the creation of the envpool, there is anum_envsargument that denotes how many envs you like to run in parallel.env=envpool.make("Pong-v5",env_type="gym",num_envs=100)The first dimension ofactionpassed to the step function should equalnum_envs.act=np.zeros(100,dtype=int)You don't need to manually reset one environment when any ofdoneis true; instead, all envs inenvpoolhave enabled auto-reset by default.Asynchronous APIimportenvpoolimportnumpyasnp# make asynchronousnum_envs=64batch_size=16env=envpool.make("Pong-v5",env_type="gym",num_envs=num_envs,batch_size=batch_size)action_num=env.action_space.nenv.async_reset()# send the initial reset signal to all envswhileTrue:obs,rew,term,trunc,info=env.recv()env_id=info["env_id"]action=np.random.randint(action_num,size=batch_size)env.send(action,env_id)In the asynchronous mode, thestepfunction is split into two parts: thesend/recvfunctions.sendtakes two arguments, a batch of action, and the correspondingenv_idthat each action should be sent to. Unlikestep,senddoes not wait for the envs to execute and return the next state, it returns immediately after the actions are fed to the envs. (The reason why it is called async mode).env.send(action,env_id)To get the "next states", we need to call therecvfunction. However,recvdoes not guarantee that you will get back the "next states" of the envs you just calledsendon. Instead, whatever envs finishes execution getsrecved first.state=env.recv()Besidesnum_envs, there is one more argumentbatch_size. Whilenum_envsdefines how many envs in total are managed by theenvpool,batch_sizespecifies the number of envs involved each time we interact withenvpool. e.g. There are 64 envs executing in theenvpool,sendandrecveach time interacts with a batch of 16 envs.envpool.make("Pong-v5",env_type="gym",num_envs=64,batch_size=16)There are other configurable arguments withenvpool.make; please check outEnvPool Python interface introduction.ContributingEnvPool is still under development. More environments will be added, and we always welcome contributions to help EnvPool better. If you would like to contribute, please check out ourcontribution guideline.LicenseEnvPool is under Apache2 license.Other third-party source-code and data are under their corresponding licenses.We do not include their source code and data in this repo.Citing EnvPoolIf you find EnvPool useful, please cite it in your publications.@inproceedings{weng2022envpool,
 author ={Weng, Jiayi and Lin, Min and Huang, Shengyi and Liu, Bo and Makoviichuk, Denys and Makoviychuk, Viktor and Liu, Zichen and Song, Yufan and Luo, Ting and Jiang, Yukun and Xu, Zhongwen and Yan, Shuicheng},
 booktitle ={Advances in Neural Information Processing Systems},
 editor ={S. Koyejo and S. Mohamed and A. Agarwal and D. Belgrave and K. Cho and A. Oh},
 pages ={22409--22421},
 publisher ={Curran Associates, Inc.},
 title ={Env{P}ool: A Highly Parallel Reinforcement Learning Environment Execution Engine},
 url ={https://proceedings.neurips.cc/paper_files/paper/2022/file/8caaf08e49ddbad6694fae067442ee21-Paper-Datasets_and_Benchmarks.pdf},
 volume ={35},
 year ={2022}}DisclaimerThis is not an official Sea Limited or Garena Online Private Limited product.