在 USD 环境中的策略推理

在 USD 环境中的策略推理#

学习了如何在 修改现有的 Direct RL 环境 中修改任务后，我们现在将了解如何在预构建的 USD 场景中运行训练好的策略。

在本教程中，我们将使用 RSL RL 库和来自 Humanoid Rough Terrain Isaac-Velocity-Rough-H1-v0 任务的训练策略，在一个简单的仓库 USD 中。

教程代码#

对于本教程，我们使用训练好的策略的检查点，导出为 jit（这是策略的离线版本）。

H1RoughEnvCfg_PLAY 配置封装了推理环境的配置值，包括要实例化的资产。

为了使用预构建的 USD 环境而不是指定的地形生成器，我们在将其传递给 ManagerBasedRLEnv 之前对配置进行以下更改。

policy_inference_in_usd.py的代码
 1# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
 2# All rights reserved.
 3#
 4# SPDX-License-Identifier: BSD-3-Clause
 5
 6"""
 7This script demonstrates policy inference in a prebuilt USD environment.
 8
 9In this example, we use a locomotion policy to control the H1 robot. The robot was trained
10using Isaac-Velocity-Rough-H1-v0. The robot is commanded to move forward at a constant velocity.
11
12.. code-block:: bash
13
14    # Run the script
15    ./isaaclab.sh -p scripts/tutorials/03_envs/policy_inference_in_usd.py --checkpoint /path/to/jit/checkpoint.pt
16
17"""
18
19"""Launch Isaac Sim Simulator first."""
20
21
22import argparse
23
24from isaaclab.app import AppLauncher
25
26# add argparse arguments
27parser = argparse.ArgumentParser(description="Tutorial on inferencing a policy on an H1 robot in a warehouse.")
28parser.add_argument("--checkpoint", type=str, help="Path to model checkpoint exported as jit.", required=True)
29
30# append AppLauncher cli args
31AppLauncher.add_app_launcher_args(parser)
32# parse the arguments
33args_cli = parser.parse_args()
34
35# launch omniverse app
36app_launcher = AppLauncher(args_cli)
37simulation_app = app_launcher.app
38
39"""Rest everything follows."""
40import io
41import os
42import torch
43
44import omni
45
46from isaaclab.envs import ManagerBasedRLEnv
47from isaaclab.terrains import TerrainImporterCfg
48from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
49
50from isaaclab_tasks.manager_based.locomotion.velocity.config.h1.rough_env_cfg import H1RoughEnvCfg_PLAY
51
52
53def main():
54    """Main function."""
55    # load the trained jit policy
56    policy_path = os.path.abspath(args_cli.checkpoint)
57    file_content = omni.client.read_file(policy_path)[2]
58    file = io.BytesIO(memoryview(file_content).tobytes())
59    policy = torch.jit.load(file)
60    env_cfg = H1RoughEnvCfg_PLAY()
61    env_cfg.scene.num_envs = 1
62    env_cfg.curriculum = None
63    env_cfg.scene.terrain = TerrainImporterCfg(
64        prim_path="/World/ground",
65        terrain_type="usd",
66        usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Simple_Warehouse/warehouse.usd",
67    )
68    env_cfg.sim.device = "cpu"
69    env_cfg.sim.use_fabric = False
70    env = ManagerBasedRLEnv(cfg=env_cfg)
71    obs, _ = env.reset()
72    while simulation_app.is_running():
73        action = policy(obs["policy"])  # run inference
74        obs, _, _, _, _ = env.step(action)
75
76
77if __name__ == "__main__":
78    main()
79    simulation_app.close()

请注意，我们已将设备设置为 CPU 并禁用了使用 Fabric 进行推理。这是因为在模拟少量环境时，CPU 模拟通常比 GPU 模拟执行得更快。

代码执行#

首先，我们需要通过运行以下命令来训练 Isaac-Velocity-Rough-H1-v0 任务:

./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/train.py --task Isaac-Velocity-Rough-H1-v0 --headless

当训练完成后，我们可以使用以下命令来可视化结果。要停止仿真，您可以关闭窗口，或者在您启动仿真的终端中按 Ctrl+C 。

./isaaclab.sh -p scripts/reinforcement_learning/rsl_rl/play.py --task Isaac-Velocity-Rough-H1-v0 --num_envs 64 --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME/POLICY_FILE.pt

运行播放脚本后，策略将被导出为 jit 和 onnx 文件，存储在实验日志目录下。请注意，并非所有学习库都支持将策略导出为`jit 或 onnx 文件。对于当前不支持此功能的库，请参考相应的 play.py 脚本，以了解如何初始化策略。

我们可以加载仓库资产，并使用导出的 jit 策略在 H1 机器人上运行推理。

./isaaclab.sh -p scripts/tutorials/03_envs/policy_inference_in_usd.py --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME/exported/policy.pt
训练 Isaac-H1-Direct-v0 任务的结果

在本教程中，我们学习了如何对现有环境配置进行一些小修改，以在预构建的 usd 环境中运行策略推理。