Training with an RL Agent

Training with an RL Agent#

In the previous tutorials, we covered how to define an RL task environment, register it into the gym registry, and interact with it using a random agent. We now move on to the next step: training an RL agent to solve the task.

Although the envs.ManagerBasedRLEnv conforms to the gymnasium.Env interface, it is not exactly a gym environment. The input and outputs of the environment are not numpy arrays, but rather based on torch tensors with the first dimension being the number of environment instances.

Additionally, most RL libraries expect their own variation of an environment interface. For example, Stable-Baselines3 expects the environment to conform to its VecEnv API which expects a list of numpy arrays instead of a single tensor. Similarly, RSL-RL, RL-Games and SKRL expect a different interface. Since there is no one-size-fits-all solution, we do not base the envs.ManagerBasedRLEnv on any particular learning library. Instead, we implement wrappers to convert the environment into the expected interface. These are specified in the isaaclab_rl module.

In this tutorial, we will use Stable-Baselines3 to train an RL agent to solve the cartpole balancing task.

Caution

Wrapping the environment with the respective learning framework’s wrapper should happen in the end, i.e. after all other wrappers have been applied. This is because the learning framework’s wrapper modifies the interpretation of environment’s APIs which may no longer be compatible with gymnasium.Env.

The Code#

For this tutorial, we use the training implementation from Stable-Baselines3 workflow in the isaaclab_rl.entrypoints.backends.train_sb3 module.

Code for train_sb3.py

# 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 Code Explained#

Most of the code above is boilerplate code to create logging directories, saving the parsed configurations, and setting up different Stable-Baselines3 components. For this tutorial, the important part is creating the environment and wrapping it with the Stable-Baselines3 wrapper.

There are three wrappers used in the code above:

gymnasium.wrappers.RecordVideo: This wrapper records a video of the environment and saves it to the specified directory. This is useful for visualizing the agent’s behavior during training.
wrappers.sb3.Sb3VecEnvWrapper: This wrapper converts the environment into a Stable-Baselines3 compatible environment.
stable_baselines3.common.vec_env.VecNormalize: This wrapper normalizes the environment’s observations and rewards.

Each of these wrappers wrap around the previous wrapper by following env = wrapper(env, *args, **kwargs) repeatedly. The final environment is then used to train the agent. For more information on how these wrappers work, please refer to the Wrapping environments documentation.

Training with an RL Agent

Contents

Training with an RL Agent#

The Code#

The Code Explained#

The Code Execution#

Training the agent#

Headless execution#

Headless execution with off-screen render#

Interactive execution#

Viewing the logs#

Playing the trained agent#