配置一个 RL Agent

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