Export Direct-Workflow Policies with LEAPP

This advanced tutorial shows how to prepare a Direct workflow policy for export with LEAPP. If your policy is manager-based, follow the standard LEAPP deployment workflow instead.

For background on LEAPP concepts, supported node patterns, state feedback, and runtime validation, see the LEAPP documentation.

Quick Start

To export a Direct workflow policy with LEAPP, you add LEAPP annotations to the environment code. During export, LEAPP traces the annotated tensors and builds an intermediate representation of the full policy pipeline. These annotations remain dormant during normal environment execution and only add a small amount of overhead until export time. They are activated by isaaclab leapp export --rl_library rsl_rl when you run the export flow.

This tutorial uses scripts/tutorials/06_deploy/anymal_c_env.py as a concrete example of adding LEAPP annotations to a Direct workflow environment. Apply the same annotation pattern to your own Direct RL environment.

This export flow requires leapp. leapp is a specialized optional extra (it is not part of --extra all). Select extras the same way as isaaclab train: add --extra leapp on every uv run, and add the backend extra that matches your task (isaacsim, ovphysx, or none for kitless Newton). See Install and run your first task and /source/policy_deployment/05_leapp/exporting_policies_with_leapp.

If you want to run the exported example with the existing IsaacContrib-Velocity-Rough-AnymalC-Direct task registration, copy the annotated tutorial environment into the task package:

Linux

Windows

cp scripts/tutorials/06_deploy/anymal_c_env.py \
   source/isaaclab_tasks/isaaclab_tasks/contrib/anymal_c_direct/anymal_c_env.py

After your environment includes the required LEAPP input, output, and state annotations, export a trained policy with the same backend extra pattern as training. For example, on Isaac Sim PhysX:

Linux

Windows

# Newton:     uv run --extra leapp isaaclab leapp export --rl_library rsl_rl ... physics=newton_mjwarp
# OV PhysX:   uv run --extra ovphysx --extra leapp isaaclab leapp export --rl_library rsl_rl ... physics=ovphysx
# Isaac Sim:
OMNI_KIT_ACCEPT_EULA=Y ACCEPT_EULA=Y uv run --extra isaacsim --extra leapp isaaclab leapp export --rl_library rsl_rl \
    --task <TASK_NAME> \
    --checkpoint <PATH_TO_CHECKPOINT> \
    --export_save_path <EXPORT_PATH> \
    physics=isaacsim_physx

The --task argument is the registered task name, such as IsaacContrib-Velocity-Rough-AnymalC-Direct. The --checkpoint argument points to the trained RSL-RL checkpoint to export. Train with uv run isaaclab train --rl_library rsl_rl --task <TASK_NAME> (see Install and run your first task); do not use the removed per-backend train.py scripts. The optional --export_save_path argument selects the output directory for the exported artifacts. If you omit it, the export is written next to the checkpoint.

Unlike with managed environments, export_method is not a command-line option. Set the export backend in the output annotation itself, for example annotate.output_tensors(..., export_with="onnx-dynamo"). See the following output annotation example for guidance.

This tutorial covers exporting direct rl policies only. direct rl policies are not currently supported by isaaclab leapp deploy.

For more information on the export arguments, see the standard LEAPP deployment workflow .

