IsaacLab | Workflow 中 rsl_rl 的 play.py 脚本精读

如是我闻： 在用IsaacLab 做强化学习实验时，回顾已训练好的模型需要调用workflow中的play.py脚本，以下是对rsl_rl的play.py脚本的逐行精读。

1. 版权声明和文件描述

# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Script to play a checkpoint if an RL agent from RSL-RL."""

"""Launch Isaac Sim Simulator first."""

2. 导入模块

import argparse
from omni.isaac.lab.app import AppLauncher

# local imports
import cli_args  # isort: skip

• 功能：
  • argparse：用于解析命令行参数，允许用户在运行脚本时指定一些选项。
  • AppLauncher：用于启动和管理 Isaac Sim 应用程序。
  • cli_args：本地导入模块，可能包含特定的命令行参数设置（如强化学习相关的参数）。

3. 命令行参数解析

# add argparse arguments
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.")
parser.add_argument(
    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
)
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
# append RSL-RL cli arguments
cli_args.add_rsl_rl_args(parser)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
args_cli = parser.parse_args()

• 功能：
  • 这里设置了一些命令行参数，用户可以在运行时指定这些参数，如是否使用 CPU、模拟环境的数量、任务名称等。
  • 通过 cli_args.add_rsl_rl_args(parser) 和 AppLauncher.add_app_launcher_args(parser) 添加了更多与 RSL-RL 和 AppLauncher 相关的参数。

4. 启动 Isaac Sim 应用

# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

• 功能：
  • 通过 AppLauncher 类，根据解析的命令行参数启动 Isaac Sim 应用，并创建一个 simulation_app 实例用于管理模拟器。

5. 导入其他库

import gymnasium as gym
import os
import torch

from rsl_rl.runners import OnPolicyRunner

import omni.isaac.lab_tasks  # noqa: F401
from omni.isaac.lab_tasks.utils import get_checkpoint_path, parse_env_cfg
from omni.isaac.lab_tasks.utils.wrappers.rsl_rl import (
    RslRlOnPolicyRunnerCfg,
    RslRlVecEnvWrapper,
    export_policy_as_jit,
    export_policy_as_onnx,
)

• 功能：
  • gymnasium：强化学习的标准库，用于创建和管理环境。
  • torch：用于深度学习的库，主要用于处理神经网络模型。
  • OnPolicyRunner：来自 RSL-RL 的一个类，用于管理强化学习代理的训练和推理过程。
  • omni.isaac.lab_tasks：Isaac Sim 相关的模块，提供了环境配置、模型检查点加载和策略导出等功能。

6. 主函数定义

def main():
    """Play with RSL-RL agent."""
    # parse configuration
    env_cfg = parse_env_cfg(
        args_cli.task, use_gpu=not args_cli.cpu, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric
    )
    agent_cfg: RslRlOnPolicyRunnerCfg = cli_args.parse_rsl_rl_cfg(args_cli.task, args_cli)

• 功能：

  • env_cfg（环境配置）

    • 用途：env_cfg 是一个环境配置对象，包含了用于设置模拟环境的参数。这些参数可能包括：
    • 任务的名称（决定要执行哪个任务）。
    • 是否使用 GPU（决定是否使用 GPU 进行加速）。
    • 模拟环境的数量（决定要同时运行多少个环境）。
    • 是否启用或禁用某些特定功能（例如 fabric 或 USD I/O 操作）。

  • 总结：env_cfg 是一个“设置包”，里面装着各种关于环境的设置，它告诉模拟器应该如何配置和运行你的任务。

  • agent_cfg（强化学习代理配置）

    • 用途：agent_cfg 是一个强化学习代理的配置对象，包含了与强化学习算法和代理相关的参数。这些参数可能包括：
    • 强化学习模型的具体设置（例如网络架构、学习率等）。
    • 从哪个检查点加载训练好的模型。
    • 日志和保存路径等其他信息。

  • 总结：agent_cfg 是另一个“设置包”，但它的重点是告诉强化学习代理应该如何行动、如何学习，以及从哪里开始。

总的来说：

• env_cfg 配置了“舞台”——也就是你的模拟环境。
• agent_cfg 配置了“演员”——也就是你的强化学习代理，它决定了演员如何在舞台上表演。

7. 创建和包装环境

    # create isaac environment
    env = gym.make(args_cli.task, cfg=env_cfg)
    # wrap around environment for rsl-rl
    env = RslRlVecEnvWrapper(env)

• 功能：
  • env_cfg 只是模拟环境，我们还需要创建并初始化一个强化学习环境 env，这个环境由 gym.make 创建，并使用 RSL-RL 的 RslRlVecEnvWrapper 进行封装，使其兼容 RSL-RL 的训练和推理流程。

8. 日志路径和模型加载

    # specify directory for logging experiments
    log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
    log_root_path = os.path.abspath(log_root_path)
    print(f"[INFO] Loading experiment from directory: {log_root_path}")
    resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)
    print(f"[INFO]: Loading model checkpoint from: {resume_path}")

• 功能：
  • 设置日志记录的路径，并根据训练运行的名称和检查点的编号加载已保存的模型检查点。

9. 加载已训练的模型

    # load previously trained model
    ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    ppo_runner.load(resume_path)
    print(f"[INFO]: Loading model checkpoint from: {resume_path}")

• 功能：
  • 使用 OnPolicyRunner 类加载之前训练好的强化学习模型，并恢复到之前的状态。

10. 获取推理策略并导出

    # obtain the trained policy for inference
    policy = ppo_runner.get_inference_policy(device=env.unwrapped.device)

    # export policy to onnx
    export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
    export_policy_as_jit(
        ppo_runner.alg.actor_critic, ppo_runner.obs_normalizer, path=export_model_dir, filename="policy.pt"
    )
    export_policy_as_onnx(ppo_runner.alg.actor_critic, path=export_model_dir, filename="policy.onnx")

• 功能：
  • 获取用于推理的策略，并将其导出为 JIT 和 ONNX 格式，以便后续使用。

11. 环境模拟循环

    # reset environment
    obs, _ = env.get_observations()
    # simulate environment
    while simulation_app.is_running():
        # run everything in inference mode
        with torch.inference_mode():
            # agent stepping
            actions = policy(obs)
            # env stepping
            obs, _, _, _ = env.step(actions)

• 功能：
  • 重置环境，获取初始观测值，然后在模拟器运行期间，不断通过加载的策略生成动作，并用这些动作推进环境状态。

12. 关闭模拟器和环境

    # close the simulator
    env.close()

• 功能：
  • 在模拟结束时关闭环境，释放资源。

13. 入口函数

if __name__ == "__main__":
    # run the main function
    main()
    # close sim app
    simulation_app.close()

• 功能：
  • 这个部分确保当脚本作为主程序运行时，main() 函数被执行，并在结束时关闭模拟器应用。

那我估计我也就是拿这个脚本来改动了，行。

非常的有品

以上