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
  1# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
  2# All rights reserved.
  3#
  4# SPDX-License-Identifier: BSD-3-Clause
  5
  6"""Stable-Baselines3 training logic for the unified reinforcement learning entrypoint."""
  7
  8from __future__ import annotations
  9
 10import argparse
 11import contextlib
 12import logging
 13import os
 14import random
 15import signal
 16import sys
 17import time
 18from datetime import datetime
 19from pathlib import Path
 20
 21from isaaclab.app import add_launcher_args
 22
 23from isaaclab_rl.entrypoints.common import (
 24    CHECKPOINT_SELECTORS,
 25    add_common_train_args,
 26    apply_env_overrides,
 27    configure_io_descriptors,
 28    create_isaaclab_env,
 29    dump_train_configs,
 30    enable_cameras_for_video,
 31    resolve_checkpoint_selector,
 32    set_hydra_args,
 33    wrap_training_capture,
 34    write_run_manifest,
 35)
 36
 37import isaaclab_tasks  # noqa: F401
 38
 39logger = logging.getLogger(__name__)
 40
 41# PLACEHOLDER: Extension template (do not remove this comment)
 42with contextlib.suppress(ImportError):
 43    import isaaclab_tasks_experimental  # noqa: F401
 44
 45
 46def _cleanup_pbar(*args):
 47    """Stop training and clean up rich progress bars on Ctrl+C."""
 48    import gc
 49
 50    tqdm_objects = [obj for obj in gc.get_objects() if "tqdm" in type(obj).__name__]
 51    for tqdm_object in tqdm_objects:
 52        if "tqdm_rich" in type(tqdm_object).__name__:
 53            tqdm_object.close()
 54    raise KeyboardInterrupt
 55
 56
 57def _parse_args(argv: list[str]) -> argparse.Namespace:
 58    """Parse Stable-Baselines3 training arguments."""
 59    from isaaclab_tasks.utils import setup_preset_cli
 60
 61    parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
 62    add_common_train_args(
 63        parser,
 64        agent_default="sb3_cfg_entry_point",
 65        agent_help="Name of the RL agent configuration entry point.",
 66        include_distributed=False,
 67    )
 68    parser.add_argument("--log_interval", type=int, default=100_000, help="Log data every n timesteps.")
 69    parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path, or latest/best.")
 70    parser.add_argument(
 71        "--keep_all_info",
 72        action="store_true",
 73        default=False,
 74        help="Use a slower SB3 wrapper but keep all the extra training info.",
 75    )
 76    add_launcher_args(parser)
 77    args_cli, hydra_args = setup_preset_cli(parser, argv)
 78    enable_cameras_for_video(args_cli)
 79    set_hydra_args(hydra_args)
 80    return args_cli
 81
 82
 83def run(argv: list[str]) -> None:
 84    """Train a Stable-Baselines3 agent."""
 85    import numpy as np
 86    from stable_baselines3 import PPO
 87    from stable_baselines3.common.callbacks import CheckpointCallback, LogEveryNTimesteps
 88    from stable_baselines3.common.vec_env import VecNormalize
 89
 90    from isaaclab.app import launch_simulation
 91    from isaaclab.envs import DirectMARLEnvCfg
 92    from isaaclab.utils.seed import configure_seed
 93
 94    from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg
 95
 96    from isaaclab_tasks.utils import resolve_task_config
 97
 98    signal.signal(signal.SIGINT, _cleanup_pbar)
 99
