DOCSEARCHbyAlgolia
Submit Docs

CleanRL

repository·master·Indexed 27 days ago

https://github.com/vwxyzjn/cleanrl

A Deep Reinforcement Learning library providing high-quality, single-file implementations of RL algorithms, including PPO, DQN, C51, SAC, DDPG, TD3, PPG, RND, and Qdagger. Designed for research-friendly use, it features easy-to-read code, Tensorboard logging, and scalability to cloud environments like AWS. It supports various environments such as Atari, Mujoco, Procgen, Envpool, and PettingZoo, with implementations available in frameworks like JAX.

Tokens
65.9K
Snippets
130
Records
281
Agent score
91%

What's inside cleanrl

  1. Overview of CleanRL

    master

    CleanRL is a Deep Reinforcement Learning library providing high-quality, single-file implementations of online DRL algorithms. It is designed for research friendliness, making it easy to understand algorithm details without navigating a complex modular library.

    Key features include:

    • Single-file implementation: Every algorithm variant is contained in a single standalone file (e.g., ppo_atari.py).
    • Benchmarked implementations: Supports 7+ algorithms and 34+ games.
    • Logging & Reproducibility: Includes Tensorboard logging, local reproducibility via seeding, and gameplay video capturing.
    • Experiment Management: Integrates with Weights and Biases and supports cloud scaling via Docker and AWS Batch.

    Note: CleanRL is not a modular library and is not intended to be imported. It is best used for understanding implementation details or prototyping advanced features where minimal lines of code provide a better debugging experience than subclassing in modular libraries.

  2. Overview of PPO Implementations

    master

    CleanRL provides several variants of the Proximal Policy Optimization (PPO) algorithm, each optimized for specific environments and hardware.

    Key variants include:

    • ppo.py: For classic control tasks (e.g., CartPole-v1) using Box observations and Discrete action spaces.
    • ppo_atari.py: For Atari games using convolutional layers and standard pre-processing.
    • ppo_continuous_action.py: For continuous action spaces with Mujoco-specific optimizations.
    • ppo_atari_envpool.py: High-speed Atari training using the Envpool vectorized environment (Linux only).
    • ppo_atari_envpool_xla_jax.py: Atari training using Envpool with JAX and XLA.
    • ppo_atari_multigpu.py: Atari training leveraging multiple GPUs.
    • ppo_pettingzoo_ma_atari.py: For multi-agent Atari environments via Pettingzoo.

  4. Overview of Robust Policy Optimization (RPO)

    master
    Robust Policy Optimization (RPO) is an enhancement to the PPO algorithm designed to encourage high-entropy actions and improve the representation of the action space. It achieves this by perturbing the mean of the action distribution using a random number drawn from a Uniform distribution.

  5. Overview of Parallel Q Network (PQN)

    master

    Parallel Q Network (PQN) is a parallelized version of the Deep Q-learning algorithm designed for higher efficiency than DQN. It achieves this by using multiple agents to interact with the environment in parallel. Key differences from DQN include:

    1. Absence of a replay buffer and target networks.
    2. Use of layer normalizations and parallel environments.

    It is suitable for both classic control tasks and Atari games.

  6. Explore implemented RL algorithm variants

    master

    CleanRL provides high-quality, single-file implementations of various Reinforcement Learning algorithms. Each implementation is documented with its usage, logged metrics, and implementation details.

    Key algorithm families include:

    • PPO (Proximal Policy Optimization): Includes variants for Atari, continuous actions, LSTM, EnvPool, Procgen, Multi-GPU, PettingZoo (multi-agent), and Isaac Gym.
    • DQN (Deep Q-Learning): Includes standard, Atari, and JAX-based versions.
    • C51 (Categorical DQN): Includes standard and Atari variants.
    • SAC (Soft Actor-Critic): Continuous action implementation.
    • DDPG (Deep Deterministic Policy Gradient): Continuous action and JAX-based versions.
    • TD3 (Twin Delayed Deep Deterministic Policy Gradient): Continuous action and JAX-based versions.
    • PPG (Phasic Policy Gradient): Procgen implementation.
    • RND (Random Network Distillation): EnvPool implementation.

  7. Deep Deterministic Policy Gradient (DDPG) Overview

    master
    DDPG is a Deep Reinforcement Learning (DRL) algorithm designed for continuous control. It extends DQN to continuous action spaces by using a deterministic actor that directly outputs continuous actions. It incorporates DQN techniques such as a replay buffer and target networks.

  8. Available Reinforcement Learning Algorithms

    master

    CleanRL provides single-file implementations for several Reinforcement Learning algorithms and their variants. Key algorithms include:

    • Proximal Policy Optimization (PPO): Includes variants for Atari, continuous actions, LSTM, EnvPool (speedup), Procgen, Multi-GPU, and PettingZoo.
    • Deep Q-Learning (DQN): Includes standard, Atari, and JAX-based implementations.
    • Categorical DQN (C51): Includes standard, Atari, and JAX-based implementations.
    • Soft Actor-Critic (SAC): Includes continuous action and Atari variants.
    • Deep Deterministic Policy Gradient (DDPG): Includes continuous action and JAX-based implementations.
    • Twin Delayed Deep Deterministic Policy Gradient (TD3): Includes continuous action and JAX-based implementations.
    • Phasic Policy Gradient (PPG): Implemented for Procgen.
    • Random Network Distillation (RND): Implemented with EnvPool.
    • Qdagger: Implemented for DQN Atari with IMPALA CNN (PyTorch and JAX).

  9. Use ppo_atari_envpool_xla_jax.py for high-performance Atari training

    master

    The ppo_atari_envpool_xla_jax.py implementation is optimized for Atari games using the following stack:

    • Envpool: A fast vectorized environment using the experimental XLA interface.
    • Jax, Flax, and Optax: Replaces torch for high-performance computation.
    • Atari Pre-processing: Uses convolutional layers and standard Atari-based techniques.
    • Observation Space: Works with pixel Box observations of shape (210, 160, 3) and Discrete action spaces.

    Important Compatibility Notes:

    • OS Support: This script does not work on Windows or macOS. It is intended for Linux.
    • EnvPool Behavior: Due to a compatibility difference between envpool and gym, the observation returned upon episode termination/truncation is the last state ($s_{last}$) rather than the new initial state ($s_{new}$). This does not appear to impact performance.

  10. Set up the CleanRL development environment

    master

    To set up a development environment, clone the repository and follow the official installation and usage documentation to configure a uv environment. To install additional dependencies required for documentation, run:

    uv pip install ".[docs]"

    To build and serve the documentation locally, use:

    uv run mkdocs serve
    uv pip install ".[docs]"