# 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
# ruff: noqa: I001
from __future__ import annotations
import gymnasium as gym
import torch
import warp as wp
from isaaclab.envs import DirectRLEnv
from isaaclab.utils import index_fill_
from .anymal_c_env_cfg import AnymalCFlatEnvCfg, AnymalCRoughEnvCfg
from leapp import annotate  # isort: skip
class AnymalCEnv(DirectRLEnv):
    cfg: AnymalCFlatEnvCfg | AnymalCRoughEnvCfg
    def __init__(self, cfg: AnymalCFlatEnvCfg | AnymalCRoughEnvCfg, render_mode: str | None = None, **kwargs):
        super().__init__(cfg, render_mode, **kwargs)
        self._robot, self._contact_sensor, self._terrain = [
            self.scene[name] for name in ("robot", "contact_sensor", "terrain")
        ]
        if isinstance(self.cfg, AnymalCRoughEnvCfg):
            self._height_scanner = self.scene["height_scanner"]
        self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)
        self._previous_actions = torch.zeros(
            self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device
        )
        self._commands = torch.zeros(self.num_envs, 3, device=self.device)
        self._episode_sums = {
            key: torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
            for key in [
                "track_lin_vel_xy_exp",
                "track_ang_vel_z_exp",
                "lin_vel_z_l2",
                "ang_vel_xy_l2",
                "dof_torques_l2",
                "dof_acc_l2",
                "action_rate_l2",
                "feet_air_time",
                "undesired_contacts",
                "flat_orientation_l2",
            ]
        }
        self._base_id, _ = self._contact_sensor.find_sensors("base")
        self._feet_ids, _ = self._contact_sensor.find_sensors(".*FOOT")
        self._undesired_contact_body_ids, _ = self._contact_sensor.find_sensors(".*THIGH")
    def _pre_physics_step(self, actions: torch.Tensor):
        self._actions = actions.clone()
        self._processed_actions = self.cfg.action_scale * self._actions + self._robot.data.default_joint_pos.torch
        # start LEAPP annotations for outputs
        annotate.update_state(self.spec.id, {"previous_actions": actions})
        annotate.output_tensors(self.spec.id, {"processed_actions": self._processed_actions}, export_with="onnx-dynamo")
        # end LEAPP annotations for outputs
    def _apply_action(self):
        self._robot.set_joint_position_target_index(target=self._processed_actions)
    def _get_observations(self) -> dict:
        self._previous_actions = self._actions.clone()
        height_data = None
        if isinstance(self.cfg, AnymalCRoughEnvCfg):
            height_data = (
                self._height_scanner.data.pos_w.torch[:, 2].unsqueeze(1)
                - self._height_scanner.data.ray_hits_w.torch[..., 2]
                - 0.5
            ).clip(-1.0, 1.0)
        # start LEAPP annotations for inputs
        root_lin_vel_b = annotate.input_tensors(self.spec.id, {"root_lin_vel_b": self._robot.data.root_lin_vel_b.torch})
        root_ang_vel_b = annotate.input_tensors(self.spec.id, {"root_ang_vel_b": self._robot.data.root_ang_vel_b.torch})
        projected_gravity_b = annotate.input_tensors(
            self.spec.id, {"projected_gravity_b": self._robot.data.projected_gravity_b.torch}
        )
        commands = annotate.input_tensors(self.spec.id, {"commands": self._commands})
        joint_pos = annotate.input_tensors(self.spec.id, {"joint_pos": self._robot.data.joint_pos.torch})
        default_joint_pos = annotate.input_tensors(
            self.spec.id, {"default_joint_pos": self._robot.data.default_joint_pos.torch}
        )
        joint_vel = annotate.input_tensors(self.spec.id, {"joint_vel": self._robot.data.joint_vel.torch})
        if height_data is not None:
            height_data = annotate.input_tensors(self.spec.id, {"height_data": height_data})
        previous_actions = annotate.state_tensors(self.spec.id, {"previous_actions": self._actions})
        # end LEAPP annotations for inputs
        obs = torch.cat(
            [
                tensor
                for tensor in (
                    root_lin_vel_b,
                    root_ang_vel_b,
                    projected_gravity_b,
                    commands,
                    joint_pos - default_joint_pos,
                    joint_vel,
                    height_data,
                    previous_actions,
                )
                if tensor is not None
            ],
            dim=-1,
        )
        observations = {"policy": obs}
        return observations
    def _get_rewards(self) -> torch.Tensor:
        lin_vel_error = torch.sum(
            torch.square(self._commands[:, :2] - self._robot.data.root_lin_vel_b.torch[:, :2]), dim=1
        )
        lin_vel_error_mapped = torch.exp(-lin_vel_error / 0.25)
        yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_ang_vel_b.torch[:, 2])
        yaw_rate_error_mapped = torch.exp(-yaw_rate_error / 0.25)
        z_vel_error = torch.square(self._robot.data.root_lin_vel_b.torch[:, 2])
        ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b.torch[:, :2]), dim=1)
        joint_torques = torch.sum(torch.square(self._robot.actuators.applied_effort.torch), dim=1)
        joint_accel = torch.sum(torch.square(self._robot.data.joint_acc.torch), dim=1)
        action_rate = torch.sum(torch.square(self._actions - self._previous_actions), dim=1)
        first_contact = self._contact_sensor.compute_first_contact(self.step_dt).torch[:, self._feet_ids]
        last_air_time = self._contact_sensor.data.last_air_time.torch[:, self._feet_ids]
        air_time = torch.sum((last_air_time - 0.5) * first_contact, dim=1) * (
            torch.linalg.norm(self._commands[:, :2], dim=1) > 0.1
        )
        net_contact_forces = self._contact_sensor.data.net_normal_forces_w_history.torch
        is_contact = (
            torch.max(torch.linalg.norm(net_contact_forces[:, :, self._undesired_contact_body_ids], dim=-1), dim=1)[0]
            > 1.0
        )
        contacts = torch.sum(is_contact, dim=1)
        flat_orientation = torch.sum(torch.square(self._robot.data.projected_gravity_b.torch[:, :2]), dim=1)
        rewards = {
            "track_lin_vel_xy_exp": lin_vel_error_mapped * self.cfg.lin_vel_reward_scale * self.step_dt,
            "track_ang_vel_z_exp": yaw_rate_error_mapped * self.cfg.yaw_rate_reward_scale * self.step_dt,
            "lin_vel_z_l2": z_vel_error * self.cfg.z_vel_reward_scale * self.step_dt,
            "ang_vel_xy_l2": ang_vel_error * self.cfg.ang_vel_reward_scale * self.step_dt,
            "dof_torques_l2": joint_torques * self.cfg.joint_torque_reward_scale * self.step_dt,
            "dof_acc_l2": joint_accel * self.cfg.joint_accel_reward_scale * self.step_dt,
            "action_rate_l2": action_rate * self.cfg.action_rate_reward_scale * self.step_dt,
            "feet_air_time": air_time * self.cfg.feet_air_time_reward_scale * self.step_dt,
            "undesired_contacts": contacts * self.cfg.undesired_contact_reward_scale * self.step_dt,
            "flat_orientation_l2": flat_orientation * self.cfg.flat_orientation_reward_scale * self.step_dt,
        }
        reward = torch.sum(torch.stack(list(rewards.values())), dim=0)
        for key, value in rewards.items():
            self._episode_sums[key] += value
        return reward
    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
        time_out = self.episode_length_buf >= self.max_episode_length - 1
        net_contact_forces = self._contact_sensor.data.net_normal_forces_w_history.torch
        died = torch.any(
            torch.max(torch.linalg.norm(net_contact_forces[:, :, self._base_id], dim=-1), dim=1)[0] > 1.0, dim=1
        )
        return died, time_out
    def _reset_idx(self, env_ids: torch.Tensor | None):
        if env_ids is None or len(env_ids) == self.num_envs:
            env_ids = wp.to_torch(self._robot._ALL_INDICES)
        assert env_ids is not None
        self._robot.reset(env_ids)
        super()._reset_idx(env_ids)
        if len(env_ids) == self.num_envs:
            self.episode_length_buf[:] = torch.randint_like(self.episode_length_buf, high=int(self.max_episode_length))
        index_fill_(self._actions, env_ids, 0.0)
        index_fill_(self._previous_actions, env_ids, 0.0)
        self._commands[env_ids] = torch.zeros_like(self._commands[env_ids]).uniform_(-1.0, 1.0)
        joint_pos = self._robot.data.default_joint_pos.torch[env_ids]
        joint_vel = self._robot.data.default_joint_vel.torch[env_ids]
        default_root_pose = self._robot.data.default_root_pose.torch[env_ids]
        default_root_vel = self._robot.data.default_root_vel.torch[env_ids]
        default_root_pose[:, :3] += self._terrain.env_origins[env_ids]
        self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids)
        self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids)
        self._robot.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids)
        self._robot.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=env_ids)
        extras = dict()
        for key in self._episode_sums.keys():
            episodic_sum_avg = torch.mean(self._episode_sums[key][env_ids])
            extras["Episode_Reward/" + key] = episodic_sum_avg / self.max_episode_length_s
            index_fill_(self._episode_sums[key], env_ids, 0.0)
        self.extras["log"] = dict()
        self.extras["log"].update(extras)
        extras = dict()
        extras["Episode_Termination/base_contact"] = torch.count_nonzero(self.reset_terminated[env_ids]).item()
        extras["Episode_Termination/time_out"] = torch.count_nonzero(self.reset_time_outs[env_ids]).item()
        self.extras["log"].update(extras)

