Skip to content

Preserve LEAPP metadata for manual deploy inputs#6185

Open
ashwinvkNV wants to merge 1 commit into
isaac-sim:developfrom
ashwinvkNV:codex/leapp-manual-input-metadata
Open

Preserve LEAPP metadata for manual deploy inputs#6185
ashwinvkNV wants to merge 1 commit into
isaac-sim:developfrom
ashwinvkNV:codex/leapp-manual-input-metadata

Conversation

@ashwinvkNV

Copy link
Copy Markdown
Contributor

Description

This PR adds a metadata-preserving manual LEAPP input path for deploy observation tensors.

Some deploy workflows need to expose a higher-level observation tensor as the LEAPP input boundary instead of tracing every simulator-side dependency used to compute that observation. For the gear assembly ROS inference task, the deploy runtime provides arm joint state and shaft pose inputs directly, while the Isaac Lab environment may compute those values from simulator-only state during training/export.

The new leapp_input_tensor(...) helper is a no-op outside LEAPP export, but during export it can annotate a computed tensor as a LEAPP input with kind, element_names, and isaaclab_connection metadata. This keeps exported YAML compatible with Isaac ROS Deploy input routing, which selects converters using the input kind.

This PR also updates the gear assembly deploy observations so the exported policy inputs are:

  • robot_joint_pos: 7 arm joints, state/joint/position
  • robot_joint_vel: 7 arm joints, state/joint/velocity
  • gear_shaft_pos: state/body/position
  • gear_shaft_quat: state/body/rotation

The action output remains arm-only and keeps its existing target/joint/position metadata.

No new dependencies are required.

Fixes #

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)

Screenshots

Not applicable.

Checklist

  • I have read and understood the contribution guidelines
  • I have run the pre-commit checks with ./isaaclab.sh --format
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the changelog and the corresponding version in the extension's config/extension.toml file
  • I have added my name to the CONTRIBUTORS.md or my name already exists there

Add a LEAPP manual input helper that can annotate computed observation
tensors with semantic metadata while remaining a no-op outside export.

Use it for gear assembly deploy observations so exported policy inputs keep
Deploy-compatible kind, element_names, and connection metadata for arm joint
state and shaft pose inputs. Cache the manually annotated joint position input
for action-side relative joint position reads to avoid duplicate LEAPP input
annotations.

Add proxy coverage for manual cached inputs and regression tests for metadata
preservation.
@github-actions github-actions Bot added the isaac-lab Related to Isaac Lab team label Jun 14, 2026
@greptile-apps

greptile-apps Bot commented Jun 14, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR introduces a leapp_input_tensor helper and a leapp_real_env unwrapper to let deploy observation terms manually annotate pre-sliced tensors as LEAPP input boundaries, preserving kind, element_names, and isaaclab_connection metadata without tracing every simulator-side dependency.

  • manual_inputs.py adds the two new helpers; _ManualTensorView adapts a plain tensor to the .torch access pattern expected by the proxy layer, and leapp_input_tensor wires the annotated result into both the shared trace cache and a new per-_DataProxy _manual_cache so the alias survives explicit cache clears.
  • proxy.py gains a _manual_cache dict on _DataProxy that is consulted (with priority) before the shared trace cache, enabling the action-side read to receive the annotated tensor after the trace cache is cleared.
  • observations.py adds joint_pos / joint_vel wrappers that slice first and annotate second, and migrates gear_shaft_pos_w / gear_shaft_quat_w to use leapp_real_env so LEAPP export sees the raw quaternion before the qw >= 0 canonicalization step.

Confidence Score: 4/5

Safe to merge for the documented gear-assembly deploy workflow; the new helper is a no-op outside LEAPP export so training is unaffected.

The core mechanism is correct and exercised by the new test. The main fragility is that _DataProxy.getattr only checks _manual_cache after confirming the property has a _leapp_semantics decorator; if a backend property lacks that decoration, a cache= call silently does nothing. All current callers use LEAPP-decorated properties, so this is not a present failure, but it is an undocumented constraint. The self.asset field left unreferenced in both gear_shaft_pos_w and gear_shaft_quat_w is minor cleanup debt.

proxy.py (_DataProxy.getattr manual-cache ordering assumption) and observations.py (dead self.asset field) are worth a second look.

Important Files Changed

Filename Overview
source/isaaclab/isaaclab/utils/leapp/manual_inputs.py New module exposing leapp_input_tensor and leapp_real_env; logic is sound for documented use cases with minor edge-case fragility in the dict-path + cache combination.
source/isaaclab/isaaclab/utils/leapp/proxy.py Adds _manual_cache to _DataProxy with priority over the shared trace cache; check is only reachable for LEAPP-semantics-annotated properties, an implicit contract not enforced at the call site.
source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py Adds joint_pos / joint_vel wrappers that slice then annotate, and migrates gear_shaft_pos_w / gear_shaft_quat_w to use leapp_real_env; self.asset is now dead code after the migration.
source/isaaclab_rl/test/export/test_leapp_proxy.py Adds a focused test verifying manual-cache survival after trace cache clear; mock extended to handle both dict and TensorSemantics signatures. joint_vel annotation path is not tested.
source/isaaclab/isaaclab/utils/leapp/init.pyi Type stub updated to export leapp_input_tensor and leapp_real_env, matching the actual init.py exports.
source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/init.pyi Type stub updated to export the new joint_pos and joint_vel observation terms.

