DOCSEARCHbyAlgolia
Submit Docs

rl_games

repository·master·Indexed 23 days ago

https://github.com/denys88/rl_games

A high-performance reinforcement learning library implemented in PyTorch (version 2.0.0). It supports advanced algorithms such as PPO and SAC and is optimized for end-to-end GPU-accelerated training pipelines using environments like NVIDIA Isaac Gym and Brax. The library features custom Triton kernels for GAE, support for torch.compile, and integration with Weights and Biases for experiment tracking.

Tokens
30.3K
Snippets
48
Records
84
Agent score
80%

What's inside rl_games

  1. How models and networks work in RL Games

    master

    In RL Games, there is a distinction between the Network (the actual neural architecture) and the Model (the wrapper that manages the network and its logic). Understanding this hierarchy is key to customizing architectures:

    1. Network Builder (NetworkBuilder): Found in algos_torch.network_builder (e.g., A2CBuilder, SACBuilder). It contains a nested Network class (derived from torch.nn.Module) which implements the forward function. This nested class takes a dictionary of tensors (like observations) and returns a tuple of tensors.
    2. Model (BaseModel): Found in algos_torch.models (e.g., ModelA2C, ModelSACContinuous). A Model contains the nested Network class and a build function to instantiate it.
    3. Algorithm usage: In standard agent/player algorithms, self.model refers to the instance of the model network class, while self.network refers to the instance of the model class.
    4. Model Builder (ModelBuilder): Located in algos_torch.model_builder, this class manages loading models via the load function based on a specified name.

  2. SMAC Training Implementation Details

    master

    When training for SMAC environments in rl_games, the following implementation details are noted for achieving high performance:

    • Algorithm: PPO (Proximal Policy Optimization).
    • Network Architecture: 4 frames + conv1d actor-critic network. Simple MLP networks are noted to perform poorly on hard environments.
    • Information Sharing: Full state information is not used for the critic; both the actor and critic receive the same agent observations.
    • Hyperparameters: miniepoch should be set to 1, as higher numbers were found to be ineffective for these environments.
    • Environment Settings: Agents are controlled independently with restricted information, typically trained at default difficulty level 7.

  3. Configure the Player for inference

    master

    The player block is used for running the trained policy in an environment (inference/evaluation mode):

    • render: Set to True to visualize the environment.
    • deterministic: Set to True to use a deterministic policy (e.g., argmax or mu) instead of a stochastic one.
    • use_vecenv: Set to True to use a vectorized environment for the player.
    • games_num: Specifies how many games to run during player mode.

  4. How Gymnasium compatibility works in RL Games

    master

    RL Games provides compatibility between the older Gym and the newer Gymnasium APIs through the gym_compat module. This module abstracts away version-specific differences to provide a unified interface.

    Key behaviors of gym_compat:

    1. Automatic Backend Detection: It detects the Python version to select the backend (Gym for Python 3.8, Gymnasium for Python 3.9+).
    2. Unified Interface: It provides a make function that handles API differences automatically.
    3. API Wrapping: It wraps Gymnasium environments to maintain backward compatibility with the legacy Gym API (specifically handling the transition from 4-value step returns to the Gymnasium standard, and 1-value reset returns).

  5. Avoid RNN/LSTM incompatibility with torch.compile

    master

    The reduce-overhead and max-autotune modes use CUDA graphs, which are incompatible with RNN/LSTM models. Using these modes with recurrent architectures will cause:

    1. Rollout buffer corruption: Sequential RNN outputs get overwritten by subsequent graph replays.
    2. Backward pass failure: LSTM internal allocations cause RuntimeError: storage data ptrs are not allocated in pool during loss.backward().

    Requirement: Always use torch_compile: "default" for any configuration involving RNN/LSTM models.

  6. Configure Triton Kernels for performance

    master

    When Triton is installed, rl_games automatically uses custom Triton kernels for performance-critical operations like GAE (Generalized Advantage Estimation). This replaces Python loops with a single fused GPU kernel.

    Disable Triton: If you need to disable Triton, set the RLG_NO_TRITON=1 environment variable:

    RLG_NO_TRITON=1 python runner.py --train --file rl_games/configs/mujoco/ant.yaml

    Benchmark Triton: To see the speedup on your hardware, run the benchmark script:

    python benchmarks/bench_triton_gae.py
    RLG_NO_TRITON=1 python runner.py --train --file rl_games/configs/mujoco/ant.yaml

  7. Configure the Network architecture

    master

    The network block defines the architecture of the agent. You can specify different types of layers and blocks:

    • CNN Block: Use type (conv2d or conv1d), activation, and initializer. You can define multiple layers using the convs list, specifying filters, kernel_size, strides, and padding for each.
    • MLP Block: Define layer sizes using units (e.g., [512, 256, 128]). Supports d2rl architecture and custom activation and initializer.
    • RNN Block: Supports lstm and gru. Specify units and layers. Use before_mlp to decide if the RNN is applied before the MLP block.
    • Space Configuration: For continuous spaces, you can configure mu_activation, sigma_activation, mu_init, and sigma_init. If using a logstd model, a sigma_init of 0 is recommended.

  8. Configure the RL training parameters

    master

    The config block controls the reinforcement learning hyperparameters and training loop behavior:

    • Optimization: Set learning_rate and lr_schedule (None, linear, or adaptive). adaptive is recommended for continuous control. Use grad_norm and truncate_grads (True) to stabilize training.
    • GAE & Rewards: Configure gamma (discount) and tau (GAE lambda). Use reward_shaper with min_val, max_val, scale_value, or shift to transform rewards.
    • Batching: Define num_actors, horizon_length, and minibatch_size. If minibatch_size_per_env is provided, it overwrites the total minibatch size as minibatch_size_per_env * num_actors.
    • Stability: Use normalize_input and normalize_value to apply running mean/std normalization. value_bootstrap is useful for locomotion environments.

  9. How the rl_games architecture works

    master

    The rl_games library is structured around four main components that work together to execute RL experiments:

    1. rl_games.torch_runner.Runner (Main Script): The entry point that instantiates algorithms and environments based on the provided configuration. It manages the execution loop (run_train() or run_play()) and handles metric logging via an observer (defaulting to DefaultAlgoObserver using Tensorboard).
    2. rl_games.common.Objectfactory() (Instantiating Algos): A factory pattern used to create algorithms and players. It uses a registry of builder functions to map string names in config files to actual class instances.
    3. RL Algorithms: Classes (e.g., inheriting from rl_games.algos_torch.A2CBase) that contain the logic for agent updates. The algorithm is responsible for instantiating the environment.
    4. Environments (vecenv & env_configurations): A two-step lookup system used to instantiate environments. The algorithm looks up an env_name in rl_games.common.env_configurations.configurations to find the vecenv_type and an env_creator function. The vecenv_type then determines which vectorized environment class (from rl_games.common.vecenv.vecenv_config) is used to wrap the environment created by the env_creator.

  10. Best practices for long-horizon manipulation training

    master

    When training long-horizon manipulation tasks (like ppo_wujihand_reorient.yaml), be aware of the following technical considerations to avoid training instability:

    • Entropy Runaway: A positive entropy bonus on a global fixed_sigma can cause sigma runaway over long runs (1B+ frames). Use state-dependent sigma with sigma_parametrization: softplus and a defined min_sigma (e.g., 0.2) to match the task's exploration floor.
    • Learning Rate (LR) Management: Use adaptive LR within a specific band (e.g., min_lr 1e-4 to max_lr 2e-4). The floor maintains the proven training rate, while the cap prevents late-training collapse caused by large negative return bursts from environment penalties.
    • Minibatch Size: Avoid small minibatches (≤10240) as they make the KL-adaptive scheduler noisy. For WujiHand, a minibatch of 16384 is recommended.
    • Optimizer Steps: Ensure sufficient optimizer steps per iteration; starving the discovery phase can occur if there are too few steps relative to the minibatch size.

  11. How Population Based Training (PBT) works

    master

    PBT in rl_games implements the DexPBT lineage. It manages a population of independent training processes, each identified by a unique policy_idx.

    Every interval_steps environment frames, each process:

    1. Saves a scored checkpoint to a shared workspace directory.
    2. Compares its objective against the population.
      • Leaders are defined as having a score > max(mean + threshold_std * std, mean + threshold_abs).
      • Underperformers are the mirror image below the mean.
    3. If a process is an underperformer, it re-executes itself from a random leader's checkpoint, applying multiplicative mutations to whitelisted hyperparameters.

    Important: Objective Configuration When PBT is enabled, you must specify an objective address. This is a dotted path used to read the score from environment info dictionaries. Because info layouts vary by backend, you must provide the correct path:

    • For Isaac Lab-style nested infos: episode.Episode_Reward/success
    • For flat info dicts: scores

    It is recommended to use a true task metric rather than raw reward if reward shaping is non-stationary.

  12. How agents and players work in RL Games

    master

    RL Games distinguishes between two primary modes of operation:

    1. Agents (Training): An agent is an instance of an algorithm (e.g., A2CAgent, SACAgent) created by an algo_factory. The core lifecycle involves calling agent.train(), which executes a loop of gathering rollout data (play_steps), preparing datasets (prepare_dataset), and performing gradient updates (train_actor_critic / calc_grad).
    2. Players (Testing/Inference): A player is an instance of a testing class (e.g., PPOPlayerContinuous, SACPlayer) created by a player_factory. The core lifecycle involves calling player.run(), which resets the environment and iteratively generates actions from observations using the model.

    Customization: You can extend the library by inheriting from BaseAlgorithm (to create new agents) or BasePlayer (to create new players), and then registering them using algo_factory.register_builder or player_factory.register_builder respectively.