Preserve LEAPP metadata for manual deploy inputs#6185
Conversation
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.
Greptile SummaryThis PR introduces a
Confidence Score: 4/5Safe 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
Sequence DiagramsequenceDiagram
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]
|
|
|
||
| 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] |
There was a problem hiding this comment.
_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.
| 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) |
There was a problem hiding this comment.
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 |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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.
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 withkind,element_names, andisaaclab_connectionmetadata. This keeps exported YAML compatible with Isaac ROS Deploy input routing, which selects converters using the inputkind.This PR also updates the gear assembly deploy observations so the exported policy inputs are:
robot_joint_pos: 7 arm joints,state/joint/positionrobot_joint_vel: 7 arm joints,state/joint/velocitygear_shaft_pos:state/body/positiongear_shaft_quat:state/body/rotationThe action output remains arm-only and keeps its existing
target/joint/positionmetadata.No new dependencies are required.
Fixes #
Type of change
Screenshots
Not applicable.
Checklist
pre-commitchecks with./isaaclab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists there