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
1# Copyright (c) 2022-2025, 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
7"""Script to train RL agent with Stable Baselines3."""
8
9"""Launch Isaac Sim Simulator first."""
10
11import argparse
12import contextlib
13import signal
14import sys
15from pathlib import Path
16
17from isaaclab.app import AppLauncher
18
19# add argparse arguments
20parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
21parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
22parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
23parser.add_argument("--video_interval", type=int, default=2000, help="Interval between video recordings (in steps).")
24parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
25parser.add_argument("--task", type=str, default=None, help="Name of the task.")
26parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
27parser.add_argument("--log_interval", type=int, default=100_000, help="Log data every n timesteps.")
28parser.add_argument("--checkpoint", type=str, default=None, help="Continue the training from checkpoint.")
29parser.add_argument("--max_iterations", type=int, default=None, help="RL Policy training iterations.")
30parser.add_argument(
31 "--keep_all_info",
32 action="store_true",
33 default=False,
34 help="Use a slower SB3 wrapper but keep all the extra training info.",
35)
36# append AppLauncher cli args
37AppLauncher.add_app_launcher_args(parser)
38# parse the arguments
39args_cli, hydra_args = parser.parse_known_args()
40# always enable cameras to record video
41if args_cli.video:
42 args_cli.enable_cameras = True
43
44# clear out sys.argv for Hydra
45sys.argv = [sys.argv[0]] + hydra_args
46
47# launch omniverse app
48app_launcher = AppLauncher(args_cli)
49simulation_app = app_launcher.app
50
51
52def cleanup_pbar(*args):
53 """
54 A small helper to stop training and
55 cleanup progress bar properly on ctrl+c
56 """
57 import gc
58
59 tqdm_objects = [obj for obj in gc.get_objects() if "tqdm" in type(obj).__name__]
60 for tqdm_object in tqdm_objects:
61 if "tqdm_rich" in type(tqdm_object).__name__:
62 tqdm_object.close()
63 raise KeyboardInterrupt
64
65
66# disable KeyboardInterrupt override
67signal.signal(signal.SIGINT, cleanup_pbar)
68
69"""Rest everything follows."""
70
71import gymnasium as gym
72import numpy as np
73import os
74import random
75from datetime import datetime
76
77from stable_baselines3 import PPO
78from stable_baselines3.common.callbacks import CheckpointCallback, LogEveryNTimesteps
79from stable_baselines3.common.vec_env import VecNormalize
80
81from isaaclab.envs import (
82 DirectMARLEnv,
83 DirectMARLEnvCfg,
84 DirectRLEnvCfg,
85 ManagerBasedRLEnvCfg,
86 multi_agent_to_single_agent,
87)
88from isaaclab.utils.dict import print_dict
89from isaaclab.utils.io import dump_pickle, dump_yaml
90
91from isaaclab_rl.sb3 import Sb3VecEnvWrapper, process_sb3_cfg
92
93import isaaclab_tasks # noqa: F401
94from isaaclab_tasks.utils.hydra import hydra_task_config
95
96# PLACEHOLDER: Extension template (do not remove this comment)
97
98
99@hydra_task_config(args_cli.task, "sb3_cfg_entry_point")
100def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: dict):
101 """Train with stable-baselines agent."""
102 # randomly sample a seed if seed = -1
103 if args_cli.seed == -1:
104 args_cli.seed = random.randint(0, 10000)
105
106 # override configurations with non-hydra CLI arguments
107 env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
108 agent_cfg["seed"] = args_cli.seed if args_cli.seed is not None else agent_cfg["seed"]
109 # max iterations for training
110 if args_cli.max_iterations is not None:
111 agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs
112
113 # set the environment seed
114 # note: certain randomizations occur in the environment initialization so we set the seed here
115 env_cfg.seed = agent_cfg["seed"]
116 env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
117
118 # directory for logging into
119 run_info = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
120 log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
121 print(f"[INFO] Logging experiment in directory: {log_root_path}")
122 # The Ray Tune workflow extracts experiment name using the logging line below, hence, do not change it (see PR #2346, comment-2819298849)
123 print(f"Exact experiment name requested from command line: {run_info}")
124 log_dir = os.path.join(log_root_path, run_info)
125 # dump the configuration into log-directory
126 dump_yaml(os.path.join(log_dir, "params", "env.yaml"), env_cfg)
127 dump_yaml(os.path.join(log_dir, "params", "agent.yaml"), agent_cfg)
128 dump_pickle(os.path.join(log_dir, "params", "env.pkl"), env_cfg)
129 dump_pickle(os.path.join(log_dir, "params", "agent.pkl"), agent_cfg)
130
131 # save command used to run the script
132 command = " ".join(sys.orig_argv)
133 (Path(log_dir) / "command.txt").write_text(command)
134
135 # post-process agent configuration
136 agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
137 # read configurations about the agent-training
138 policy_arch = agent_cfg.pop("policy")
139 n_timesteps = agent_cfg.pop("n_timesteps")
140
141 # create isaac environment
142 env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)
143
144 # convert to single-agent instance if required by the RL algorithm
145 if isinstance(env.unwrapped, DirectMARLEnv):
146 env = multi_agent_to_single_agent(env)
147
148 # wrap for video recording
149 if args_cli.video:
150 video_kwargs = {
151 "video_folder": os.path.join(log_dir, "videos", "train"),
152 "step_trigger": lambda step: step % args_cli.video_interval == 0,
153 "video_length": args_cli.video_length,
154 "disable_logger": True,
155 }
156 print("[INFO] Recording videos during training.")
157 print_dict(video_kwargs, nesting=4)
158 env = gym.wrappers.RecordVideo(env, **video_kwargs)
159
160 # wrap around environment for stable baselines
161 env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)
162
163 norm_keys = {"normalize_input", "normalize_value", "clip_obs"}
164 norm_args = {}
165 for key in norm_keys:
166 if key in agent_cfg:
167 norm_args[key] = agent_cfg.pop(key)
168
169 if norm_args and norm_args.get("normalize_input"):
170 print(f"Normalizing input, {norm_args=}")
171 env = VecNormalize(
172 env,
173 training=True,
174 norm_obs=norm_args["normalize_input"],
175 norm_reward=norm_args.get("normalize_value", False),
176 clip_obs=norm_args.get("clip_obs", 100.0),
177 gamma=agent_cfg["gamma"],
178 clip_reward=np.inf,
179 )
180
181 # create agent from stable baselines
182 agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
183 if args_cli.checkpoint is not None:
184 agent = agent.load(args_cli.checkpoint, env, print_system_info=True)
185
186 # callbacks for agent
187 checkpoint_callback = CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2)
188 callbacks = [checkpoint_callback, LogEveryNTimesteps(n_steps=args_cli.log_interval)]
189
190 # train the agent
191 with contextlib.suppress(KeyboardInterrupt):
192 agent.learn(
193 total_timesteps=n_timesteps,
194 callback=callbacks,
195 progress_bar=True,
196 log_interval=None,
197 )
198 # save the final model
199 agent.save(os.path.join(log_dir, "model"))
200 print("Saving to:")
201 print(os.path.join(log_dir, "model.zip"))
202
203 if isinstance(env, VecNormalize):
204 print("Saving normalization")
205 env.save(os.path.join(log_dir, "model_vecnormalize.pkl"))
206
207 # close the simulator
208 env.close()
209
210
211if __name__ == "__main__":
212 # run the main function
213 main()
214 # close sim app
215 simulation_app.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.
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#
If the --headless
flag is set, the simulation is not rendered during training. This is useful
when training on a remote server or when you do not want to see the simulation. Typically, it speeds
up the training process since only physics simulation step is performed.
./isaaclab.sh -p scripts/reinforcement_learning/sb3/train.py --task Isaac-Cartpole-v0 --num_envs 64 --headless
Headless execution with off-screen render#
Since the above command does not render the simulation, it is not possible to visualize the agent’s
behavior during training. To visualize the agent’s behavior, we pass the --enable_cameras
which
enables off-screen rendering. Additionally, we pass the flag --video
which records a video of the
agent’s behavior during training.
./isaaclab.sh -p scripts/reinforcement_learning/sb3/train.py --task Isaac-Cartpole-v0 --num_envs 64 --headless --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, you can ignore the --headless
flag and run the
training script as follows:
./isaaclab.sh -p scripts/reinforcement_learning/sb3/train.py --task Isaac-Cartpole-v0 --num_envs 64
This will open the Isaac Sim window and you can see the agent training in the environment. However, this
will slow down the training process since the simulation is rendered on the screen. 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
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.