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_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 diff --git a/source/isaaclab_rl/test/export/test_leapp_proxy.py b/source/isaaclab_rl/test/export/test_leapp_proxy.py new file mode 100644 index 000000000000..d57cbe1580b7 --- /dev/null +++ b/source/isaaclab_rl/test/export/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"}