配置一个 RL Agent

配置一个 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()。