Fix LEAPP projected gravity export

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.

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.
 

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"}