Adds action term for inverse kinematics (#182)

# Description

This MR simplifies the `DifferentialInverseKinematics` to expect that
the user gives the right frame for the end-effector pose as inputs.
Additionally, it adds the action term corresponding to the controller.

## Type of change

- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./orbit.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] 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

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

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

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

 Changelog
 ---------
 
+0.9.9 (2023-10-12)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added the property :attr:`omni.isaac.orbit.assets.Articulation.is_fixed_base` to the articulation class to
+  check if the base of the articulation is fixed or floating.
+* Added the task-space action term corresponding to the differential inverse-kinematics controller.
+
+Fixed
+^^^^^
+
+* Simplified the :class:`omni.isaac.orbit.controllers.DifferentialIKController` to assume that user provides the
+  correct end-effector poses and Jacobians. Earlier it was doing internal frame transformations which made the
+  code more complicated and error-prone.
+
+
 0.9.8 (2023-09-30)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/articulation/articulation.py

@@ -17,6 +17,7 @@ import omni.physics.tensors.impl.api as physx
 from omni.isaac.core.articulations import ArticulationView
 from omni.isaac.core.prims import RigidPrimView
 from omni.isaac.core.utils.types import ArticulationActions
+from pxr import Usd, UsdPhysics
 
 import omni.isaac.orbit.utils.math as math_utils
 import omni.isaac.orbit.utils.string as string_utils
@@ -56,6 +57,11 @@ class Articulation(RigidObject):
         """Data related to articulation."""
         return self._data
 
+    @property
+    def is_fixed_base(self) -> bool:
+        """Whether the articulation is a fixed-base or floating-base system."""
+        return self._is_fixed_base
+
     @property
     def num_joints(self) -> int:
         """Number of joints in articulation."""
@@ -397,6 +403,18 @@ class Articulation(RigidObject):
         # check that initialization was successful
         if len(self.body_names) != self.num_bodies:
             raise RuntimeError("Failed to initialize all bodies properly in the articulation.")
+        # -- fixed base based on root joint
+        self._is_fixed_base = False
+        for prim in Usd.PrimRange(self._root_view.prims[0]):
+            joint_prim = UsdPhysics.FixedJoint(prim)
+            # we check all joints under the root prim and classify the asset as fixed base if there exists
+            # a fixed joint that has only one target (i.e. the root link).
+            if joint_prim and joint_prim.GetJointEnabledAttr().Get():
+                body_0_exist = joint_prim.GetBody0Rel.GetTargets() != []
+                body_1_exist = joint_prim.GetBody1Rel.GetTargets() != []
+                if not (body_0_exist and body_1_exist):
+                    self._is_fixed_base = True
+                    break
         # log information about the articulation
         carb.log_info(f"Articulation initialized at: {self.cfg.prim_path}")
         carb.log_info(f"Number of bodies: {self.num_bodies}")

source/extensions/omni.isaac.orbit/omni/isaac/orbit/assets/config/universal_robots.py

