Skip to navigation

Policy Inference in USD Environment

View as Markdown

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.

# 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.
~~~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()

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:

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.

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

H1 policy running in a warehouse USD scene

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