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.


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",
        "rsl_rl_with_symmetry_cfg_entry_point": (
            f"{agents.__name__}.rsl_rl_ppo_cfg:CartpolePPORunnerWithSymmetryCfg"
        ),
        "skrl_cfg_entry_point": f"{agents.__name__}:skrl_manager_ppo_cfg.yaml",
        "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
    },

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:

Specifying the configuration entry point as a string
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",
   },
)
Specifying the configuration entry point as a class
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:

Code for train_sb3.py with SB3
  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 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:

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

  • Training with the PPO configuration with symmetry augmentation:

    # PPO training with symmetry augmentation
./isaaclab.sh train --rl_library rsl_rl --task Isaac-Cartpole \
  --agent rsl_rl_with_symmetry_cfg_entry_point \
  --run_name ppo_with_symmetry_data_augmentation

# you can use hydra to disable symmetry augmentation but enable mirror loss computation
./isaaclab.sh train --rl_library rsl_rl --task Isaac-Cartpole \
  --agent rsl_rl_with_symmetry_cfg_entry_point \
  --run_name ppo_without_symmetry_data_augmentation \
  agent.algorithm.symmetry_cfg.use_data_augmentation=false

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.