Exporting Direct Workflow Policies with LEAPP

Exporting Direct Workflow Policies with LEAPP#

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

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

Overview#

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 scripts/reinforcement_learning/leapp/rsl_rl/export.py 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>=0.5.2. Before exporting, install LEAPP into the Isaac Lab Python environment:

./isaaclab.sh -p -m pip install leapp

isaaclab.bat -p -m pip install leapp

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

cp scripts/tutorials/06_deploy/anymal_c_env.py \
   source/isaaclab_tasks/isaaclab_tasks/direct/anymal_c/anymal_c_env.py

copy scripts\tutorials\06_deploy\anymal_c_env.py ^
   source\isaaclab_tasks\isaaclab_tasks\direct\anymal_c\anymal_c_env.py

After your environment includes the required LEAPP input, output, and state annotations, export a trained policy with:

./isaaclab.sh -p scripts/reinforcement_learning/leapp/rsl_rl/export.py \
    --task <TASK_NAME> \
    --checkpoint <PATH_TO_CHECKPOINT> \
    --export_save_path <EXPORT_PATH>

isaaclab.bat -p scripts\reinforcement_learning\leapp\rsl_rl\export.py ^
    --task <TASK_NAME> ^
    --checkpoint <PATH_TO_CHECKPOINT> ^
    --export_save_path <EXPORT_PATH>

The --task argument is the registered task name, such as Isaac-Velocity-Rough-Anymal-C-Direct-v0. The --checkpoint argument points to the trained RSL-RL checkpoint to export. 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.

Warning

This tutorial covers exporting direct rl policies only. direct rl policies are not currently supported by scripts/reinforcement_learning/leapp/deploy.py.

For more information on the export arguments, see the manager-based LEAPP export guide.

Full example script
  1# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
  2# All rights reserved.
  3#
  4# SPDX-License-Identifier: BSD-3-Clause
  5
  6# ruff: noqa: I001
  7
  8from __future__ import annotations
  9
 10import gymnasium as gym
 11import torch
 12import warp as wp
 13
 14import isaaclab.sim as sim_utils
 15from isaaclab.assets import Articulation
 16from isaaclab.envs import DirectRLEnv
 17from isaaclab.sensors import ContactSensor, RayCaster
 18
 19from .anymal_c_env_cfg import AnymalCFlatEnvCfg, AnymalCRoughEnvCfg
 20from leapp import annotate  # isort: skip
 21
 22
 23class AnymalCEnv(DirectRLEnv):
 24    cfg: AnymalCFlatEnvCfg | AnymalCRoughEnvCfg
 25
 26    def __init__(self, cfg: AnymalCFlatEnvCfg | AnymalCRoughEnvCfg, render_mode: str | None = None, **kwargs):
 27        super().__init__(cfg, render_mode, **kwargs)
 28
 29        self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)
 30        self._previous_actions = torch.zeros(
 31            self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device
 32        )
 33
 34        self._commands = torch.zeros(self.num_envs, 3, device=self.device)
 35
 36        self._episode_sums = {
 37            key: torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
 38            for key in [
 39                "track_lin_vel_xy_exp",
 40                "track_ang_vel_z_exp",
 41                "lin_vel_z_l2",
 42                "ang_vel_xy_l2",
 43                "dof_torques_l2",
 44                "dof_acc_l2",
 45                "action_rate_l2",
 46                "feet_air_time",
 47                "undesired_contacts",
 48                "flat_orientation_l2",
 49            ]
 50        }
 51        self._base_id, _ = self._contact_sensor.find_sensors("base")
 52        self._feet_ids, _ = self._contact_sensor.find_sensors(".*FOOT")
 53        self._undesired_contact_body_ids, _ = self._contact_sensor.find_sensors(".*THIGH")
 54
 55    def _setup_scene(self):
 56        self._robot = Articulation(self.cfg.robot)
 57        self.scene.articulations["robot"] = self._robot
 58        self._contact_sensor = ContactSensor(self.cfg.contact_sensor)
 59        self.scene.sensors["contact_sensor"] = self._contact_sensor
 60        if isinstance(self.cfg, AnymalCRoughEnvCfg):
 61            self._height_scanner = RayCaster(self.cfg.height_scanner)
 62            self.scene.sensors["height_scanner"] = self._height_scanner
 63        self.cfg.terrain.num_envs = self.scene.cfg.num_envs
 64        self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing
 65        self._terrain = self.cfg.terrain.class_type(self.cfg.terrain)
 66        self.scene.clone_environments(copy_from_source=False)
 67        if self.device == "cpu":
 68            self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path])
 69        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
 70        light_cfg.func("/World/Light", light_cfg)
 71
 72    def _pre_physics_step(self, actions: torch.Tensor):
 73        self._actions = actions.clone()
 74        self._processed_actions = self.cfg.action_scale * self._actions + self._robot.data.default_joint_pos.torch
 75        # start LEAPP annotations for outputs
 76        annotate.update_state(self.spec.id, {"previous_actions": actions})
 77        annotate.output_tensors(self.spec.id, {"processed_actions": self._processed_actions}, export_with="onnx-dynamo")
 78        # end LEAPP annotations for outputs
 79
 80    def _apply_action(self):
 81        self._robot.set_joint_position_target_index(target=self._processed_actions)
 82
 83    def _get_observations(self) -> dict:
 84        self._previous_actions = self._actions.clone()
 85        height_data = None
 86        if isinstance(self.cfg, AnymalCRoughEnvCfg):
 87            height_data = (
 88                self._height_scanner.data.pos_w.torch[:, 2].unsqueeze(1)
 89                - self._height_scanner.data.ray_hits_w.torch[..., 2]
 90                - 0.5
 91            ).clip(-1.0, 1.0)
 92        # start LEAPP annotations for inputs
 93        root_lin_vel_b = annotate.input_tensors(self.spec.id, {"root_lin_vel_b": self._robot.data.root_lin_vel_b.torch})
 94        root_ang_vel_b = annotate.input_tensors(self.spec.id, {"root_ang_vel_b": self._robot.data.root_ang_vel_b.torch})
 95        projected_gravity_b = annotate.input_tensors(
 96            self.spec.id, {"projected_gravity_b": self._robot.data.projected_gravity_b.torch}
 97        )
 98        commands = annotate.input_tensors(self.spec.id, {"commands": self._commands})
 99        joint_pos = annotate.input_tensors(self.spec.id, {"joint_pos": self._robot.data.joint_pos.torch})
