USD 环境中的策略推理

USD 环境中的策略推理#

本部署示例在预先构建的 USD 场景中运行已训练策略，沿用训练任务的观测、动作和机器人配置。

本教程使用 RSL RL 库，将人形机器人崎岖地形任务 Isaac-Velocity-Rough-H1 的已训练策略运行在一个简单的仓库 USD 场景中。

教程代码#

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

脚本通过 parse_env_cfg 解析 H1RoughEnvCfg ，包括命令行传入的物理预设。调用配置的 play_mode 方法，会在训练配置之上应用回放和推理覆盖项，例如减少环境数量、禁用观测噪声。

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

policy_inference_in_usd.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

"""
This script demonstrates policy inference in a prebuilt USD environment.

In this example, we use a locomotion policy to control the H1 robot. The robot was trained
using Isaac-Velocity-Rough-H1. The robot is commanded to move forward at a constant velocity.

.. code-block:: bash

    # Run the script
    uv run python scripts/tutorials/03_envs/policy_inference_in_usd.py --checkpoint /path/to/jit/checkpoint.pt

"""

import argparse

from isaaclab.app import add_launcher_args, launch_simulation
from isaaclab.utils import instantiate

# add argparse arguments
parser = argparse.ArgumentParser(description="Tutorial on inferencing a policy on an H1 robot in a warehouse.")
parser.add_argument("--checkpoint", type=str, help="Path to model checkpoint exported as jit.", required=True)

add_launcher_args(parser)
# parse the arguments, forwarding unrecognized ones as Hydra-style task config overrides
args_cli, hydra_overrides = parser.parse_known_args()

import os

import torch

from isaaclab.terrains import TerrainImporterCfg
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, read_file

from isaaclab_tasks.utils import parse_env_cfg


def main():
    """Main function."""
    # load the trained jit policy
    policy_path = os.path.abspath(args_cli.checkpoint)
    file = read_file(policy_path)
    policy = torch.jit.load(file, map_location=args_cli.device)

    # setup environment
    env_cfg = parse_env_cfg("Isaac-Velocity-Rough-H1", device=args_cli.device, num_envs=1, overrides=hydra_overrides)
    env_cfg.play_mode()
    env_cfg.curriculum = None
    env_cfg.scene.terrain = TerrainImporterCfg(
        prim_path="/World/ground",
        terrain_type="usd",
        usd_path=f"{ISAAC_NUCLEUS_DIR}/Environments/Simple_Warehouse/warehouse.usd",
    )
    # The warehouse is enclosed: start the height-scan rays below its roof so they hit the floor.
    env_cfg.scene.height_scanner.offset.pos = (0.0, 0.0, 2.0)
    if args_cli.device == "cpu":
        env_cfg.sim.use_fabric = False

    with launch_simulation(env_cfg, args_cli):
        # create environment
        env = instantiate(env_cfg)

        # run inference with the policy
        obs, _ = env.reset()
        with torch.inference_mode():
            while env.sim.is_running():
                action = policy(obs["policy"])
                obs, _, _, _, _ = env.step(action)
        env.close()


if __name__ == "__main__":
    main()

脚本对策略加载和仿真都使用 --device 指定的设备。只有显式选择 --device cpu 时，才禁用 Fabric。高度扫描器从仓库屋顶下方发射射线，使向下的射线测量地面。

训练、导出和推理应使用相同的物理预设。以下命令使用 Newton MJWarp 和 newton_gl 可视化器，无需 Isaac Sim。对于使用 Isaac Sim 训练的 PhysX 检查点，应在整个流程中使用 physics=isaacsim_physx 并安装 Isaac Sim；仍可使用 Newton GL 查看器。跨后端策略迁移需要额外验证，参见在 PhysX 和 Newton 之间迁移策略。

代码执行#

首先运行以下命令，训练 Isaac-Velocity-Rough-H1 任务：

uv（推荐）

uv run isaaclab train --rl_library rsl_rl --task Isaac-Velocity-Rough-H1 physics=newton_mjwarp

isaaclab.sh / isaaclab.bat

./isaaclab.sh train --rl_library rsl_rl --task Isaac-Velocity-Rough-H1 physics=newton_mjwarp

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

uv（推荐）

uv run isaaclab play --rl_library rsl_rl --task Isaac-Velocity-Rough-H1 physics=newton_mjwarp --num_envs 64 --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME/POLICY_FILE.pt --viz newton_gl

isaaclab.sh / isaaclab.bat

./isaaclab.sh play --rl_library rsl_rl --task Isaac-Velocity-Rough-H1 physics=newton_mjwarp --num_envs 64 --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME/POLICY_FILE.pt --viz newton_gl

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

我们可以加载仓库资产，并使用导出的 jit 策略在 H1 机器人上运行推理 (policy.pt 文件在 exported/ 路径下)。

uv（推荐）

uv run python scripts/tutorials/03_envs/policy_inference_in_usd.py --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME/exported/policy.pt physics=newton_mjwarp --viz newton_gl

isaaclab.sh / isaaclab.bat

./isaaclab.sh -p scripts/tutorials/03_envs/policy_inference_in_usd.py --checkpoint logs/rsl_rl/h1_rough/EXPERIMENT_NAME/exported/policy.pt physics=newton_mjwarp --viz newton_gl

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