From 30d66c07b4721545653316ba1b1312d38543f4ff Mon Sep 17 00:00:00 2001 From: Lionel Gulich Date: Sat, 6 Jun 2026 17:05:00 +0200 Subject: [PATCH 1/5] Fix LEAPP projected gravity export Expose root orientation as the LEAPP graph input for projected gravity observations and compute the projected gravity tensor inside the exported graph. This keeps the policy observation contract unchanged while matching deployment inputs. --- .../lgulich-leapp-root-quat-gravity.rst | 6 + .../isaaclab/utils/leapp/export_annotator.py | 30 +++++ .../isaaclab/test/utils/test_leapp_proxy.py | 104 ++++++++++++++++++ 3 files changed, 140 insertions(+) create mode 100644 source/isaaclab/changelog.d/lgulich-leapp-root-quat-gravity.rst create mode 100644 source/isaaclab/test/utils/test_leapp_proxy.py diff --git a/source/isaaclab/changelog.d/lgulich-leapp-root-quat-gravity.rst b/source/isaaclab/changelog.d/lgulich-leapp-root-quat-gravity.rst new file mode 100644 index 000000000000..92af9a52251c --- /dev/null +++ b/source/isaaclab/changelog.d/lgulich-leapp-root-quat-gravity.rst @@ -0,0 +1,6 @@ +Fixed +^^^^^ + +* Fixed LEAPP export of :func:`isaaclab.envs.mdp.projected_gravity` to expose + root orientation as the graph input and compute projected gravity inside the + exported graph. diff --git a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py index fb335811535a..3ca82a66c082 100644 --- a/source/isaaclab/isaaclab/utils/leapp/export_annotator.py +++ b/source/isaaclab/isaaclab/utils/leapp/export_annotator.py @@ -291,6 +291,8 @@ def _patch_observation_manager(self, obs_manager, proxy_env): term_cfg.func = self._wrap_last_action(original_func) elif func_name == "generated_commands": term_cfg.func = self._wrap_generated_commands(original_func, term_cfg) + elif func_name == "projected_gravity": + term_cfg.func = self._wrap_projected_gravity(original_func, proxy_env) else: term_cfg.func = self._wrap_with_proxy(original_func, proxy_env) @@ -445,6 +447,34 @@ def wrapped(*args, **kwargs): wrapped.__name__ = getattr(original_func, "__name__", "unknown") return wrapped + @staticmethod + def _wrap_projected_gravity(original_func, proxy_env): + """Wrap projected gravity as root-quaternion input plus fixed gravity projection. + + Deployment backends generally provide body orientation, not an already + projected gravity vector. During export, keep the policy observation as + projected gravity while exposing ``root_quat_w`` at the LEAPP graph + boundary. + """ + + def wrapped(*args, **kwargs): + """Compute projected gravity from an annotated root quaternion input.""" + kwargs.pop("inspect", None) + asset_cfg = kwargs.get("asset_cfg") + if asset_cfg is None and len(args) > 1: + asset_cfg = args[1] + asset_name = getattr(asset_cfg, "name", "robot") + root_quat_w = proxy_env.scene[asset_name].data.root_quat_w.torch + gravity_w = torch.zeros((*root_quat_w.shape[:-1], 3), dtype=root_quat_w.dtype, device=root_quat_w.device) + gravity_w[..., 2] = -1.0 + quat_xyz = root_quat_w[..., :3] + quat_w = root_quat_w[..., 3:4] + t = torch.cross(quat_xyz, gravity_w, dim=-1) * 2.0 + return gravity_w - quat_w * t + torch.cross(quat_xyz, t, dim=-1) + + wrapped.__name__ = getattr(original_func, "__name__", "unknown") + return wrapped + def _wrap_last_action(self, original_func): """Wrap ``last_action`` as a LEAPP state tensor. diff --git a/source/isaaclab/test/utils/test_leapp_proxy.py b/source/isaaclab/test/utils/test_leapp_proxy.py new file mode 100644 index 000000000000..d57cbe1580b7 --- /dev/null +++ b/source/isaaclab/test/utils/test_leapp_proxy.py @@ -0,0 +1,104 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import math +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("leapp") + +from isaaclab.envs import mdp +from isaaclab.test.mock_interfaces.assets.mock_articulation import MockArticulationData +from isaaclab.utils import math as math_utils +from isaaclab.utils.leapp import utils as leapp_utils +from isaaclab.utils.leapp.export_annotator import ExportPatcher +from isaaclab.utils.leapp.leapp_semantics import InputKindEnum +from isaaclab.utils.leapp.proxy import _DataProxy, _EnvProxy + + +class _TestScene(dict): + """Minimal scene mapping for LEAPP proxy tests.""" + + sensors = {} + + +def _make_articulation_data() -> tuple[MockArticulationData, torch.Tensor]: + """Create mock articulation data with a non-identity root orientation.""" + data = MockArticulationData(num_instances=2, num_joints=0, num_bodies=1, device="cpu") + root_pose_w = torch.zeros(2, 7, dtype=torch.float32) + root_pose_w[:, 6] = 1.0 + root_pose_w[1, 3] = math.sin(math.pi / 4.0) + root_pose_w[1, 6] = math.cos(math.pi / 4.0) + data.set_root_link_pose_w(root_pose_w) + return data, root_pose_w + + +def _capture_leapp_inputs(monkeypatch: pytest.MonkeyPatch) -> list: + """Capture LEAPP input annotations while returning their tensor references.""" + annotated_inputs = [] + + def _record_input_tensor(task_name, semantics): + annotated_inputs.append((task_name, semantics)) + return semantics.ref + + monkeypatch.setattr(leapp_utils.annotate, "input_tensors", _record_input_tensor) + return annotated_inputs + + +def test_direct_projected_gravity_b_read_preserves_vector3d_input(monkeypatch: pytest.MonkeyPatch): + """Test direct data proxy reads keep projected gravity as its own semantic input.""" + annotated_inputs = _capture_leapp_inputs(monkeypatch) + data, _ = _make_articulation_data() + + proxy = _DataProxy( + data, + entity_name="robot", + task_name="Isaac-Velocity-Flat-G1-v0", + property_resolution_cache={}, + cache={}, + input_name_resolver=lambda property_name: f"robot_{property_name}", + ) + + assert proxy.projected_gravity_b.torch.shape == (2, 3) + + assert len(annotated_inputs) == 1 + task_name, semantics = annotated_inputs[0] + assert task_name == "Isaac-Velocity-Flat-G1-v0" + assert semantics.name == "robot_projected_gravity_b" + assert semantics.kind == InputKindEnum.VECTOR3D + assert semantics.extra == {"isaaclab_connection": "state:robot:projected_gravity_b"} + + +def test_projected_gravity_observation_exports_root_quat_w_input(monkeypatch: pytest.MonkeyPatch): + """Test the projected-gravity observation is export-lowered through root quaternion.""" + annotated_inputs = _capture_leapp_inputs(monkeypatch) + data, root_pose_w = _make_articulation_data() + scene = _TestScene({"robot": SimpleNamespace(data=data)}) + env = SimpleNamespace(scene=scene) + proxy_env = _EnvProxy(env, "Isaac-Velocity-Flat-G1-v0", {}, {}) + + term_cfg = SimpleNamespace(func=mdp.projected_gravity, noise="noise") + obs_manager = SimpleNamespace(_group_obs_term_cfgs={"policy": [term_cfg]}, compute=lambda *args, **kwargs: None) + patcher = ExportPatcher(export_method="onnx-dynamo", required_obs_groups={"policy"}) + patcher.task_name = "Isaac-Velocity-Flat-G1-v0" + patcher._patch_observation_manager(obs_manager, proxy_env) + + projected_gravity_b = term_cfg.func(env) + + expected = math_utils.quat_apply_inverse( + root_pose_w[:, 3:7], + torch.tensor([[0.0, 0.0, -1.0]], dtype=torch.float32).expand(2, 3), + ) + assert torch.allclose(projected_gravity_b, expected) + assert term_cfg.noise is None + + assert len(annotated_inputs) == 1 + task_name, semantics = annotated_inputs[0] + assert task_name == "Isaac-Velocity-Flat-G1-v0" + assert semantics.name == "robot_root_quat_w" + assert semantics.kind == InputKindEnum.BODY_ROTATION + assert semantics.extra == {"isaaclab_connection": "state:robot:root_quat_w"} From 57bcf60e276c4e07ad3265e85c1d3c43127aecf7 Mon Sep 17 00:00:00 2001 From: Lionel Gulich Date: Sun, 7 Jun 2026 22:52:29 +0200 Subject: [PATCH 2/5] Move LEAPP proxy tests to export test directory --- .../test/utils => isaaclab_rl/test/export}/test_leapp_proxy.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename source/{isaaclab/test/utils => isaaclab_rl/test/export}/test_leapp_proxy.py (100%) diff --git a/source/isaaclab/test/utils/test_leapp_proxy.py b/source/isaaclab_rl/test/export/test_leapp_proxy.py similarity index 100% rename from source/isaaclab/test/utils/test_leapp_proxy.py rename to source/isaaclab_rl/test/export/test_leapp_proxy.py From d30f2b737fded81a835c0558fdd7deb0b9a3bcc1 Mon Sep 17 00:00:00 2001 From: Lionel Gulich Date: Sun, 7 Jun 2026 22:54:57 +0200 Subject: [PATCH 3/5] Mark isaaclab_rl test move as changelog skip --- .../isaaclab_rl/changelog.d/lgulich-leapp-root-quat-gravity.skip | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 source/isaaclab_rl/changelog.d/lgulich-leapp-root-quat-gravity.skip diff --git a/source/isaaclab_rl/changelog.d/lgulich-leapp-root-quat-gravity.skip b/source/isaaclab_rl/changelog.d/lgulich-leapp-root-quat-gravity.skip new file mode 100644 index 000000000000..e69de29bb2d1 From 2bd4a4b6eb08f4395820cc6694a80668549fe5b5 Mon Sep 17 00:00:00 2001 From: Lionel Gulich Date: Mon, 8 Jun 2026 16:52:43 +0200 Subject: [PATCH 4/5] Fix recurrent LEAPP export dropping the LSTM hidden state The export wrote the traced recurrent state to `memory.hidden_state`, but rsl_rl's Memory module reads/writes `hidden_states` (plural). The registered state tensors were therefore never fed into the LSTM: LEAPP flagged the `actor_state_*` inputs as unused and dropped them, producing a stateless ONNX (the LSTM reset to zeros every call). Write to `hidden_states` so the hidden state is threaded as graph I/O and the recurrent policy exports as stateful. Co-Authored-By: Claude Opus 4.8 --- scripts/reinforcement_learning/leapp/rsl_rl/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py index fedda81cb961..064f79c2530f 100644 --- a/scripts/reinforcement_learning/leapp/rsl_rl/export.py +++ b/scripts/reinforcement_learning/leapp/rsl_rl/export.py @@ -199,7 +199,7 @@ def ensure_actor_hidden_state_initialized(policy_nn, batch_size: int, device, dt actor_state = (zeros.clone(), zeros.clone()) else: actor_state = zeros - memory.hidden_state = actor_state + memory.hidden_states = actor_state return actor_state @@ -329,7 +329,7 @@ def export_rsl_rl_agent( ) actor_memory = get_actor_memory_module(policy_nn) if actor_memory is not None: - actor_memory.hidden_state = actor_hidden_from_registered(registered_state, actor_hidden) + actor_memory.hidden_states = actor_hidden_from_registered(registered_state, actor_hidden) actions = policy(obs) From 6f18d1cfa4e984770c0a519dd3e82441204f40d3 Mon Sep 17 00:00:00 2001 From: Lionel Gulich Date: Mon, 8 Jun 2026 19:01:02 +0200 Subject: [PATCH 5/5] Revert recurrent hidden_state attribute change My earlier commit changed the actor memory write to `hidden_states` (plural), but rsl_rl 5.0.1 (which Isaac Lab pins) stores the RNN state in `hidden_state` (singular, rsl_rl/modules/rnn.py). The plural form only exists in older/vendored rsl_rl (e.g. AGILE's 2.3.3 fork) and as the algorithm-level (actor, critic) tuple. The original `hidden_state` is correct for 5.0.1, so reverting; the AGILE recurrent-export issue is a vendored-rsl_rl version mismatch to be fixed by upgrading AGILE to rsl_rl 5.0.1. Co-Authored-By: Claude Opus 4.8 --- scripts/reinforcement_learning/leapp/rsl_rl/export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/reinforcement_learning/leapp/rsl_rl/export.py b/scripts/reinforcement_learning/leapp/rsl_rl/export.py index 064f79c2530f..fedda81cb961 100644 --- a/scripts/reinforcement_learning/leapp/rsl_rl/export.py +++ b/scripts/reinforcement_learning/leapp/rsl_rl/export.py @@ -199,7 +199,7 @@ def ensure_actor_hidden_state_initialized(policy_nn, batch_size: int, device, dt actor_state = (zeros.clone(), zeros.clone()) else: actor_state = zeros - memory.hidden_states = actor_state + memory.hidden_state = actor_state return actor_state @@ -329,7 +329,7 @@ def export_rsl_rl_agent( ) actor_memory = get_actor_memory_module(policy_nn) if actor_memory is not None: - actor_memory.hidden_states = actor_hidden_from_registered(registered_state, actor_hidden) + actor_memory.hidden_state = actor_hidden_from_registered(registered_state, actor_hidden) actions = policy(obs)