@@ -28,6 +28,10 @@ from ..articulation import ArticulationCfg
 UR10_CFG = ArticulationCfg(
     spawn=sim_utils.UsdFileCfg(
         usd_path=f"{ISAAC_ORBIT_NUCLEUS_DIR}/Robots/UniversalRobots/UR10/ur10_instanceable.usd",
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=False,
+            max_depenetration_velocity=5.0,
+        ),
         activate_contact_sensors=False,
     ),
     init_state=ArticulationCfg.InitialStateCfg(

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

@@ -5,15 +5,7 @@
 
 from __future__ import annotations
 
-# enable motion generation extension
-from omni.isaac.core.utils.extensions import enable_extension
+from .differential_ik import DifferentialIKController
+from .differential_ik_cfg import DifferentialIKControllerCfg
 
-# TODO: Maybe keep this extension enabled by default? -- Fix the app experience.
-enable_extension("omni.isaac.motion_generation")
-
-
-# get module libraries
-from .differential_inverse_kinematics import DifferentialInverseKinematics
-from .rmp_flow import RmpFlowController
-
-__all__ = ["DifferentialInverseKinematics", "RmpFlowController"]
+__all__ = ["DifferentialIKController", "DifferentialIKControllerCfg"]

source/extensions/omni.isaac.orbit/omni/isaac/orbit/controllers/differential_ik_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 typing_extensions import Literal
+
+from omni.isaac.orbit.utils import configclass
+
+from .differential_ik import DifferentialIKController
+
+
+@configclass
+class DifferentialIKControllerCfg:
+    """Configuration for differential inverse kinematics controller."""
+
+    class_type: type = DifferentialIKController
+    """The associated controller class."""
+
+    command_type: Literal["position", "pose"] = MISSING
+    """Type of task-space command to control the articulation's body.
+
+    If "position", then the controller only controls the position of the articulation's body.
+    Otherwise, the controller controls the pose of the articulation's body.
+    """
+
+    use_relative_mode: bool = False
+    """Whether to use relative mode for the controller. Defaults to False.
+
+    If True, then the controller treats the input command as a delta change in the position/pose.
+    Otherwise, the controller treats the input command as the absolute position/pose.
+    """
+
+    ik_method: Literal["pinv", "svd", "trans", "dls"] = MISSING
+    """Method for computing inverse of Jacobian."""
+
+    ik_params: dict[str, float] | None = None
+    """Parameters for the inverse-kinematics method. Defaults to None, in which case the default
+    parameters for the method are used.
+
+    - Moore-Penrose pseudo-inverse ("pinv"):
+        - "k_val": Scaling of computed delta-joint positions (default: 1.0).
+    - Adaptive Singular Value Decomposition ("svd"):
+        - "k_val": Scaling of computed delta-joint positions (default: 1.0).
+        - "min_singular_value": Single values less than this are suppressed to zero (default: 1e-5).
+    - Jacobian transpose ("trans"):
+        - "k_val": Scaling of computed delta-joint positions (default: 1.0).
+    - Damped Moore-Penrose pseudo-inverse ("dls"):
+        - "lambda_val": Damping coefficient (default: 0.1).
+    """
+
+    def __post_init__(self):
+        # check valid input
+        if self.command_type not in ["position", "pose"]:
+            raise ValueError(f"Unsupported inverse-kinematics command: {self.command_type}.")
+        if self.ik_method not in ["pinv", "svd", "trans", "dls"]:
+            raise ValueError(f"Unsupported inverse-kinematics method: {self.ik_method}.")
+        # default parameters for different inverse kinematics approaches.
+        default_ik_params = {
+            "pinv": {"k_val": 1.0},
+            "svd": {"k_val": 1.0, "min_singular_value": 1e-5},
+            "trans": {"k_val": 1.0},
+            "dls": {"lambda_val": 0.1},
+        }
+        # update parameters for IK-method if not provided
+        ik_params = default_ik_params[self.ik_method].copy()
+        if self.ik_params is not None:
+            ik_params.update(self.ik_params)
+        self.ik_params = ik_params

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/actions_cfg.py

@@ -9,8 +9,9 @@ from dataclasses import MISSING
 
 from omni.isaac.orbit.managers.action_manager import ActionTerm, ActionTermCfg
 from omni.isaac.orbit.utils import configclass
+from omni.isaac.orbit.controllers import DifferentialIKControllerCfg
 
-from . import binary_joint_actions, joint_actions, non_holonomic_actions
+from . import binary_joint_actions, joint_actions, non_holonomic_actions, task_space_actions
 
 ##
 # Joint actions.
@@ -142,3 +143,27 @@ class NonHolonomicActionCfg(ActionTermCfg):
     """Scale factor for the action. Defaults to (1.0, 1.0)."""
     offset: tuple[float, float] = (0.0, 0.0)
     """Offset factor for the action. Defaults to (0.0, 0.0)."""
+
+
+##
+# Task-space Actions.
+##
+
+
+@configclass
+class DifferentialInverseKinematicsActionCfg(ActionTermCfg):
+    """Configuration for inverse differential kinematics action term.
+
+    See :class:`DifferentialInverseKinematicsAction` for more details.
+    """
+
+    class_type: type[ActionTerm] = task_space_actions.DifferentialInverseKinematicsAction
+
+    joint_names: list[str] = MISSING
+    """List of joint names or regex expressions that the action will be mapped to."""
+    body_name: str = MISSING
+    """Name of the body or frame for which IK is performed."""
+    scale: float | tuple[float, ...] = 1.0
+    """Scale factor for the action. Defaults to 1.0."""
+    controller: DifferentialIKControllerCfg = MISSING
+    """The configuration for the differential IK controller."""

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/joint_actions.py

@@ -74,7 +74,7 @@ class JointAction(ActionTerm):
         if isinstance(cfg.scale, float):
             self._scale = cfg.scale
         elif isinstance(cfg.scale, dict):
-            self._scale = torch.ones(1, self.action_dim, device=self.device)
+            self._scale = torch.ones(1.0, self.action_dim, device=self.device)
             # resolve the dictionary config
             index_list, _, value_list = string_utils.resolve_matching_names_values(self.cfg.scale, self._joint_names)
             self._scale[:, index_list] = torch.tensor(value_list, device=self.device)

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/non_holonomic_actions.py

@@ -95,7 +95,7 @@ class NonHolonomicAction(ActionTerm):
         )
 
         # create tensors for raw and processed actions
