使用 LEAPP 导出 Direct 工作流策略

使用 LEAPP 导出 Direct 工作流策略#

本进阶教程介绍如何为 Direct 工作流策略添加 LEAPP 导出支持。如果策略采用基于管理器的工作流，请参阅 标准 LEAPP 部署工作流 。

LEAPP 概念、支持的节点模式、状态反馈及运行时验证等背景知识，见 LEAPP 文档 。

快速开始#

要通过 LEAPP 导出 Direct 工作流策略，需要在环境代码中添加 LEAPP 标注。导出时，LEAPP 会跟踪已标注的张量，构建完整策略流水线的中间表示。这些标注在环境正常运行时不启用，仅带来少量开销，直到导出阶段才被激活。运行导出流程时，由 isaaclab leapp export --rl_library rsl_rl 激活它们。

本教程以 scripts/tutorials/06_deploy/anymal_c_env.py 为例，说明如何在 Direct 工作流环境中添加 LEAPP 标注。您可以将同样的标注方法用于自己的 Direct RL 环境。

这一导出流程需要 leapp 。它是专用的可选额外依赖，不包含在 --extra all 中。额外依赖的选择方式与 isaaclab train 相同：每次 uv run 都添加 --extra leapp ，并添加与任务匹配的后端额外依赖（ isaacsim 、 ovphysx ，不依赖 Kit 的 Newton 则无需添加）。见 安装并运行第一个任务 及 使用 LEAPP 导出策略 。

如果希望使用现有的 IsaacContrib-Velocity-Rough-AnymalC-Direct 任务注册来运行导出的示例，请将已标注的教程环境复制到任务包中：

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

copy scripts\tutorials\06_deploy\anymal_c_env.py ^
   source\isaaclab_tasks\isaaclab_tasks\contrib\anymal_c_direct\anymal_c_env.py

环境具备所需的 LEAPP 输入、输出及状态标注后，可按与训练相同的后端额外依赖设置导出已训练策略。例如，使用 Isaac Sim PhysX 时：

# 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

:: 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:
set OMNI_KIT_ACCEPT_EULA=Y
set 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

--task 参数是已注册的任务名，如 IsaacContrib-Velocity-Rough-AnymalC-Direct 。 --checkpoint 指向要导出的已训练 RSL-RL 检查点。训练请使用 uv run isaaclab train --rl_library rsl_rl --task <TASK_NAME> （见 安装并运行第一个任务 ），不要使用已移除的各后端 train.py 脚本。可选的 --export_save_path 指定导出产物的输出目录；省略时，导出到检查点所在目录。

备注

与受管理的环境不同，这里的 export_method 不是命令行选项。请在输出标注中指定导出后端，例如 annotate.output_tensors(..., export_with="onnx-dynamo") 。具体用法见下文的输出标注示例。

警告

本教程仅介绍 Direct RL 策略的导出。 isaaclab leapp deploy 目前不支持 Direct RL 策略。

导出参数的更多信息，见 标准 LEAPP 部署工作流 。

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

标注的工作原理#

主要工作是识别环境中的输入、输出和持久状态，并向 LEAPP 注册。本示例使用四个标注辅助函数：

  • annotate.input_tensors() 标记进入策略流水线的张量。

  • annotate.output_tensors() 标记离开流水线中环境侧部分的张量。

  • annotate.state_tensors() 标记作为持久状态的张量。

  • annotate.update_state() 在每步之后更新该持久状态。

输入标注#

输入标注通常放在 _get_observations() 中，因为该方法负责收集传给策略的张量。

# 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() 包装一个张量，使 LEAPP 能跟踪依赖该张量的所有下游操作。该函数有两个重要参数：

  • self.spec.id 标识张量所属的节点。运行 isaaclab leapp export 时，该 ID 与导出的策略节点对应。

  • 第二个参数是字典，将唯一的张量名映射到张量本身。LEAPP 在导出元数据和调试中使用这些名称。

为便于阅读，本例逐个注册观测张量； annotate.input_tensors() 也支持一次注册多个张量。

备注

未显式标注的输入会自动作为常量内联。对于固定变换或默认值等数据，这可能正是所需行为。

输出标注#

输出标注应放在环境已完成机器人指令准备的位置。本例中，该位置在 _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() 标记离开流水线中环境侧部分的张量。与输入标注一样，调用时使用 self.spec.id ，以及将张量名映射到张量的字典。

export_with 参数将输出标注限定为特定导出后端。支持的后端名有 onnx-dynamo 、 onnx-torchscript 、 jit-script 和 jit-trace 。必须提供该参数，才能根据跟踪结果实际生成中间表示（IR）。

与 annotate.input_tensors() 不同，输出标注应只对流水线阶段的最终输出执行一次。本例中，要导出的张量是 processed_actions 。调用 annotate.output_tensors() 后，无需使用返回值。

备注

传入 annotate.output_tensors() 的所有张量都必须是已跟踪的张量，即由输入或输入派生的张量创建。

警告

不要在 _apply_action() 中添加输出标注。根据控制降采样倍数设置，该方法可能在每个环境步内被调用多次，从而导致跟踪出的流水线不正确。

状态标注#

如果策略依赖内部状态或反馈环路，请用 annotate.state_tensors() 显式注册这些数据，并通过 annotate.update_state() 更新。

本例将前一步动作作为观测的一部分，因此 previous_actions 是反馈状态：

  • 在 _get_observations() 中调用 annotate.state_tensors() ，使状态参与被跟踪的观测流水线。

  • 在 _pre_physics_step() 中调用 annotate.update_state() ，为下一步更新存储值。

两次调用的状态名必须一致。本例中，两者都使用 previous_actions ，使 LEAPP 能正确传递反馈张量。

语义标注#

本例介绍跟踪流水线所需的最少标注。更复杂的导出工作流还可以附加语义元数据，让下游运行时了解每个张量的含义。要部署到下游部署库，必须提供语义标注。

对于 Direct 环境，语义标注是可选的，需由用户显式编写。与基于管理器的导出方式不同，Isaac Lab 不会自动推断 Direct 环境的张量语义，用户需自行提供这些数据。LEAPP 通过 TensorSemantics 支持这一功能。可用它更精确地描述张量含义，使导出流水线更易于检查、验证和集成到部署系统。

备注

语义标注的编写方法，见 LEAPP 语义标注指南 及 LEAPP API 参考 。

目录