Sequence Diagram

sequenceDiagram
    participant OT as Observation Term
    participant LIT as leapp_input_tensor
    participant SP as _SceneProxy / _DataProxy
    participant AN as annotate.input_tensors
    participant AC as Action Term

    Note over OT,AC: LEAPP Export path only (env is _EnvProxy)

    OT->>LIT: leapp_input_tensor(env, name, sliced_tensor, kind, cache)
    LIT->>AN: input_tensors(task_id, TensorSemantics)
    AN-->>LIT: annotated_tensor
    LIT->>SP: proxy_cache write
    LIT->>SP: data_proxy._manual_cache write
    LIT-->>OT: annotated_tensor

    Note over SP: Trace cache cleared externally

    AC->>SP: proxy_env.scene[robot].data.joint_pos
    SP->>SP: _resolve_annotated_property → non-None
    SP->>SP: check _manual_cache → HIT
    SP-->>AC: _ManualTensorView(annotated_tensor)
    AC->>AC: .torch[:, joint_ids]
Loading

Comments Outside Diff (1)

  1. source/isaaclab_tasks/isaaclab_tasks/contrib/deploy/mdp/observations.py, line 106-118 (link)

    P2 self.asset is now dead code in both gear_shaft_pos_w and gear_shaft_quat_w

    Both __init__ methods still assign self.asset = env.scene[self.asset_cfg.name], but after this PR neither __call__ reads self.asset anymore — both now do real_env.scene[self.asset_cfg.name] on every invocation. The cached field and its type annotation are dead code and may confuse future readers.

    Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Reviews (1): Last reviewed commit: "Preserve LEAPP metadata for manual deplo..." | Re-trigger Greptile

Comment on lines 162 to 169

cache = object.__getattribute__(self, "_cache")
cache_key = (id(real_data), name)
manual_cache = object.__getattribute__(self, "_manual_cache")
if cache_key in manual_cache:
return manual_cache[cache_key]
if cache_key in cache:
return cache[cache_key]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 _manual_cache silently skipped for non-semantic properties

The _manual_cache check is only reached when _resolve_annotated_property returns a non-None resolution — i.e., when the property has a _leapp_semantics decorator somewhere in the MRO. If a caller passes cache=("robot", "some_prop") for a property that is NOT LEAPP-annotated, leapp_input_tensor writes to the manual cache successfully, but every subsequent read falls through to getattr(real_data, name) and returns the raw value, silently ignoring the cached annotated tensor. Consider adding a guard or a comment documenting this assumption.

Comment on lines +76 to +86
if kind is None and element_names is None and extra is None:
annotated = annotate.input_tensors(env.unwrapped.spec.id, {name: tensor})
else:
semantics = TensorSemantics(
name=name,
ref=tensor,
kind=kind,
element_names=element_names,
extra=extra,
)
annotated = annotate.input_tensors(env.unwrapped.spec.id, semantics)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Dict-path result type is not validated before being stored in the manual cache

When kind, element_names, and extra are all None, annotate.input_tensors is called with a plain {name: tensor} dict. If cache is also non-None, annotated is passed directly to _ManualTensorView(annotated), which expects a torch.Tensor. The real LEAPP library could return a different type from its dict-signature overload. In practice all current callers that supply cache= also supply kind, so the risk is contained; still, worth documenting the assumption or adding a guard before _ManualTensorView(annotated).

# the real env and annotate it before canonicalization so the ONNX graph
# owns the ``qw >= 0`` sign convention.
real_env = leapp_real_env(env)
base_quat = real_env.scene[self.asset_cfg.name].data.root_quat_w.torch

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah here can just be
base_quat = annotate.input_tensor('policy', TensorSemantics('gear_shaft_quat', base_quat, kind=InputKindEnum.BODY_ROTATION, element_names=QUAT_XYZW_ELEMENT_NAMES)

similar to everything above

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why was this needed?
this seems like a re-implementation of leapp tensors.

this seems to also hook onto the env proxy which I don't think is needed. That is only for automatic annotation at the source.

@frlai frlai left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the AI was trying to do this to preserve the ability to run everything in simulation using deploy.py. I'm not sure if that's a goal for us but if you're able to run everything in deploy.py using this method that would actually be pretty interesting and might be a useful change. But for a MVP, most of these changes can be cut down i believe.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

isaac-lab Related to Isaac Lab team

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants