Configuring an RL Agent

In the previous tutorial, we saw how to train an RL agent to solve the cartpole balancing task using the Stable-Baselines3 library. In this tutorial, we will see how to configure the training process to use different RL libraries and different training algorithms.

In the directory scripts/reinforcement_learning, you will find the scripts for different RL libraries. These are organized into subdirectories named after the library name. Each subdirectory contains the training and playing scripts for the library.

To configure a learning library with a specific task, you need to create a configuration file for the learning agent. This configuration file is used to create an instance of the learning agent and is used to configure the training process. Similar to the environment registration shown in the Registering an Environment tutorial, you can register the learning agent with the gymnasium.register method.

The Code

As an example, we will look at the configuration included for the task Isaac-Cartpole-v0 in the isaaclab_tasks package. This is the same task that we used in the Training with an RL Agent tutorial.

gym.register(
    id="Isaac-Cartpole-v0",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg",
        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg",
        "rsl_rl_with_symmetry_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerWithSymmetryCfg",
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_ppo_cfg.yaml",
        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
    },

The Code Explained

Under the attribute kwargs, we can see the configuration for the different learning libraries. The key is the name of the library and the value is the path to the configuration instance. This configuration instance can be a string, a class, or an instance of the class. For example, the value of the key "rl_games_cfg_entry_point" is a string that points to the configuration YAML file for the RL-Games library. Meanwhile, the value of the key "rsl_rl_cfg_entry_point" points to the configuration class for the RSL-RL library.

The pattern used for specifying an agent configuration class follows closely to that used for specifying the environment configuration entry point. This means that while the following are equivalent:

Specifying the configuration entry point as a string

from . import agents

gym.register(
   id="Isaac-Cartpole-v0",
   entry_point="isaaclab.envs:ManagerBasedRLEnv",
   disable_env_checker=True,
   kwargs={
      "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg",
      "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg",
   },
)

Specifying the configuration entry point as a class

from . import agents

gym.register(
   id="Isaac-Cartpole-v0",
   entry_point="isaaclab.envs:ManagerBasedRLEnv",
   disable_env_checker=True,
   kwargs={
      "env_cfg_entry_point": f"{__name__}.cartpole_env_cfg:CartpoleEnvCfg",
      "rsl_rl_cfg_entry_point": agents.rsl_rl_ppo_cfg.CartpolePPORunnerCfg,
   },
)

The first code block is the preferred way to specify the configuration entry point. The second code block is equivalent to the first one, but it leads to import of the configuration class which slows down the import time. This is why we recommend using strings for the configuration entry point.

The reinforcement learning entrypoints are configured by default to read the <library_name>_cfg_entry_point from the kwargs dictionary to retrieve the configuration instance.

For instance, the following code block shows how the Stable-Baselines3 training implementation reads the configuration instance:

