Adds in frame transformer and demo script (#183)

# Description
Adds in a new `Sensor` that performs frame transforms.

No tests as part of this initial PR as this feature is needed by
developers. Tests will be added in with the improvements:
https://github.com/isaac-orbit/orbit/issues/202.

Fixes #133 

## Type of change
- New feature (non-breaking change which adds functionality)

## Screenshots

![image](https://github.com/isaac-orbit/orbit/assets/142246516/d4e2ebde-377f-4e22-9c14-8ec19b7b7498)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file

---------
Signed-off-by: jsmith-bdai <142246516+jsmith-bdai@users.noreply.github.com>

source/extensions/omni.isaac.orbit/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.9.25"
+version = "0.9.26"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.9.26 (2023-10-31)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added the sensor implementation for :class:`omni.isaac.orbit.sensors.FrameTransformer` class. Currently,
+  it handles obtaining the transformation between two frames in the same articulation.
+
+
 0.9.25 (2023-10-27)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/scene/interactive_scene.py

@@ -19,7 +19,7 @@ from omni.isaac.version import get_version
 from pxr import PhysxSchema
 
 from omni.isaac.orbit.assets import Articulation, ArticulationCfg, AssetBaseCfg, RigidObject, RigidObjectCfg
-from omni.isaac.orbit.sensors import SensorBase, SensorBaseCfg
+from omni.isaac.orbit.sensors import FrameTransformerCfg, SensorBase, SensorBaseCfg
 from omni.isaac.orbit.terrains import TerrainImporter, TerrainImporterCfg
 
 from .interactive_scene_cfg import InteractiveSceneCfg
@@ -342,6 +342,14 @@ class InteractiveScene:
             elif isinstance(asset_cfg, RigidObjectCfg):
                 self.rigid_objects[asset_name] = asset_cfg.class_type(asset_cfg)
             elif isinstance(asset_cfg, SensorBaseCfg):
+                # Update target frame path(s)' regex name space for FrameTransformer
+                if isinstance(asset_cfg, FrameTransformerCfg):
+                    updated_target_frames = []
+                    for target_frame in asset_cfg.target_frames:
+                        target_frame.prim_path = target_frame.prim_path.format(ENV_REGEX_NS=self.env_regex_ns)
+                        updated_target_frames.append(target_frame)
+                    asset_cfg.target_frames = updated_target_frames
+
                 self.sensors[asset_name] = asset_cfg.class_type(asset_cfg)
             elif isinstance(asset_cfg, AssetBaseCfg):
                 # manually spawn asset

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/__init__.py

@@ -5,16 +5,29 @@
 
 
 """
-This subpackage contains the sensor classes that are compatible with the Isaac Sim. We include both
+This subpackage contains the sensor classes that are compatible with Isaac Sim. We include both
 USD-based and custom sensors. The USD-based sensors are the ones that are available in Omniverse and
 require creating a USD prim for them. Custom sensors, on the other hand, are the ones that are
 implemented in Python and do not require creating a USD prim for them.
+
+A prim path (or expression) is still set for each sensor based on the following schema:
+
++-------------------+--------------------------+---------------------------------------------------------------+
+| Sensor Type       | Example Prim Path        | Pre-check                                                     |
++===================+==========================+===============================================================+
+| Camera            | /World/robot/base/camera | Leaf is available, and it will spawn a USD camera             |
+| Contact Sensor    | /World/robot/feet_*      | Leaf is available and checks if the schema exists             |
+| Ray Casters       | /World/robot/base        | Leaf exists and is a physics body (Articulation / Rigid Body) |
+| Frame Transformer | /World/robot/base        | Leaf exists and is a physics body (Articulation / Rigid Body) |
++-------------------+--------------------------+---------------------------------------------------------------+
+
 """
 
 from __future__ import annotations
 
 from .camera import *  # noqa: F401, F403
 from .contact_sensor import *  # noqa: F401, F403
+from .frame_transformer import *  # noqa: F401
 from .ray_caster import *  # noqa: F401, F403
 from .sensor_base import SensorBase  # noqa: F401
 from .sensor_base_cfg import SensorBaseCfg  # noqa: F401

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/frame_transformer/__init__.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Frame transform sensor for calculating frame transform of articulations.
+"""
+
+from __future__ import annotations
+
+from .frame_transformer import FrameTransformer
+from .frame_transformer_cfg import FrameTransformerCfg, OffsetCfg
+from .frame_transformer_data import FrameTransformerData
+
+__all__ = ["FrameTransformer", "FrameTransformerCfg", "FrameTransformerData", "OffsetCfg"]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/frame_transformer/frame_transformer_cfg.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+
+from __future__ import annotations
+
+from dataclasses import MISSING
+
+from omni.isaac.orbit.utils import configclass
+
+from ..sensor_base_cfg import SensorBaseCfg
+from .frame_transformer import FrameTransformer
+
+
+@configclass
+class OffsetCfg:
+    """The offset pose of one frame relative to another frame."""
+
+    pos: tuple[float, float, float] = (0.0, 0.0, 0.0)
+    """Translation w.r.t. the parent frame. Defaults to (0.0, 0.0, 0.0)."""
+    rot: tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0)
+    """Quaternion rotation ``(w, x, y, z)`` w.r.t. the parent frame. Defaults to (1.0, 0.0, 0.0, 0.0)."""
+
+
+@configclass
+class FrameTransformerCfg(SensorBaseCfg):
+    """Configuration for the frame transformer sensor."""
+
+    @configclass
+    class FrameCfg:
+        """Information specific to a coordinate frame."""
+
+        prim_path: str = MISSING
+        """The prim path corresponding to the parent rigid body.
+
+        This prim should be part of the same articulation as :attr:`FrameTransformerCfg.prim_path`.
+        """
+        name: str | None = None
+        """User-defined name for the new coordinate frame. Defaults to None.
+
+        If None, then the name is extracted from the leaf of the prim path.
+        """
+
+        offset: OffsetCfg = OffsetCfg()
+        """The pose offset from the parent prim frame."""
+
+    class_type: type = FrameTransformer
+
+    prim_path: str = MISSING
+    """The prim path of the body to transform from (source frame)."""
+
+    source_frame_offset: OffsetCfg = OffsetCfg()
+    """The pose offset from the source prim frame."""
+
+    target_frames: list[FrameCfg] = MISSING
+    """A list of the target frames.
+
+    This allows a single FrameTransformer to handle multiple target prims. For example, in a quadruped,
+    we can use a single FrameTransformer to track each foot's position and orientation in the body
+    frame using four frame offsets.
+    """

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/frame_transformer/frame_transformer_data.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+from dataclasses import dataclass
+
+
+@dataclass
+class FrameTransformerData:
+    """Data container for the frame transformer sensor."""
+
+    target_frame_names: list[str] = None
+    """Target frame names (this denotes the order that frame data will be ordered).
+
+    The frame names are resolved from the :attr:`FrameTransformerCfg.FrameCfg.name` field.
+    This usually follows the order in which the frames are defined in the config. However, in the case of
+    regex matching, the order may be different.
+    """
+
+    target_pos_source: torch.Tensor = None
+    """Position of the target frame(s) relative to source frame.
+
+    Shape is (N, M, 3), where N is the number of environments, and M is the number of target frames.
+    """
+    target_rot_source: torch.Tensor = None
+    """Orientation of the target frame(s) relative to source frame quaternion ``(w, x, y, z)``.
+
+    Shape is (N, M, 4), where N is the number of environments, and M is the number of target frames.
+    """
+    target_pos_w: torch.Tensor = None
+    """Position of the target frame(s) after offset (in world frame).
+
+    Shape is (N, M, 3), where N is the number of environments, and M is the number of target frames.
+    """
+    target_rot_w: torch.Tensor = None
+    """Orientation of the target frame(s) after offset (in world frame) quaternion ``(w, x, y, z)``.
+
+    Shape is (N, M, 4), where N is the number of environments, and M is the number of target frames.
+    """
+    source_pos_w: torch.Tensor = None
+    """Position of the source frame after offset (in world frame).
+
+    Shape is (N, 3), where N is the number of environments.
+    """
+    source_rot_w: torch.Tensor = None
+    """Orientation of the source frame after offset (in world frame) quaternion ``(w, x, y, z)``.
+
+    Shape is (N, 4), where N is the number of environments.
+    """

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/math.py

@@ -53,6 +53,7 @@ __all__ = [
     "quat_rotate",
     "quat_rotate_inverse",
     # Transformations
+    "is_identity_pose",
     "combine_frame_transforms",
     "subtract_frame_transforms",
     "compute_pose_error",
@@ -524,6 +525,27 @@ Transformations
 """
 
 
+def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool:
+    """Checks if input poses are identity transforms.
+
+    The function checks if the input position and orientation are close to zero and
+    identity respectively using L2-norm. It does NOT check the error in the orientation.
+
+    Args:
+        pos: The cartesian position. Shape is (N, 3).
+        rot: The quaternion orientation in (w, x, y, z). Shape is (N, 4).
+
+    Returns:
+        True if all the input poses result in identity transform. Otherwise, False.
+    """
+    # create identity transformations
+    pos_identity = torch.zeros_like(pos)
+    rot_identity = torch.zeros_like(rot)
+    rot_identity[..., 0] = 1
+    # compare input to identity
+    return torch.allclose(pos, pos_identity) and torch.allclose(rot, rot_identity)
+
+
 # @torch.jit.script
 def combine_frame_transforms(
     t01: torch.Tensor, q01: torch.Tensor, t12: torch.Tensor = None, q12: torch.Tensor = None

source/extensions/omni.isaac.orbit/test/sensors/check_frame_transformer.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+This script checks the FrameTransformer sensor by visualizing the frames that it creates.
+"""
+
+from __future__ import annotations
+
+"""Launch Isaac Sim Simulator first."""
+
+
+import argparse
+import math
+from scipy.spatial.transform import Rotation
+
+from omni.isaac.orbit.app import AppLauncher
+
+# add argparse arguments
+parser = argparse.ArgumentParser(
+    description="This script checks the FrameTransformer sensor by visualizing the frames that it creates."
+)
+parser.add_argument(
+    "--mode",
+    type=str,
+    choices=["feet", "all"],
+    default="all",
+    help="Dictates the the type of frames to use for FrameTransformer.",
+)
+parser.add_argument(
+    "--visualize",
+    action="store_true",
+    help=(
+        "Whether to enable FrameTransformer's debug_vis (True)                     or visualize each frame one at a"
+        " time and print to console (False)."
+    ),
+)
+parser.add_argument("--num_envs", type=int, default=2, help="Number of environments to spawn.")
+AppLauncher.add_app_launcher_args(parser)
+args_cli = parser.parse_args()
+
+# launch omniverse app
+app_launcher = AppLauncher(headless=args_cli.headless)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import omni.isaac.orbit.sim as sim_utils
+from omni.isaac.orbit.assets import AssetBaseCfg
+from omni.isaac.orbit.assets.config.anymal import ANYMAL_C_CFG
+from omni.isaac.orbit.markers import VisualizationMarkers
+from omni.isaac.orbit.markers.config import FRAME_MARKER_CFG
+from omni.isaac.orbit.scene import InteractiveScene, InteractiveSceneCfg
+from omni.isaac.orbit.sensors import FrameTransformerCfg, OffsetCfg
+from omni.isaac.orbit.sim import SimulationContext
+from omni.isaac.orbit.terrains import TerrainImporterCfg
+from omni.isaac.orbit.utils import configclass
+from omni.isaac.orbit.utils.timer import Timer
+
+
+def quat_from_euler_rpy(roll, pitch, yaw, degrees=False):
+    """Converts Euler XYZ to Quaternion (w, x, y, z)."""
+    quat = Rotation.from_euler("xyz", (roll, pitch, yaw), degrees=degrees).as_quat()
+    return tuple(quat[[3, 0, 1, 2]].tolist())
+
+
+def euler_rpy_apply(rpy, xyz, degrees=False):
+    """Applies rotation from Euler XYZ on position vector."""
+    rot = Rotation.from_euler("xyz", rpy, degrees=degrees)
+    return tuple(rot.apply(xyz).tolist())
+
+
+@configclass
+class MySceneCfg(InteractiveSceneCfg):
+    """Example scene configuration."""
+
+    # terrain - flat terrain plane
+    terrain = TerrainImporterCfg(
+        prim_path="/World/ground",
+        terrain_type="plane",
+    )
+
+    # articulation - robot
+    robot = ANYMAL_C_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+
+    if args_cli.mode == "feet":
+        # Example where only feet position frames are created
+        frame_transformer = FrameTransformerCfg(
+            prim_path="{ENV_REGEX_NS}/Robot/base",
+            target_frames=[
+                FrameTransformerCfg.FrameCfg(
+                    name="LF_FOOT",
+                    prim_path="{ENV_REGEX_NS}/Robot/LF_SHANK",
+                    offset=OffsetCfg(
+                        pos=euler_rpy_apply(rpy=(0, 0, -math.pi / 2), xyz=(0.08795, 0.01305, -0.33797)),
+                        rot=quat_from_euler_rpy(0, 0, -math.pi / 2),
+                    ),
+                ),
+                FrameTransformerCfg.FrameCfg(
+                    name="RF_FOOT",
+                    prim_path="{ENV_REGEX_NS}/Robot/RF_SHANK",
+                    offset=OffsetCfg(
+                        pos=euler_rpy_apply(rpy=(0, 0, math.pi / 2), xyz=(0.08795, -0.01305, -0.33797)),
+                        rot=quat_from_euler_rpy(0, 0, math.pi / 2),
+                    ),
+                ),
+                FrameTransformerCfg.FrameCfg(
+                    name="LH_FOOT",
+                    prim_path="{ENV_REGEX_NS}/Robot/LH_SHANK",
+                    offset=OffsetCfg(
+                        pos=euler_rpy_apply(rpy=(0, 0, -math.pi / 2), xyz=(-0.08795, 0.01305, -0.33797)),
+                        rot=quat_from_euler_rpy(0, 0, -math.pi / 2),
+                    ),
+                ),
+                FrameTransformerCfg.FrameCfg(
+                    name="RH_FOOT",
+                    prim_path="{ENV_REGEX_NS}/Robot/RH_SHANK",
+                    offset=OffsetCfg(
+                        pos=euler_rpy_apply(rpy=(0, 0, math.pi / 2), xyz=(-0.08795, -0.01305, -0.33797)),
+                        rot=quat_from_euler_rpy(0, 0, math.pi / 2),
+                    ),
+                ),
+            ],
+            debug_vis=args_cli.visualize,
+        )
+
+    elif args_cli.mode == "all":
+        # Example using .* to get full body + LF_FOOT to ensure these work together
+        frame_transformer = FrameTransformerCfg(
+            prim_path="{ENV_REGEX_NS}/Robot/base",
+            target_frames=[
+                FrameTransformerCfg.FrameCfg(prim_path="{ENV_REGEX_NS}/Robot/.*"),
+                FrameTransformerCfg.FrameCfg(
+                    prim_path="{ENV_REGEX_NS}/Robot/LF_SHANK",
+                    name="LF_FOOT",
+                    offset=OffsetCfg(
+                        pos=euler_rpy_apply(rpy=(0, 0, -math.pi / 2), xyz=(0.08795, 0.01305, -0.33797)),
+                        rot=quat_from_euler_rpy(0, 0, -math.pi / 2),
+                    ),
+                ),
+            ],
+            debug_vis=args_cli.visualize,
+        )
+
+    # extras - light
+    light = AssetBaseCfg(
+        prim_path="/World/light",
+        spawn=sim_utils.DistantLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)),
+        init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, 500.0)),
+    )
+
+
+def main():
+    """Main function."""
+
+    # Load kit helper
+    sim = SimulationContext(sim_utils.SimulationCfg(dt=0.005))
+    # Set main camera
+    sim.set_camera_view(eye=[5, 5, 5], target=[0.0, 0.0, 0.0])
+
+    # Spawn things into stage
+    with Timer("Setup scene"):
+        scene = InteractiveScene(MySceneCfg(num_envs=args_cli.num_envs, env_spacing=5.0, lazy_sensor_update=False))
+
+    # Play the simulator
+    with Timer("Time taken to play the simulator"):
+        sim.reset()
+
+    # Now we are ready!
+    print("[INFO]: Setup complete...")
+
+    # default joint targets
+    robot_actions = scene.articulations["robot"].data.default_joint_pos.clone()
+    # Define simulation stepping
+    sim_dt = sim.get_physics_dt()
+    sim_time = 0.0
+    count = 0
+
+    # We only want one visualization at a time. This visualizer will be used
+    # to step through each frame so the user can verify that the correct frame
+    # is being visualized as the frame names are printing to console
+    if not args_cli.visualize:
+        cfg = FRAME_MARKER_CFG.replace(prim_path="/Visuals/FrameTransformer")
+        cfg.markers["frame"].scale = (0.05, 0.05, 0.05)
+        transform_visualizer = VisualizationMarkers(cfg)
+    else:
+        transform_visualizer = None
+
+    frame_index = 0
+    # Simulate physics
+    while simulation_app.is_running():
+        # If simulation is stopped, then exit.
+        if sim.is_stopped():
+            break
+        # If simulation is paused, then skip.
+        if not sim.is_playing():
+            sim.step()
+            continue
+        # # reset
+        if count % 50 == 0:
+            # reset counters
+            sim_time = 0.0
+            count = 0
+            # reset root state
+            root_state = scene.articulations["robot"].data.default_root_state_w.clone()
+            root_state[:, :3] += scene.env_origins
+            joint_pos = scene.articulations["robot"].data.default_joint_pos
+            joint_vel = scene.articulations["robot"].data.default_joint_vel
+            # -- set root state
+            # -- robot
+            scene.articulations["robot"].write_root_state_to_sim(root_state)
+            scene.articulations["robot"].write_joint_state_to_sim(joint_pos, joint_vel)
+            root_state[:, 1] += 1.0
+            # reset buffers
+            scene.reset()
+            print(">>>>>>>> Reset!")
+        # perform this loop at policy control freq (50 Hz)
+        for _ in range(4):
+            # set joint targets
+            scene.articulations["robot"].set_joint_position_target(robot_actions)
+            # write data to sim
+            scene.write_data_to_sim()
+            # perform step
+            sim.step()
+            # read data from sim
+            scene.update(sim_dt)
+
+        # Change the frame that we are visualizing to ensure that frame names
+        # are correctly associated with the frames
+        if not args_cli.visualize:
+            if count % 50 == 0:
+                frame_names = scene["frame_transformer"].data.target_frame_names
+
+                frame_name = frame_names[frame_index]
+                print(f"Displaying {frame_index}: {frame_name}")
+                frame_index += 1
+                frame_index = frame_index % len(frame_names)
+            # visualize frame
+            pos = scene["frame_transformer"].data.target_pos_w[:, frame_index]
+            rot = scene["frame_transformer"].data.target_rot_w[:, frame_index]
+            transform_visualizer.visualize(pos, rot)
+
+        # update sim-time
+        sim_time += sim_dt * 4
+        count += 1
+
+
+if __name__ == "__main__":
+    # Run the main function
+    main()
+    # Close the simulator
+    simulation_app.close()

source/extensions/omni.isaac.orbit/test/utils/test_math.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import torch
+import unittest
+
+"""Launch Isaac Sim Simulator first.
+
+This is only needed because of warp dependency.
+"""
+
+from omni.isaac.orbit.app import AppLauncher
+
+# launch omniverse app in headless mode
+app_launcher = AppLauncher(headless=True)
+
+import omni.isaac.orbit.utils.math as math_utils
+
+
+class TestMathUtilities(unittest.TestCase):
+    """Test fixture for checking math utilities in Orbit."""
+
+    def test_is_identity_pose(self):
+        """Test is_identity_pose method."""
+        identity_pos_one_row = torch.zeros(3)
+        identity_rot_one_row = torch.tensor((1.0, 0.0, 0.0, 0.0))
+
+        self.assertTrue(math_utils.is_identity_pose(identity_pos_one_row, identity_rot_one_row))
+
+        identity_pos_one_row[0] = 1.0
+        identity_rot_one_row[1] = 1.0
+
+        self.assertFalse(math_utils.is_identity_pose(identity_pos_one_row, identity_rot_one_row))
+
+        identity_pos_multi_row = torch.zeros(3, 3)
+        identity_rot_multi_row = torch.zeros(3, 4)
+        identity_rot_multi_row[:, 0] = 1.0
+
+        self.assertTrue(math_utils.is_identity_pose(identity_pos_multi_row, identity_rot_multi_row))
+
+        identity_pos_multi_row[0, 0] = 1.0
+        identity_rot_multi_row[0, 1] = 1.0
+
+        self.assertFalse(math_utils.is_identity_pose(identity_pos_multi_row, identity_rot_multi_row))
+
+
+if __name__ == "__main__":
+    unittest.main()