-        self._raw_actions = torch.zeros(self.num_envs, 2, device=self.device)
+        self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device)
         self._processed_actions = torch.zeros_like(self.raw_actions)
         self._joint_vel_command = torch.zeros(self.num_envs, 3, device=self.device)
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/envs/mdp/actions/task_space_actions.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 typing import TYPE_CHECKING
+
+import carb
+
+from omni.isaac.orbit.assets.articulation import Articulation
+from omni.isaac.orbit.controllers.differential_ik import DifferentialIKController
+from omni.isaac.orbit.managers.action_manager import ActionTerm
+from omni.isaac.orbit.utils.math import subtract_frame_transforms
+
+if TYPE_CHECKING:
+    from omni.isaac.orbit.envs import BaseEnv
+
+    from . import actions_cfg
+
+
+class DifferentialInverseKinematicsAction(ActionTerm):
+    r"""Inverse Kinematics action term.
+
+    This action term performs pre-processing of the raw actions using scaling transformation.
+
+    .. math::
+        \text{action} = \text{scaling} \times \text{input action}
+        \text{joint position} = J^{-} \times \text{action}
+
+    where :math:`\text{scaling}` is the scaling applied to the input action, and :math:`\text{input action}`
+    is the input action from the user, :math:`J` is the Jacobian over the articulation's actuated joints,
+    and \text{joint position} is the desired joint position command for the articulation's joints.
+    """
+
+    cfg: actions_cfg.DifferentialInverseKinematicsActionCfg
+    """The configuration of the action term."""
+    _asset: Articulation
+    """The articulation asset on which the action term is applied."""
+    _scale: torch.Tensor
+    """The scaling factor applied to the input action. Shape is (1, action_dim)."""
+
+    def __init__(self, cfg: actions_cfg.DifferentialInverseKinematicsActionCfg, env: BaseEnv):
+        # initialize the action term
+        super().__init__(cfg, env)
+
+        # resolve the joints over which the action term is applied
+        self._joint_ids, self._joint_names = self._asset.find_joints(self.cfg.joint_names)
+        self._num_joints = len(self._joint_ids)
+        # parse the body index
+        body_ids, body_names = self._asset.find_bodies(self.cfg.body_name)
+        if len(body_ids) != 1:
+            raise ValueError(
+                f"Expected one match for the body name: {self.cfg.body_name}. Found {len(body_ids)}: {body_names}."
+            )
+        # save only the first body index
+        self._body_idx = body_ids[0]
+        self._body_name = body_names[0]
+        # check if articulation is fixed-base
+        # if fixed-base then the jacobian for the base is not computed
+        # this means that number of bodies is one less than the articulation's number of bodies
+        if self._asset.is_fixed_base:
+            self._jacobi_body_idx = self._body_idx - 1
+        else:
+            self._jacobi_body_idx = self._body_idx
+
+        # log info for debugging
+        carb.log_info(
+            f"Resolved joint names for the action term {self.__class__.__name__}: {self._joint_names} [{self._joint_ids}]"
+        )
+        carb.log_info(
+            f"Resolved body name for the action term {self.__class__.__name__}: {self._body_name} [{self._body_idx}]"
+        )
+        # Avoid indexing across all joints for efficiency
+        if self._num_joints == self._asset.num_joints:
+            self._joint_ids = ...
+
+        # create the differential IK controller
+        self._controller = DifferentialIKController(cfg=self.cfg.controller, num_envs=self.num_envs, device=self.device)
+
+        # create tensors for raw and processed actions
+        self._raw_actions = torch.zeros(self.num_envs, self.action_dim, device=self.device)
+        self._processed_actions = torch.zeros_like(self.raw_actions)
+
+        # save the scale as tensors
+        self._scale = torch.zeros((self.num_envs, self.action_dim), device=self.device)
+        self._scale[:] = torch.tensor(self.cfg.scale, device=self.device)
+
+    """
+    Properties.
+    """
+
+    @property
+    def action_dim(self) -> int:
+        return self._controller.action_dim
+
+    @property
+    def raw_actions(self) -> torch.Tensor:
+        return self._raw_actions
+
+    @property
+    def processed_actions(self) -> torch.Tensor:
+        return self._processed_actions
+
+    """
+    Operations.
+    """
+
+    def process_actions(self, actions: torch.Tensor):
+        # store the raw actions
+        self._raw_actions[:] = actions
+        self._processed_actions[:] = self.raw_actions * self._scale
+        # obtain quantities from simulation
+        ee_pos_curr, ee_quat_curr = self._compute_body_pose()
+        # set command into controller
+        self._ik_controller.set_command(self._processed_actions, ee_pos_curr, ee_quat_curr)
+
+    def apply_actions(self):
+        # obtain quantities from simulation
+        jacobian = self._asset.root_physx_view.get_jacobians()[:, self._jacobi_body_idx, :, self._joint_ids]
+        ee_pos_curr, ee_quat_curr = self._compute_body_pose()
+        joint_pos = self._asset.data.joint_pos[:, self._joint_ids]
+        # compute the delta in joint-space
+        joint_pos_des = self._ik_controller.compute(ee_pos_curr, ee_quat_curr, jacobian, joint_pos)
+        # set the joint position command
+        self._asset.set_joint_position_target(joint_pos_des, self._joint_ids)
+
+    """
+    Helper functions.
+    """
+
+    def _compute_body_pose(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Computes the pose of the body in the root frame.
+
+        Returns:
+            A tuple of the body's position and orientation in the root frame.
+        """
+        # obtain quantities from simulation
+        ee_pose_w = self._asset.data.body_state_w[:, self._body_idx, :7]
+        root_pose_w = self._asset.data.root_state_w[:, :7]
+        # compute the pose of the body in the root frame
+        return subtract_frame_transforms(root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])

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

@@ -588,7 +588,7 @@ def compute_pose_error(
     q01: torch.Tensor,
     t02: torch.Tensor,
     q02: torch.Tensor,
-    rot_error_type: Literal["quat", "axis_angle"],
+    rot_error_type: Literal["quat", "axis_angle"] = "axis_angle",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Compute the position and orientation error between source and target frames.
 
@@ -598,6 +598,7 @@ def compute_pose_error(
         t02: Position of target frame.
         q02: Quaternion orientation of target frame.
         rot_error_type: The rotation error type to return: "quat", "axis_angle".
+            Defaults to "axis_angle".
 
     Returns:
         A tuple containing position and orientation error.

source/extensions/omni.isaac.orbit/test/assets/test_articulation.py

@@ -67,6 +67,8 @@ class TestArticulation(unittest.TestCase):
         self.sim.reset()
         # Check if robot is initialized
         self.assertTrue(robot._is_initialized)
+        # Check that floating base
+        self.assertFalse(robot.is_fixed_base)
         # Check buffers that exists and have correct shapes
         self.assertTrue(robot.data.root_pos_w.shape == (1, 3))
         self.assertTrue(robot.data.root_quat_w.shape == (1, 4))
@@ -94,6 +96,8 @@ class TestArticulation(unittest.TestCase):
         self.sim.reset()
         # Check if robot is initialized
         self.assertTrue(robot._is_initialized)
+        # Check that fixed base
+        self.assertTrue(robot.is_fixed_base)
         # Check buffers that exists and have correct shapes
         self.assertTrue(robot.data.root_pos_w.shape == (1, 3))
         self.assertTrue(robot.data.root_quat_w.shape == (1, 4))

source/extensions/omni.isaac.orbit/test/controllers/test_differential_ik.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
+
+"""Launch Isaac Sim Simulator first."""
+
+from omni.isaac.orbit.app import AppLauncher
+
+# launch omniverse app
+simulation_app = AppLauncher(headless=True).app
+
+"""Rest everything follows."""
+
+import gc
+import sys
+import torch
+import traceback
+import unittest
+
+import carb
+import omni.isaac.core.utils.prims as prim_utils
+import omni.isaac.core.utils.stage as stage_utils
+from omni.isaac.cloner import GridCloner
+
+import omni.isaac.orbit.sim as sim_utils
+from omni.isaac.orbit.assets import Articulation
+from omni.isaac.orbit.assets.config import FRANKA_PANDA_ARM_WITH_PANDA_HAND_CFG, UR10_CFG
+from omni.isaac.orbit.controllers import DifferentialIKController, DifferentialIKControllerCfg
+from omni.isaac.orbit.utils.math import compute_pose_error, sample_uniform, subtract_frame_transforms
+
+
+class TestDifferentialIKController(unittest.TestCase):
+    """Test fixture for checking that differential IK controller tracks commands properly."""
+
+    def setUp(self) -> None:
+        """Create a blank new stage for each test."""
+        # Wait for spawning
+        stage_utils.create_new_stage()
+        # Constants
+        self.num_envs = 128
+        # Load kit helper
+        sim_cfg = sim_utils.SimulationCfg(dt=0.01, shutdown_app_on_stop=False)
+        self.sim = sim_utils.SimulationContext(sim_cfg)
+
+        # Create a ground plane
+        cfg = sim_utils.GroundPlaneCfg()
+        cfg.func("/World/GroundPlane", cfg)
+
+        # Create interface to clone the scene
+        cloner = GridCloner(spacing=2.0)
+        cloner.define_base_env("/World/envs")
+        self.env_prim_paths = cloner.generate_paths(f"/World/envs/env", self.num_envs)
+        # create source prim
+        prim_utils.define_prim(self.env_prim_paths[0], "Xform")
+        # clone the env xform
+        self.env_origins = cloner.clone(
+            source_prim_path=self.env_prim_paths[0],
+            prim_paths=self.env_prim_paths,
+            replicate_physics=True,
+        )
+
+        # Define goals for the arm
+        ee_goals_set = [
+            [0.5, 0.5, 0.7, 0.707, 0, 0.707, 0],
+            [0.5, -0.4, 0.6, 0.707, 0.707, 0.0, 0.0],
+            [0.5, 0, 0.5, 0.0, 1.0, 0.0, 0.0],
+        ]
+        self.ee_pose_b_des_set = torch.tensor(ee_goals_set, device=self.sim.device)
+
+    def tearDown(self) -> None:
+        """Stops simulator after each test."""
+        # stop simulation
+        self.sim.stop()
+        self.sim.clear()
+        self.sim.clear_all_callbacks()
+        self.sim.clear_instance()
+
+    """
+    Test fixtures.
+    """
+
+    def test_franka_ik_pose_abs(self):
+        """Test IK controller for Franka arm with Franka hand."""
+        # Create robot instance
+        robot_cfg = FRANKA_PANDA_ARM_WITH_PANDA_HAND_CFG.replace(prim_path="/World/envs/env_.*/Robot")
+        robot_cfg.spawn.rigid_props.disable_gravity = True
+        robot = Articulation(cfg=robot_cfg)
+
+        # Create IK controller
+        diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls")
+        diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        # Run the controller and check that it converges to the goal
+        self._run_ik_controller(robot, diff_ik_controller, "panda_hand", ["panda_joint.*"])
+
+    def test_ur10_ik_pose_abs(self):
+        """Test IK controller for UR10 arm."""
+        # Create robot instance
+        robot_cfg = UR10_CFG.replace(prim_path="/World/envs/env_.*/Robot")
+        robot_cfg.spawn.rigid_props.disable_gravity = True
+        robot = Articulation(cfg=robot_cfg)
+
+        # Create IK controller
+        diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls")
+        diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        # Run the controller and check that it converges to the goal
+        self._run_ik_controller(robot, diff_ik_controller, "ee_link", [".*"])
+
+    """
+    Helper functions.
+    """
+
+    def _run_ik_controller(
+        self,
+        robot: Articulation,
+        diff_ik_controller: DifferentialIKController,
+        ee_frame_name: str,
+        arm_joint_names: list[str],
+    ):
+        # Define simulation stepping
+        sim_dt = self.sim.get_physics_dt()
+        # Play the simulator
+        self.sim.reset()
+
+        # Obtain the frame index of the end-effector
+        ee_frame_idx = robot.find_bodies(ee_frame_name)[0][0]
+        ee_jacobi_idx = ee_frame_idx - 1
+        # Obtain joint indices
+        arm_joint_ids = robot.find_joints(arm_joint_names)[0]
+        # Update existing buffers
+        # Note: We need to update buffers before the first step for the controller.
+        robot.update(dt=sim_dt)
+
+        # Track the given command
+        current_goal_idx = 0
+        # Current goal for the arm
+        ee_pose_b_des = torch.zeros(self.num_envs, diff_ik_controller.action_dim, device=self.sim.device)
+        ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx]
+        # Compute current pose of the end-effector
+        ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
+        root_pose_w = robot.data.root_state_w[:, 0:7]
+        ee_pos_b, ee_quat_b = subtract_frame_transforms(
+            root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
+        )
+
+        # Now we are ready!
+        for count in range(1500):
+            # reset every 150 steps
+            if count % 250 == 0:
+                # check that we converged to the goal
+                if count > 0:
+                    pos_error, rot_error = compute_pose_error(
+                        ee_pos_b, ee_quat_b, ee_pose_b_des[:, 0:3], ee_pose_b_des[:, 3:7]
+                    )
+                    pos_error_norm = torch.norm(pos_error, dim=-1)
+                    rot_error_norm = torch.norm(rot_error, dim=-1)
+                    # desired error (zer)
+                    des_error = torch.zeros_like(pos_error_norm)
+                    # check convergence
+                    torch.testing.assert_close(pos_error_norm, des_error, rtol=0.0, atol=1e-3)
+                    torch.testing.assert_close(rot_error_norm, des_error, rtol=0.0, atol=1e-3)
+                # reset joint state
+                joint_pos = robot.data.default_joint_pos.clone()
+                joint_vel = robot.data.default_joint_vel.clone()
+                # joint_pos *= sample_uniform(0.9, 1.1, joint_pos.shape, joint_pos.device)
+                robot.write_joint_state_to_sim(joint_pos, joint_vel)
+                robot.set_joint_position_target(joint_pos)
+                robot.write_data_to_sim()
+                robot.reset()
+                # reset actions
+                ee_pose_b_des[:] = self.ee_pose_b_des_set[current_goal_idx]
+                joint_pos_des = joint_pos[:, arm_joint_ids].clone()
+                # update goal for next iteration
+                current_goal_idx = (current_goal_idx + 1) % len(self.ee_pose_b_des_set)
+                # set the controller commands
+                diff_ik_controller.reset()
+                diff_ik_controller.set_command(ee_pose_b_des)
+            else:
+                # at reset, the jacobians are not updated to the latest state
+                # so we MUST skip the first step
+                # obtain quantities from simulation
+                jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids]
+                ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 0:7]
+                root_pose_w = robot.data.root_state_w[:, 0:7]
+                joint_pos = robot.data.joint_pos[:, arm_joint_ids]
+                # compute frame in root frame
+                ee_pos_b, ee_quat_b = subtract_frame_transforms(
+                    root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
+                )
+                # compute the joint commands
+                joint_pos_des = diff_ik_controller.compute(ee_pos_b, ee_quat_b, jacobian, joint_pos)
+
+            # apply actions
+            robot.set_joint_position_target(joint_pos_des, arm_joint_ids)
+            robot.write_data_to_sim()
+            # perform step
+            self.sim.step(render=False)
+            # update buffers
+            robot.update(sim_dt)
+
+        return None
+
+
+if __name__ == "__main__":
+    try:
+        unittest.main()
+    except Exception as err:
+        carb.log_error(err)
+        carb.log_error(traceback.format_exc())
+        raise
+    finally:
+        # close sim app
+        simulation_app.close()