Code for train_sb3.py with SB3

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Stable-Baselines3 training logic for the unified reinforcement learning entrypoint."""

from __future__ import annotations

import argparse
import contextlib
import logging
import os
import random
import signal
import sys
import time
from datetime import datetime
from pathlib import Path

from common import (
    add_common_train_args,
    add_isaaclab_launcher_args,
    apply_env_overrides,
    configure_io_descriptors,
    create_isaaclab_env,
    dump_train_configs,
    enable_cameras_for_video,
    set_hydra_args,
    wrap_record_video,
)

import isaaclab_tasks  # noqa: F401

logger = logging.getLogger(__name__)

# PLACEHOLDER: Extension template (do not remove this comment)
with contextlib.suppress(ImportError):
    import isaaclab_tasks_experimental  # noqa: F401


def _cleanup_pbar(*args):
    """Stop training and clean up rich progress bars on Ctrl+C."""
    import gc

    tqdm_objects = [obj for obj in gc.get_objects() if "tqdm" in type(obj).__name__]
    for tqdm_object in tqdm_objects:
        if "tqdm_rich" in type(tqdm_object).__name__:
            tqdm_object.close()
    raise KeyboardInterrupt


def _parse_args(argv: list[str]) -> argparse.Namespace:
    """Parse Stable-Baselines3 training arguments."""
    parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
    add_common_train_args(
        parser,
        agent_default="sb3_cfg_entry_point",
        agent_help="Name of the RL agent configuration entry point.",
        include_distributed=False,
    )
    parser.add_argument("--log_interval", type=int, default=100_000, help="Log data every n timesteps.")
    parser.add_argument("--checkpoint", type=str, default=None, help="Continue the training from checkpoint.")
    parser.add_argument(
        "--keep_all_info",
        action="store_true",
        default=False,
        help="Use a slower SB3 wrapper but keep all the extra training info.",
    )
    from isaaclab_tasks.utils import setup_preset_cli

    add_isaaclab_launcher_args(parser)
    # setup_preset_cli registers preset-selection help text + runs parse_known_args; the
    # physics=/renderer=/presets= tokens pass through the remainder for hydra to parse later.
    args_cli, hydra_args = setup_preset_cli(parser, argv)
    enable_cameras_for_video(args_cli)
    set_hydra_args(hydra_args)
    return args_cli


def run(argv: list[str]) -> None:
    """Train a Stable-Baselines3 agent."""
    import numpy as np
    from stable_baselines3 import PPO
    from stable_baselines3.common.callbacks import CheckpointCallback, LogEveryNTimesteps
    from stable_baselines3.common.vec_env import VecNormalize

    from isaaclab.envs import DirectMARLEnvCfg

    from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg

    from isaaclab_tasks.utils import launch_simulation, resolve_task_config

    signal.signal(signal.SIGINT, _cleanup_pbar)

    args_cli = _parse_args(argv)
    env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)

    with launch_simulation(env_cfg, args_cli):
        if args_cli.seed == -1:
            args_cli.seed = random.randint(0, 10000)

        apply_env_overrides(args_cli, env_cfg)
        agent_cfg["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["seed"]
        if args_cli.max_iterations is not None:
            agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs

        env_cfg.seed = agent_cfg["seed"]

        run_info = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
        print(f"[INFO] Logging experiment in directory: {log_root_path}")
        print(f"Exact experiment name requested from command line: {run_info}")
        log_dir = os.path.join(log_root_path, run_info)
        dump_train_configs(log_dir, env_cfg, agent_cfg)

        command = " ".join(sys.orig_argv)
        (Path(log_dir) / "command.txt").write_text(command)

        agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
        policy_arch = agent_cfg.pop("policy")
        n_timesteps = agent_cfg.pop("n_timesteps")

        configure_io_descriptors(env_cfg, args_cli, logger)
        env_cfg.log_dir = log_dir

        env = create_isaaclab_env(
            args_cli.task,
            env_cfg,
            args_cli,
            convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg),
        )
        env = wrap_record_video(env, log_dir, args_cli)

        start_time = time.time()
        env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)

        norm_keys = {"normalize_input", "normalize_value", "clip_obs"}
        norm_args = {}
        for key in norm_keys:
            if key in agent_cfg:
                norm_args[key] = agent_cfg.pop(key)

        if norm_args and norm_args.get("normalize_input"):
            print(f"Normalizing input, {norm_args=}")
            env = VecNormalize(
                env,
                training=True,
                norm_obs=norm_args["normalize_input"],
                norm_reward=norm_args.get("normalize_value", False),
                clip_obs=norm_args.get("clip_obs", 100.0),
                gamma=agent_cfg["gamma"],
                clip_reward=np.inf,
            )

        agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
        if args_cli.checkpoint is not None:
            agent = agent.load(args_cli.checkpoint, env, print_system_info=True)

        checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2)
        callbacks = [checkpoint_callback, LogEveryNTimesteps(n_steps=args_cli.log_interval)]

        with contextlib.suppress(KeyboardInterrupt):
            agent.learn(
                total_timesteps=n_timesteps,
                callback=callbacks,
                progress_bar=True,
                log_interval=None,
            )

        agent.save(os.path.join(log_dir, "model"))
        print("Saving to:")
        print(os.path.join(log_dir, "model.zip"))

        if isinstance(env, VecNormalize):
            print("Saving normalization")
            env.save(os.path.join(log_dir, "model_vecnormalize.pkl"))

        print(f"Training time: {round(time.time() - start_time, 2)} seconds")
        env.close()

The argument --rl_library selects the reinforcement learning library. The --agent argument selects the library-specific configuration entry point from the kwargs dictionary, so you can manually specify alternate configuration instances.

The Code Execution

Since for the cartpole balancing task, RSL-RL library offers two configuration instances, we can use the --agent argument to specify the configuration instance to use.

• Training with the standard PPO configuration:

  # standard PPO training
  ./isaaclab.sh train --rl_library rsl_rl --task Isaac-Cartpole --headless \
    --run_name ppo
  

• Training with the PPO configuration with symmetry augmentation:

  # PPO training with symmetry augmentation
  ./isaaclab.sh train --rl_library rsl_rl --task Isaac-Cartpole --headless \
    --agent rsl_rl_with_symmetry_cfg_entry_point \
    --run_name ppo_with_symmetry_data_augmentation
  
  # you can use hydra to disable symmetry augmentation but enable mirror loss computation
  ./isaaclab.sh train --rl_library rsl_rl --task Isaac-Cartpole --headless \
    --agent rsl_rl_with_symmetry_cfg_entry_point \
    --run_name ppo_without_symmetry_data_augmentation \
    agent.algorithm.symmetry_cfg.use_data_augmentation=false
  

The --run_name argument is used to specify the name of the run. This is used to create a directory for the run in the logs/rsl_rl/cartpole directory.