配置一个 RL Agent#
每个任务为其支持的每个强化学习库注册智能体配置,用于选择算法、网络及训练设置。
本教程在 注册环境 的基础上,创建智能体配置并通过 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 logging
14import os
15import signal
16import sys
17import time
18from datetime import datetime
19from pathlib import Path
20
21import numpy as np
22from stable_baselines3 import PPO
23from stable_baselines3.common.callbacks import CheckpointCallback, LogEveryNTimesteps
24from stable_baselines3.common.vec_env import VecNormalize
25
26from isaaclab.app import add_launcher_args, launch_simulation, report_activity
27from isaaclab.envs import DirectMARLEnvCfg
28from isaaclab.utils.seed import configure_seed
29
30import isaaclab_tasks # noqa: F401
31from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli
32
33from ...sb3 import Sb3VecEnvWrapper, process_sb3_cfg
34from ..common import (
35 CHECKPOINT_SELECTORS,
36 add_common_train_args,
37 apply_env_overrides,
38 apply_video_recording,
39 close_env,
40 create_isaaclab_env,
41 dump_train_configs,
42 enable_cameras_for_video,
43 pre_launch_video_config,
44 resolve_checkpoint_selector,
45 set_hydra_args,
46 show_run_summary,
47 startup_screen,
48 wrap_sensor_capture,
49 write_run_manifest,
50)
51from . import cli_args_sb3 as cli_args
52
53logger = logging.getLogger(__name__)
54
55# PLACEHOLDER: Extension template (do not remove this comment)
56with contextlib.suppress(ImportError):
57 import isaaclab_tasks_experimental # noqa: F401
58
59
60def _cleanup_pbar(*args) -> None:
61 """Stop training and close the rich progress bars on Ctrl+C."""
62 for obj in gc.get_objects():
63 if "tqdm_rich" in type(obj).__name__:
64 obj.close()
65 raise KeyboardInterrupt
66
67
68def _parse_args(argv: list[str]) -> argparse.Namespace:
69 """Parse Stable-Baselines3 training arguments."""
70 parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
71 add_common_train_args(
72 parser,
73 agent_default="sb3_cfg_entry_point",
74 agent_help="Name of the RL agent configuration entry point.",
75 include_distributed=False,
76 )
77 parser.add_argument("--log_interval", type=int, default=100_000, help="Log data every n timesteps.")
78 parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path, or latest/best.")
79 parser.add_argument(
80 "--keep_all_info",
81 action="store_true",
82 default=False,
83 help="Use a slower SB3 wrapper but keep all the extra training info.",
84 )
85 add_launcher_args(parser)
86 args_cli, hydra_args = setup_preset_cli(parser, argv)
87 enable_cameras_for_video(args_cli)
88 set_hydra_args(hydra_args)
89 return args_cli
90
91
92def _resolve_checkpoint(args_cli: argparse.Namespace, log_root_path: str) -> str | None:
93 """Resolve the checkpoint to resume from, or None when training starts from scratch."""
94 if args_cli.checkpoint in CHECKPOINT_SELECTORS:
95 return resolve_checkpoint_selector(
96 log_root_path,
97 args_cli.checkpoint,
98 library="sb3",
99 task=args_cli.task,
100 checkpoint_pattern=r"model(?:_.*)?\.zip",
101 preferred_checkpoint_pattern=r"model\.zip",
102 metadata={"agent": args_cli.agent},
103 )
104 return args_cli.checkpoint
105
106
107def run(argv: list[str]) -> None:
108 """Train a Stable-Baselines3 agent."""
109 signal.signal(signal.SIGINT, _cleanup_pbar)
110 args_cli = _parse_args(argv)
111 with startup_screen(args_cli, num_stages=3) as screen:
112 env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
113 pre_launch_video_config(env_cfg, args_cli)
114 screen.stage("Launching simulation")
115 with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
116 show_run_summary(screen, args_cli, env_cfg, library="sb3", action="train")
117 apply_env_overrides(args_cli, env_cfg)
118 agent_cfg = cli_args.update_sb3_cfg(agent_cfg, args_cli)
119 if args_cli.max_iterations is not None:
120 agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs
121 env_cfg.seed = agent_cfg["seed"]
122
123 log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
124 logger.info(f"Logging experiment in directory: {log_root_path}")
125 run_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
126 logger.info(f"Exact experiment name requested from command line: {run_name}")
127 log_dir = os.path.join(log_root_path, run_name)
128 write_run_manifest(log_dir, library="sb3", task=args_cli.task, metadata={"agent": args_cli.agent})
129 dump_train_configs(log_dir, env_cfg, agent_cfg)
130 (Path(log_dir) / "command.txt").write_text(" ".join(sys.orig_argv))
131
132 checkpoint_path = _resolve_checkpoint(args_cli, log_root_path)
133 agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
134 policy_arch = agent_cfg.pop("policy")
135 n_timesteps = agent_cfg.pop("n_timesteps")
136 norm_args = {}
137 for key in ("normalize_input", "normalize_value", "clip_obs"):
138 if key in agent_cfg:
139 norm_args[key] = agent_cfg.pop(key)
140
141 env_cfg.log_dir = log_dir
142 apply_video_recording(env_cfg, log_dir, args_cli)
143
144 screen.stage("Creating environment")
145 env = create_isaaclab_env(
146 args_cli.task,
147 env_cfg,
148 args_cli,
149 convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg),
150 )
151 cleanup.callback(lambda: close_env(env))
152 env = wrap_sensor_capture(env, log_dir, args_cli)
153
154 screen.stage("Preparing agent")
155 start_time = time.time()
156 report_activity("Wrapping environment")
157 env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)
158 report_activity(None)
159 if norm_args.get("normalize_input"):
160 print(f"Normalizing input, {norm_args=}")
161 env = VecNormalize(
162 env,
163 training=True,
164 norm_obs=norm_args["normalize_input"],
165 norm_reward=norm_args.get("normalize_value", False),
166 clip_obs=norm_args.get("clip_obs", 100.0),
167 gamma=agent_cfg["gamma"],
168 clip_reward=np.inf,
169 )
170
171 report_activity("Building policy")
172 agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
173 report_activity(None)
174 if checkpoint_path is not None:
175 agent = agent.load(checkpoint_path, env, print_system_info=True)
176
177 # configure_seed must run after PPO construction and loading so torch determinism does not disturb
178 # SB3's initialization
179 if args_cli.deterministic:
180 configure_seed(env_cfg.seed, torch_deterministic=True)
181
182 callbacks = [
183 CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2),
184 LogEveryNTimesteps(n_steps=args_cli.log_interval),
185 ]
186
187 screen.close()
188 with contextlib.suppress(KeyboardInterrupt):
189 agent.learn(total_timesteps=n_timesteps, callback=callbacks, progress_bar=True, log_interval=None)
190
191 agent.save(os.path.join(log_dir, "model"))
192 print(f"Saving to:\n{os.path.join(log_dir, 'model.zip')}")
193 if isinstance(env, VecNormalize):
194 print("Saving normalization")
195 env.save(os.path.join(log_dir, "model_vecnormalize.pkl"))
196 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
使用具有对称增强的 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
--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()。