配置一个 RL Agent#
在之前的教程中,我们看到了如何使用 Stable-Baselines3 库训练 RL 智能体以解决 cartpole 平衡任务。在本教程中,我们将看到如何配置训练过程以使用不同的 RL 库和不同的训练算法。
在目录 scripts/reinforcement_learning 中,您将找到不同 RL 库的脚本。这些按照库名称命名的子目录进行组织。每个子目录包含该库的训练和回放脚本。
要让某个强化学习库运行特定任务,需要为学习智能体创建配置文件,用于实例化智能体并配置训练过程。与 注册环境 教程中的环境注册类似,可以通过 gymnasium.register 注册学习智能体。
代码#
以 isaaclab_tasks 包中 Isaac-Cartpole 任务的配置为例。这也是 与 RL Agent 进行训练 教程所使用的任务。
##
# 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": (
代码解释#
在 kwargs 属性下,我们可以看到不同学习库的配置。键是库的名称,值是配置实例的路径。这个配置实例可以是字符串、类或类的实例。例如,键 "rl_games_cfg_entry_point" 的值是指向 RL-Games 库配置 YAML 文件的字符串。同时,键 "rsl_rl_cfg_entry_point" 的值指向 RSL-RL 库的配置类。
用于指定代理配置类的模式与用于指定环境配置入口点的模式非常相似。这意味着以下内容是等效的:
指定配置入口点为字符串
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,
},
)
第一个代码块是指定配置入口的首选方式。第二个代码块等同于第一个代码块,但会导致配置类的导入,从而减慢导入时间。这就是为什么我们建议使用字符串作为配置入口的原因。
强化学习入口默认从 kwargs 字典读取 <library_name>_cfg_entry_point ,以获取配置实例。
例如,以下代码展示了 Stable-Baselines3 的训练实现如何读取配置实例:
使用 SB3 的 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 backend of the unified reinforcement learning entrypoint."""
7
8from __future__ import annotations
9
10import argparse
11import contextlib
12import gc
13import os
14import signal
15import sys
16import time
17from datetime import datetime
18from pathlib import Path
19
20import numpy as np
21from stable_baselines3 import PPO
22from stable_baselines3.common.callbacks import CheckpointCallback, LogEveryNTimesteps
23from stable_baselines3.common.vec_env import VecNormalize
24
25from isaaclab.app import add_launcher_args, launch_simulation, report_activity
26from isaaclab.envs import DirectMARLEnvCfg
27from isaaclab.utils.seed import configure_seed
28
29import isaaclab_tasks # noqa: F401
30from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli
31
32from ...sb3 import Sb3VecEnvWrapper, process_sb3_cfg
33from ..common import (
34 CHECKPOINT_SELECTORS,
35 add_common_train_args,
36 apply_env_overrides,
37 apply_video_recording,
38 close_env,
39 create_isaaclab_env,
40 dump_train_configs,
41 enable_cameras_for_video,
42 pre_launch_video_config,
43 resolve_checkpoint_selector,
44 resolve_seed,
45 set_hydra_args,
46 show_run_summary,
47 startup_screen,
48 wrap_sensor_capture,
49 write_run_manifest,
50)
51
52# PLACEHOLDER: Extension template (do not remove this comment)
53with contextlib.suppress(ImportError):
54 import isaaclab_tasks_experimental # noqa: F401
55
56
57def _cleanup_pbar(*args) -> None:
58 """Stop training and close the rich progress bars on Ctrl+C."""
59 for obj in gc.get_objects():
60 if "tqdm_rich" in type(obj).__name__:
61 obj.close()
62 raise KeyboardInterrupt
63
64
65def _parse_args(argv: list[str]) -> argparse.Namespace:
66 """Parse Stable-Baselines3 training arguments."""
67 parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
68 add_common_train_args(
69 parser,
70 agent_default="sb3_cfg_entry_point",
71 agent_help="Name of the RL agent configuration entry point.",
72 include_distributed=False,
73 )
74 parser.add_argument("--log_interval", type=int, default=100_000, help="Log data every n timesteps.")
75 parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path, or latest/best.")
76 parser.add_argument(
77 "--keep_all_info",
78 action="store_true",
79 default=False,
80 help="Use a slower SB3 wrapper but keep all the extra training info.",
81 )
82 add_launcher_args(parser)
83 args_cli, hydra_args = setup_preset_cli(parser, argv)
84 enable_cameras_for_video(args_cli)
85 set_hydra_args(hydra_args)
86 return args_cli
87
88
89def _resolve_checkpoint(args_cli: argparse.Namespace, log_root_path: str) -> str | None:
90 """Resolve the checkpoint to resume from, or None when training starts from scratch."""
91 if args_cli.checkpoint in CHECKPOINT_SELECTORS:
92 return resolve_checkpoint_selector(
93 log_root_path,
94 args_cli.checkpoint,
95 library="sb3",
96 task=args_cli.task,
97 checkpoint_pattern=r"model(?:_.*)?\.zip",
98 preferred_checkpoint_pattern=r"model\.zip",
99 metadata={"agent": args_cli.agent},
100 )
101 return args_cli.checkpoint
102
103
104def run(argv: list[str]) -> None:
105 """Train a Stable-Baselines3 agent."""
106 signal.signal(signal.SIGINT, _cleanup_pbar)
107 args_cli = _parse_args(argv)
108 with startup_screen(args_cli, num_stages=3) as screen:
109 env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
110 pre_launch_video_config(env_cfg, args_cli)
111 show_run_summary(screen, args_cli, env_cfg, library="sb3", action="train")
112 screen.stage("Launching simulation")
113 with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
114 apply_env_overrides(args_cli, env_cfg)
115 args_cli.seed = resolve_seed(args_cli.seed)
116 if args_cli.seed is not None:
117 agent_cfg["seed"] = args_cli.seed
118 if args_cli.max_iterations is not None:
119 agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs
120 env_cfg.seed = agent_cfg["seed"]
121
122 log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
123 print(f"[INFO] Logging experiment in directory: {log_root_path}")
124 run_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
125 print(f"Exact experiment name requested from command line: {run_name}")
126 log_dir = os.path.join(log_root_path, run_name)
127 write_run_manifest(log_dir, library="sb3", task=args_cli.task, metadata={"agent": args_cli.agent})
128 dump_train_configs(log_dir, env_cfg, agent_cfg)
129 (Path(log_dir) / "command.txt").write_text(" ".join(sys.orig_argv))
130
131 checkpoint_path = _resolve_checkpoint(args_cli, log_root_path)
132 agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
133 policy_arch = agent_cfg.pop("policy")
134 n_timesteps = agent_cfg.pop("n_timesteps")
135 norm_args = {}
136 for key in ("normalize_input", "normalize_value", "clip_obs"):
137 if key in agent_cfg:
138 norm_args[key] = agent_cfg.pop(key)
139
140 env_cfg.log_dir = log_dir
141 apply_video_recording(env_cfg, log_dir, args_cli)
142
143 screen.stage("Creating environment")
144 env = create_isaaclab_env(
145 args_cli.task,
146 env_cfg,
147 args_cli,
148 convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg),
149 )
150 cleanup.callback(lambda: close_env(env))
151 env = wrap_sensor_capture(env, log_dir, args_cli)
152
153 screen.stage("Preparing agent")
154 start_time = time.time()
155 report_activity("Wrapping environment")
156 env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)
157 report_activity(None)
158 if norm_args.get("normalize_input"):
159 print(f"Normalizing input, {norm_args=}")
160 env = VecNormalize(
161 env,
162 training=True,
163 norm_obs=norm_args["normalize_input"],
164 norm_reward=norm_args.get("normalize_value", False),
165 clip_obs=norm_args.get("clip_obs", 100.0),
166 gamma=agent_cfg["gamma"],
167 clip_reward=np.inf,
168 )
169
170 report_activity("Building policy")
171 agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
172 report_activity(None)
173 if checkpoint_path is not None:
174 agent = agent.load(checkpoint_path, env, print_system_info=True)
175
176 # configure_seed must run after PPO construction and loading so torch determinism does not disturb
177 # SB3's initialization
178 if args_cli.deterministic:
179 configure_seed(env_cfg.seed, torch_deterministic=True)
180
181 callbacks = [
182 CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2),
183 LogEveryNTimesteps(n_steps=args_cli.log_interval),
184 ]
185
186 screen.close()
187 with contextlib.suppress(KeyboardInterrupt):
188 agent.learn(total_timesteps=n_timesteps, callback=callbacks, progress_bar=True, log_interval=None)
189
190 agent.save(os.path.join(log_dir, "model"))
191 print(f"Saving to:\n{os.path.join(log_dir, 'model.zip')}")
192 if isinstance(env, VecNormalize):
193 print("Saving normalization")
194 env.save(os.path.join(log_dir, "model_vecnormalize.pkl"))
195 print(f"Training time: {round(time.time() - start_time, 2)} seconds")
--rl_library 参数选择强化学习库; --agent 参数从 kwargs 字典选择对应库的配置入口,因此可以手动指定其他配置实例。
代码执行#
由于 cartpole 平衡任务,RSL-RL 库提供两个配置实例,我们可以使用 --agent 参数来指定要使用的配置实例。
使用标准 PPO 配置进行训练:
# standard PPO training uv run isaaclab train --rl_library rsl_rl --task Isaac-Cartpole \ --run_name ppo
# standard PPO training ./isaaclab.sh train --rl_library rsl_rl --task Isaac-Cartpole \ --run_name ppo
使用具有对称增强的 PPO 配置进行训练:
# PPO training with symmetry augmentation uv run isaaclab 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 uv run isaaclab 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
# 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
--run_name 参数用于指定运行的名称。这用于在 logs/rsl_rl/cartpole 目录中为运行创建一个目录。
继续 RSL-RL 训练时,如果需要修改策略学习率,可使用 --checkpoint /path/to/model.pt --no-resume_optimizer_state。它会恢复模型权重和迭代次数,但保留新配置的策略优化器及其学习率。默认的 --resume_optimizer_state 则照常恢复检查点中的优化器状态。RND 模型及其独立优化器仍会恢复检查点状态。蒸馏时,此选项要求学生检查点同时包含学生和教师;从 PPO 教师检查点初始化新学生时,应省略此选项。
已安装的下游任务通过 isaaclab.tasks 入口点组发现。自定义启动器可像内置 CLI 和直接运行的强化学习工作进程一样,在查询任务 ID 前调用 isaaclab.cli.load_external_tasks()。