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 torchfrom rsl_rl.runners import OnPolicyRunnerimport 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 configurationenv_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 environmentenv = gym.make(args_cli.task, cfg=env_cfg)# wrap around environment for rsl-rlenv = RslRlVecEnvWrapper(env)
- 功能:
env_cfg
只是模拟环境,我们还需要创建并初始化一个强化学习环境env
,这个环境由gym.make
创建,并使用 RSL-RL 的RslRlVecEnvWrapper
进行封装,使其兼容 RSL-RL 的训练和推理流程。
8. 日志路径和模型加载
# specify directory for logging experimentslog_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 modelppo_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 inferencepolicy = ppo_runner.get_inference_policy(device=env.unwrapped.device)# export policy to onnxexport_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 environmentobs, _ = env.get_observations()# simulate environmentwhile simulation_app.is_running():# run everything in inference modewith torch.inference_mode():# agent steppingactions = policy(obs)# env steppingobs, _, _, _ = env.step(actions)
- 功能:
- 重置环境,获取初始观测值,然后在模拟器运行期间,不断通过加载的策略生成动作,并用这些动作推进环境状态。
12. 关闭模拟器和环境
# close the simulatorenv.close()
- 功能:
- 在模拟结束时关闭环境,释放资源。
13. 入口函数
if __name__ == "__main__":# run the main functionmain()# close sim appsimulation_app.close()
- 功能:
- 这个部分确保当脚本作为主程序运行时,
main()
函数被执行,并在结束时关闭模拟器应用。
- 这个部分确保当脚本作为主程序运行时,
那我估计我也就是拿这个脚本来改动了,行。
非常的有品
以上