100        default_joint_pos = annotate.input_tensors(
101            self.spec.id, {"default_joint_pos": self._robot.data.default_joint_pos.torch}
102        )
103        joint_vel = annotate.input_tensors(self.spec.id, {"joint_vel": self._robot.data.joint_vel.torch})
104        if height_data is not None:
105            height_data = annotate.input_tensors(self.spec.id, {"height_data": height_data})
106        previous_actions = annotate.state_tensors(self.spec.id, {"previous_actions": self._actions})
107        # end LEAPP annotations for inputs
108
109        obs = torch.cat(
110            [
111                tensor
112                for tensor in (
113                    root_lin_vel_b,
114                    root_ang_vel_b,
115                    projected_gravity_b,
116                    commands,
117                    joint_pos - default_joint_pos,
118                    joint_vel,
119                    height_data,
120                    previous_actions,
121                )
122                if tensor is not None
123            ],
124            dim=-1,
125        )
126        observations = {"policy": obs}
127        return observations
128
129    def _get_rewards(self) -> torch.Tensor:
130        lin_vel_error = torch.sum(
131            torch.square(self._commands[:, :2] - self._robot.data.root_lin_vel_b.torch[:, :2]), dim=1
132        )
133        lin_vel_error_mapped = torch.exp(-lin_vel_error / 0.25)
134        yaw_rate_error = torch.square(self._commands[:, 2] - self._robot.data.root_ang_vel_b.torch[:, 2])
135        yaw_rate_error_mapped = torch.exp(-yaw_rate_error / 0.25)
136        z_vel_error = torch.square(self._robot.data.root_lin_vel_b.torch[:, 2])
137        ang_vel_error = torch.sum(torch.square(self._robot.data.root_ang_vel_b.torch[:, :2]), dim=1)
138        joint_torques = torch.sum(torch.square(self._robot.data.applied_torque.torch), dim=1)
139        joint_accel = torch.sum(torch.square(self._robot.data.joint_acc.torch), dim=1)
140        action_rate = torch.sum(torch.square(self._actions - self._previous_actions), dim=1)
141        first_contact = self._contact_sensor.compute_first_contact(self.step_dt).torch[:, self._feet_ids]
142        last_air_time = self._contact_sensor.data.last_air_time.torch[:, self._feet_ids]
143        air_time = torch.sum((last_air_time - 0.5) * first_contact, dim=1) * (
144            torch.linalg.norm(self._commands[:, :2], dim=1) > 0.1
145        )
146        net_contact_forces = self._contact_sensor.data.net_forces_w_history.torch
147        is_contact = (
148            torch.max(torch.linalg.norm(net_contact_forces[:, :, self._undesired_contact_body_ids], dim=-1), dim=1)[0]
149            > 1.0
150        )
151        contacts = torch.sum(is_contact, dim=1)
152        flat_orientation = torch.sum(torch.square(self._robot.data.projected_gravity_b.torch[:, :2]), dim=1)
153
154        rewards = {
155            "track_lin_vel_xy_exp": lin_vel_error_mapped * self.cfg.lin_vel_reward_scale * self.step_dt,
156            "track_ang_vel_z_exp": yaw_rate_error_mapped * self.cfg.yaw_rate_reward_scale * self.step_dt,
157            "lin_vel_z_l2": z_vel_error * self.cfg.z_vel_reward_scale * self.step_dt,
158            "ang_vel_xy_l2": ang_vel_error * self.cfg.ang_vel_reward_scale * self.step_dt,
159            "dof_torques_l2": joint_torques * self.cfg.joint_torque_reward_scale * self.step_dt,
160            "dof_acc_l2": joint_accel * self.cfg.joint_accel_reward_scale * self.step_dt,
161            "action_rate_l2": action_rate * self.cfg.action_rate_reward_scale * self.step_dt,
162            "feet_air_time": air_time * self.cfg.feet_air_time_reward_scale * self.step_dt,
163            "undesired_contacts": contacts * self.cfg.undesired_contact_reward_scale * self.step_dt,
164            "flat_orientation_l2": flat_orientation * self.cfg.flat_orientation_reward_scale * self.step_dt,
165        }
166        reward = torch.sum(torch.stack(list(rewards.values())), dim=0)
167        for key, value in rewards.items():
168            self._episode_sums[key] += value
169        return reward
170
171    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
172        time_out = self.episode_length_buf >= self.max_episode_length - 1
173        net_contact_forces = self._contact_sensor.data.net_forces_w_history.torch
174        died = torch.any(
175            torch.max(torch.linalg.norm(net_contact_forces[:, :, self._base_id], dim=-1), dim=1)[0] > 1.0, dim=1
176        )
177        return died, time_out
178
179    def _reset_idx(self, env_ids: torch.Tensor | None):
180        if env_ids is None or len(env_ids) == self.num_envs:
181            env_ids = wp.to_torch(self._robot._ALL_INDICES)
182        assert env_ids is not None
183        self._robot.reset(env_ids)
184        super()._reset_idx(env_ids)
185        if len(env_ids) == self.num_envs:
186            self.episode_length_buf[:] = torch.randint_like(self.episode_length_buf, high=int(self.max_episode_length))
187        self._actions[env_ids] = 0.0
188        self._previous_actions[env_ids] = 0.0
189        self._commands[env_ids] = torch.zeros_like(self._commands[env_ids]).uniform_(-1.0, 1.0)
190        joint_pos = self._robot.data.default_joint_pos.torch[env_ids]
191        joint_vel = self._robot.data.default_joint_vel.torch[env_ids]
192        default_root_pose = self._robot.data.default_root_pose.torch[env_ids]
193        default_root_vel = self._robot.data.default_root_vel.torch[env_ids]
194        default_root_pose[:, :3] += self._terrain.env_origins[env_ids]
195        self._robot.write_root_pose_to_sim_index(root_pose=default_root_pose, env_ids=env_ids)
196        self._robot.write_root_velocity_to_sim_index(root_velocity=default_root_vel, env_ids=env_ids)
197        self._robot.write_joint_position_to_sim_index(position=joint_pos, env_ids=env_ids)
198        self._robot.write_joint_velocity_to_sim_index(velocity=joint_vel, env_ids=env_ids)
199        extras = dict()
200        for key in self._episode_sums.keys():
201            episodic_sum_avg = torch.mean(self._episode_sums[key][env_ids])
202            extras["Episode_Reward/" + key] = episodic_sum_avg / self.max_episode_length_s
203            self._episode_sums[key][env_ids] = 0.0
204        self.extras["log"] = dict()
205        self.extras["log"].update(extras)
206        extras = dict()
207        extras["Episode_Termination/base_contact"] = torch.count_nonzero(self.reset_terminated[env_ids]).item()
208        extras["Episode_Termination/time_out"] = torch.count_nonzero(self.reset_time_outs[env_ids]).item()
209        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 use export.py, 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.

Note

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.

Note

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

Warning

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.

For direct environments, semantic annotations are optional and should 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.

Note

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