100    args_cli = _parse_args(argv)
101    env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
102
103    with launch_simulation(env_cfg, args_cli):
104        if args_cli.seed == -1:
105            args_cli.seed = random.randint(0, 10000)
106
107        apply_env_overrides(args_cli, env_cfg)
108        agent_cfg["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["seed"]
109        if args_cli.max_iterations is not None:
110            agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs
111
112        env_cfg.seed = agent_cfg["seed"]
113
114        run_info = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
115        log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
116        print(f"[INFO] Logging experiment in directory: {log_root_path}")
117        print(f"Exact experiment name requested from command line: {run_info}")
118        log_dir = os.path.join(log_root_path, run_info)
119        write_run_manifest(
120            log_dir,
121            library="sb3",
122            task=args_cli.task,
123            metadata={"agent": args_cli.agent},
124        )
125        dump_train_configs(log_dir, env_cfg, agent_cfg)
126
127        command = " ".join(sys.orig_argv)
128        (Path(log_dir) / "command.txt").write_text(command)
129
130        agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
131        policy_arch = agent_cfg.pop("policy")
132        n_timesteps = agent_cfg.pop("n_timesteps")
133
134        configure_io_descriptors(env_cfg, args_cli, logger)
135        env_cfg.log_dir = log_dir
136
137        env = create_isaaclab_env(
138            args_cli.task,
139            env_cfg,
140            args_cli,
141            convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg),
142        )
143        env = wrap_training_capture(env, log_dir, args_cli)
144
145        start_time = time.time()
146        env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)
147
148        norm_keys = {"normalize_input", "normalize_value", "clip_obs"}
149        norm_args = {}
150        for key in norm_keys:
151            if key in agent_cfg:
152                norm_args[key] = agent_cfg.pop(key)
153
154        if norm_args and norm_args.get("normalize_input"):
155            print(f"Normalizing input, {norm_args=}")
156            env = VecNormalize(
157                env,
158                training=True,
159                norm_obs=norm_args["normalize_input"],
160                norm_reward=norm_args.get("normalize_value", False),
161                clip_obs=norm_args.get("clip_obs", 100.0),
162                gamma=agent_cfg["gamma"],
163                clip_reward=np.inf,
164            )
165
166        agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
167        if args_cli.checkpoint in CHECKPOINT_SELECTORS:
168            checkpoint_path = resolve_checkpoint_selector(
169                log_root_path,
170                args_cli.checkpoint,
171                library="sb3",
172                task=args_cli.task,
173                checkpoint_pattern=r"model(?:_.*)?\.zip",
174                preferred_checkpoint_pattern=r"model\.zip",
175                metadata={"agent": args_cli.agent},
176            )
177            agent = agent.load(checkpoint_path, env, print_system_info=True)
178        elif args_cli.checkpoint is not None:
179            agent = agent.load(args_cli.checkpoint, env, print_system_info=True)
180
181        # configure_seed must run after PPO construction/load so torch determinism does not disturb SB3's initialization
182        if args_cli.deterministic:
183            configure_seed(env_cfg.seed, torch_deterministic=True)
184
185        checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2)
186        callbacks = [checkpoint_callback, LogEveryNTimesteps(n_steps=args_cli.log_interval)]
187
188        with contextlib.suppress(KeyboardInterrupt):
189            agent.learn(
190                total_timesteps=n_timesteps,
191                callback=callbacks,
192                progress_bar=True,
193                log_interval=None,
194            )
195
196        agent.save(os.path.join(log_dir, "model"))
197        print("Saving to:")
198        print(os.path.join(log_dir, "model.zip"))
199
200        if isinstance(env, VecNormalize):
201            print("Saving normalization")
202            env.save(os.path.join(log_dir, "model_vecnormalize.pkl"))
203
204        print(f"Training time: {round(time.time() - start_time, 2)} seconds")
205        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:

  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 train --rl_library sb3 --task Isaac-Cartpole --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 train --rl_library sb3 --task Isaac-Cartpole --num_envs 64 --video

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

For tasks with on-scene cameras, you can also save the sensor image outputs directly during training with --capture_env_sensors. See Capturing sensor frames during training for the available options and output formats.

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 command as follows:

./isaaclab.sh train --rl_library sb3 --task Isaac-Cartpole --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

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 play --rl_library sb3 --task Isaac-Cartpole --num_envs 32 --viz kit

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