Policy Inference in USD Environment

This deployment example runs a trained policy in a prebuilt USD scene using the training task’s observations, actions, and robot configuration.

In this tutorial, we will use the RSL RL library and the trained policy from the Humanoid Rough Terrain Isaac-Velocity-Rough-H1 task in a simple warehouse USD.

The Tutorial Code

For this tutorial, we use the trained policy’s checkpoint exported as jit (which is an offline version of the policy).

The script resolves H1RoughEnvCfg with parse_env_cfg, including any physics preset passed on the command line. Calling its play_mode method applies the play/inference overrides (such as a reduced number of environments and disabled observation noise) on top of the training configuration.

In order to use a prebuilt USD environment instead of the terrain generator specified, we make the following changes to the config before passing it to the ManagerBasedRLEnv.

Code for 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

# 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.envs import ManagerBasedRLEnv
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 = ManagerBasedRLEnv(cfg=env_cfg)

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


if __name__ == "__main__":
    main()

The script uses --device for both policy loading and simulation. It disables Fabric only when --device cpu is explicitly selected. The height scanner starts below the warehouse roof so its downward rays measure the floor.

Keep the same physics preset for training, export, and inference. The commands below use Newton MJWarp with the newton_gl visualizer and do not require Isaac Sim. For a PhysX checkpoint trained with Isaac Sim, use physics=isaacsim_physx throughout and install Isaac Sim; the Newton GL viewer can still be used. Cross-backend policy transfer needs additional validation; see Transfer Policies Between PhysX and Newton.

The Code Execution

First, we need to train the Isaac-Velocity-Rough-H1 task by running the following:

uv (Recommended)

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

When the training is finished, we can visualize the result with the following command. To stop the simulation, you can either close the window, or press Ctrl+C in the terminal where you started the simulation.

uv (Recommended)

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

After running the play script, the policy will be exported to jit and onnx files under the experiment logs directory. Note that not all learning libraries support exporting the policy to a jit or onnx file. For libraries that don’t currently support this functionality, please refer to the corresponding play.py script for the library to learn about how to initialize the policy.

We can then load the warehouse asset and run inference on the H1 robot using the exported jit policy (policy.pt file in the exported/ directory).

uv (Recommended)

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

In this tutorial, we learnt how to make minor modifications to an existing environment config to run policy inference in a prebuilt usd environment.