Skip to navigation

Configuring an RL Agent

View as Markdown

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:

from . import agents
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",
"rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerCfg",
},
)
from . import agents
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",
"rsl_rl_cfg_entry_point": agents.rsl_rl_ppo_cfg.CartpolePPORunnerCfg,
},
)

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:

# 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 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,
resolve_seed,
set_hydra_args,
show_run_summary,
startup_screen,
wrap_sensor_capture,
write_run_manifest,
)
# 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)
show_run_summary(screen, args_cli, env_cfg, library="sb3", action="train")
screen.stage("Launching simulation")
with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
apply_env_overrides(args_cli, env_cfg)
args_cli.seed = resolve_seed(args_cli.seed)
if args_cli.seed is not None:
agent_cfg["seed"] = args_cli.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"]
log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
print(f"[INFO] Logging experiment in directory: {log_root_path}")
run_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
print(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:
  • Training with the PPO configuration with symmetry augmentation:

The --run_name argument is used to specify the name of the run. This is used to create a directory for the run in the logs/rsl_rl/cartpole directory.

When continuing RSL-RL training with a different policy learning rate, use --checkpoint /path/to/model.pt --no-resume_optimizer_state. This restores model weights and the iteration count while keeping the newly configured policy optimizer and its learning rate. By default, --resume_optimizer_state restores the checkpoint’s optimizer state as usual. RND models and their separate optimizer retain their checkpoint state. For distillation, this option requires a student checkpoint containing both student and teacher; omit it when initializing a new student from a PPO teacher checkpoint.

Installed downstream tasks are discovered through the isaaclab.tasks entry-point group. Custom launchers can call isaaclab.cli.load_external_tasks() before looking up task IDs, as the built-in CLI and direct RL workers do.