Configuring an RL Agent

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 in the isaaclab_tasks package. This is the same task that we used in the Training with an RL Agent tutorial.

##
# Register Gym environments -- manager-based workflow.
##

gym.register(
    id="Isaac-Cartpole",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,
    kwargs={
        "env_cfg_entry_point": f"{__name__}.cartpole_manager_env_cfg:CartpoleEnvCfg",
        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_manager_ppo_cfg.yaml",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg",
        "torchrl_cfg_entry_point": f"{agents.__name__}.torchrl_ppo_cfg:CartpolePPOCfg",
        "default_agent": "rsl_rl",
        "rsl_rl_with_symmetry_cfg_entry_point": (

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:

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 backend of the unified reinforcement learning entrypoint."""

from __future__ import annotations

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

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.app import add_launcher_args, launch_simulation, report_activity
from isaaclab.envs import DirectMARLEnvCfg
from isaaclab.utils.seed import configure_seed

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli

from ...sb3 import Sb3VecEnvWrapper, process_sb3_cfg
from ..common import (
    CHECKPOINT_SELECTORS,
    add_common_train_args,
    apply_env_overrides,
    apply_video_recording,
    close_env,
    create_isaaclab_env,
    dump_train_configs,
    enable_cameras_for_video,
    pre_launch_video_config,
    resolve_checkpoint_selector,
    set_hydra_args,
    show_run_summary,
    startup_screen,
    wrap_sensor_capture,
    write_run_manifest,
)
from . import cli_args_sb3 as cli_args

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) -> None:
    """Stop training and close the rich progress bars on Ctrl+C."""
    for obj in gc.get_objects():
        if "tqdm_rich" in type(obj).__name__:
            obj.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="Checkpoint path, or latest/best.")
    parser.add_argument(
        "--keep_all_info",
        action="store_true",
        default=False,
        help="Use a slower SB3 wrapper but keep all the extra training info.",
    )
    add_launcher_args(parser)
    args_cli, hydra_args = setup_preset_cli(parser, argv)
    enable_cameras_for_video(args_cli)
    set_hydra_args(hydra_args)
    return args_cli


def _resolve_checkpoint(args_cli: argparse.Namespace, log_root_path: str) -> str | None:
    """Resolve the checkpoint to resume from, or None when training starts from scratch."""
    if args_cli.checkpoint in CHECKPOINT_SELECTORS:
        return resolve_checkpoint_selector(
            log_root_path,
            args_cli.checkpoint,
            library="sb3",
            task=args_cli.task,
            checkpoint_pattern=r"model(?:_.*)?\.zip",
            preferred_checkpoint_pattern=r"model\.zip",
            metadata={"agent": args_cli.agent},
        )
    return args_cli.checkpoint


def run(argv: list[str]) -> None:
    """Train a Stable-Baselines3 agent."""
    signal.signal(signal.SIGINT, _cleanup_pbar)
    args_cli = _parse_args(argv)
    with startup_screen(args_cli, num_stages=3) as screen:
        env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
        pre_launch_video_config(env_cfg, args_cli)
        screen.stage("Launching simulation")
        with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
            show_run_summary(screen, args_cli, env_cfg, library="sb3", action="train")
            apply_env_overrides(args_cli, env_cfg)
            agent_cfg = cli_args.update_sb3_cfg(agent_cfg, args_cli)
            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"]

            log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
            logger.info(f"Logging experiment in directory: {log_root_path}")
            run_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            logger.info(f"Exact experiment name requested from command line: {run_name}")
            log_dir = os.path.join(log_root_path, run_name)
            write_run_manifest(log_dir, library="sb3", task=args_cli.task, metadata={"agent": args_cli.agent})
            dump_train_configs(log_dir, env_cfg, agent_cfg)
            (Path(log_dir) / "command.txt").write_text(" ".join(sys.orig_argv))

            checkpoint_path = _resolve_checkpoint(args_cli, log_root_path)
            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")
            norm_args = {}
            for key in ("normalize_input", "normalize_value", "clip_obs"):
                if key in agent_cfg:
                    norm_args[key] = agent_cfg.pop(key)

            env_cfg.log_dir = log_dir
            apply_video_recording(env_cfg, log_dir, args_cli)

            screen.stage("Creating environment")
            env = create_isaaclab_env(
                args_cli.task,
                env_cfg,
                args_cli,
                convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg),
            )
            cleanup.callback(lambda: close_env(env))
            env = wrap_sensor_capture(env, log_dir, args_cli)

            screen.stage("Preparing agent")
            start_time = time.time()
            report_activity("Wrapping environment")
            env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)
            report_activity(None)
            if 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,
                )

            report_activity("Building policy")
            agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
            report_activity(None)
            if checkpoint_path is not None:
                agent = agent.load(checkpoint_path, env, print_system_info=True)

            # configure_seed must run after PPO construction and loading so torch determinism does not disturb
            # SB3's initialization
            if args_cli.deterministic:
                configure_seed(env_cfg.seed, torch_deterministic=True)

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

            screen.close()
            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(f"Saving to:\n{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")

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:

uv (Recommended)

# standard PPO training
uv run isaaclab train --rl_library rsl_rl --task Isaac-Cartpole \
  --run_name ppo

Training with the PPO configuration with symmetry augmentation: