与 RL Agent 进行训练

与 RL Agent 进行训练#

在之前的教程中，我们介绍了如何定义一个 RL 任务环境、将其注册到 gym 注册表中，并使用一个随机 agent 与其交互。现在我们继续进行下一步: 训练一个 RL agent 来解决这个任务。

尽管 envs.ManagerBasedRLEnv 符合 gymnasium.Env 接口，但它并不完全是一个 gym 环境。环境的输入和输出不是 numpy 数组，而是基于 torch tensors，其中第一个维度是环境实例的数量。

此外，大多数 RL 库都期望其自己的环境接口变体。例如， Stable-Baselines3 期望环境符合其 VecEnv API ，该 API 期望接收一个 numpy 数组列表而不是一个单一的张量。类似地， RSL-RL 、 RL-Games 和 SKRL 也预期另一个接口。由于没有一种适合所有情况的解决方案，我们不将 envs.ManagerBasedRLEnv 基于任何特定的学习库。相反，我们实现了包装器来将环境转换为所期望的接口。这些包装器在 isaaclab_rl 模块中指定。

在本教程中，我们将使用 Stable-Baselines3 来训练一个 RL agent 来解决 cartpole 平衡任务。

小心

在应用所有其他包装器后，应在最后使用相应学习框架的包装器包装环境。这是因为学习框架的包装器会修改环境 API 的解释，这可能不再与 gymnasium.Env 兼容。

代码#

本教程使用 Stable-Baselines3 工作流的训练实现，位于 isaaclab_rl.entrypoints.backends.train_sb3 模块。

train_sb3.py 代码

# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Stable-Baselines3 training backend of the unified reinforcement learning entrypoint."""

from __future__ import annotations

import argparse
import contextlib
import gc
import os
import signal
import sys
import time
from datetime import datetime
from pathlib import Path

import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback, LogEveryNTimesteps
from stable_baselines3.common.vec_env import VecNormalize

from isaaclab.app import add_launcher_args, launch_simulation, report_activity
from isaaclab.envs import DirectMARLEnvCfg
from isaaclab.utils.seed import configure_seed

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli

from ...sb3 import Sb3VecEnvWrapper, process_sb3_cfg
from ..common import (
    CHECKPOINT_SELECTORS,
    add_common_train_args,
    apply_env_overrides,
    apply_video_recording,
    close_env,
    create_isaaclab_env,
    dump_train_configs,
    enable_cameras_for_video,
    pre_launch_video_config,
    resolve_checkpoint_selector,
    resolve_seed,
    set_hydra_args,
    show_run_summary,
    startup_screen,
    wrap_sensor_capture,
    write_run_manifest,
)

# PLACEHOLDER: Extension template (do not remove this comment)
with contextlib.suppress(ImportError):
    import isaaclab_tasks_experimental  # noqa: F401


def _cleanup_pbar(*args) -> None:
    """Stop training and close the rich progress bars on Ctrl+C."""
    for obj in gc.get_objects():
        if "tqdm_rich" in type(obj).__name__:
            obj.close()
    raise KeyboardInterrupt


def _parse_args(argv: list[str]) -> argparse.Namespace:
    """Parse Stable-Baselines3 training arguments."""
    parser = argparse.ArgumentParser(description="Train an RL agent with Stable-Baselines3.")
    add_common_train_args(
        parser,
        agent_default="sb3_cfg_entry_point",
        agent_help="Name of the RL agent configuration entry point.",
        include_distributed=False,
    )
    parser.add_argument("--log_interval", type=int, default=100_000, help="Log data every n timesteps.")
    parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path, or latest/best.")
    parser.add_argument(
        "--keep_all_info",
        action="store_true",
        default=False,
        help="Use a slower SB3 wrapper but keep all the extra training info.",
    )
    add_launcher_args(parser)
    args_cli, hydra_args = setup_preset_cli(parser, argv)
    enable_cameras_for_video(args_cli)
    set_hydra_args(hydra_args)
    return args_cli


def _resolve_checkpoint(args_cli: argparse.Namespace, log_root_path: str) -> str | None:
    """Resolve the checkpoint to resume from, or None when training starts from scratch."""
    if args_cli.checkpoint in CHECKPOINT_SELECTORS:
        return resolve_checkpoint_selector(
            log_root_path,
            args_cli.checkpoint,
            library="sb3",
            task=args_cli.task,
            checkpoint_pattern=r"model(?:_.*)?\.zip",
            preferred_checkpoint_pattern=r"model\.zip",
            metadata={"agent": args_cli.agent},
        )
    return args_cli.checkpoint


def run(argv: list[str]) -> None:
    """Train a Stable-Baselines3 agent."""
    signal.signal(signal.SIGINT, _cleanup_pbar)
    args_cli = _parse_args(argv)
    with startup_screen(args_cli, num_stages=3) as screen:
        env_cfg, agent_cfg = resolve_task_config(args_cli.task, args_cli.agent)
        pre_launch_video_config(env_cfg, args_cli)
        show_run_summary(screen, args_cli, env_cfg, library="sb3", action="train")
        screen.stage("Launching simulation")
        with launch_simulation(env_cfg, args_cli), contextlib.ExitStack() as cleanup:
            apply_env_overrides(args_cli, env_cfg)
            args_cli.seed = resolve_seed(args_cli.seed)
            if args_cli.seed is not None:
                agent_cfg["seed"] = args_cli.seed
            if args_cli.max_iterations is not None:
                agent_cfg["n_timesteps"] = args_cli.max_iterations * agent_cfg["n_steps"] * env_cfg.scene.num_envs
            env_cfg.seed = agent_cfg["seed"]

            log_root_path = os.path.abspath(os.path.join("logs", "sb3", args_cli.task))
            print(f"[INFO] Logging experiment in directory: {log_root_path}")
            run_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            print(f"Exact experiment name requested from command line: {run_name}")
            log_dir = os.path.join(log_root_path, run_name)
            write_run_manifest(log_dir, library="sb3", task=args_cli.task, metadata={"agent": args_cli.agent})
            dump_train_configs(log_dir, env_cfg, agent_cfg)
            (Path(log_dir) / "command.txt").write_text(" ".join(sys.orig_argv))

            checkpoint_path = _resolve_checkpoint(args_cli, log_root_path)
            agent_cfg = process_sb3_cfg(agent_cfg, env_cfg.scene.num_envs)
            policy_arch = agent_cfg.pop("policy")
            n_timesteps = agent_cfg.pop("n_timesteps")
            norm_args = {}
            for key in ("normalize_input", "normalize_value", "clip_obs"):
                if key in agent_cfg:
                    norm_args[key] = agent_cfg.pop(key)

            env_cfg.log_dir = log_dir
            apply_video_recording(env_cfg, log_dir, args_cli)

            screen.stage("Creating environment")
            env = create_isaaclab_env(
                args_cli.task,
                env_cfg,
                args_cli,
                convert_marl_to_single_agent=isinstance(env_cfg, DirectMARLEnvCfg),
            )
            cleanup.callback(lambda: close_env(env))
            env = wrap_sensor_capture(env, log_dir, args_cli)

            screen.stage("Preparing agent")
            start_time = time.time()
            report_activity("Wrapping environment")
            env = Sb3VecEnvWrapper(env, fast_variant=not args_cli.keep_all_info)
            report_activity(None)
            if norm_args.get("normalize_input"):
                print(f"Normalizing input, {norm_args=}")
                env = VecNormalize(
                    env,
                    training=True,
                    norm_obs=norm_args["normalize_input"],
                    norm_reward=norm_args.get("normalize_value", False),
                    clip_obs=norm_args.get("clip_obs", 100.0),
                    gamma=agent_cfg["gamma"],
                    clip_reward=np.inf,
                )

            report_activity("Building policy")
            agent = PPO(policy_arch, env, verbose=1, tensorboard_log=log_dir, **agent_cfg)
            report_activity(None)
            if checkpoint_path is not None:
                agent = agent.load(checkpoint_path, env, print_system_info=True)

            # configure_seed must run after PPO construction and loading so torch determinism does not disturb
            # SB3's initialization
            if args_cli.deterministic:
                configure_seed(env_cfg.seed, torch_deterministic=True)

            callbacks = [
                CheckpointCallback(save_freq=1000, save_path=log_dir, name_prefix="model", verbose=2),
                LogEveryNTimesteps(n_steps=args_cli.log_interval),
            ]

            screen.close()
            with contextlib.suppress(KeyboardInterrupt):
                agent.learn(total_timesteps=n_timesteps, callback=callbacks, progress_bar=True, log_interval=None)

            agent.save(os.path.join(log_dir, "model"))
            print(f"Saving to:\n{os.path.join(log_dir, 'model.zip')}")
            if isinstance(env, VecNormalize):
                print("Saving normalization")
                env.save(os.path.join(log_dir, "model_vecnormalize.pkl"))
            print(f"Training time: {round(time.time() - start_time, 2)} seconds")

代码解释#

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

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

gymnasium.wrappers.RecordVideo: 这个包装器记录环境的视频并将其保存到指定目录。这对于在训练过程中可视化 agent 的行为非常有用。
wrappers.sb3.Sb3VecEnvWrapper: 这个包装器将环境转换为 Stable-Baselines3 兼容的环境。
stable_baselines3.common.vec_env.VecNormalize: 这个包装器对环境的观测和奖励进行标准化。

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

与 RL Agent 进行训练

目录

与 RL Agent 进行训练#

代码#

代码解释#

代码执行#

训练 agent#

无界面执行#

无界面执行与离屏渲染#

交互式执行#

查看日志#

播放经过训练的 agent#