uv run mkdocs serve

  11. Use `dqn_atari_jax.py` for Atari games

    master

    dqn_atari_jax.py is a JAX-based implementation of Deep Q-Learning optimized for Atari games. It uses JAX, Flax, and Optax instead of PyTorch, offering a 25%-50% speedup. It is designed for Atari's pixel Box observation space (210, 160, 3) and Discrete action spaces.

    Note: JAX does not work on Windows. Use Windows Subsystem for Linux (WSL) if you are on a Windows machine.

    Installation (using uv/poetry)

    uv pip install ".[atari, jax]"
uv pip install --upgrade "jax[cuda11_cudnn82]==0.4.8" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

    Installation (using pip)

    pip install -r requirements/requirements-atari.txt
pip install -r requirements/requirements-jax.txt
pip install --upgrade "jax[cuda]==0.3.17" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
    uv run python cleanrl/dqn_atari_jax.py --env-id BreakoutNoFrameskip-v4
uv run python cleanrl/dqn_atari_jax.py --env-id PongNoFrameskip-v4

  12. Use JAX-based algorithm variants for faster training

    master

    CleanRL supports JAX-based implementations of several algorithms, which are typically significantly faster than their PyTorch equivalents.

    Available JAX variants and their performance benefits:

    • dqn_atari_jax.py: ~25% faster than dqn_atari.py.
    • td3_continuous_action_jax.py: ~2.5-4x faster than td3_continuous_action.py.
    • ddpg_continuous_action_jax.py: ~2.5-4x faster than ddpg_continuous_action.py.
    • ppo_atari_envpool_xla_jax.py: ~3x faster than OpenAI Baselines' PPO.