DOCSEARCHbyAlgolia
Submit Docs

EnvPool Documentation

repository·main·Indexed 23 days ago

https://github.com/sail-sg/envpool

A high-performance C++-based batched environment pool for Reinforcement Learning. EnvPool provides massive throughput for environment parallelization, supporting Gymnasium and dm_env APIs across Linux, macOS, and Windows. It features synchronous and asynchronous execution modes, built-in batched rendering, XLA support for JAX, and integration with RL libraries such as Stable-Baselines3, Tianshou, ACME, CleanRL, and rl_games.

Tokens
30.3K
Snippets
41
Records
145
Agent score
80%

What's inside EnvPool

  1. Overview of EnvPool

    main

    EnvPool is a high-performance, C++-based batched environment pool designed for Reinforcement Learning (RL). It uses pybind11 and a thread pool to achieve massive throughput (e.g., ~1M Atari FPS or ~3M MuJoCo FPS on high-end hardware).

    Key Features:

    • Multi-platform: Supports Linux, macOS, and Windows.
    • API Compatibility: Supports both gymnasium and DeepMind dm_env APIs.
    • Execution Modes: Supports both synchronous and asynchronous execution, as well as single and multi-player environments.
    • Rendering: Built-in rendering with batched rgb_array output and human display mode.
    • Integration: Compatible with major RL libraries like Stable-Baselines3, Tianshou, ACME, CleanRL, and rl_games.
    • XLA Support: Provides an XLA interface for JAX jit functions.

  2. Overview of EnvPool features and performance

    main

    EnvPool is a high-performance, C++-based batched environment pool designed for Reinforcement Learning (RL). It uses a thread pool and pybind11 to provide high throughput for environment simulations.

    Key Features:

    • API Compatibility: Supports both OpenAI gym and DeepMind dm_env APIs.
    • Execution Modes: Supports both synchronous and asynchronous execution.
    • Environment Types: Supports single-player and multi-player environments.
    • High Throughput: Capable of ~1M Atari frames/s or ~3M Mujoco steps/s on high-end hardware (e.g., DGX-A100 with 256 CPU cores).
    • Efficiency: Provides significantly higher throughput compared to Python subprocess-based vector environments (up to ~20x).
    • Integration: Compatible with XLA (via JAX jit) and major RL libraries like Stable-Baselines3, Tianshou, ACME, CleanRL, and rl_games.
    • General Purpose: Unlike purely GPU-based solutions (e.g., Brax or Isaac-gym), EnvPool is a general solution for speeding up various RL environment parallelizations.

  3. Compare EnvPool speedup against Python baseline

    main

    EnvPool provides significant speedups over standard Python environments for reinforcement learning tasks. The speedup is measured in FPS (Frames Per Second) across different systems (Laptop, Workstation, TPU-VM, DGX-A100) and environments (Atari Pong-v5, Mujoco Ant-v3, dm_control cheetah run).

    Typical speedup factors observed:

    • Laptop: ~1.27x to 1.87x
    • Workstation: ~1.32x to 1.85x
    • TPU-VM: ~1.38x to 1.88x
    • DGX-A100: ~1.45x to 2.07x

  4. Use Sudoku puzzle databases in EnvPool

    main

    EnvPool includes packaged Sudoku puzzle databases from the Jumanji project. These assets are used for reset sampling in the Sudoku-v0 and Sudoku-very-easy-v0 environments.

    Available puzzle sets:

    • Sudoku-v0: Uses the 10000_mixed_puzzles.npy dataset.
    • Sudoku-very-easy-v0: Uses the 1000_very_easy_puzzles.npy dataset.

  5. Performance Benchmarks: Atari Throughput

    main

    EnvPool demonstrates high throughput for Atari environments, scaling significantly better than standard For-loop or Subprocess methods.

    Key observations:

    • Sync vs Async: EnvPool (async) generally provides higher throughput than EnvPool (sync) as the number of environments increases.
    • Scaling: On high-end hardware like DGX-A100, EnvPool (async) can achieve throughput exceeding 800,000 FPS for Atari environments.
    • NUMA Optimization: On certain systems (e.g., TPU-VM, DGX-A100), EnvPool (numa+async) can provide even higher performance by optimizing for Non-Uniform Memory Access architectures.

  6. Performance Benchmarks: Single Environment Speedup

    main
    EnvPool provides significant speedups compared to standard Python environments for single-environment execution. Across various hardware systems (Laptop, Workstation, TPU-VM, DGX-A100), EnvPool consistently outperforms Python for tasks like Atari Pong-v5, Mujoco Ant-v3, and dm_control cheetah. Typical speedup factors range from approximately 1.27x to over 2.0x depending on the environment and hardware.

  7. Available Toy Text Environments

    main

    EnvPool provides several 'toy text' environments that are largely compatible with Gymnasium and bsuite. These environments are useful for testing reinforcement learning algorithms on discrete, low-dimensional state spaces.

    Supported environments include:

    • Catch-v0: Intercept falling balls with a paddle.
    • FrozenLake-v1 and FrozenLake8x8-v1: Navigate a grid world with walkable tiles and uncertain movement.
    • Taxi-v3: Pick up and drop off passengers at specific locations in a grid.
    • NChain-v0: Move along a linear chain to reach a large reward at the end.
    • CliffWalking-v1 and CliffWalkingSlippery-v1: Navigate a grid while avoiding a cliff.
    • Blackjack-v1: A card game aiming for a sum near 21.

  8. Performance Benchmarks: Mujoco Throughput

    main

    EnvPool provides massive throughput improvements for Mujoco environments compared to standard Python implementations.

    Key observations:

    • Async Advantage: EnvPool (async) scales exceptionally well. On a DGX-A100, it can reach over 2,300,000 FPS.
    • NUMA Optimization: For Mujoco on TPU-VM and DGX-A100, using EnvPool (numa+async) offers a substantial performance boost over standard EnvPool (async), particularly as the number of environments scales up.

  9. Understand Go observation shapes and scoring

    main

    For PGX-compatible Go tasks (Go9x9-v1, Go13x13-v1, Go19x19-v1):

    • Observation Shape: Boolean shape (N, N, 17) using AlphaGo Zero history planes.
    • Action Space: N * N + 1 discrete actions, where the final action is pass.
    • Scoring: Uses Tromp-Taylor scoring.
    • Legal Actions: Uses SSK (Simple Seki Knowledge) legal-action filtering.

    For Chinese-rule variants:

    • Scoring: Chinese area scoring (stones plus empty regions bordered by exactly one color; neutral regions are not counted).
    • Superko: Positional superko moves are masked as illegal actions rather than being accepted and then resulting in a terminal loss.

  10. Compare FrozenLake-v1 and FrozenLake8x8-v1

    main

    Both environments involve controlling a character in a grid world where movement is uncertain. The primary differences are the grid size and episode constraints:

    FeatureFrozenLake-v1FrozenLake8x8-v1
    Size48
    Max Episode Steps100200
    Reward Threshold0.70.85

  11. Understand Auto Reset behavior in EnvPool

    main

    EnvPool enables auto-reset by default. When an environment reaches its termination condition (e.g., max_episode_steps), the next call to env.step(action) will automatically trigger an env.reset() instead of executing the action in the finished episode.

    Key behaviors of auto-reset:

    • If a call triggers a reset, the provided action is discarded.
    • The call returns done = False and reward = 0 (representing the start of a new episode).
    • The elapsed step count resets to 0.

    Example Scenario: If an environment has max_episode_steps = 3, calling env.step(action) five times results in the following sequence:

    #User CallActualElapsedMisc
    1env.step(a)env.reset()0
    2env.step(a)env.step(a)1
    3env.step(a)env.step(a)2
    4env.step(a)env.step(a)3Hit max_episode_steps
    5env.step(a)env.reset()0

  12. Use CliffWalking-v1 and CliffWalkingSlippery-v1

    main

    The CliffWalking environments use a 4x12 matrix. The agent starts at [3, 0] and must reach the goal at [3, 11]. Stepping into the cliff ([3, 1..10]) results in a -100 reward and a reset to the start. Each step incurs a -1 reward.

    • CliffWalking-v1: Deterministic movement.
    • CliffWalkingSlippery-v1: Stochastic movement where each action samples uniformly from the intended direction and its two perpendicular alternatives.

    Note: CliffWalking-v0 is a legacy alias for CliffWalking-v1, and tabular/CliffWalking-v0 is also available.