与 RL Agent 进行训练

与 RL Agent 进行训练#

使用 注册环境 中的任务，训练 Stable-Baselines3 策略来平衡 Cartpole。

Isaac Lab 环境返回批量 PyTorch 张量，各强化学习库则要求各自的环境接口。isaaclab_rl 中的包装器将观测、动作和回合信息适配为相应接口。本教程使用 Stable-Baselines3 的 VecEnv 包装器；其他包装器见 包装环境，训练命令参考见 强化学习。

小心

Gymnasium 包装器应先于所用强化学习库的 Isaac Lab 包装器应用。RL 包装器会改变环境接口，之后可能不再支持 Gymnasium 包装器。

代码#

本教程使用 Stable-Baselines3 工作流的训练实现，位于 isaaclab_rl.entrypoints.backends.train_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")

代码解释#

上面的大部分代码是创建日志目录、保存解析的配置和设置不同的 Stable-Baselines3 组件的样板代码。对于本教程，重要的部分是创建环境并使用 Stable-Baselines3 包装器对其进行包装。

代码中使用了三个包装器:

  1. gymnasium.wrappers.RecordVideo: 这个包装器记录环境的视频并将其保存到指定目录。这对于在训练过程中可视化 agent 的行为非常有用。

  2. wrappers.sb3.Sb3VecEnvWrapper: 这个包装器将环境转换为 Stable-Baselines3 兼容的环境。

  3. stable_baselines3.common.vec_env.VecNormalize: 这个包装器对环境的观测和奖励进行标准化。

这些包装器中的每一个都通过反复执行 env = wrapper(env, *args, **kwargs) 来包装前一个包装器。然后使用最终的环境来训练 agent。有关这些包装器如何工作的更多信息，请参考 包装环境 文档。

代码执行#

我们训练一个从 Stable-Baselines3 学习的 PPO agent 来解决 cartpole 平衡任务。

训练 agent#

通过以下选项，可无可视化器训练、实时观察仿真或录制视频。

无界面执行#

未请求可视化器时，训练期间不会打开交互式可视化窗口。这适用于远程服务器训练，或不需要实时画面反馈的情况，因为实时反馈会增加一些计算开销。工作流启用传感器或相机数据采集后，渲染仍可保持运行。

uv run --extra sb3 isaaclab train --rl_library sb3 --task Isaac-Cartpole --num_envs 64

无界面执行与离屏渲染#

若需在不打开交互窗口的情况下检查训练行为，请传入 --video。默认从无窗口 Newton GL 可视化器录制场景。选择其他可视化器或相机传感器输出，见 录制视频。

uv run --extra sb3,video isaaclab train --rl_library sb3 --task Isaac-Cartpole --num_envs 64 --video

视频保存在 logs/sb3/Isaac-Cartpole/<run-dir>/videos/train 目录，可以用任意视频播放器打开。

对于场景中带有相机的任务，还可以通过 --capture_env_sensors 在训练期间直接保存传感器图像输出。可用选项和输出格式见 在训练期间采集传感器帧 。

交互式执行#

上述两种方式可用于训练智能体，但无法与仿真交互、观察实时情况。如果需要交互，可按以下方式运行训练命令：

uv run --extra sb3 isaaclab train --rl_library sb3 --task Isaac-Cartpole --num_envs 64 --viz kit

这会打开 Kit 可视化窗口，可以看到智能体在环境中训练。不过，启用交互式画面反馈可能降低训练速度。可以在屏幕右下角停靠的 "Isaac Lab" 窗口中切换不同渲染模式。更多信息见 sim.SimulationContext.RenderMode 。

查看日志#

在单独的终端中，您可以通过执行以下命令监视训练进度:

# execute from the root directory of the repository
uv run python -m tensorboard.main --logdir logs/sb3/Isaac-Cartpole

播放经过训练的 agent#

一旦训练完成，您可以通过执行以下命令来可视化经过训练的 agent:

# execute from the root directory of the repository
uv run --extra sb3 isaaclab play --rl_library sb3 --task Isaac-Cartpole --num_envs 32 --viz kit

上述命令加载 logs/sb3/Isaac-Cartpole 目录中的最新检查点。也可以用 --checkpoint 指定某个检查点。

目录