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 script from Stable-Baselines3 workflow in the scripts/reinforcement_learning/sb3 directory.

Code for train.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


"""Script to train RL agent with Stable Baselines3."""

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

import gymnasium as gym
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, ManagerBasedRLEnvCfg
from isaaclab.utils.dict import print_dict
from isaaclab.utils.io import dump_yaml

from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import add_launcher_args, launch_simulation, resolve_task_config

logger = logging.getLogger(__name__)

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

# -- argparse ----------------------------------------------------------------
parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument(
    "--agent", type=str, default="sb3_cfg_entry_point", help="Name of the RL agent configuration entry point."
)
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
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("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
parser.add_argument("--export_io_descriptors", action="store_true", default=False, help="Export IO descriptors.")
parser.add_argument(
    "--keep_all_info",
    action="store_true",
    default=False,
    help="Use a slower SB3 wrapper but keep all the extra training info.",
)
parser.add_argument(
    "--ray-proc-id", "-rid", type=int, default=None, help="Automatically configured by Ray integration, otherwise None."
)
add_launcher_args(parser)
args_cli, hydra_args = parser.parse_known_args()

if args_cli.video:
    args_cli.enable_cameras = True

sys.argv = [sys.argv[0]] + hydra_args


def cleanup_pbar(*args):
    """
    A small helper to stop training and
    cleanup progress bar properly 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


signal.signal(signal.SIGINT, cleanup_pbar)


def main():
    """Train with stable-baselines agent."""
    env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
    with launch_simulation(env_cfg, args_cli):
        # randomly sample a seed if seed = -1
        if args_cli.seed == -1:
            args_cli.seed = random.randint(0, 10000)

        # override configurations with non-hydra CLI arguments
        env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
        agent_cfg["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["seed"]
        # max iterations for training
        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

        # set the environment seed
        env_cfg.seed = agent_cfg["seed"]
        env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device

        # directory for logging into
        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 the configuration into log-directory
        dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
        dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)

        # save command used to run the script
        command = " ".join(sys.orig_argv)
        (Path(log_dir) / "command.txt").write_text(command)

        # post-process agent configuration
        agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
        # read configurations about the agent-training
        policy_arch = agent_cfg.pop("policy")
        n_timesteps = agent_cfg.pop("n_timesteps")

        # set the IO descriptors export flag if requested
        if isinstance(env_cfg, ManagerBasedRLEnvCfg):
            env_cfg.export_io_descriptors = args_cli.export_io_descriptors
        else:
            logger.warning(
                "IO descriptors are only supported for manager based RL environments."
                " No IO descriptors will be exported."
            )

        # set the log directory for the environment
        env_cfg.log_dir = log_dir

        # create isaac environment
        env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)

        # convert to single-agent instance if required by the RL algorithm
        if isinstance(env.unwrapped.cfg, DirectMARLEnvCfg):
            from isaaclab.envs import multi_agent_to_single_agent

            env = multi_agent_to_single_agent(env)

        # wrap for video recording
        if args_cli.video:
            video_kwargs = {
                "video_folder": os.path.join(log_dir, "videos", "train"),
                "step_trigger": lambda step: step % args_cli.video_interval == 0,
                "video_length": args_cli.video_length,
                "disable_logger": True,
            }
            print("[INFO] Recording videos during training.")
            print_dict(video_kwargs, nesting=4)
            env = gym.wrappers.RecordVideo(env, **video_kwargs)

        start_time = time.time()

        # wrap around environment for stable baselines
        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,
            )

        # create agent from stable baselines
        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)

        # callbacks for agent
        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)]

        # train the agent
        with contextlib.suppress(KeyboardInterrupt):
            agent.learn(
                total_timesteps=n_timesteps,
                callback=callbacks,
                progress_bar=True,
                log_interval=None,
            )
        # save the final model
        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")

        # close the simulator
        env.close()


if __name__ == "__main__":
    main()

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:

1. 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.

2. wrappers.sb3.Sb3VecEnvWrapper: This wrapper converts the environment into a Stable-Baselines3 compatible environment.

3. 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.

The Code Execution

We train a PPO agent from Stable-Baselines3 to solve the cartpole balancing task.

Training the agent

There are three main ways to train the agent. Each of them has their own advantages and disadvantages. It is up to you to decide which one you prefer based on your use case.

Headless execution

When no visualizer is requested, no interactive visualizer window is opened during training. This is useful when training on a remote server or when you do not need live visual feedback, which can add some compute cost. Rendering can still be active for sensor/camera data capture when enabled by the workflow.

./isaaclab.sh -p scripts/reinforcement_learning/sb3/train.py --task Isaac-Cartpole-v0 --num_envs 64

Headless execution with off-screen render

Since the above command does not open an interactive visualizer, it is not possible to monitor behavior live in a viewport window. To capture visual output during training, enable camera/sensor rendering in the workflow and pass --video to record the agent behavior.

./isaaclab.sh -p scripts/reinforcement_learning/sb3/train.py --task Isaac-Cartpole-v0 --num_envs 64 --video

The videos are saved to the logs/sb3/Isaac-Cartpole-v0/<run-dir>/videos/train directory. You can open these videos using any video player.

Interactive execution

While the above two methods are useful for training the agent, they don’t allow you to interact with the simulation to see what is happening. In this case, run the training script as follows:

./isaaclab.sh -p scripts/reinforcement_learning/sb3/train.py --task Isaac-Cartpole-v0 --num_envs 64 --viz kit

This will open the Kit visualizer window and you can see the agent training in the environment. However, this can slow down the training process because interactive visual feedback is enabled. As a workaround, you can switch between different render modes in the "Isaac Lab" window that is docked on the bottom-right corner of the screen. To learn more about these render modes, please check the sim.SimulationContext.RenderMode class.

Viewing the logs

On a separate terminal, you can monitor the training progress by executing the following command:

# execute from the root directory of the repository
./isaaclab.sh -p -m tensorboard.main --logdir logs/sb3/Isaac-Cartpole-v0

Playing the trained agent

Once the training is complete, you can visualize the trained agent by executing the following command:

# execute from the root directory of the repository
./isaaclab.sh -p scripts/reinforcement_learning/sb3/play.py --task Isaac-Cartpole-v0 --num_envs 32 --use_last_checkpoint --viz kit

The above command will load the latest checkpoint from the logs/sb3/Isaac-Cartpole-v0 directory. You can also specify a specific checkpoint by passing the --checkpoint flag.