How the Annotations Work

The main task is to identify the inputs, outputs, and persistent state in the environment and register them with LEAPP. In this example, the script uses four annotation helpers:

annotate.input_tensors marks tensors that enter the policy pipeline.
annotate.output_tensors marks tensors that leave the environment-side part of the pipeline.
annotate.state_tensors marks tensors that behave like persistent state.
annotate.update_state updates that persistent state after each step.

Input Annotations

Input annotations usually belong in _get_observations(), because that method collects the tensors that are passed to the policy.

# start LEAPP annotations for inputs
root_lin_vel_b = annotate.input_tensors(self.spec.id, {"root_lin_vel_b": self._robot.data.root_lin_vel_b.torch})
root_ang_vel_b = annotate.input_tensors(self.spec.id, {"root_ang_vel_b": self._robot.data.root_ang_vel_b.torch})
projected_gravity_b = annotate.input_tensors(
    self.spec.id, {"projected_gravity_b": self._robot.data.projected_gravity_b.torch}
)
commands = annotate.input_tensors(self.spec.id, {"commands": self._commands})
joint_pos = annotate.input_tensors(self.spec.id, {"joint_pos": self._robot.data.joint_pos.torch})
default_joint_pos = annotate.input_tensors(
    self.spec.id, {"default_joint_pos": self._robot.data.default_joint_pos.torch}
)
joint_vel = annotate.input_tensors(self.spec.id, {"joint_vel": self._robot.data.joint_vel.torch})
if height_data is not None:
    height_data = annotate.input_tensors(self.spec.id, {"height_data": height_data})
