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.

        "env_cfg_entry_point": f"{__name__}.cartpole_direct_camera_env_cfg:CartpoleCameraEnvCfg",
        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml",
        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpoleCameraDirectPPORunnerCfg",
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_direct_camera_ppo_cfg.yaml",
    },
)

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

gym.register(
    id="Isaac-Cartpole",
    entry_point="isaaclab.envs:ManagerBasedRLEnv",
    disable_env_checker=True,

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 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 isaaclab.app import add_launcher_args, report_activity

from isaaclab_rl.entrypoints.common import (
    CHECKPOINT_SELECTORS,
    add_common_train_args,
    apply_env_overrides,
    apply_video_recording,
    configure_io_descriptors,
    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_training_capture,
    write_run_manifest,
)

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."""
    from isaaclab_tasks.utils import setup_preset_cli

    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, agent_library="sb3")
    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.app import launch_simulation
    from isaaclab.envs import DirectMARLEnvCfg
    from isaaclab.utils.seed import configure_seed

    from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg

    from isaaclab_tasks.utils import resolve_task_config

    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)
        show_run_summary(screen, args_cli, env_cfg, library="sb3", action="train")
        pre_launch_video_config(env_cfg, args_cli=args_cli)
        screen.stage("Launching simulation")
        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)
            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)

            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
            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),
            )
            env = wrap_training_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)

            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,
                )

            report_activity("Building policy")
            agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
            report_activity(None)
            if args_cli.checkpoint in CHECKPOINT_SELECTORS:
                checkpoint_path = 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},
                )
                agent = agent.load(checkpoint_path, env, print_system_info=True)
            elif args_cli.checkpoint is not None:
                agent = agent.load(args_cli.checkpoint, env, print_system_info=True)

            # configure_seed must run after PPO construction/load so torch determinism does not disturb
            # SB3's initialization
            if args_cli.deterministic:
                configure_seed(env_cfg.seed, torch_deterministic=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)]

            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("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:

uv (Recommended)

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

isaaclab.sh / isaaclab.bat

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

Training with the PPO configuration with symmetry augmentation: