[Draft] Added export with leapp to isaac lab. by frlai · Pull Request #1 · frlai/IsaacLab

frlai · 2026-01-12T19:21:45Z

DO NOT MERGE
Exports isaac lab managed environment policies with LEAPP. This implementation focuses on performing exports using only monkey patching. All edits functional edits are localized to scripts/reinforcement_learning/rsl_rl/annotate_functions_for_export.py. where many articulation functions, observation manager and action manager functions are patched to annotate traced tensors.

The final expected output includes 3 model fliles: observation_manager, action_manager, and policy.

This currently works for all managed environment policies for rslrl built into isaac lab. The generated yaml also incorporates info from the io_descriptors such as joint names and units.

Limitations:
The computation for projected_gravity_b cannot be traced without some modifications to the source code since key functions are scripted and strictly only take torch tensors, the current TracedTensor strategy fails. This can be fixed if we modify how some functions in math.py work

pipeline now correctly reports, removed all modifications to original moved patching logic to seperate file and expanded capability applied workaround for reused input functions bug

lgulich

Looks good overall. Left some small comments.

Two bigger questions:

Do you already have an idea of how to add additional metadata to the exported LEAPP graph? I'm thinking mainly about the joint order but there might also be other required data.
It would be nice if we could also export the kp and kd gains. Either they could be just a constant output from the resulting LEAPP graph or in more complex cases they might be a dynamic output from the policy. Is there a way we can do this?

lgulich · 2026-01-13T09:01:01Z

+
+    # start annotation tracing
+    # Note: all patching is done at module/class level before isaaclab_tasks import
+    annotate.start(task_name)


do I understand correctly that this file is the same as the play.py script, just with the annotate commands added here?

yes. I'm not sure what we want to do with the entry point. I also have the limitation where I can only support --num_envs 1. Using this for my testing just for now.

…s that copy data using [:]

…g the same observation multiple times

frlai · 2026-01-16T21:58:48Z

u

newest version addresses both of these
for semantic data captured by io_descriptors I am only keeping:

joint_names
units

there is also a field that maps the semantic data to the main Leapp graph. Tell me if you need more information from io_descriptors.

kellyguo11 · 2026-01-29T22:31:50Z

+    _annotated_tensor_cache: dict[str, torch.Tensor] = field(default_factory=dict, repr=False)
+
+    # Articulation properties to annotate
+    OBSERVATION_PROPERTIES: frozenset[str] = frozenset({


do we anticipate this will be a fixed list? would it make sense to parse this from properties available in ArticulationData?

I'm actually unable to use this method to parse all the functions in ArticulationData. the reason is because some processing functions are decorated with torch.jit.script which will cause an error if you try to pass a wrapped function in. Currently these functions are the only ones used by the ObservationManager. I'm curious about all the possible use cases and I'm wondering if it is possible for a user to have a managed environment and used some ArticulationData that is outside of this list.

kellyguo11 · 2026-01-30T01:06:07Z

+            tensors[term_name] = term.processed_actions
+
+            # Handle variable impedance (dynamic gains)
+            osc = getattr(term, "_osc", None)


there seems to be many places where we are checking for specific strings, which feels like could be fragile. are these limitations because of how lab is exposing data?

for this case you're commenting on I needed to check for specific attributes by name because I needed a way to detect if the gains were dynamic or static. Not familiar with how the kp and kd gains are configured and this was just the first thing that works. Open to revising this part.

kellyguo11 · 2026-01-30T01:08:02Z

+            for term_cfg in term_cfgs:
+                term_cfg.noise = None
+
+    def _patch_observation_functions(self):


are these patches the workaround for injecting the annotations into the code so that we don't have these in the main source code? I wonder if it makes sense to have a set of specialized classes for deployment that has these injections in there. if something ends up changing in the framework, we might need to be careful about impacting these patches.

Either that or we can just leave the annotations in the source code. leapp is built to noop unless the start flag has been set. For example if you're ok with one line additions to articulation functions or mdp observation functions that will add one function call and one if statement check we can skip adding this entirely. Similar story on the action manager side. I'm slightly against actually writing a separate environment wrapper or something in case the logic might end up diverging.

rafaelcathomen · 2026-01-30T10:56:44Z

+                if func_name == "last_action":
+                    self._original_obs_funcs[(group_name, term_idx)] = original_func
+                    term_cfg.func = self._make_patched_last_action(original_func)
+
+                elif func_name == "generated_commands":


This assumes that the user uses the names last_action and generated_commands which might not be the case. Is there a better way to check whether an observation function is a command (could also have multiple command functions) or a last action?

The function that the example tasks use to generate commands is stored in source/isaaclab/isaaclab/envs/mdp/observations.py. This function previously has a @generic_io_descriptor decorator on it so I assumed targeting this function was enough. Similar story with last action. Are you saying you have a situation where in a managed environment it is possible for some user to define their own command function?

Yes, a user might implement their own command, for example a height tracking command. Then they might create observations velocity_tracking_command and seperately height_tracking_command

if the command is self implemented a user might need to write their own annotations for compatibility.
I'm not sure if there is a good way to detect this.

rafaelcathomen · 2026-01-30T14:57:16Z

+        static_values = {}
+
+        for term_name, term in action_manager._terms.items():
+            tensors[term_name] = term.processed_actions


In IsaacLab processed_actions can be a joint position target, a relative joint position target (logic is in apply_actions), a joint velocity target or a joint effort target. How do we handle these action types? What if a user implements their own action which uses multiple of these? The most generic action would be [joint_pos_target, joint_vel_target, feed_forward_torque, kp, kd], no?

it should handle multiple actions since it just iterates through all the action terms in the manager.
but you're right not capturing the final output after apply_actions was an oversight. I will look fixing this. should be a relatively simple patch.

rafaelcathomen · 2026-02-02T16:22:37Z

+        """Temporarily apply annotating properties."""
+        if not self._annotations_active:
+            # Clear the tensor cache at the start of each compute_group call
+            self._annotated_tensor_cache.clear()


Did you try it with an environment with multiple observation groups with shared observation terms? I had to move the cache clearing to avoid duplicates across group.

…er side observation patching

…tives

added initial version

78bcc18

pipeline now correctly reports, removed all modifications to original moved patching logic to seperate file and expanded capability applied workaround for reused input functions bug

github-actions Bot added the isaac-lab label Jan 12, 2026

lgulich reviewed Jan 13, 2026

View reviewed changes

frlai added 11 commits January 14, 2026 16:41

added static outputs to graph. addressd comments

39f5dbc

added precommit

6194308

added static buffers to graph. we can now trace class tensor variable…

a69f55f

…s that copy data using [:]

removed random observation corruption for during export

74dbb8f

added io descriptor info into leapp generated yaml

a83a600

ugly implementation for automatic mapping from io descriptors to leapp

d658429

added action mapping for io_descriptor output

7a1a046

cleaned up utility to annotate environment, added protection for usin…

cc69ff0

…g the same observation multiple times

added documentation

fa8ffd9

added sample generated yaml file

619b4b3

fixed bug in action_manager

d05a2d2

frlai added 3 commits January 28, 2026 10:16

switched to export with onnx

b10a247

exports as one single policy

f85131b

fixed readme

847adbc

kellyguo11 reviewed Jan 30, 2026

View reviewed changes

rafaelcathomen reviewed Jan 30, 2026

View reviewed changes

rafaelcathomen reviewed Feb 2, 2026

View reviewed changes

frlai added 8 commits March 9, 2026 14:35

changes for export to work after v0.5.0 update

b84b312

refactor and added semantic data

e2d2038

updated annotations to use local patching again

6d4e3b8

refactor for unified patching class. new class also does action manag…

e6cc033

…er side observation patching

new setup adds kp and kd semantics

e0e97af

implemented deployment environment. need to validate against more models

971e38b

added a test for export. expanded annotations to include sensor primi…

b6f5a61

…tives

fixed issue with callables in observation terms

93025a7

Conversation

frlai commented Jan 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lgulich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

frlai commented Jan 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

frlai commented Jan 12, 2026 •

edited

Loading