previous_actions = annotate.state_tensors(self.spec.id, {"previous_actions": self._actions})
# end LEAPP annotations for inputs

annotate.input_tensors() wraps a tensor so LEAPP can trace all downstream operations that depend on it. The function takes two important arguments:

self.spec.id identifies the node that owns the tensor. When you run isaaclab leapp export, this ID matches the exported policy node.
The second argument is a dictionary that maps a unique tensor name to the tensor itself. LEAPP uses these names in the exported metadata and for debugging.

In this example, the observation tensors are registered one by one for readability, but annotate.input_tensors() can also register multiple tensors in a single call.

Any inputs not explicitly annotated will be automatically inlined as a constant. This may be desired for certain values such as constant transforms or default values.

Output Annotations

Output annotations should be placed where the environment has finished preparing the command that will be applied to the robot. In this example, that happens in _pre_physics_step().

# start LEAPP annotations for outputs
annotate.update_state(self.spec.id, {"previous_actions": actions})
annotate.output_tensors(self.spec.id, {"processed_actions": self._processed_actions}, export_with="onnx-dynamo")
# end LEAPP annotations for outputs

annotate.output_tensors() marks the tensors that leave the environment-side part of the pipeline. As with input annotations, the call uses self.spec.id together with a dictionary that maps tensor names to tensors.

The export_with argument restricts an output annotation to specific export backends. The supported backend names are onnx-dynamo, onnx-torchscript, jit-script, and jit-trace. This argument is needed to actually generate the IR based on the tracing.

Unlike annotate.input_tensors(), output annotation should happen once for the final outputs of the pipeline stage. In this example, processed_actions is the tensor that should be exported. After calling annotate.output_tensors(), you do not need to use a return value.

All tensors passed to annotate.output_tensors() must be traced tensors. These tensors are created from inputs or tensors derived from inputs.

Do not place output annotations in _apply_action(). That method may be called multiple times per environment step, depending on the decimation setting, which would make the traced pipeline incorrect.

State Annotations

If your policy depends on internal state or feedback loops, register that data explicitly with annotate.state_tensors() and update it with annotate.update_state().

In this example, the environment uses the previous action as part of the observation. That makes previous_actions a feedback state:

annotate.state_tensors() is called in _get_observations() so the state can participate in the traced observation pipeline.
annotate.update_state() is called in _pre_physics_step() so the stored value is updated for the next step.

The state name must match across both calls. Here, both functions use the name previous_actions, which lets LEAPP route the feedback tensor correctly.

Semantic Annotations

This example covers the minimum annotations needed to trace the pipeline. In more advanced export workflows, you may also want to attach semantic metadata so downstream runtimes know what each tensor represents. Semantic Annotations is a prerequisite for being able to deploy to a downstream deployment library.

For direct environments, semantic annotations are optional and need to be authored explicitly by the user. Unlike the manager-based export path, Isaac Lab does not infer tensor semantics automatically for direct environments, instead it is up to the user to provide this data. LEAPP provides this through TensorSemantics. You can use it to describe the meaning of tensors more precisely and make the exported pipeline easier to inspect, validate, and integrate into deployment systems.

Refer to the LEAPP semantic annotation guide and LEAPP API reference for details on authoring semantic annotations.