Adds `OperationSpaceController` to docs and tests and implement corresponding...

Adds `OperationSpaceController` to docs and tests and implement corresponding action/action_cfg classes (#913)

# Description

This PR adds the `OperationalSpaceController` to the docs and provides
some tests for its parametric features. Moreover, it implements the
corresponding `OperationalSpaceControllerAction` and
`OperationalSpaceControllerActionCfg` classes so they can be used with
manager-based environments.

Fixes #873 

## Type of change

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [x] 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
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there
``

---------
Signed-off-by: Özhan Özen <41010165+ozhanozen@users.noreply.github.com>
Co-authored-by: James Tigue <jtigue@theaiinstitute.com>
Co-authored-by: jtigue-bdai <166445701+jtigue-bdai@users.noreply.github.com>
Co-authored-by: Kelly Guo <kellyguo123@hotmail.com>

docs/source/api/lab/omni.isaac.lab.controllers.rst

@@ -9,6 +9,8 @@
 
     DifferentialIKController
     DifferentialIKControllerCfg
+    OperationalSpaceController
+    OperationalSpaceControllerCfg
 
 Differential Inverse Kinematics
 -------------------------------
@@ -23,3 +25,17 @@ Differential Inverse Kinematics
     :inherited-members:
     :show-inheritance:
     :exclude-members: __init__, class_type
+
+Operational Space controllers
+-----------------------------
+
+.. autoclass:: OperationalSpaceController
+    :members:
+    :inherited-members:
+    :show-inheritance:
+
+.. autoclass:: OperationalSpaceControllerCfg
+    :members:
+    :inherited-members:
+    :show-inheritance:
+    :exclude-members: __init__, class_type

docs/source/overview/environments.rst

@@ -91,7 +91,7 @@ Manipulation
 Environments based on fixed-arm manipulation tasks.
 
 For many of these tasks, we include configurations with different arm action spaces. For example,
-for the reach environment:
+for the lift-cube environment:
 
 * |lift-cube-link|: Franka arm with joint position control
 * |lift-cube-ik-abs-link|: Franka arm with absolute IK control
@@ -421,6 +421,10 @@ Comprehensive List of Environments
       -
       - Manager Based
       -
+    * - Isaac-Reach-Franka-OSC-v0
+      - Isaac-Reach-Franka-OSC-Play-v0
+      - Manager Based
+      - **rsl_rl** (PPO)
     * - Isaac-Reach-Franka-v0
       - Isaac-Reach-Franka-Play-v0
       - Manager Based

docs/source/tutorials/05_controllers/run_osc.rst

+Using an operational space controller
+=====================================
+
+.. currentmodule:: omni.isaac.lab
+
+Sometimes, controlling the end-effector pose of the robot using a differential IK controller is not sufficient.
+For example, we might want to enforce a very specific pose tracking error dynamics in the task space, actuate the robot
+with joint effort/torque commands, or apply a contact force at a specific direction while controlling the motion of
+the other directions (e.g., washing the surface of the table with a cloth). In such tasks, we can use an
+operational space controller (OSC).
+
+.. rubric:: References for the operational space control:
+
+1. O Khatib. A unified approach for motion and force control of robot manipulators:
+   The operational space formulation. IEEE Journal of Robotics and Automation, 3(1):43–53, 1987. URL http://dx.doi.org/10.1109/JRA.1987.1087068.
+
+2. Robot Dynamics Lecture Notes by Marco Hutter (ETH Zurich). URL https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2017/RD_HS2017script.pdf
+
+In this tutorial, we will learn how to use an OSC to control the robot.
+We will use the :class:`controllers.OperationalSpaceController` class to apply a constant force perpendicular to a
+tilted wall surface while tracking a desired end-effector pose in all the other directions.
+
+The Code
+~~~~~~~~
+
+The tutorial corresponds to the ``run_osc.py`` script in the
+``source/standalone/tutorials/05_controllers`` directory.
+
+
+.. dropdown:: Code for run_osc.py
+   :icon: code
+
+   .. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+      :language: python
+      :linenos:
+
+
+Creating an Operational Space Controller
+----------------------------------------
+
+The :class:`~controllers.OperationalSpaceController` class computes the joint
+efforts/torques for a robot to do simultaneous motion and force control in task space.
+
+The reference frame of this task space could be an arbitrary coordinate frame in Euclidean space. By default,
+it is the robot's base frame. However, in certain cases, it could be easier to define target coordinates w.r.t. a
+different frame. In such cases, the pose of this task reference frame, w.r.t. to the robot's base frame, should be
+provided in the ``set_command`` method's ``current_task_frame_pose_b`` argument. For example, in this tutorial, it
+makes sense to define the target commands w.r.t. a frame that is parallel to the wall surface, as the force control
+direction would be then only nonzero in the z-axis of this frame. The target pose, which is set to have the same
+orientation as the wall surface, is such a candidate and is used as the task frame in this tutorial. Therefore, all
+the arguments to the :class:`~controllers.OperationalSpaceControllerCfg` should be set with this task reference frame
+in mind.
+
+For the motion control, the task space targets could be given as absolute (i.e., defined w.r.t. the robot base,
+``target_types: "pose_abs"``) or relative the the end-effector's current pose (i.e., ``target_types: "pose_rel"``).
+For the force control, the task space targets could be given as absolute (i.e., defined w.r.t. the robot base,
+``target_types: "force_abs"``). If it is desired to apply pose and force control simultaneously, the ``target_types``
+should be a list such as ``["pose_abs", "wrench_abs"]`` or ``["pose_rel", "wrench_abs"]``.
+
+The axes that the motion and force control will be applied can be specified using the ``motion_control_axes_task`` and
+``force_control_axes_task`` arguments, respectively. These lists should consist of 0/1 for all six axes (position and
+rotation) and be complementary to each other (e.g., for the x-axis, if the ``motion_control_axes_task`` is ``0``, the
+``force_control_axes_task`` should be ``1``).
+
+For the motion control axes, desired stiffness, and damping ratio values can be specified using the
+``motion_control_stiffness`` and ``motion_damping_ratio_task`` arguments, which can be a scalar (same value for all
+axes) or a list of six scalars, one value corresponding to each axis. If desired, the stiffness and damping ratio
+values could be a command parameter (e.g., to learn the values using RL or change them on the go). For this,
+``impedance_mode`` should be either ``"variable_kp"`` to include the stiffness values within the command or
+``"variable"`` to include both the stiffness and damping ratio values. In these cases, ``motion_stiffness_limits_task``
+and ``motion_damping_limits_task`` should be set as well, which puts bounds on the stiffness and damping ratio values.
+
+For contact force control, it is possible to apply an open-loop force control by not setting the
+``contact_wrench_stiffness_task``, or apply a closed-loop force control (with the feed-forward term) by setting
+the desired stiffness values using the ``contact_wrench_stiffness_task`` argument, which can be a scalar or a list
+of six scalars. Please note that, currently, only the linear part of the contact wrench (hence the first three
+elements of the ``contact_wrench_stiffness_task``) is considered in the closed-loop control, as the rotational part
+cannot be measured with the contact sensors.
+
+For the motion control, ``inertial_dynamics_decoupling`` should be set to ``True`` to use the robot's inertia matrix
+to decouple the desired accelerations in the task space. This is important for the motion control to be accurate,
+especially for rapid movements. This inertial decoupling accounts for the coupling between all the six motion axes.
+If desired, the inertial coupling between the translational and rotational axes could be ignored by setting the
+``partial_inertial_dynamics_decoupling`` to ``True``.
+
+If it is desired to include the gravity compensation in the operational space command, the ``gravity_compensation``
+should be set to ``True``.
+
+The included OSC implementation performs the computation in a batched format and uses PyTorch operations.
+
+In this tutorial, we will use ``"pose_abs"`` for controlling the motion in all axes except the z-axis and
+``"wrench_abs"`` for controlling the force in the z-axis. Moreover, we will include the full inertia decoupling in
+the motion control and not include the gravity compensation, as the gravity is disabled from the robot configuration.
+Finally, we set the impedance mode to ``"variable_kp"`` to dynamically change the stiffness values
+(``motion_damping_ratio_task`` is set to ``1``: the kd values adapt according to kp values to maintain a critically
+damped response).
+
+.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+   :language: python
+   :start-at: # Create the OSC
+   :end-at: osc = OperationalSpaceController(osc_cfg, num_envs=scene.num_envs, device=sim.device)
+
+Updating the states of the robot
+--------------------------------------------
+
+The OSC implementation is a computation-only class. Thus, it expects the user to provide the necessary information
+about the robot. This includes the robot's Jacobian matrix, mass/inertia matrix, end-effector pose, velocity, and
+contact force, all in the root frame. Moreover, the user should provide gravity compensation vector, if desired.
+
+.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+   :language: python
+   :start-at: # Update robot states
+   :end-before: return jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b
+
+
+Computing robot command
+-----------------------
+
+The OSC separates the operation of setting the desired command and computing the desired joint positions.
+To set the desired command, the user should provide command vector, which  includes the target commands
+(i.e., in the order they appear in the ``target_types`` argument of the OSC configuration),
+and the desired stiffness and damping ratio values if the impedance_mode is set to ``"variable_kp"`` or ``"variable"``.
+They should be all in the same coordinate frame as the task frame (e.g., indicated with ``_task`` subscript) and
+concatanated together.
+
+In this tutorial, the desired wrench is already defined w.r.t. the task frame, and the desired pose is transformed
+to the task frame as the following:
+
+.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+   :language: python
+   :start-at: # Convert the target commands to the task frame
+   :end-at: return command, task_frame_pose_b
+
+The OSC command is set with the command vector in the task frame, the end-effector pose in the base frame, and the
+task (reference) frame pose in the base frame as the following. This information is needed, as the internal
+computations are done in the base frame.
+
+.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+   :language: python
+   :start-at: # set the osc command
+   :end-at: osc.set_command(command=command, current_ee_pose_b=ee_pose_b, current_task_frame_pose_b=task_frame_pose_b)
+
+The joint effort/torque values are computed using the provided robot states and the desired command as the following:
+
+.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+   :language: python
+   :start-at: # compute the joint commands
+   :end-at: )
+
+
+The computed joint effort/torque targets can then be applied on the robot.
+
+.. literalinclude:: ../../../../source/standalone/tutorials/05_controllers/run_osc.py
+   :language: python
+   :start-at: # apply actions
+   :end-at: robot.write_data_to_sim()
+
+
+The Code Execution
+~~~~~~~~~~~~~~~~~~
+
+You can now run the script and see the result:
+
+.. code-block:: bash
+
+   ./isaaclab.sh -p source/standalone/tutorials/05_controllers/run_osc.py --num_envs 128
+
+The script will start a simulation with 128 robots. The robots will be controlled using the OSC.
+The current and desired end-effector poses should be displayed using frame markers in addition to the red tilted wall.
+You should see that the robot reaches the desired pose while applying a constant force perpendicular to the wall
+surface.
+
+.. figure:: ../../_static/tutorials/tutorial_operational_space_controller.jpg
+    :align: center
+    :figwidth: 100%
+    :alt: result of run_osc.py
+
+To stop the simulation, you can either close the window or press ``Ctrl+C`` in the terminal.

docs/source/tutorials/index.rst

@@ -101,3 +101,4 @@ tutorials show you how to use motion generators to control the robots at the tas
     :titlesonly:
 
     05_controllers/run_diff_ik
+    05_controllers/run_osc

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.29.1"
+version = "0.29.2"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.29.2 (2024-12-16)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed errors within the calculations of :class:`omni.isaac.lab.controllers.OperationalSpaceController`.
+
+Added
+^^^^^
+
+* Added :class:`omni.isaac.lab.controllers.OperationalSpaceController` to API documentation.
+* Added test cases for :class:`omni.isaac.lab.controllers.OperationalSpaceController`.
+* Added a tutorial for :class:`omni.isaac.lab.controllers.OperationalSpaceController`.
+* Added the implementation of :class:`omni.isaac.lab.envs.mdp.actions.OperationalSpaceControllerAction` class.
+
+
 0.29.1 (2024-12-15)
 ~~~~~~~~~~~~~~~~~~~
 

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

@@ -13,3 +13,5 @@ commands to be sent to the robot.
 
 from .differential_ik import DifferentialIKController
 from .differential_ik_cfg import DifferentialIKControllerCfg
+from .operational_space import OperationalSpaceController
+from .operational_space_cfg import OperationalSpaceControllerCfg

source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space.py

@@ -3,92 +3,40 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-import torch
-from collections.abc import Sequence
-from dataclasses import MISSING
-
-from omni.isaac.lab.utils import configclass
-from omni.isaac.lab.utils.math import apply_delta_pose, compute_pose_error
-
-
-@configclass
-class OperationSpaceControllerCfg:
-    """Configuration for operation-space controller."""
-
-    command_types: Sequence[str] = MISSING
-    """Type of command.
-
-    It has two sub-strings joined by underscore:
-        - type of command mode: "position", "pose", "force"
-        - type of command resolving: "abs" (absolute), "rel" (relative)
-    """
-
-    impedance_mode: str = MISSING
-    """Type of gains for motion control: "fixed", "variable", "variable_kp"."""
-
-    uncouple_motion_wrench: bool = False
-    """Whether to decouple the wrench computation from task-space pose (motion) error."""
-
-    motion_control_axes: Sequence[int] = (1, 1, 1, 1, 1, 1)
-    """Motion direction to control. Mark as 0/1 for each axis."""
-    force_control_axes: Sequence[int] = (0, 0, 0, 0, 0, 0)
-    """Force direction to control. Mark as 0/1 for each axis."""
-
-    inertial_compensation: bool = False
-    """Whether to perform inertial compensation for motion control (inverse dynamics)."""
-
-    gravity_compensation: bool = False
-    """Whether to perform gravity compensation."""
-
-    stiffness: float | Sequence[float] = MISSING
-    """The positional gain for determining wrenches based on task-space pose error."""
-
-    damping_ratio: float | Sequence[float] | None = None
-    """The damping ratio is used in-conjunction with positional gain to compute wrenches
-    based on task-space velocity error.
-
-    The following math operation is performed for computing velocity gains:
-        :math:`d_gains = 2 * sqrt(p_gains) * damping_ratio`.
-    """
-
-    stiffness_limits: tuple[float, float] = (0, 300)
-    """Minimum and maximum values for positional gains.
-
-    Note: Used only when :obj:`impedance_mode` is "variable" or "variable_kp".
-    """
-
-    damping_ratio_limits: tuple[float, float] = (0, 100)
-    """Minimum and maximum values for damping ratios used to compute velocity gains.
+from __future__ import annotations
 
-    Note: Used only when :obj:`impedance_mode` is "variable".
-    """
+import torch
+from typing import TYPE_CHECKING
 
-    force_stiffness: float | Sequence[float] = None
-    """The positional gain for determining wrenches for closed-loop force control.
+from omni.isaac.lab.utils.math import (
+    apply_delta_pose,
+    combine_frame_transforms,
+    compute_pose_error,
+    matrix_from_quat,
+    subtract_frame_transforms,
+)
 
-    If obj:`None`, then open-loop control of desired forces is performed.
-    """
+if TYPE_CHECKING:
+    from .operational_space_cfg import OperationalSpaceControllerCfg
 
-    position_command_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
-    """Scaling of the position command received. Used only in relative mode."""
-    rotation_command_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
-    """Scaling of the rotation command received. Used only in relative mode."""
 
-
-class OperationSpaceController:
-    """Operation-space controller.
+class OperationalSpaceController:
+    """Operational-space controller.
 
     Reference:
-        [1] https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2017/RD_HS2017script.pdf
+
+    1. `A unified approach for motion and force control of robot manipulators: The operational space formulation <http://dx.doi.org/10.1109/JRA.1987.1087068>`_
+       by Oussama Khatib (Stanford University)
+    2. `Robot Dynamics Lecture Notes <https://ethz.ch/content/dam/ethz/special-interest/mavt/robotics-n-intelligent-systems/rsl-dam/documents/RobotDynamics2017/RD_HS2017script.pdf>`_
+       by Marco Hutter (ETH Zurich)
     """
 
-    def __init__(self, cfg: OperationSpaceControllerCfg, num_robots: int, num_dof: int, device: str):
-        """Initialize operation-space controller.
+    def __init__(self, cfg: OperationalSpaceControllerCfg, num_envs: int, device: str):
+        """Initialize operational-space controller.
 
         Args:
-            cfg: The configuration for operation-space controller.
-            num_robots: The number of robots to control.
-            num_dof: The number of degrees of freedom of the robot.
+            cfg: The configuration for operational-space controller.
+            num_envs: The number of environments.
             device: The device to use for computations.
 
         Raises:
@@ -96,79 +44,105 @@ class OperationSpaceController:
         """
         # store inputs
         self.cfg = cfg
-        self.num_robots = num_robots
-        self.num_dof = num_dof
+        self.num_envs = num_envs
         self._device = device
 
         # resolve tasks-pace target dimensions
         self.target_list = list()
-        for command_type in self.cfg.command_types:
-            if "position" in command_type:
-                self.target_list.append(3)
-            elif command_type == "pose_rel":
+        for command_type in self.cfg.target_types:
+            if command_type == "pose_rel":
                 self.target_list.append(6)
             elif command_type == "pose_abs":
                 self.target_list.append(7)
-            elif command_type == "force_abs":
+            elif command_type == "wrench_abs":
                 self.target_list.append(6)
             else:
                 raise ValueError(f"Invalid control command: {command_type}.")
         self.target_dim = sum(self.target_list)
 
         # create buffers
-        # -- selection matrices
-        self._selection_matrix_motion = torch.diag(
-            torch.tensor(self.cfg.motion_control_axes, dtype=torch.float, device=self._device)
+        # -- selection matrices, which might be defined in the task reference frame different from the root frame
+        self._selection_matrix_motion_task = torch.diag_embed(
+            torch.tensor(self.cfg.motion_control_axes_task, dtype=torch.float, device=self._device)
+            .unsqueeze(0)
+            .repeat(self.num_envs, 1)
         )
-        self._selection_matrix_force = torch.diag(
-            torch.tensor(self.cfg.force_control_axes, dtype=torch.float, device=self._device)
+        self._selection_matrix_force_task = torch.diag_embed(
+            torch.tensor(self.cfg.contact_wrench_control_axes_task, dtype=torch.float, device=self._device)
+            .unsqueeze(0)
+            .repeat(self.num_envs, 1)
         )
+        # -- selection matrices in root frame
+        self._selection_matrix_motion_b = torch.zeros_like(self._selection_matrix_motion_task)
+        self._selection_matrix_force_b = torch.zeros_like(self._selection_matrix_force_task)
         # -- commands
-        self._task_space_target = torch.zeros(self.num_robots, self.target_dim, device=self._device)
-        # -- scaling of command
-        self._position_command_scale = torch.diag(torch.tensor(self.cfg.position_command_scale, device=self._device))
-        self._rotation_command_scale = torch.diag(torch.tensor(self.cfg.rotation_command_scale, device=self._device))
+        self._task_space_target_task = torch.zeros(self.num_envs, self.target_dim, device=self._device)
+        # -- buffers for motion/force control
+        self.desired_ee_pose_task = None
+        self.desired_ee_pose_b = None
+        self.desired_ee_wrench_task = None
+        self.desired_ee_wrench_b = None
         # -- motion control gains
-        self._p_gains = torch.zeros(self.num_robots, 6, device=self._device)
-        self._p_gains[:] = torch.tensor(self.cfg.stiffness, device=self._device)
-        self._d_gains = 2 * torch.sqrt(self._p_gains) * torch.tensor(self.cfg.damping_ratio, device=self._device)
+        self._motion_p_gains_task = torch.diag_embed(
+            torch.ones(self.num_envs, 6, device=self._device)
+            * torch.tensor(self.cfg.motion_stiffness_task, dtype=torch.float, device=self._device)
+        )
+        # -- -- zero out the axes that are not motion controlled, as keeping them non-zero will cause other axes
+        # -- -- to act due to coupling
+        self._motion_p_gains_task[:] = self._selection_matrix_motion_task @ self._motion_p_gains_task[:]
+        self._motion_d_gains_task = torch.diag_embed(
+            2
+            * torch.diagonal(self._motion_p_gains_task, dim1=-2, dim2=-1).sqrt()
+            * torch.as_tensor(self.cfg.motion_damping_ratio_task, dtype=torch.float, device=self._device).reshape(1, -1)
+        )
+        # -- -- motion control gains in root frame
+        self._motion_p_gains_b = torch.zeros_like(self._motion_p_gains_task)
+        self._motion_d_gains_b = torch.zeros_like(self._motion_d_gains_task)
         # -- force control gains
-        if self.cfg.force_stiffness is not None:
-            self._p_wrench_gains = torch.zeros(self.num_robots, 6, device=self._device)
-            self._p_wrench_gains[:] = torch.tensor(self.cfg.force_stiffness, device=self._device)
+        if self.cfg.contact_wrench_stiffness_task is not None:
+            self._contact_wrench_p_gains_task = torch.diag_embed(
+                torch.ones(self.num_envs, 6, device=self._device)
+                * torch.tensor(self.cfg.contact_wrench_stiffness_task, dtype=torch.float, device=self._device)
+            )
+            self._contact_wrench_p_gains_task[:] = (
+                self._selection_matrix_force_task @ self._contact_wrench_p_gains_task[:]
+            )
+            # -- -- force control gains in root frame
+            self._contact_wrench_p_gains_b = torch.zeros_like(self._contact_wrench_p_gains_task)
         else:
-            self._p_wrench_gains = None
+            self._contact_wrench_p_gains_task = None
+            self._contact_wrench_p_gains_b = None
         # -- position gain limits
-        self._p_gains_limits = torch.zeros(self.num_robots, 6, device=self._device)
-        self._p_gains_limits[..., 0], self._p_gains_limits[..., 1] = (
-            self.cfg.stiffness_limits[0],
-            self.cfg.stiffness_limits[1],
+        self._motion_p_gains_limits = torch.zeros(self.num_envs, 6, 2, device=self._device)
+        self._motion_p_gains_limits[..., 0], self._motion_p_gains_limits[..., 1] = (
+            self.cfg.motion_stiffness_limits_task[0],
+            self.cfg.motion_stiffness_limits_task[1],
         )
         # -- damping ratio limits
-        self._damping_ratio_limits = torch.zeros_like(self._p_gains_limits)
-        self._damping_ratio_limits[..., 0], self._damping_ratio_limits[..., 1] = (
-            self.cfg.damping_ratio_limits[0],
-            self.cfg.damping_ratio_limits[1],
+        self._motion_damping_ratio_limits = torch.zeros_like(self._motion_p_gains_limits)
+        self._motion_damping_ratio_limits[..., 0], self._motion_damping_ratio_limits[..., 1] = (
+            self.cfg.motion_damping_ratio_limits_task[0],
+            self.cfg.motion_damping_ratio_limits_task[1],
         )
-        # -- storing outputs
-        self._desired_torques = torch.zeros(self.num_robots, self.num_dof, self.num_dof, device=self._device)
+        # -- end-effector contact wrench
+        self._ee_contact_wrench_b = torch.zeros(self.num_envs, 6, device=self._device)
 
     """
     Properties.
     """
 
     @property
-    def num_actions(self) -> int:
+    def action_dim(self) -> int:
         """Dimension of the action space of controller."""
         # impedance mode
         if self.cfg.impedance_mode == "fixed":
-            # task-space pose
+            # task-space targets
             return self.target_dim
         elif self.cfg.impedance_mode == "variable_kp":
-            # task-space pose + stiffness
+            # task-space targets + stiffness
             return self.target_dim + 6
         elif self.cfg.impedance_mode == "variable":
-            # task-space pose + stiffness + damping
+            # task-space targets + stiffness + damping
             return self.target_dim + 6 + 6
         else:
             raise ValueError(f"Invalid impedance mode: {self.cfg.impedance_mode}.")
@@ -177,184 +151,301 @@ class OperationSpaceController:
     Operations.
     """
 
-    def initialize(self):
-        """Initialize the internals."""
-        pass
-
-    def reset_idx(self, robot_ids: torch.Tensor = None):
+    def reset(self):
         """Reset the internals."""
-        pass
+        self.desired_ee_pose_b = None
+        self.desired_ee_pose_task = None
+        self.desired_ee_wrench_b = None
+        self.desired_ee_wrench_task = None
 
-    def set_command(self, command: torch.Tensor):
-        """Set target end-effector pose or force command.
+    def set_command(
+        self,
+        command: torch.Tensor,
+        current_ee_pose_b: torch.Tensor | None = None,
+        current_task_frame_pose_b: torch.Tensor | None = None,
+    ):
+        """Set the task-space targets and impedance parameters.
 
         Args:
-            command: The target end-effector pose or force command.
+            command (torch.Tensor): A concatenated tensor of shape (``num_envs``, ``action_dim``) containing task-space
+                targets (i.e., pose/wrench) and impedance parameters.
+            current_ee_pose_b (torch.Tensor, optional): Current end-effector pose, in root frame, of shape
+                (``num_envs``, 7), containing position and quaternion ``(w, x, y, z)``. Required for relative
+                commands. Defaults to None.
+            current_task_frame_pose_b: Current pose of the task frame, in root frame, in which the targets and the
+                (motion/wrench) control axes are defined. It is a tensor of shape (``num_envs``, 7),
+                containing position and the quaternion ``(w, x, y, z)``. Defaults to None.
+
+        Format:
+            Task-space targets, ordered according to 'command_types':
+
+                Absolute pose: shape (``num_envs``, 7), containing position and quaternion ``(w, x, y, z)``.
+                Relative pose: shape (``num_envs``, 6), containing delta position and rotation in axis-angle form.
+                Absolute wrench: shape (``num_envs``, 6), containing force and torque.
+
+            Impedance parameters: stiffness for ``variable_kp``, or stiffness, followed by damping ratio for
+            ``variable``:
+
+                Stiffness: shape (``num_envs``, 6)
+                Damping ratio: shape (``num_envs``, 6)
+
+        Raises:
+            ValueError: When the command dimensions are invalid.
+            ValueError: When an invalid impedance mode is provided.
+            ValueError: When the current end-effector pose is not provided for the ``pose_rel`` command.
+            ValueError: When an invalid control command is provided.
         """
-        # check input size
-        if command.shape != (self.num_robots, self.num_actions):
+        # Check the input dimensions
+        if command.shape != (self.num_envs, self.action_dim):
             raise ValueError(
-                f"Invalid command shape '{command.shape}'. Expected: '{(self.num_robots, self.num_actions)}'."
+                f"Invalid command shape '{command.shape}'. Expected: '{(self.num_envs, self.action_dim)}'."
             )
-        # impedance mode
+
+        # Resolve the impedance parameters
         if self.cfg.impedance_mode == "fixed":
-            # joint positions
-            self._task_space_target[:] = command
+            # task space targets (i.e., pose/wrench)
+            self._task_space_target_task[:] = command
         elif self.cfg.impedance_mode == "variable_kp":
             # split input command
-            task_space_command, stiffness = torch.tensor_split(command, (self.target_dim, 6), dim=-1)
+            task_space_command, stiffness = torch.split(command, [self.target_dim, 6], dim=-1)
             # format command
-            stiffness = stiffness.clip_(min=self._p_gains_limits[0], max=self._p_gains_limits[1])
-            # joint positions + stiffness
-            self._task_space_target[:] = task_space_command.squeeze(dim=-1)
-            self._p_gains[:] = stiffness
-            self._d_gains[:] = 2 * torch.sqrt(self._p_gains)  # critically damped
+            stiffness = stiffness.clip_(
+                min=self._motion_p_gains_limits[..., 0], max=self._motion_p_gains_limits[..., 1]
+            )
+            # task space targets + stiffness
+            self._task_space_target_task[:] = task_space_command.squeeze(dim=-1)
+            self._motion_p_gains_task[:] = torch.diag_embed(stiffness)
+            self._motion_p_gains_task[:] = self._selection_matrix_motion_task @ self._motion_p_gains_task[:]
+            self._motion_d_gains_task = torch.diag_embed(
+                2
+                * torch.diagonal(self._motion_p_gains_task, dim1=-2, dim2=-1).sqrt()
+                * torch.as_tensor(self.cfg.motion_damping_ratio_task, dtype=torch.float, device=self._device).reshape(
+                    1, -1
+                )
+            )
         elif self.cfg.impedance_mode == "variable":
             # split input command
-            task_space_command, stiffness, damping_ratio = torch.tensor_split(command, 3, dim=-1)
+            task_space_command, stiffness, damping_ratio = torch.split(command, [self.target_dim, 6, 6], dim=-1)
             # format command
-            stiffness = stiffness.clip_(min=self._p_gains_limits[0], max=self._p_gains_limits[1])
-            damping_ratio = damping_ratio.clip_(min=self._damping_ratio_limits[0], max=self._damping_ratio_limits[1])
-            # joint positions + stiffness + damping
-            self._task_space_target[:] = task_space_command
-            self._p_gains[:] = stiffness
-            self._d_gains[:] = 2 * torch.sqrt(self._p_gains) * damping_ratio
+            stiffness = stiffness.clip_(
+                min=self._motion_p_gains_limits[..., 0], max=self._motion_p_gains_limits[..., 1]
+            )
+            damping_ratio = damping_ratio.clip_(
+                min=self._motion_damping_ratio_limits[..., 0], max=self._motion_damping_ratio_limits[..., 1]
+            )
+            # task space targets + stiffness + damping
+            self._task_space_target_task[:] = task_space_command
+            self._motion_p_gains_task[:] = torch.diag_embed(stiffness)
+            self._motion_p_gains_task[:] = self._selection_matrix_motion_task @ self._motion_p_gains_task[:]
+            self._motion_d_gains_task[:] = torch.diag_embed(
+                2 * torch.diagonal(self._motion_p_gains_task, dim1=-2, dim2=-1).sqrt() * damping_ratio
+            )
         else:
             raise ValueError(f"Invalid impedance mode: {self.cfg.impedance_mode}.")
 
+        if current_task_frame_pose_b is None:
+            current_task_frame_pose_b = torch.tensor(
+                [[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]] * self.num_envs, device=self._device
+            )
+
+        # Resolve the target commands
+        target_groups = torch.split(self._task_space_target_task, self.target_list, dim=1)
+        for command_type, target in zip(self.cfg.target_types, target_groups):
+            if command_type == "pose_rel":
+                # check input is provided
+                if current_ee_pose_b is None:
+                    raise ValueError("Current pose is required for 'pose_rel' command.")
+                # Transform the current pose from base/root frame to task frame
+                current_ee_pos_task, current_ee_rot_task = subtract_frame_transforms(
+                    current_task_frame_pose_b[:, :3],
+                    current_task_frame_pose_b[:, 3:],
+                    current_ee_pose_b[:, :3],
+                    current_ee_pose_b[:, 3:],
+                )
+                # compute targets in task frame
+                desired_ee_pos_task, desired_ee_rot_task = apply_delta_pose(
+                    current_ee_pos_task, current_ee_rot_task, target
+                )
+                self.desired_ee_pose_task = torch.cat([desired_ee_pos_task, desired_ee_rot_task], dim=-1)
+            elif command_type == "pose_abs":
+                # compute targets
+                self.desired_ee_pose_task = target.clone()
+            elif command_type == "wrench_abs":
+                # compute targets
+                self.desired_ee_wrench_task = target.clone()
+            else:
+                raise ValueError(f"Invalid control command: {command_type}.")
+
+        # Rotation of task frame wrt root frame, converts a coordinate from task frame to root frame.
+        R_task_b = matrix_from_quat(current_task_frame_pose_b[:, 3:])
+        # Rotation of root frame wrt task frame, converts a coordinate from root frame to task frame.
+        R_b_task = R_task_b.mT
+
+        # Transform motion control stiffness gains from task frame to root frame
+        self._motion_p_gains_b[:, 0:3, 0:3] = R_task_b @ self._motion_p_gains_task[:, 0:3, 0:3] @ R_b_task
+        self._motion_p_gains_b[:, 3:6, 3:6] = R_task_b @ self._motion_p_gains_task[:, 3:6, 3:6] @ R_b_task
+
+        # Transform motion control damping gains from task frame to root frame
+        self._motion_d_gains_b[:, 0:3, 0:3] = R_task_b @ self._motion_d_gains_task[:, 0:3, 0:3] @ R_b_task
+        self._motion_d_gains_b[:, 3:6, 3:6] = R_task_b @ self._motion_d_gains_task[:, 3:6, 3:6] @ R_b_task
+
+        # Transform contact wrench gains from task frame to root frame (if applicable)
+        if self._contact_wrench_p_gains_task is not None and self._contact_wrench_p_gains_b is not None:
+            self._contact_wrench_p_gains_b[:, 0:3, 0:3] = (
+                R_task_b @ self._contact_wrench_p_gains_task[:, 0:3, 0:3] @ R_b_task
+            )
+            self._contact_wrench_p_gains_b[:, 3:6, 3:6] = (
+                R_task_b @ self._contact_wrench_p_gains_task[:, 3:6, 3:6] @ R_b_task
+            )
+
+        # Transform selection matrices from target frame to base frame
+        self._selection_matrix_motion_b[:, 0:3, 0:3] = (
+            R_task_b @ self._selection_matrix_motion_task[:, 0:3, 0:3] @ R_b_task
+        )
+        self._selection_matrix_motion_b[:, 3:6, 3:6] = (
+            R_task_b @ self._selection_matrix_motion_task[:, 3:6, 3:6] @ R_b_task
+        )
+        self._selection_matrix_force_b[:, 0:3, 0:3] = (
+            R_task_b @ self._selection_matrix_force_task[:, 0:3, 0:3] @ R_b_task
+        )
+        self._selection_matrix_force_b[:, 3:6, 3:6] = (
+            R_task_b @ self._selection_matrix_force_task[:, 3:6, 3:6] @ R_b_task
+        )
+
+        # Transform desired pose from task frame to root frame
+        if self.desired_ee_pose_task is not None:
+            self.desired_ee_pose_b = torch.zeros_like(self.desired_ee_pose_task)
+            self.desired_ee_pose_b[:, :3], self.desired_ee_pose_b[:, 3:] = combine_frame_transforms(
+                current_task_frame_pose_b[:, :3],
+                current_task_frame_pose_b[:, 3:],
+                self.desired_ee_pose_task[:, :3],
+                self.desired_ee_pose_task[:, 3:],
+            )
+
+        # Transform desired wrenches to root frame
+        if self.desired_ee_wrench_task is not None:
+            self.desired_ee_wrench_b = torch.zeros_like(self.desired_ee_wrench_task)
+            self.desired_ee_wrench_b[:, :3] = (R_task_b @ self.desired_ee_wrench_task[:, :3].unsqueeze(-1)).squeeze(-1)
+            self.desired_ee_wrench_b[:, 3:] = (R_task_b @ self.desired_ee_wrench_task[:, 3:].unsqueeze(-1)).squeeze(
+                -1
+            ) + torch.cross(current_task_frame_pose_b[:, :3], self.desired_ee_wrench_b[:, :3], dim=-1)
+
     def compute(
         self,
-        jacobian: torch.Tensor,
-        ee_pose: torch.Tensor | None = None,
-        ee_vel: torch.Tensor | None = None,
-        ee_force: torch.Tensor | None = None,
+        jacobian_b: torch.Tensor,
+        current_ee_pose_b: torch.Tensor | None = None,
+        current_ee_vel_b: torch.Tensor | None = None,
+        current_ee_force_b: torch.Tensor | None = None,
         mass_matrix: torch.Tensor | None = None,
         gravity: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Performs inference with the controller.
 
         Args:
-            jacobian: The Jacobian matrix of the end-effector.
-            ee_pose: The current end-effector pose. It is a tensor of shape
-                (num_robots, 7), which contains the position and quaternion (w, x, y, z). Defaults to None.
-            ee_vel: The current end-effector velocity. It is a tensor of shape
-                (num_robots, 6), which contains the linear and angular velocities. Defaults to None.
-            ee_force: The current external force on the end-effector.
-                It is a tensor of shape (num_robots, 3), which contains the linear force. Defaults to None.
-            mass_matrix: The joint-space inertial matrix. Defaults to None.
-            gravity: The joint-space gravity vector. Defaults to None.
+            jacobian_b: The Jacobian matrix of the end-effector in root frame. It is a tensor of shape
+                (``num_envs``, 6, ``num_DoF``).
+            current_ee_pose_b: The current end-effector pose in root frame. It is a tensor of shape
+                (``num_envs``, 7), which contains the position and quaternion ``(w, x, y, z)``. Defaults to ``None``.
+            current_ee_vel_b: The current end-effector velocity in root frame. It is a tensor of shape
+                (``num_envs``, 6), which contains the linear and angular velocities. Defaults to None.
+            current_ee_force_b: The current external force on the end-effector in root frame.
+                It is a tensor of shape (``num_envs``, 3), which contains the linear force. Defaults to ``None``.
+            mass_matrix: The joint-space mass/inertia matrix. It is a tensor of shape (``num_envs``, ``num_DoF``,
+            ``num_DoF``). Defaults to ``None``.
+            gravity: The joint-space gravity vector. It is a tensor of shape (``num_envs``, ``num_DoF``). Defaults
+            to ``None``.
 
         Raises:
-            ValueError: When the end-effector pose is not provided for the 'position_rel' command.
-            ValueError: When the end-effector pose is not provided for the 'position_abs' command.
-            ValueError: When the end-effector pose is not provided for the 'pose_rel' command.
-            ValueError: When an invalid command type is provided.
-            ValueError: When motion-control is enabled but the end-effector pose or velocity is not provided.
-            ValueError: When force-control is enabled but the end-effector force is not provided.
-            ValueError: When inertial compensation is enabled but the mass matrix  is not provided.
+            ValueError: When motion-control is enabled but the current end-effector pose or velocity is not provided.
+            ValueError: When inertial dynamics decoupling is enabled but the mass matrix is not provided.
+            ValueError: When the current end-effector pose is not provided for the ``pose_rel`` command.
+            ValueError: When closed-loop force control is enabled but the current end-effector force is not provided.
             ValueError: When gravity compensation is enabled but the gravity vector is not provided.
 
         Returns:
-            The target joint torques commands.
+            Tensor: The joint efforts computed by the controller. It is a tensor of shape (``num_envs``, ``num_DoF``).
         """
-        # buffers for motion/force control
-        desired_ee_pos = None
-        desired_ee_rot = None
-        desired_ee_force = None
-        # resolve the commands
-        target_groups = torch.split(self._task_space_target, self.target_list)
-        for command_type, target in zip(self.cfg.command_types, target_groups):
-            if command_type == "position_rel":
-                # check input is provided
-                if ee_pose is None:
-                    raise ValueError("End-effector pose is required for 'position_rel' command.")
-                # scale command
-                target @= self._position_command_scale
-                # compute targets
-                desired_ee_pos = ee_pose[:, :3] + target
-                desired_ee_rot = ee_pose[:, 3:]
-            elif command_type == "position_abs":
-                # check input is provided
-                if ee_pose is None:
-                    raise ValueError("End-effector pose is required for 'position_abs' command.")
-                # compute targets
-                desired_ee_pos = target
-                desired_ee_rot = ee_pose[:, 3:]
-            elif command_type == "pose_rel":
-                # check input is provided
-                if ee_pose is None:
-                    raise ValueError("End-effector pose is required for 'pose_rel' command.")
-                # scale command
-                target[:, 0:3] @= self._position_command_scale
-                target[:, 3:6] @= self._rotation_command_scale
-                # compute targets
-                desired_ee_pos, desired_ee_rot = apply_delta_pose(ee_pose[:, :3], ee_pose[:, 3:], target)
-            elif command_type == "pose_abs":
-                # compute targets
-                desired_ee_pos = target[:, 0:3]
-                desired_ee_rot = target[:, 3:7]
-            elif command_type == "force_abs":
-                # compute targets
-                desired_ee_force = target
-            else:
-                raise ValueError(f"Invalid control command: {self.cfg.command_type}.")
 
-        # reset desired joint torques
-        self._desired_torques[:] = 0
-        # compute for motion-control
-        if desired_ee_pos is not None:
+        # deduce number of DoF
+        num_DoF = jacobian_b.shape[2]
+        # create joint effort vector
+        joint_efforts = torch.zeros(self.num_envs, num_DoF, device=self._device)
+
+        # compute joint efforts for motion-control
+        if self.desired_ee_pose_b is not None:
             # check input is provided
-            if ee_pose is None or ee_vel is None:
-                raise ValueError("End-effector pose and velocity are required for motion control.")
+            if current_ee_pose_b is None or current_ee_vel_b is None:
+                raise ValueError("Current end-effector pose and velocity are required for motion control.")
             # -- end-effector tracking error
-            pose_error = compute_pose_error(
-                ee_pose[:, :3], ee_pose[:, 3:], desired_ee_pos, desired_ee_rot, rot_error_type="axis_angle"
+            pose_error_b = torch.cat(
+                compute_pose_error(
+                    current_ee_pose_b[:, :3],
+                    current_ee_pose_b[:, 3:],
+                    self.desired_ee_pose_b[:, :3],
+                    self.desired_ee_pose_b[:, 3:],
+                    rot_error_type="axis_angle",
+                ),
+                dim=-1,
             )
-            velocity_error = -ee_vel  # zero target velocity
-            # -- desired end-effector acceleration (spring damped system)
-            des_ee_acc = self._p_gains * pose_error + self._d_gains * velocity_error
-            # -- inertial compensation
-            if self.cfg.inertial_compensation:
+            velocity_error_b = -current_ee_vel_b  # zero target velocity. The target is assumed to be stationary.
+            # -- desired end-effector acceleration (spring-damper system)
+            des_ee_acc_b = self._motion_p_gains_b @ pose_error_b.unsqueeze(
+                -1
+            ) + self._motion_d_gains_b @ velocity_error_b.unsqueeze(-1)
+            # -- Inertial dynamics decoupling
+            if self.cfg.inertial_dynamics_decoupling:
                 # check input is provided
                 if mass_matrix is None:
-                    raise ValueError("Mass matrix is required for inertial compensation.")
-                # compute task-space dynamics quantities
-                # wrench = (J M^(-1) J^T)^(-1) * \ddot(x_des)
+                    raise ValueError("Mass matrix is required for inertial decoupling.")
+                # Compute operational space mass matrix
                 mass_matrix_inv = torch.inverse(mass_matrix)
-                if self.cfg.uncouple_motion_wrench:
-                    # decoupled-mass matrices
-                    lambda_pos = torch.inverse(jacobian[:, 0:3] @ mass_matrix_inv * jacobian[:, 0:3].T)
-                    lambda_ori = torch.inverse(jacobian[:, 3:6] @ mass_matrix_inv * jacobian[:, 3:6].T)
-                    # desired end-effector wrench (from pseudo-dynamics)
-                    decoupled_force = lambda_pos @ des_ee_acc[:, 0:3]
-                    decoupled_torque = lambda_ori @ des_ee_acc[:, 3:6]
-                    des_motion_wrench = torch.cat(decoupled_force, decoupled_torque)
+                if self.cfg.partial_inertial_dynamics_decoupling:
+                    # Create a zero tensor representing the mass matrix, to fill in the non-zero elements later
+                    os_mass_matrix = torch.zeros(self.num_envs, 6, 6, device=self._device)
+                    # Fill in the translational and rotational parts of the inertia separately, ignoring their coupling
+                    os_mass_matrix[:, 0:3, 0:3] = torch.inverse(
+                        jacobian_b[:, 0:3] @ mass_matrix_inv @ jacobian_b[:, 0:3].mT
+                    )
+                    os_mass_matrix[:, 3:6, 3:6] = torch.inverse(
+                        jacobian_b[:, 3:6] @ mass_matrix_inv @ jacobian_b[:, 3:6].mT
+                    )
                 else:
-                    # coupled dynamics
-                    lambda_full = torch.inverse(jacobian @ mass_matrix_inv * jacobian.T)
-                    # desired end-effector wrench (from pseudo-dynamics)
-                    des_motion_wrench = lambda_full @ des_ee_acc
+                    # Calculate the operational space mass matrix fully accounting for the couplings
+                    os_mass_matrix = torch.inverse(jacobian_b @ mass_matrix_inv @ jacobian_b.mT)
+                # (Generalized) operational space command forces
+                # F = (J M^(-1) J^T)^(-1) * \ddot(x_des) = M_task * \ddot(x_des)
+                os_command_forces_b = os_mass_matrix @ des_ee_acc_b
             else:
-                # task-space impedance control
-                # wrench = \ddot(x_des)
-                des_motion_wrench = des_ee_acc
-            # -- joint-space wrench
-            self._desired_torques += jacobian.T @ self._selection_matrix_motion @ des_motion_wrench
-
-        # compute for force control
-        if desired_ee_force is not None:
-            # -- task-space wrench
-            if self.cfg.stiffness is not None:
+                # Task-space impedance control: command forces = \ddot(x_des).
+                # Please note that the definition of task-space impedance control varies in literature.
+                # This implementation ignores the inertial term. For inertial decoupling,
+                # use inertial_dynamics_decoupling=True.
+                os_command_forces_b = des_ee_acc_b
+            # -- joint-space commands
+            joint_efforts += (jacobian_b.mT @ self._selection_matrix_motion_b @ os_command_forces_b).squeeze(-1)
+
+        # compute joint efforts for contact wrench/force control
+        if self.desired_ee_wrench_b is not None:
+            # -- task-space contact wrench
+            if self.cfg.contact_wrench_stiffness_task is not None:
                 # check input is provided
-                if ee_force is None:
-                    raise ValueError("End-effector force is required for closed-loop force control.")
-                # closed-loop control
-                des_force_wrench = desired_ee_force + self._p_wrench_gains * (desired_ee_force - ee_force)
+                if current_ee_force_b is None:
+                    raise ValueError("Current end-effector force is required for closed-loop force control.")
+                # We can only measure the force component at the contact, so only apply the feedback for only the force
+                # component, keep the control of moment components open loop
+                self._ee_contact_wrench_b[:, 0:3] = current_ee_force_b
+                self._ee_contact_wrench_b[:, 3:6] = self.desired_ee_wrench_b[:, 3:6]
+                # closed-loop control with feedforward term
+                os_contact_wrench_command_b = self.desired_ee_wrench_b.unsqueeze(
+                    -1
+                ) + self._contact_wrench_p_gains_b @ (self.desired_ee_wrench_b - self._ee_contact_wrench_b).unsqueeze(
+                    -1
+                )
             else:
                 # open-loop control
-                des_force_wrench = desired_ee_force
-            # -- joint-space wrench
-            self._desired_torques += jacobian.T @ self._selection_matrix_force @ des_force_wrench
+                os_contact_wrench_command_b = self.desired_ee_wrench_b.unsqueeze(-1)
+            # -- joint-space commands
+            joint_efforts += (jacobian_b.mT @ self._selection_matrix_force_b @ os_contact_wrench_command_b).squeeze(-1)
 
         # add gravity compensation (bias correction)
         if self.cfg.gravity_compensation:
@@ -362,6 +453,6 @@ class OperationSpaceController:
             if gravity is None:
                 raise ValueError("Gravity vector is required for gravity compensation.")
             # add gravity compensation
-            self._desired_torques += gravity
+            joint_efforts += gravity
 
-        return self._desired_torques
+        return joint_efforts

source/extensions/omni.isaac.lab/omni/isaac/lab/controllers/operational_space_cfg.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from collections.abc import Sequence
+from dataclasses import MISSING
+
+from omni.isaac.lab.utils import configclass
+
+from .operational_space import OperationalSpaceController
+
+
+@configclass
+class OperationalSpaceControllerCfg:
+    """Configuration for operational-space controller."""
+
+    class_type: type = OperationalSpaceController
+    """The associated controller class."""
+
+    target_types: Sequence[str] = MISSING
+    """Type of task-space targets.
+
+    It has two sub-strings joined by underscore:
+        - type of task-space target: ``"pose"``, ``"wrench"``
+        - reference for the task-space targets: ``"abs"`` (absolute), ``"rel"`` (relative, only for pose)
+    """
+
+    motion_control_axes_task: Sequence[int] = (1, 1, 1, 1, 1, 1)
+    """Motion direction to control in task reference frame. Mark as ``0/1`` for each axis."""
+
+    contact_wrench_control_axes_task: Sequence[int] = (0, 0, 0, 0, 0, 0)
+    """Contact wrench direction to control in task reference frame. Mark as 0/1 for each axis."""
+
+    inertial_dynamics_decoupling: bool = False
+    """Whether to perform inertial dynamics decoupling for motion control (inverse dynamics)."""
+
+    partial_inertial_dynamics_decoupling: bool = False
+    """Whether to ignore the inertial coupling between the translational & rotational motions."""
+
+    gravity_compensation: bool = False
+    """Whether to perform gravity compensation."""
+
+    impedance_mode: str = "fixed"
+    """Type of gains for motion control: ``"fixed"``, ``"variable"``, ``"variable_kp"``."""
+
+    motion_stiffness_task: float | Sequence[float] = (100.0, 100.0, 100.0, 100.0, 100.0, 100.0)
+    """The positional gain for determining operational space command forces based on task-space pose error."""
+
+    motion_damping_ratio_task: float | Sequence[float] = (1.0, 1.0, 1.0, 1.0, 1.0, 1.0)
+    """The damping ratio is used in-conjunction with positional gain to compute operational space command forces
+    based on task-space velocity error.
+
+    The following math operation is performed for computing velocity gains:
+        :math:`d_gains = 2 * sqrt(p_gains) * damping_ratio`.
+    """
+
+    motion_stiffness_limits_task: tuple[float, float] = (0, 1000)
+    """Minimum and maximum values for positional gains.
+
+    Note: Used only when :obj:`impedance_mode` is ``"variable"`` or ``"variable_kp"``.
+    """
+
+    motion_damping_ratio_limits_task: tuple[float, float] = (0, 100)
+    """Minimum and maximum values for damping ratios used to compute velocity gains.
+
+    Note: Used only when :obj:`impedance_mode` is ``"variable"``.
+    """
+
+    contact_wrench_stiffness_task: float | Sequence[float] | None = None
+    """The proportional gain for determining operational space command forces for closed-loop contact force control.
+
+    If ``None``, then open-loop control of desired contact wrench is performed.
+
+    Note: since only the linear forces could be measured at the moment,
+    only the first three elements are used for the feedback loop.
+    """

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

@@ -5,7 +5,7 @@
 
 from dataclasses import MISSING
 
-from omni.isaac.lab.controllers import DifferentialIKControllerCfg
+from omni.isaac.lab.controllers import DifferentialIKControllerCfg, OperationalSpaceControllerCfg
 from omni.isaac.lab.managers.action_manager import ActionTerm, ActionTermCfg
 from omni.isaac.lab.utils import configclass
 
@@ -248,3 +248,58 @@ class DifferentialInverseKinematicsActionCfg(ActionTermCfg):
     """Scale factor for the action. Defaults to 1.0."""
     controller: DifferentialIKControllerCfg = MISSING
     """The configuration for the differential IK controller."""
+
+
+@configclass
+class OperationalSpaceControllerActionCfg(ActionTermCfg):
+    """Configuration for operational space controller action term.
+
+    See :class:`OperationalSpaceControllerAction` for more details.
+    """
+
+    @configclass
+    class OffsetCfg:
+        """The offset pose from parent frame to child frame.
+
+        On many robots, end-effector frames are fictitious frames that do not have a corresponding
+        rigid body. In such cases, it is easier to define this transform w.r.t. their parent rigid body.
+        For instance, for the Franka Emika arm, the end-effector is defined at an offset to the the
+        "panda_hand" 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)."""
+
+    class_type: type[ActionTerm] = task_space_actions.OperationalSpaceControllerAction
+
+    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 motion/force control is performed."""
+
+    body_offset: OffsetCfg | None = None
+    """Offset of target frame w.r.t. to the body frame. Defaults to None, in which case no offset is applied."""
+
+    task_frame_rel_path: str = None
+    """The path of a ``RigidObject``, relative to the sub-environment, representing task frame. Defaults to None."""
+
+    controller_cfg: OperationalSpaceControllerCfg = MISSING
+    """The configuration for the operational space controller."""
+
+    position_scale: float = 1.0
+    """Scale factor for the position targets. Defaults to 1.0."""
+
+    orientation_scale: float = 1.0
+    """Scale factor for the orientation (quad for ``pose_abs`` or axis-angle for ``pose_rel``). Defaults to 1.0."""
+
+    wrench_scale: float = 1.0
+    """Scale factor for the wrench targets. Defaults to 1.0."""
+
+    stiffness_scale: float = 1.0
+    """Scale factor for the stiffness commands. Defaults to 1.0."""
+
+    damping_ratio_scale: float = 1.0
+    """Scale factor for the damping ratio commands. Defaults to 1.0."""

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/actions/task_space_actions.py

@@ -10,12 +10,16 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 import omni.log
+from pxr import UsdPhysics
 
 import omni.isaac.lab.utils.math as math_utils
 import omni.isaac.lab.utils.string as string_utils
 from omni.isaac.lab.assets.articulation import Articulation
 from omni.isaac.lab.controllers.differential_ik import DifferentialIKController
+from omni.isaac.lab.controllers.operational_space import OperationalSpaceController
 from omni.isaac.lab.managers.action_manager import ActionTerm
+from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg, FrameTransformer, FrameTransformerCfg
+from omni.isaac.lab.sim.utils import find_matching_prims
 
 if TYPE_CHECKING:
     from omni.isaac.lab.envs import ManagerBasedEnv
@@ -223,3 +227,423 @@ class DifferentialInverseKinematicsAction(ActionTerm):
             jacobian[:, 3:, :] = torch.bmm(math_utils.matrix_from_quat(self._offset_rot), jacobian[:, 3:, :])
 
         return jacobian
+
+
+class OperationalSpaceControllerAction(ActionTerm):
+    r"""Operational space controller action term.
+
+    This action term performs pre-processing of the raw actions for operational space control.
+
+    """
+
+    cfg: actions_cfg.OperationalSpaceControllerActionCfg
+    """The configuration of the action term."""
+    _asset: Articulation
+    """The articulation asset on which the action term is applied."""
+    _contact_sensor: ContactSensor = None
+    """The contact sensor for the end-effector body."""
+    _task_frame_transformer: FrameTransformer = None
+    """The frame transformer for the task frame."""
+
+    def __init__(self, cfg: actions_cfg.OperationalSpaceControllerActionCfg, env: ManagerBasedEnv):
+        # initialize the action term
+        super().__init__(cfg, env)
+
+        self._sim_dt = env.sim.get_physics_dt()
+
+        # 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_DoF = len(self._joint_ids)
+        # parse the ee 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 ee body name: {self.cfg.body_name}. Found {len(body_ids)}: {body_names}."
+            )
+        # save only the first ee body index
+        self._ee_body_idx = body_ids[0]
+        self._ee_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_ee_body_idx = self._ee_body_idx - 1
+            self._jacobi_joint_idx = self._joint_ids
+        else:
+            self._jacobi_ee_body_idx = self._ee_body_idx
+            self._jacobi_joint_idx = [i + 6 for i in self._joint_ids]
+
+        # log info for debugging
+        omni.log.info(
+            f"Resolved joint names for the action term {self.__class__.__name__}:"
+            f" {self._joint_names} [{self._joint_ids}]"
+        )
+        omni.log.info(
+            f"Resolved ee body name for the action term {self.__class__.__name__}:"
+            f" {self._ee_body_name} [{self._ee_body_idx}]"
+        )
+        # Avoid indexing across all joints for efficiency
+        if self._num_DoF == self._asset.num_joints:
+            self._joint_ids = slice(None)
+
+        # convert the fixed offsets to torch tensors of batched shape
+        if self.cfg.body_offset is not None:
+            self._offset_pos = torch.tensor(self.cfg.body_offset.pos, device=self.device).repeat(self.num_envs, 1)
+            self._offset_rot = torch.tensor(self.cfg.body_offset.rot, device=self.device).repeat(self.num_envs, 1)
+        else:
+            self._offset_pos, self._offset_rot = None, None
+
+        # create contact sensor if any of the command is wrench_abs, and if stiffness is provided
+        if (
+            "wrench_abs" in self.cfg.controller_cfg.target_types
+            and self.cfg.controller_cfg.contact_wrench_stiffness_task is not None
+        ):
+            self._contact_sensor_cfg = ContactSensorCfg(prim_path=self._asset.cfg.prim_path + "/" + self._ee_body_name)
+            self._contact_sensor = ContactSensor(self._contact_sensor_cfg)
+            if not self._contact_sensor.is_initialized:
+                self._contact_sensor._initialize_impl()
+                self._contact_sensor._is_initialized = True
+
+        # Initialize the task frame transformer if a relative path for the RigidObject, representing the task frame,
+        # is provided.
+        if self.cfg.task_frame_rel_path is not None:
+            # The source RigidObject can be any child of the articulation asset (we will not use it),
+            # hence, we will use the first RigidObject child.
+            root_rigidbody_path = self._first_RigidObject_child_path()
+            task_frame_transformer_path = "/World/envs/env_.*/" + self.cfg.task_frame_rel_path
+            task_frame_transformer_cfg = FrameTransformerCfg(
+                prim_path=root_rigidbody_path,
+                target_frames=[
+                    FrameTransformerCfg.FrameCfg(
+                        name="task_frame",
+                        prim_path=task_frame_transformer_path,
+                    ),
+                ],
+            )
+            self._task_frame_transformer = FrameTransformer(task_frame_transformer_cfg)
+            if not self._task_frame_transformer.is_initialized:
+                self._task_frame_transformer._initialize_impl()
+                self._task_frame_transformer._is_initialized = True
+            # create tensor for task frame pose wrt the root frame
+            self._task_frame_pose_b = torch.zeros(self.num_envs, 7, device=self.device)
+        else:
+            # create an empty reference for task frame pose
+            self._task_frame_pose_b = None
+
+        # create the operational space controller
+        self._osc = OperationalSpaceController(cfg=self.cfg.controller_cfg, 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)
+
+        # create tensors for the dynamic-related quantities
+        self._jacobian_b = torch.zeros(self.num_envs, 6, self._num_DoF, device=self.device)
+        self._mass_matrix = torch.zeros(self.num_envs, self._num_DoF, self._num_DoF, device=self.device)
+        self._gravity = torch.zeros(self.num_envs, self._num_DoF, device=self.device)
+
+        # create tensors for the ee states
+        self._ee_pose_w = torch.zeros(self.num_envs, 7, device=self.device)
+        self._ee_pose_b = torch.zeros(self.num_envs, 7, device=self.device)
+        self._ee_pose_b_no_offset = torch.zeros(self.num_envs, 7, device=self.device)  # The original ee without offset
+        self._ee_vel_w = torch.zeros(self.num_envs, 6, device=self.device)
+        self._ee_vel_b = torch.zeros(self.num_envs, 6, device=self.device)
+        self._ee_force_w = torch.zeros(self.num_envs, 3, device=self.device)  # Only the forces are used for now
+        self._ee_force_b = torch.zeros(self.num_envs, 3, device=self.device)  # Only the forces are used for now
+
+        # create the joint effort tensor
+        self._joint_efforts = torch.zeros(self.num_envs, self._num_DoF, device=self.device)
+
+        # save the scale as tensors
+        self._position_scale = torch.tensor(self.cfg.position_scale, device=self.device)
+        self._orientation_scale = torch.tensor(self.cfg.orientation_scale, device=self.device)
+        self._wrench_scale = torch.tensor(self.cfg.wrench_scale, device=self.device)
+        self._stiffness_scale = torch.tensor(self.cfg.stiffness_scale, device=self.device)
+        self._damping_ratio_scale = torch.tensor(self.cfg.damping_ratio_scale, device=self.device)
+
+        # indexes for the various command elements (e.g., pose_rel, stifness, etc.) within the command tensor
+        self._pose_abs_idx = None
+        self._pose_rel_idx = None
+        self._wrench_abs_idx = None
+        self._stiffness_idx = None
+        self._damping_ratio_idx = None
+        self._resolve_command_indexes()
+
+    """
+    Properties.
+    """
+
+    @property
+    def action_dim(self) -> int:
+        """Dimension of the action space of operational space control."""
+        return self._osc.action_dim
+
+    @property
+    def raw_actions(self) -> torch.Tensor:
+        """Raw actions for operational space control."""
+        return self._raw_actions
+
+    @property
+    def processed_actions(self) -> torch.Tensor:
+        """Processed actions for operational space control."""
+        return self._processed_actions
+
+    @property
+    def jacobian_w(self) -> torch.Tensor:
+        return self._asset.root_physx_view.get_jacobians()[:, self._jacobi_ee_body_idx, :, self._jacobi_joint_idx]
+
+    @property
+    def jacobian_b(self) -> torch.Tensor:
+        jacobian = self.jacobian_w
+        base_rot = self._asset.data.root_quat_w
+        base_rot_matrix = math_utils.matrix_from_quat(math_utils.quat_inv(base_rot))
+        jacobian[:, :3, :] = torch.bmm(base_rot_matrix, jacobian[:, :3, :])
+        jacobian[:, 3:, :] = torch.bmm(base_rot_matrix, jacobian[:, 3:, :])
+        return jacobian
+
+    """
+    Operations.
+    """
+
+    def process_actions(self, actions: torch.Tensor):
+        """Pre-processes the raw actions and sets them as commands for for operational space control.
+
+        Args:
+            actions (torch.Tensor): The raw actions for operational space control. It is a tensor of
+                shape (``num_envs``, ``action_dim``).
+        """
+
+        # Update ee pose, which would be used by relative targets (i.e., pose_rel)
+        self._compute_ee_pose()
+
+        # Update task frame pose w.r.t. the root frame.
+        self._compute_task_frame_pose()
+
+        # Pre-process the raw actions for operational space control.
+        self._preprocess_actions(actions)
+
+        # set command into controller
+        self._osc.set_command(
+            command=self._processed_actions,
+            current_ee_pose_b=self._ee_pose_b,
+            current_task_frame_pose_b=self._task_frame_pose_b,
+        )
+
+    def apply_actions(self):
+        """Computes the joint efforts for operational space control and applies them to the articulation."""
+
+        # Update the relevant states and dynamical quantities
+        self._compute_dynamic_quantities()
+        self._compute_ee_jacobian()
+        self._compute_ee_pose()
+        self._compute_ee_velocity()
+        self._compute_ee_force()
+        # Calculate the joint efforts
+        self._joint_efforts[:] = self._osc.compute(
+            jacobian_b=self._jacobian_b,
+            current_ee_pose_b=self._ee_pose_b,
+            current_ee_vel_b=self._ee_vel_b,
+            current_ee_force_b=self._ee_force_b,
+            mass_matrix=self._mass_matrix,
+            gravity=self._gravity,
+        )
+        self._asset.set_joint_effort_target(self._joint_efforts, joint_ids=self._joint_ids)
+
+    def reset(self, env_ids: Sequence[int] | None = None) -> None:
+        """Resets the raw actions and the sensors if available.
+
+        Args:
+            env_ids (Sequence[int] | None): The environment indices to reset. If ``None``, all environments are reset.
+        """
+        self._raw_actions[env_ids] = 0.0
+        if self._contact_sensor is not None:
+            self._contact_sensor.reset(env_ids)
+        if self._task_frame_transformer is not None:
+            self._task_frame_transformer.reset(env_ids)
+
+    """
+    Helper functions.
+
+    """
+
+    def _first_RigidObject_child_path(self):
+        """Finds the first ``RigidObject`` child under the articulation asset.
+
+        Raises:
+            ValueError: If no child ``RigidObject`` is found under the articulation asset.
+
+        Returns:
+            str: The path to the first ``RigidObject`` child under the articulation asset.
+        """
+        child_prims = find_matching_prims(self._asset.cfg.prim_path + "/.*")
+        rigid_child_prim = None
+        # Loop through the list and stop at the first RigidObject found
+        for prim in child_prims:
+            if prim.HasAPI(UsdPhysics.RigidBodyAPI):
+                rigid_child_prim = prim
+                break
+        if rigid_child_prim is None:
+            raise ValueError("No child rigid body found under the expression: '{self._asset.cfg.prim_path}'/.")
+        rigid_child_prim_path = rigid_child_prim.GetPath().pathString
+        # Remove the specific env index from the path string
+        rigid_child_prim_path = self._asset.cfg.prim_path + "/" + rigid_child_prim_path.split("/")[-1]
+        return rigid_child_prim_path
+
+    def _resolve_command_indexes(self):
+        """Resolves the indexes for the various command elements within the command tensor.
+
+        Raises:
+            ValueError: If any command index is left unresolved.
+        """
+        # First iterate over the target types to find the indexes of the different command elements
+        cmd_idx = 0
+        for target_type in self.cfg.controller_cfg.target_types:
+            if target_type == "pose_abs":
+                self._pose_abs_idx = cmd_idx
+                cmd_idx += 7
+            elif target_type == "pose_rel":
+                self._pose_rel_idx = cmd_idx
+                cmd_idx += 6
+            elif target_type == "wrench_abs":
+                self._wrench_abs_idx = cmd_idx
+                cmd_idx += 6
+            else:
+                raise ValueError("Undefined target_type for OSC within OperationalSpaceControllerAction.")
+        # Then iterate over the impedance parameters depending on the impedance mode
+        if (
+            self.cfg.controller_cfg.impedance_mode == "variable_kp"
+            or self.cfg.controller_cfg.impedance_mode == "variable"
+        ):
+            self._stiffness_idx = cmd_idx
+            cmd_idx += 6
+            if self.cfg.controller_cfg.impedance_mode == "variable":
+                self._damping_ratio_idx = cmd_idx
+                cmd_idx += 6
+
+        # Check if any command is left unresolved
+        if self.action_dim != cmd_idx:
+            raise ValueError("Not all command indexes have been resolved.")
+
+    def _compute_dynamic_quantities(self):
+        """Computes the dynamic quantities for operational space control."""
+
+        self._mass_matrix[:] = self._asset.root_physx_view.get_mass_matrices()[:, self._joint_ids, :][
+            :, :, self._joint_ids
+        ]
+        self._gravity[:] = self._asset.root_physx_view.get_generalized_gravity_forces()[:, self._joint_ids]
+
+    def _compute_ee_jacobian(self):
+        """Computes the geometric Jacobian of the ee body frame in root frame.
+
+        This function accounts for the target frame offset and applies the necessary transformations to obtain
+        the right Jacobian from the parent body Jacobian.
+        """
+        # Get the Jacobian in root frame
+        self._jacobian_b[:] = self.jacobian_b
+
+        # account for the offset
+        if self.cfg.body_offset is not None:
+            # Modify the jacobian to account for the offset
+            # -- translational part
+            # v_link = v_ee + w_ee x r_link_ee = v_J_ee * q + w_J_ee * q x r_link_ee
+            #        = (v_J_ee + w_J_ee x r_link_ee ) * q
+            #        = (v_J_ee - r_link_ee_[x] @ w_J_ee) * q
+            self._jacobian_b[:, 0:3, :] += torch.bmm(-math_utils.skew_symmetric_matrix(self._offset_pos), self._jacobian_b[:, 3:, :])  # type: ignore
+            # -- rotational part
+            # w_link = R_link_ee @ w_ee
+            self._jacobian_b[:, 3:, :] = torch.bmm(math_utils.matrix_from_quat(self._offset_rot), self._jacobian_b[:, 3:, :])  # type: ignore
+
+    def _compute_ee_pose(self):
+        """Computes the pose of the ee frame in root frame."""
+        # Obtain quantities from simulation
+        self._ee_pose_w[:] = self._asset.data.body_state_w[:, self._ee_body_idx, :7]
+        # Compute the pose of the ee body in the root frame
+        self._ee_pose_b_no_offset[:, 0:3], self._ee_pose_b_no_offset[:, 3:7] = math_utils.subtract_frame_transforms(
+            self._asset.data.root_pos_w, self._asset.data.root_quat_w, self._ee_pose_w[:, 0:3], self._ee_pose_w[:, 3:7]
+        )
+        # Account for the offset
+        if self.cfg.body_offset is not None:
+            self._ee_pose_b[:, 0:3], self._ee_pose_b[:, 3:7] = math_utils.combine_frame_transforms(
+                self._ee_pose_b_no_offset[:, 0:3], self._ee_pose_b_no_offset[:, 3:7], self._offset_pos, self._offset_rot
+            )
+        else:
+            self._ee_pose_b[:] = self._ee_pose_b_no_offset
+
+    def _compute_ee_velocity(self):
+        """Computes the velocity of the ee frame in root frame."""
+        # Extract end-effector velocity in the world frame
+        self._ee_vel_w[:] = self._asset.data.body_vel_w[:, self._ee_body_idx, :]
+        # Compute the relative velocity in the world frame
+        relative_vel_w = self._ee_vel_w - self._asset.data.root_vel_w
+
+        # Convert ee velocities from world to root frame
+        self._ee_vel_b[:, 0:3] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, relative_vel_w[:, 0:3])
+        self._ee_vel_b[:, 3:6] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, relative_vel_w[:, 3:6])
+
+        # Account for the offset
+        if self.cfg.body_offset is not None:
+            # Compute offset vector in root frame
+            r_offset_b = math_utils.quat_rotate(self._ee_pose_b_no_offset[:, 3:7], self._offset_pos)
+            # Adjust the linear velocity to account for the offset
+            self._ee_vel_b[:, :3] += torch.cross(self._ee_vel_b[:, 3:], r_offset_b, dim=-1)
+            # Angular velocity is not affected by the offset
+
+    def _compute_ee_force(self):
+        """Computes the contact forces on the ee frame in root frame."""
+        # Obtain contact forces only if the contact sensor is available
+        if self._contact_sensor is not None:
+            self._contact_sensor.update(self._sim_dt)
+            self._ee_force_w[:] = self._contact_sensor.data.net_forces_w[:, 0, :]  # type: ignore
+            # Rotate forces and torques into root frame
+            self._ee_force_b[:] = math_utils.quat_rotate_inverse(self._asset.data.root_quat_w, self._ee_force_w)
+
+    def _compute_task_frame_pose(self):
+        """Computes the pose of the task frame in root frame."""
+        # Update task frame pose if task frame rigidbody is provided
+        if self._task_frame_transformer is not None and self._task_frame_pose_b is not None:
+            self._task_frame_transformer.update(self._sim_dt)
+            # Calculate the pose of the task frame in the root frame
+            self._task_frame_pose_b[:, :3], self._task_frame_pose_b[:, 3:] = math_utils.subtract_frame_transforms(
+                self._asset.data.root_pos_w,
+                self._asset.data.root_quat_w,
+                self._task_frame_transformer.data.target_pos_w[:, 0, :],
+                self._task_frame_transformer.data.target_quat_w[:, 0, :],
+            )
+
+    def _preprocess_actions(self, actions: torch.Tensor):
+        """Pre-processes the raw actions for operational space control.
+
+        Args:
+            actions (torch.Tensor): The raw actions for operational space control. It is a tensor of
+                shape (``num_envs``, ``action_dim``).
+        """
+        # Store the raw actions. Please note that the actions contain task space targets
+        # (in the order of the target_types), and possibly the impedance parameters depending on impedance_mode.
+        self._raw_actions[:] = actions
+        # Initialize the processed actions with raw actions.
+        self._processed_actions[:] = self._raw_actions
+        # Go through the command types one by one, and apply the pre-processing if needed.
+        if self._pose_abs_idx is not None:
+            self._processed_actions[:, self._pose_abs_idx : self._pose_abs_idx + 3] *= self._position_scale
+            self._processed_actions[:, self._pose_abs_idx + 3 : self._pose_abs_idx + 7] *= self._orientation_scale
+        if self._pose_rel_idx is not None:
+            self._processed_actions[:, self._pose_rel_idx : self._pose_rel_idx + 3] *= self._position_scale
+            self._processed_actions[:, self._pose_rel_idx + 3 : self._pose_rel_idx + 6] *= self._orientation_scale
+        if self._wrench_abs_idx is not None:
+            self._processed_actions[:, self._wrench_abs_idx : self._wrench_abs_idx + 6] *= self._wrench_scale
+        if self._stiffness_idx is not None:
+            self._processed_actions[:, self._stiffness_idx : self._stiffness_idx + 6] *= self._stiffness_scale
+            self._processed_actions[:, self._stiffness_idx : self._stiffness_idx + 6] = torch.clamp(
+                self._processed_actions[:, self._stiffness_idx : self._stiffness_idx + 6],
+                min=self.cfg.controller_cfg.motion_stiffness_limits_task[0],
+                max=self.cfg.controller_cfg.motion_stiffness_limits_task[1],
+            )
+        if self._damping_ratio_idx is not None:
+            self._processed_actions[
+                :, self._damping_ratio_idx : self._damping_ratio_idx + 6
+            ] *= self._damping_ratio_scale
+            self._processed_actions[:, self._damping_ratio_idx : self._damping_ratio_idx + 6] = torch.clamp(
+                self._processed_actions[:, self._damping_ratio_idx : self._damping_ratio_idx + 6],
+                min=self.cfg.controller_cfg.motion_damping_ratio_limits_task[0],
+                max=self.cfg.controller_cfg.motion_damping_ratio_limits_task[1],
+            )

source/extensions/omni.isaac.lab/test/controllers/test_operational_space.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Launch Isaac Sim Simulator first."""
+
+from omni.isaac.lab.app import AppLauncher, run_tests
+
+# launch omniverse app
+simulation_app = AppLauncher(headless=True).app
+
+"""Rest everything follows."""
+
+import torch
+import unittest
+
+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.lab.sim as sim_utils
+from omni.isaac.lab.assets import Articulation
+from omni.isaac.lab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg
+from omni.isaac.lab.markers import VisualizationMarkers
+from omni.isaac.lab.markers.config import FRAME_MARKER_CFG
+from omni.isaac.lab.sensors import ContactSensor, ContactSensorCfg
+from omni.isaac.lab.utils.math import (
+    apply_delta_pose,
+    combine_frame_transforms,
+    compute_pose_error,
+    matrix_from_quat,
+    quat_inv,
+    quat_rotate_inverse,
+    subtract_frame_transforms,
+)
+
+##
+# Pre-defined configs
+##
+from omni.isaac.lab_assets import FRANKA_PANDA_CFG  # isort:skip
+
+
+class TestOperationalSpaceController(unittest.TestCase):
+    """Test fixture for checking that Operational Space controller tracks commands properly."""
+
+    def setUp(self):
+        """Create a blank new stage for each test."""
+        # Wait for spawning
+        stage_utils.create_new_stage()
+        # Constants
+        self.num_envs = 16
+        # Load kit helper
+        sim_cfg = sim_utils.SimulationCfg(dt=0.01)
+        self.sim = sim_utils.SimulationContext(sim_cfg)
+        # TODO: Remove this once we have a better way to handle this.
+        self.sim._app_control_on_stop_handle = None
+
+        # Create a ground plane
+        cfg = sim_utils.GroundPlaneCfg()
+        cfg.func("/World/GroundPlane", cfg)
+
+        # Markers
+        frame_marker_cfg = FRAME_MARKER_CFG.copy()
+        frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
+        self.ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current"))
+        self.goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal"))
+
+        light_cfg = sim_utils.DistantLightCfg(intensity=5.0, exposure=10.0)
+        light_cfg.func(
+            "/Light",
+            light_cfg,
+            translation=[0, 0, 1],
+        )
+
+        # Create interface to clone the scene
+        cloner = GridCloner(spacing=2.0)
+        cloner.define_base_env("/World/envs")
+        self.env_prim_paths = cloner.generate_paths("/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,
+        )
+
+        self.robot_cfg = FRANKA_PANDA_CFG.replace(prim_path="/World/envs/env_.*/Robot")
+        self.robot_cfg.actuators["panda_shoulder"].stiffness = 0.0
+        self.robot_cfg.actuators["panda_shoulder"].damping = 0.0
+        self.robot_cfg.actuators["panda_forearm"].stiffness = 0.0
+        self.robot_cfg.actuators["panda_forearm"].damping = 0.0
+        self.robot_cfg.spawn.rigid_props.disable_gravity = True
+
+        # Define the ContactSensor
+        self.contact_forces = None
+
+        # Define the target sets
+        ee_goal_abs_pos_set_b = torch.tensor(
+            [
+                [0.5, 0.5, 0.7],
+                [0.5, -0.4, 0.6],
+                [0.5, 0, 0.5],
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_abs_quad_set_b = torch.tensor(
+            [
+                [0.707, 0.0, 0.707, 0.0],
+                [0.707, 0.707, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0],
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_rel_pos_set = torch.tensor(
+            [
+                [0.2, 0.0, 0.0],
+                [0.2, 0.2, 0.0],
+                [0.2, 0.2, -0.2],
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_rel_axisangle_set = torch.tensor(
+            [
+                [0.0, torch.pi / 2, 0.0],  # for [0.707, 0, 0.707, 0]
+                [torch.pi / 2, 0.0, 0.0],  # for [0.707, 0.707, 0, 0]
+                [torch.pi, 0.0, 0.0],  # for [0.0, 1.0, 0, 0]
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_abs_wrench_set_b = torch.tensor(
+            [
+                [0.0, 0.0, 10.0, 0.0, -1.0, 0.0],
+                [0.0, 10.0, 0.0, 0.0, 0.0, 0.0],
+                [10.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            ],
+            device=self.sim.device,
+        )
+        kp_set = torch.tensor(
+            [
+                [200.0, 200.0, 200.0, 200.0, 200.0, 200.0],
+                [240.0, 240.0, 240.0, 240.0, 240.0, 240.0],
+                [160.0, 160.0, 160.0, 160.0, 160.0, 160.0],
+            ],
+            device=self.sim.device,
+        )
+        d_ratio_set = torch.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+                [1.1, 1.1, 1.1, 1.1, 1.1, 1.1],
+                [0.9, 0.9, 0.9, 0.9, 0.9, 0.9],
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_hybrid_set_b = torch.tensor(
+            [
+                [0.6, 0.2, 0.5, 0.0, 0.707, 0.0, 0.707, 10.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                [0.6, -0.29, 0.6, 0.0, 0.707, 0.0, 0.707, 10.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                [0.6, 0.1, 0.8, 0.0, 0.5774, 0.0, 0.8165, 4.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_pose_set_tilted_b = torch.tensor(
+            [
+                [0.6, 0.15, 0.3, 0.0, 0.92387953, 0.0, 0.38268343],
+                [0.6, -0.3, 0.3, 0.0, 0.92387953, 0.0, 0.38268343],
+                [0.8, 0.0, 0.5, 0.0, 0.92387953, 0.0, 0.38268343],
+            ],
+            device=self.sim.device,
+        )
+        ee_goal_wrench_set_tilted_task = torch.tensor(
+            [
+                [0.0, 0.0, 10.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 10.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 10.0, 0.0, 0.0, 0.0],
+            ],
+            device=self.sim.device,
+        )
+
+        # Define goals for the arm [xyz]
+        self.target_abs_pos_set_b = ee_goal_abs_pos_set_b.clone()
+        # Define goals for the arm [xyz + quat_wxyz]
+        self.target_abs_pose_set_b = torch.cat([ee_goal_abs_pos_set_b, ee_goal_abs_quad_set_b], dim=-1)
+        # Define goals for the arm [xyz]
+        self.target_rel_pos_set = ee_goal_rel_pos_set.clone()
+        # Define goals for the arm [xyz + axis-angle]
+        self.target_rel_pose_set_b = torch.cat([ee_goal_rel_pos_set, ee_goal_rel_axisangle_set], dim=-1)
+        # Define goals for the arm [force_xyz + torque_xyz]
+        self.target_abs_wrench_set = ee_goal_abs_wrench_set_b.clone()
+        # Define goals for the arm [xyz + quat_wxyz] and variable kp [kp_xyz + kp_rot_xyz]
+        self.target_abs_pose_variable_kp_set = torch.cat([self.target_abs_pose_set_b, kp_set], dim=-1)
+        # Define goals for the arm [xyz + quat_wxyz] and the variable imp. [kp_xyz + kp_rot_xyz + d_xyz + d_rot_xyz]
+        self.target_abs_pose_variable_set = torch.cat([self.target_abs_pose_set_b, kp_set, d_ratio_set], dim=-1)
+        # Define goals for the arm pose [xyz + quat_wxyz] and wrench [force_xyz + torque_xyz]
+        self.target_hybrid_set_b = ee_goal_hybrid_set_b.clone()
+        # Define goals for the arm pose, and wrench, and kp
+        self.target_hybrid_variable_kp_set = torch.cat([self.target_hybrid_set_b, kp_set], dim=-1)
+        # Define goals for the arm pose [xyz + quat_wxyz] in root and and wrench [force_xyz + torque_xyz] in task frame
+        self.target_hybrid_set_tilted = torch.cat([ee_goal_pose_set_tilted_b, ee_goal_wrench_set_tilted_task], dim=-1)
+
+        # Reference frame for targets
+        self.frame = "root"
+
+    def tearDown(self):
+        """Stops simulator after each test."""
+        # stop simulation
+        self.sim.stop()
+        # self.sim.clear()  # FIXME: This hangs the test for some reason when LIVESTREAM is not enabled.
+        self.sim.clear_all_callbacks()
+        self.sim.clear_instance()
+        # Make contact_forces None after relevant tests otherwise other tests give warning
+        self.contact_forces = None
+
+    """
+    Test fixtures.
+    """
+
+    def test_franka_pose_abs_without_inertial_decoupling(self):
+        """Test absolute pose control with fixed impedance and without inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_stiffness_task=[400.0, 400.0, 400.0, 100.0, 100.0, 100.0],
+            motion_damping_ratio_task=[5.0, 5.0, 5.0, 0.001, 0.001, 0.001],
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b)
+
+    def test_franka_pose_abs_with_partial_inertial_decoupling(self):
+        """Test absolute pose control with fixed impedance and partial inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=True,
+            gravity_compensation=False,
+            motion_stiffness_task=1000.0,
+            motion_damping_ratio_task=1.0,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b)
+
+    def test_franka_pose_abs_fixed_impedance_with_gravity_compensation(self):
+        """Test absolute pose control with fixed impedance, gravity compensation, and inertial dynamics decoupling."""
+        self.robot_cfg.spawn.rigid_props.disable_gravity = False
+        robot = Articulation(cfg=self.robot_cfg)
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=True,
+            motion_stiffness_task=500.0,
+            motion_damping_ratio_task=1.0,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b)
+
+    def test_franka_pose_abs(self):
+        """Test absolute pose control with fixed impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_stiffness_task=500.0,
+            motion_damping_ratio_task=1.0,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b)
+
+    def test_franka_pose_rel(self):
+        """Test relative pose control with fixed impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_rel"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_stiffness_task=500.0,
+            motion_damping_ratio_task=1.0,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_rel_pose_set_b)
+
+    def test_franka_pose_abs_variable_impedance(self):
+        """Test absolute pose control with variable impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs"],
+            impedance_mode="variable",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_variable_set)
+
+    def test_franka_wrench_abs_open_loop(self):
+        """Test open loop absolute force control."""
+        robot = Articulation(cfg=self.robot_cfg)
+
+        obstacle_spawn_cfg = sim_utils.CuboidCfg(
+            size=(0.7, 0.7, 0.01),
+            collision_props=sim_utils.CollisionPropertiesCfg(),
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+            activate_contact_sensors=True,
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle1",
+            obstacle_spawn_cfg,
+            translation=(0.2, 0.0, 0.93),
+            orientation=(0.9848, 0.0, -0.1736, 0.0),
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle2",
+            obstacle_spawn_cfg,
+            translation=(0.2, 0.35, 0.7),
+            orientation=(0.707, 0.707, 0.0, 0.0),
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle3",
+            obstacle_spawn_cfg,
+            translation=(0.55, 0.0, 0.7),
+            orientation=(0.707, 0.0, 0.707, 0.0),
+        )
+        contact_forces_cfg = ContactSensorCfg(
+            prim_path="/World/envs/env_.*/obstacle.*",
+            update_period=0.0,
+            history_length=50,
+            debug_vis=False,
+            force_threshold=0.1,
+        )
+        self.contact_forces = ContactSensor(contact_forces_cfg)
+
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["wrench_abs"],
+            motion_control_axes_task=[0, 0, 0, 0, 0, 0],
+            contact_wrench_control_axes_task=[1, 1, 1, 1, 1, 1],
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_wrench_set)
+
+    def test_franka_wrench_abs_closed_loop(self):
+        """Test closed loop absolute force control."""
+        robot = Articulation(cfg=self.robot_cfg)
+
+        obstacle_spawn_cfg = sim_utils.CuboidCfg(
+            size=(0.7, 0.7, 0.01),
+            collision_props=sim_utils.CollisionPropertiesCfg(),
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+            activate_contact_sensors=True,
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle1",
+            obstacle_spawn_cfg,
+            translation=(0.2, 0.0, 0.93),
+            orientation=(0.9848, 0.0, -0.1736, 0.0),
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle2",
+            obstacle_spawn_cfg,
+            translation=(0.2, 0.35, 0.7),
+            orientation=(0.707, 0.707, 0.0, 0.0),
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle3",
+            obstacle_spawn_cfg,
+            translation=(0.55, 0.0, 0.7),
+            orientation=(0.707, 0.0, 0.707, 0.0),
+        )
+        contact_forces_cfg = ContactSensorCfg(
+            prim_path="/World/envs/env_.*/obstacle.*",
+            update_period=0.0,
+            history_length=2,
+            debug_vis=False,
+            force_threshold=0.1,
+        )
+        self.contact_forces = ContactSensor(contact_forces_cfg)
+
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["wrench_abs"],
+            contact_wrench_stiffness_task=[
+                0.2,
+                0.2,
+                0.2,
+                0.0,
+                0.0,
+                0.0,
+            ],  # Zero torque feedback as we cannot contact torque
+            motion_control_axes_task=[0, 0, 0, 0, 0, 0],
+            contact_wrench_control_axes_task=[1, 1, 1, 1, 1, 1],
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_wrench_set)
+
+    def test_franka_hybrid_decoupled_motion(self):
+        """Test hybrid control with fixed impedance and partial inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+
+        obstacle_spawn_cfg = sim_utils.CuboidCfg(
+            size=(1.0, 1.0, 0.01),
+            collision_props=sim_utils.CollisionPropertiesCfg(),
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+            activate_contact_sensors=True,
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle1",
+            obstacle_spawn_cfg,
+            translation=(self.target_hybrid_set_b[0, 0] + 0.05, 0.0, 0.7),
+            orientation=(0.707, 0.0, 0.707, 0.0),
+        )
+        contact_forces_cfg = ContactSensorCfg(
+            prim_path="/World/envs/env_.*/obstacle.*",
+            update_period=0.0,
+            history_length=2,
+            debug_vis=False,
+            force_threshold=0.1,
+        )
+        self.contact_forces = ContactSensor(contact_forces_cfg)
+
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs", "wrench_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=True,
+            gravity_compensation=False,
+            motion_stiffness_task=300.0,
+            motion_damping_ratio_task=1.0,
+            contact_wrench_stiffness_task=[0.1, 0.0, 0.0, 0.0, 0.0, 0.0],
+            motion_control_axes_task=[0, 1, 1, 1, 1, 1],
+            contact_wrench_control_axes_task=[1, 0, 0, 0, 0, 0],
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_leftfinger", ["panda_joint.*"], self.target_hybrid_set_b)
+
+    def test_franka_hybrid_variable_kp_impedance(self):
+        """Test hybrid control with variable kp impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+
+        obstacle_spawn_cfg = sim_utils.CuboidCfg(
+            size=(1.0, 1.0, 0.01),
+            collision_props=sim_utils.CollisionPropertiesCfg(),
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+            activate_contact_sensors=True,
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle1",
+            obstacle_spawn_cfg,
+            translation=(self.target_hybrid_set_b[0, 0] + 0.05, 0.0, 0.7),
+            orientation=(0.707, 0.0, 0.707, 0.0),
+        )
+        contact_forces_cfg = ContactSensorCfg(
+            prim_path="/World/envs/env_.*/obstacle.*",
+            update_period=0.0,
+            history_length=2,
+            debug_vis=False,
+            force_threshold=0.1,
+        )
+        self.contact_forces = ContactSensor(contact_forces_cfg)
+
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs", "wrench_abs"],
+            impedance_mode="variable_kp",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_damping_ratio_task=0.8,
+            contact_wrench_stiffness_task=[0.1, 0.0, 0.0, 0.0, 0.0, 0.0],
+            motion_control_axes_task=[0, 1, 1, 1, 1, 1],
+            contact_wrench_control_axes_task=[1, 0, 0, 0, 0, 0],
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(
+            robot, osc, "panda_leftfinger", ["panda_joint.*"], self.target_hybrid_variable_kp_set
+        )
+
+    def test_franka_taskframe_pose_abs(self):
+        """Test absolute pose control in task frame with fixed impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        self.frame = "task"
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_stiffness_task=500.0,
+            motion_damping_ratio_task=1.0,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_abs_pose_set_b)
+
+    def test_franka_taskframe_pose_rel(self):
+        """Test relative pose control in task frame with fixed impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        self.frame = "task"
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_rel"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_stiffness_task=500.0,
+            motion_damping_ratio_task=1.0,
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_hand", ["panda_joint.*"], self.target_rel_pose_set_b)
+
+    def test_franka_taskframe_hybrid(self):
+        """Test hybrid control in task frame with fixed impedance and inertial dynamics decoupling."""
+        robot = Articulation(cfg=self.robot_cfg)
+        self.frame = "task"
+
+        obstacle_spawn_cfg = sim_utils.CuboidCfg(
+            size=(2.0, 1.5, 0.01),
+            collision_props=sim_utils.CollisionPropertiesCfg(),
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+            activate_contact_sensors=True,
+        )
+        obstacle_spawn_cfg.func(
+            "/World/envs/env_.*/obstacle1",
+            obstacle_spawn_cfg,
+            translation=(self.target_hybrid_set_tilted[0, 0] + 0.085, 0.0, 0.3),
+            orientation=(0.9238795325, 0.0, -0.3826834324, 0.0),
+        )
+        contact_forces_cfg = ContactSensorCfg(
+            prim_path="/World/envs/env_.*/obstacle.*",
+            update_period=0.0,
+            history_length=2,
+            debug_vis=False,
+            force_threshold=0.1,
+        )
+        self.contact_forces = ContactSensor(contact_forces_cfg)
+
+        osc_cfg = OperationalSpaceControllerCfg(
+            target_types=["pose_abs", "wrench_abs"],
+            impedance_mode="fixed",
+            inertial_dynamics_decoupling=True,
+            partial_inertial_dynamics_decoupling=False,
+            gravity_compensation=False,
+            motion_stiffness_task=400.0,
+            motion_damping_ratio_task=1.0,
+            contact_wrench_stiffness_task=[0.0, 0.0, 0.1, 0.0, 0.0, 0.0],
+            motion_control_axes_task=[1, 1, 0, 1, 1, 1],
+            contact_wrench_control_axes_task=[0, 0, 1, 0, 0, 0],
+        )
+        osc = OperationalSpaceController(osc_cfg, num_envs=self.num_envs, device=self.sim.device)
+
+        self._run_op_space_controller(robot, osc, "panda_leftfinger", ["panda_joint.*"], self.target_hybrid_set_tilted)
+
+    """
+    Helper functions
+    """
+
+    def _run_op_space_controller(
+        self,
+        robot: Articulation,
+        osc: OperationalSpaceController,
+        ee_frame_name: str,
+        arm_joint_names: list[str],
+        target_set: torch.tensor,
+    ):
+        """Run the operational space controller with the given parameters.
+
+        Args:
+            robot (Articulation): The robot to control.
+            osc (OperationalSpaceController): The operational space controller.
+            ee_frame_name (str): The name of the end-effector frame.
+            arm_joint_names (list[str]): The names of the arm joints.
+            target_set (torch.tensor): The target set to track.
+        """
+        # Initialize the masks for evaluating target convergence according to selection matrices
+        self.pos_mask = torch.tensor(osc.cfg.motion_control_axes_task[:3], device=self.sim.device).view(1, 3)
+        self.rot_mask = torch.tensor(osc.cfg.motion_control_axes_task[3:], device=self.sim.device).view(1, 3)
+        self.wrench_mask = torch.tensor(osc.cfg.contact_wrench_control_axes_task, device=self.sim.device).view(1, 6)
+        self.force_mask = self.wrench_mask[:, 0:3]  # Take only the force components as we can measure only these
+
+        # 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]
+        # 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)
+
+        # get the updated states
+        jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = self._update_states(
+            robot, ee_frame_idx, arm_joint_ids
+        )
+
+        # Track the given target command
+        current_goal_idx = 0  # Current goal index for the arm
+        command = torch.zeros(
+            self.num_envs, osc.action_dim, device=self.sim.device
+        )  # Generic target command, which can be pose, position, force, etc.
+        ee_target_pose_b = torch.zeros(self.num_envs, 7, device=self.sim.device)  # Target pose in the body frame
+        ee_target_pose_w = torch.zeros(
+            self.num_envs, 7, device=self.sim.device
+        )  # Target pose in the world frame (for marker)
+
+        # Set joint efforts to zero
+        zero_joint_efforts = torch.zeros(self.num_envs, robot.num_joints, device=self.sim.device)
+        joint_efforts = torch.zeros(self.num_envs, len(arm_joint_ids), device=self.sim.device)
+
+        # Now we are ready!
+        for count in range(1501):
+            # reset every 500 steps
+            if count % 500 == 0:
+                # check that we converged to the goal
+                if count > 0:
+                    self._check_convergence(osc, ee_pose_b, ee_target_pose_b, ee_force_b, command)
+                # reset joint state to default
+                default_joint_pos = robot.data.default_joint_pos.clone()
+                default_joint_vel = robot.data.default_joint_vel.clone()
+                robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel)
+                robot.set_joint_effort_target(zero_joint_efforts)  # Set zero torques in the initial step
+                robot.write_data_to_sim()
+                robot.reset()
+                # reset contact sensor
+                if self.contact_forces is not None:
+                    self.contact_forces.reset()
+                # reset target pose
+                robot.update(sim_dt)
+                _, _, _, ee_pose_b, _, _, _, _ = self._update_states(
+                    robot, ee_frame_idx, arm_joint_ids
+                )  # at reset, the jacobians are not updated to the latest state
+                command, ee_target_pose_b, ee_target_pose_w, current_goal_idx = self._update_target(
+                    osc, root_pose_w, ee_pose_b, target_set, current_goal_idx
+                )
+                # set the osc command
+                osc.reset()
+                command, task_frame_pose_b = self._convert_to_task_frame(
+                    osc, command=command, ee_target_pose_b=ee_target_pose_b
+                )
+                osc.set_command(
+                    command=command, current_ee_pose_b=ee_pose_b, current_task_frame_pose_b=task_frame_pose_b
+                )
+            else:
+                # get the updated states
+                jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = (
+                    self._update_states(robot, ee_frame_idx, arm_joint_ids)
+                )
+                # compute the joint commands
+                joint_efforts = osc.compute(
+                    jacobian_b=jacobian_b,
+                    current_ee_pose_b=ee_pose_b,
+                    current_ee_vel_b=ee_vel_b,
+                    current_ee_force_b=ee_force_b,
+                    mass_matrix=mass_matrix,
+                    gravity=gravity,
+                )
+                robot.set_joint_effort_target(joint_efforts, joint_ids=arm_joint_ids)
+                robot.write_data_to_sim()
+
+            # update marker positions
+            self.ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])
+            self.goal_marker.visualize(ee_target_pose_w[:, 0:3], ee_target_pose_w[:, 3:7])
+
+            # perform step
+            self.sim.step(render=False)
+            # update buffers
+            robot.update(sim_dt)
+
+    def _update_states(
+        self,
+        robot: Articulation,
+        ee_frame_idx: int,
+        arm_joint_ids: list[int],
+    ):
+        """Update the states of the robot and obtain the relevant quantities for the operational space controller.
+
+        Args:
+            robot (Articulation): The robot to control.
+            ee_frame_idx (int): The index of the end-effector frame.
+            arm_joint_ids (list[int]): The indices of the arm joints.
+
+        Returns:
+            jacobian_b (torch.tensor): The Jacobian in the root frame.
+            mass_matrix (torch.tensor): The mass matrix.
+            gravity (torch.tensor): The gravity vector.
+            ee_pose_b (torch.tensor): The end-effector pose in the root frame.
+            ee_vel_b (torch.tensor): The end-effector velocity in the root frame.
+            root_pose_w (torch.tensor): The root pose in the world frame.
+            ee_pose_w (torch.tensor): The end-effector pose in the world frame.
+            ee_force_b (torch.tensor): The end-effector force in the root frame.
+        """
+        # obtain dynamics related quantities from simulation
+        ee_jacobi_idx = ee_frame_idx - 1
+        jacobian_w = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids]
+        mass_matrix = robot.root_physx_view.get_mass_matrices()[:, arm_joint_ids, :][:, :, arm_joint_ids]
+        gravity = robot.root_physx_view.get_generalized_gravity_forces()[:, arm_joint_ids]
+        # Convert the Jacobian from world to root frame
+        jacobian_b = jacobian_w.clone()
+        root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_state_w[:, 3:7]))
+        jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :])
+        jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :])
+
+        # Compute current pose of the end-effector
+        root_pose_w = robot.data.root_state_w[:, 0:7]
+        ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 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]
+        )
+        ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1)
+
+        # Compute the current velocity of the end-effector
+        ee_vel_w = robot.data.body_vel_w[:, ee_frame_idx, :]  # Extract end-effector velocity in the world frame
+        root_vel_w = robot.data.root_vel_w  # Extract root velocity in the world frame
+        relative_vel_w = ee_vel_w - root_vel_w  # Compute the relative velocity in the world frame
+        ee_lin_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 0:3])  # From world to root frame
+        ee_ang_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 3:6])
+        ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1)
+
+        # Calculate the contact force
+        ee_force_w = torch.zeros(self.num_envs, 3, device=self.sim.device)
+        if self.contact_forces is not None:  # Only modify if it exist
+            sim_dt = self.sim.get_physics_dt()
+            self.contact_forces.update(sim_dt)  # update contact sensor
+            # Calculate the contact force by averaging over last four time steps (i.e., to smoothen) and
+            # taking the max of three surfaces as only one should be the contact of interest
+            ee_force_w, _ = torch.max(torch.mean(self.contact_forces.data.net_forces_w_history, dim=1), dim=1)
+
+        # This is a simplification, only for the sake of testing.
+        ee_force_b = ee_force_w
+
+        return jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b
+
+    def _update_target(
+        self,
+        osc: OperationalSpaceController,
+        root_pose_w: torch.tensor,
+        ee_pose_b: torch.tensor,
+        target_set: torch.tensor,
+        current_goal_idx: int,
+    ):
+        """Update the target for the operational space controller.
+
+        Args:
+            osc (OperationalSpaceController): The operational space controller.
+            root_pose_w (torch.tensor): The root pose in the world frame.
+            ee_pose_b (torch.tensor): The end-effector pose in the body frame.
+            target_set (torch.tensor): The target set to track.
+            current_goal_idx (int): The current goal index.
+
+        Returns:
+            command (torch.tensor): The target command.
+            ee_target_pose_b (torch.tensor): The end-effector target pose in the body frame.
+            ee_target_pose_w (torch.tensor): The end-effector target pose in the world frame.
+            next_goal_idx (int): The next goal index.
+
+        Raises:
+            ValueError: If the target type is undefined.
+        """
+        # update the ee desired command
+        command = torch.zeros(self.num_envs, osc.action_dim, device=self.sim.device)
+        command[:] = target_set[current_goal_idx]
+
+        # update the ee desired pose
+        ee_target_pose_b = torch.zeros(self.num_envs, 7, device=self.sim.device)
+        for target_type in osc.cfg.target_types:
+            if target_type == "pose_abs":
+                ee_target_pose_b[:] = command[:, :7]
+            elif target_type == "pose_rel":
+                ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7] = apply_delta_pose(
+                    ee_pose_b[:, :3], ee_pose_b[:, 3:], command[:, :7]
+                )
+            elif target_type == "wrench_abs":
+                pass  # ee_target_pose_b could stay at the root frame for force control, what matters is ee_target_b
+            else:
+                raise ValueError("Undefined target_type within _update_target().")
+
+        # update the target desired pose in world frame (for marker)
+        ee_target_pos_w, ee_target_quat_w = combine_frame_transforms(
+            root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7]
+        )
+        ee_target_pose_w = torch.cat([ee_target_pos_w, ee_target_quat_w], dim=-1)
+
+        next_goal_idx = (current_goal_idx + 1) % len(target_set)
+
+        return command, ee_target_pose_b, ee_target_pose_w, next_goal_idx
+
+    def _convert_to_task_frame(
+        self, osc: OperationalSpaceController, command: torch.tensor, ee_target_pose_b: torch.tensor
+    ):
+        """Convert the target command to the task frame if required.
+
+        Args:
+            osc (OperationalSpaceController): The operational space controller.
+            command (torch.tensor): The target command to convert.
+            ee_target_pose_b (torch.tensor): The end-effector target pose in the body frame.
+
+        Returns:
+            command (torch.tensor): The converted target command.
+            task_frame_pose_b (torch.tensor): The task frame pose in the body frame.
+
+        Raises:
+            ValueError: If the frame is invalid.
+        """
+        command = command.clone()
+        task_frame_pose_b = None
+        if self.frame == "root":
+            # No need to transform anything if they are already in root frame
+            pass
+        elif self.frame == "task":
+            # Convert target commands from base to the task frame
+            command = command.clone()
+            task_frame_pose_b = ee_target_pose_b.clone()
+
+            cmd_idx = 0
+            for target_type in osc.cfg.target_types:
+                if target_type == "pose_abs":
+                    command[:, :3], command[:, 3:7] = subtract_frame_transforms(
+                        task_frame_pose_b[:, :3], task_frame_pose_b[:, 3:], command[:, :3], command[:, 3:7]
+                    )
+                    cmd_idx += 7
+                elif target_type == "pose_rel":
+                    # Compute rotation matrices
+                    R_task_b = matrix_from_quat(task_frame_pose_b[:, 3:])  # Task frame to base frame
+                    R_b_task = R_task_b.mT  # Base frame to task frame
+                    # Transform the delta position and orientation from base to task frame
+                    command[:, :3] = (R_b_task @ command[:, :3].unsqueeze(-1)).squeeze(-1)
+                    command[:, 3:7] = (R_b_task @ command[:, 3:7].unsqueeze(-1)).squeeze(-1)
+                    cmd_idx += 6
+                elif target_type == "wrench_abs":
+                    # These are already defined in target frame for ee_goal_wrench_set_tilted_task (since it is
+                    # easier), so not transforming
+                    cmd_idx += 6
+                else:
+                    raise ValueError("Undefined target_type within _convert_to_task_frame().")
+        else:
+            # Raise error for invalid frame
+            raise ValueError("Invalid frame selection for target setting inside the test_operational_space.")
+
+        return command, task_frame_pose_b
+
+    def _check_convergence(
+        self,
+        osc: OperationalSpaceController,
+        ee_pose_b: torch.tensor,
+        ee_target_pose_b: torch.tensor,
+        ee_force_b: torch.tensor,
+        ee_target_b: torch.tensor,
+    ):
+        """Check the convergence to the target.
+
+        Args:
+            osc (OperationalSpaceController): The operational space controller.
+            ee_pose_b (torch.tensor): The end-effector pose in the body frame.
+            ee_target_pose_b (torch.tensor): The end-effector target pose in the body frame.
+            ee_force_b (torch.tensor): The end-effector force in the body frame.
+            ee_target_b (torch.tensor): The end-effector target in the body frame.
+
+        Raises:
+            AssertionError: If the convergence is not achieved.
+            ValueError: If the target type is undefined.
+        """
+        cmd_idx = 0
+        for target_type in osc.cfg.target_types:
+            if target_type == "pose_abs":
+                pos_error, rot_error = compute_pose_error(
+                    ee_pose_b[:, 0:3], ee_pose_b[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7]
+                )
+                pos_error_norm = torch.norm(pos_error * self.pos_mask, dim=-1)
+                rot_error_norm = torch.norm(rot_error * self.rot_mask, 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=0.1)
+                torch.testing.assert_close(rot_error_norm, des_error, rtol=0.0, atol=0.1)
+                cmd_idx += 7
+            elif target_type == "pose_rel":
+                pos_error, rot_error = compute_pose_error(
+                    ee_pose_b[:, 0:3], ee_pose_b[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7]
+                )
+                pos_error_norm = torch.norm(pos_error * self.pos_mask, dim=-1)
+                rot_error_norm = torch.norm(rot_error * self.rot_mask, 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=0.1)
+                torch.testing.assert_close(rot_error_norm, des_error, rtol=0.0, atol=0.1)
+                cmd_idx += 6
+            elif target_type == "wrench_abs":
+                force_target_b = ee_target_b[:, cmd_idx : cmd_idx + 3].clone()
+                # Convert to base frame if the target was defined in task frame
+                if self.frame == "task":
+                    task_frame_pose_b = ee_target_pose_b.clone()
+                    R_task_b = matrix_from_quat(task_frame_pose_b[:, 3:])
+                    force_target_b[:] = (R_task_b @ force_target_b[:].unsqueeze(-1)).squeeze(-1)
+                force_error = ee_force_b - force_target_b
+                force_error_norm = torch.norm(
+                    force_error * self.force_mask, dim=-1
+                )  # ignore torque part as we cannot measure it
+                des_error = torch.zeros_like(force_error_norm)
+                # check convergence: big threshold here as the force control is not precise when the robot moves
+                torch.testing.assert_close(force_error_norm, des_error, rtol=0.0, atol=1.0)
+                cmd_idx += 6
+            else:
+                raise ValueError("Undefined target_type within _check_convergence().")
+
+
+if __name__ == "__main__":
+    run_tests()

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.10.14"
+version = "0.10.15"
 
 # Description
 title = "Isaac Lab Environments"

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

 Changelog
 ---------
 
+0.10.15 (2024-12-16)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+* Added ``Isaac-Reach-Franka-OSC-v0`` and ``Isaac-Reach-Franka-OSC-Play-v0``
+  variations of the manager based reach environment that uses
+  :class:`omni.isaac.lab.envs.mdp.actions.OperationalSpaceControllerAction`.
+
+
 0.10.14 (2024-12-03)
 ~~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/__init__.py

@@ -65,3 +65,27 @@ gym.register(
     },
     disable_env_checker=True,
 )
+
+##
+# Operational Space Control
+##
+
+gym.register(
+    id="Isaac-Reach-Franka-OSC-v0",
+    entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.osc_env_cfg:FrankaReachEnvCfg",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg",
+    },
+)
+
+gym.register(
+    id="Isaac-Reach-Franka-OSC-Play-v0",
+    entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": f"{__name__}.osc_env_cfg:FrankaReachEnvCfg_PLAY",
+        "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaReachPPORunnerCfg",
+    },
+)

source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/manipulation/reach/config/franka/osc_env_cfg.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from omni.isaac.lab.controllers.operational_space_cfg import OperationalSpaceControllerCfg
+from omni.isaac.lab.envs.mdp.actions.actions_cfg import OperationalSpaceControllerActionCfg
+from omni.isaac.lab.utils import configclass
+
+from . import joint_pos_env_cfg
+
+##
+# Pre-defined configs
+##
+from omni.isaac.lab_assets.franka import FRANKA_PANDA_CFG  # isort: skip
+
+
+@configclass
+class FrankaReachEnvCfg(joint_pos_env_cfg.FrankaReachEnvCfg):
+    def __post_init__(self):
+        # post init of parent
+        super().__post_init__()
+
+        # Set Franka as robot
+        # We remove stiffness and damping for the shoulder and forearm joints for effort control
+        self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+        self.scene.robot.actuators["panda_shoulder"].stiffness = 0.0
+        self.scene.robot.actuators["panda_shoulder"].damping = 0.0
+        self.scene.robot.actuators["panda_forearm"].stiffness = 0.0
+        self.scene.robot.actuators["panda_forearm"].damping = 0.0
+        self.scene.robot.spawn.rigid_props.disable_gravity = True
+
+        # If closed-loop contact force control is desired, contact sensors should be enabled for the robot
+        # self.scene.robot.spawn.activate_contact_sensors = True
+
+        self.actions.arm_action = OperationalSpaceControllerActionCfg(
+            asset_name="robot",
+            joint_names=["panda_joint.*"],
+            body_name="panda_hand",
+            # If a task frame different from articulation root/base is desired, a RigidObject, e.g., "task_frame",
+            # can be added to the scene and its relative path could provided as task_frame_rel_path
+            # task_frame_rel_path="task_frame",
+            controller_cfg=OperationalSpaceControllerCfg(
+                target_types=["pose_abs"],
+                impedance_mode="variable_kp",
+                inertial_dynamics_decoupling=True,
+                partial_inertial_dynamics_decoupling=False,
+                gravity_compensation=False,
+                motion_stiffness_task=100.0,
+                motion_damping_ratio_task=1.0,
+                motion_stiffness_limits_task=(50.0, 200.0),
+            ),
+            position_scale=1.0,
+            orientation_scale=1.0,
+            stiffness_scale=100.0,
+        )
+        # Removing these observations as they are not needed for OSC and we want keep the observation space small
+        self.observations.policy.joint_pos = None
+        self.observations.policy.joint_vel = None
+
+
+@configclass
+class FrankaReachEnvCfg_PLAY(FrankaReachEnvCfg):
+    def __post_init__(self):
+        # post init of parent
+        super().__post_init__()
+        # make a smaller scene for play
+        self.scene.num_envs = 16
+        self.scene.env_spacing = 2.5
+        # disable randomization for play
+        self.observations.policy.enable_corruption = False

source/standalone/tutorials/05_controllers/run_osc.py

+# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+This script demonstrates how to use the operational space controller (OSC) with the simulator.
+
+The OSC controller can be configured in different modes. It uses the dynamical quantities such as Jacobians and
+mass matricescomputed by PhysX.
+
+.. code-block:: bash
+
+    # Usage
+    ./isaaclab.sh -p source/standalone/tutorials/05_controllers/run_osc.py
+
+"""
+
+"""Launch Isaac Sim Simulator first."""
+
+import argparse
+
+from omni.isaac.lab.app import AppLauncher
+
+# add argparse arguments
+parser = argparse.ArgumentParser(description="Tutorial on using the operational space controller.")
+parser.add_argument("--num_envs", type=int, default=128, help="Number of environments to spawn.")
+# append AppLauncher cli args
+AppLauncher.add_app_launcher_args(parser)
+# parse the arguments
+args_cli = parser.parse_args()
+
+# launch omniverse app
+app_launcher = AppLauncher(args_cli)
+simulation_app = app_launcher.app
+
+"""Rest everything follows."""
+
+import torch
+
+import omni.isaac.lab.sim as sim_utils
+from omni.isaac.lab.assets import Articulation, AssetBaseCfg
+from omni.isaac.lab.controllers import OperationalSpaceController, OperationalSpaceControllerCfg
+from omni.isaac.lab.markers import VisualizationMarkers
+from omni.isaac.lab.markers.config import FRAME_MARKER_CFG
+from omni.isaac.lab.scene import InteractiveScene, InteractiveSceneCfg
+from omni.isaac.lab.sensors import ContactSensorCfg
+from omni.isaac.lab.utils import configclass
+from omni.isaac.lab.utils.math import (
+    combine_frame_transforms,
+    matrix_from_quat,
+    quat_inv,
+    quat_rotate_inverse,
+    subtract_frame_transforms,
+)
+
+##
+# Pre-defined configs
+##
+from omni.isaac.lab_assets import FRANKA_PANDA_HIGH_PD_CFG  # isort:skip
+
+
+@configclass
+class SceneCfg(InteractiveSceneCfg):
+    """Configuration for a simple scene with a tilted wall."""
+
+    # ground plane
+    ground = AssetBaseCfg(
+        prim_path="/World/defaultGroundPlane",
+        spawn=sim_utils.GroundPlaneCfg(),
+    )
+
+    # lights
+    dome_light = AssetBaseCfg(
+        prim_path="/World/Light", spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75))
+    )
+
+    # Tilted wall
+    tilted_wall = AssetBaseCfg(
+        prim_path="{ENV_REGEX_NS}/TiltedWall",
+        spawn=sim_utils.CuboidCfg(
+            size=(2.0, 1.5, 0.01),
+            collision_props=sim_utils.CollisionPropertiesCfg(),
+            visual_material=sim_utils.PreviewSurfaceCfg(diffuse_color=(1.0, 0.0, 0.0), opacity=0.1),
+            rigid_props=sim_utils.RigidBodyPropertiesCfg(kinematic_enabled=True),
+            activate_contact_sensors=True,
+        ),
+        init_state=AssetBaseCfg.InitialStateCfg(
+            pos=(0.6 + 0.085, 0.0, 0.3), rot=(0.9238795325, 0.0, -0.3826834324, 0.0)
+        ),
+    )
+
+    contact_forces = ContactSensorCfg(
+        prim_path="/World/envs/env_.*/TiltedWall",
+        update_period=0.0,
+        history_length=2,
+        debug_vis=False,
+    )
+
+    robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
+    robot.actuators["panda_shoulder"].stiffness = 0.0
+    robot.actuators["panda_shoulder"].damping = 0.0
+    robot.actuators["panda_forearm"].stiffness = 0.0
+    robot.actuators["panda_forearm"].damping = 0.0
+    robot.spawn.rigid_props.disable_gravity = True
+
+
+def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
+    """Runs the simulation loop.
+
+    Args:
+        sim: (SimulationContext) Simulation context.
+        scene: (InteractiveScene) Interactive scene.
+    """
+
+    # Extract scene entities for readability.
+    robot = scene["robot"]
+    contact_forces = scene["contact_forces"]
+
+    # Obtain indices for the end-effector and arm joints
+    ee_frame_name = "panda_leftfinger"
+    arm_joint_names = ["panda_joint.*"]
+    ee_frame_idx = robot.find_bodies(ee_frame_name)[0][0]
+    arm_joint_ids = robot.find_joints(arm_joint_names)[0]
+
+    # Create the OSC
+    osc_cfg = OperationalSpaceControllerCfg(
+        target_types=["pose_abs", "wrench_abs"],
+        impedance_mode="variable_kp",
+        inertial_dynamics_decoupling=True,
+        partial_inertial_dynamics_decoupling=False,
+        gravity_compensation=False,
+        motion_damping_ratio_task=1.0,
+        contact_wrench_stiffness_task=[0.0, 0.0, 0.1, 0.0, 0.0, 0.0],
+        motion_control_axes_task=[1, 1, 0, 1, 1, 1],
+        contact_wrench_control_axes_task=[0, 0, 1, 0, 0, 0],
+    )
+    osc = OperationalSpaceController(osc_cfg, num_envs=scene.num_envs, device=sim.device)
+
+    # Markers
+    frame_marker_cfg = FRAME_MARKER_CFG.copy()
+    frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
+    ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current"))
+    goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal"))
+
+    # Define targets for the arm
+    ee_goal_pose_set_tilted_b = torch.tensor(
+        [
+            [0.6, 0.15, 0.3, 0.0, 0.92387953, 0.0, 0.38268343],
+            [0.6, -0.3, 0.3, 0.0, 0.92387953, 0.0, 0.38268343],
+            [0.8, 0.0, 0.5, 0.0, 0.92387953, 0.0, 0.38268343],
+        ],
+        device=sim.device,
+    )
+    ee_goal_wrench_set_tilted_task = torch.tensor(
+        [
+            [0.0, 0.0, 10.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 10.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 10.0, 0.0, 0.0, 0.0],
+        ],
+        device=sim.device,
+    )
+    kp_set_task = torch.tensor(
+        [
+            [360.0, 360.0, 360.0, 360.0, 360.0, 360.0],
+            [420.0, 420.0, 420.0, 420.0, 420.0, 420.0],
+            [320.0, 320.0, 320.0, 320.0, 320.0, 320.0],
+        ],
+        device=sim.device,
+    )
+    ee_target_set = torch.cat([ee_goal_pose_set_tilted_b, ee_goal_wrench_set_tilted_task, kp_set_task], dim=-1)
+
+    # Define simulation stepping
+    sim_dt = sim.get_physics_dt()
+
+    # Update existing buffers
+    # Note: We need to update buffers before the first step for the controller.
+    robot.update(dt=sim_dt)
+
+    # get the updated states
+    jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = update_states(
+        sim, scene, robot, ee_frame_idx, arm_joint_ids, contact_forces
+    )
+
+    # Track the given target command
+    current_goal_idx = 0  # Current goal index for the arm
+    command = torch.zeros(
+        scene.num_envs, osc.action_dim, device=sim.device
+    )  # Generic target command, which can be pose, position, force, etc.
+    ee_target_pose_b = torch.zeros(scene.num_envs, 7, device=sim.device)  # Target pose in the body frame
+    ee_target_pose_w = torch.zeros(scene.num_envs, 7, device=sim.device)  # Target pose in the world frame (for marker)
+
+    # Set joint efforts to zero
+    zero_joint_efforts = torch.zeros(scene.num_envs, robot.num_joints, device=sim.device)
+    joint_efforts = torch.zeros(scene.num_envs, len(arm_joint_ids), device=sim.device)
+
+    count = 0
+    # Simulation loop
+    while simulation_app.is_running():
+        # reset every 500 steps
+        if count % 500 == 0:
+            # reset joint state to default
+            default_joint_pos = robot.data.default_joint_pos.clone()
+            default_joint_vel = robot.data.default_joint_vel.clone()
+            robot.write_joint_state_to_sim(default_joint_pos, default_joint_vel)
+            robot.set_joint_effort_target(zero_joint_efforts)  # Set zero torques in the initial step
+            robot.write_data_to_sim()
+            robot.reset()
+            # reset contact sensor
+            contact_forces.reset()
+            # reset target pose
+            robot.update(sim_dt)
+            _, _, _, ee_pose_b, _, _, _, _ = update_states(
+                sim, scene, robot, ee_frame_idx, arm_joint_ids, contact_forces
+            )  # at reset, the jacobians are not updated to the latest state
+            command, ee_target_pose_b, ee_target_pose_w, current_goal_idx = update_target(
+                sim, scene, osc, root_pose_w, ee_target_set, current_goal_idx
+            )
+            # set the osc command
+            osc.reset()
+            command, task_frame_pose_b = convert_to_task_frame(osc, command=command, ee_target_pose_b=ee_target_pose_b)
+            osc.set_command(command=command, current_ee_pose_b=ee_pose_b, current_task_frame_pose_b=task_frame_pose_b)
+        else:
+            # get the updated states
+            jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b = update_states(
+                sim, scene, robot, ee_frame_idx, arm_joint_ids, contact_forces
+            )
+            # compute the joint commands
+            joint_efforts = osc.compute(
+                jacobian_b=jacobian_b,
+                current_ee_pose_b=ee_pose_b,
+                current_ee_vel_b=ee_vel_b,
+                current_ee_force_b=ee_force_b,
+                mass_matrix=mass_matrix,
+                gravity=gravity,
+            )
+            # apply actions
+            robot.set_joint_effort_target(joint_efforts, joint_ids=arm_joint_ids)
+            robot.write_data_to_sim()
+
+        # update marker positions
+        ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])
+        goal_marker.visualize(ee_target_pose_w[:, 0:3], ee_target_pose_w[:, 3:7])
+
+        # perform step
+        sim.step(render=True)
+        # update robot buffers
+        robot.update(sim_dt)
+        # update buffers
+        scene.update(sim_dt)
+        # update sim-time
+        count += 1
+
+
+# Update robot states
+def update_states(
+    sim: sim_utils.SimulationContext,
+    scene: InteractiveScene,
+    robot: Articulation,
+    ee_frame_idx: int,
+    arm_joint_ids: list[int],
+    contact_forces,
+):
+    """Update the robot states.
+
+    Args:
+        sim: (SimulationContext) Simulation context.
+        scene: (InteractiveScene) Interactive scene.
+        robot: (Articulation) Robot articulation.
+        ee_frame_idx: (int) End-effector frame index.
+        arm_joint_ids: (list[int]) Arm joint indices.
+        contact_forces: (ContactSensor) Contact sensor.
+
+    Returns:
+        jacobian_b (torch.tensor): Jacobian in the body frame.
+        mass_matrix (torch.tensor): Mass matrix.
+        gravity (torch.tensor): Gravity vector.
+        ee_pose_b (torch.tensor): End-effector pose in the body frame.
+        ee_vel_b (torch.tensor): End-effector velocity in the body frame.
+        root_pose_w (torch.tensor): Root pose in the world frame.
+        ee_pose_w (torch.tensor): End-effector pose in the world frame.
+        ee_force_b (torch.tensor): End-effector force in the body frame.
+
+    Raises:
+        ValueError: Undefined target_type.
+    """
+    # obtain dynamics related quantities from simulation
+    ee_jacobi_idx = ee_frame_idx - 1
+    jacobian_w = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, arm_joint_ids]
+    mass_matrix = robot.root_physx_view.get_mass_matrices()[:, arm_joint_ids, :][:, :, arm_joint_ids]
+    gravity = robot.root_physx_view.get_generalized_gravity_forces()[:, arm_joint_ids]
+    # Convert the Jacobian from world to root frame
+    jacobian_b = jacobian_w.clone()
+    root_rot_matrix = matrix_from_quat(quat_inv(robot.data.root_state_w[:, 3:7]))
+    jacobian_b[:, :3, :] = torch.bmm(root_rot_matrix, jacobian_b[:, :3, :])
+    jacobian_b[:, 3:, :] = torch.bmm(root_rot_matrix, jacobian_b[:, 3:, :])
+
+    # Compute current pose of the end-effector
+    root_pose_w = robot.data.root_state_w[:, 0:7]
+    ee_pose_w = robot.data.body_state_w[:, ee_frame_idx, 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]
+    )
+    ee_pose_b = torch.cat([ee_pos_b, ee_quat_b], dim=-1)
+
+    # Compute the current velocity of the end-effector
+    ee_vel_w = robot.data.body_vel_w[:, ee_frame_idx, :]  # Extract end-effector velocity in the world frame
+    root_vel_w = robot.data.root_vel_w  # Extract root velocity in the world frame
+    relative_vel_w = ee_vel_w - root_vel_w  # Compute the relative velocity in the world frame
+    ee_lin_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 0:3])  # From world to root frame
+    ee_ang_vel_b = quat_rotate_inverse(robot.data.root_quat_w, relative_vel_w[:, 3:6])
+    ee_vel_b = torch.cat([ee_lin_vel_b, ee_ang_vel_b], dim=-1)
+
+    # Calculate the contact force
+    ee_force_w = torch.zeros(scene.num_envs, 3, device=sim.device)
+    sim_dt = sim.get_physics_dt()
+    contact_forces.update(sim_dt)  # update contact sensor
+    # Calculate the contact force by averaging over last four time steps (i.e., to smoothen) and
+    # taking the max of three surfaces as only one should be the contact of interest
+    ee_force_w, _ = torch.max(torch.mean(contact_forces.data.net_forces_w_history, dim=1), dim=1)
+
+    # This is a simplification, only for the sake of testing.
+    ee_force_b = ee_force_w
+
+    return jacobian_b, mass_matrix, gravity, ee_pose_b, ee_vel_b, root_pose_w, ee_pose_w, ee_force_b
+
+
+# Update the target commands
+def update_target(
+    sim: sim_utils.SimulationContext,
+    scene: InteractiveScene,
+    osc: OperationalSpaceController,
+    root_pose_w: torch.tensor,
+    ee_target_set: torch.tensor,
+    current_goal_idx: int,
+):
+    """Update the targets for the operational space controller.
+
+    Args:
+        sim: (SimulationContext) Simulation context.
+        scene: (InteractiveScene) Interactive scene.
+        osc: (OperationalSpaceController) Operational space controller.
+        root_pose_w: (torch.tensor) Root pose in the world frame.
+        ee_target_set: (torch.tensor) End-effector target set.
+        current_goal_idx: (int) Current goal index.
+
+    Returns:
+        command (torch.tensor): Updated target command.
+        ee_target_pose_b (torch.tensor): Updated target pose in the body frame.
+        ee_target_pose_w (torch.tensor): Updated target pose in the world frame.
+        next_goal_idx (int): Next goal index.
+
+    Raises:
+        ValueError: Undefined target_type.
+    """
+
+    # update the ee desired command
+    command = torch.zeros(scene.num_envs, osc.action_dim, device=sim.device)
+    command[:] = ee_target_set[current_goal_idx]
+
+    # update the ee desired pose
+    ee_target_pose_b = torch.zeros(scene.num_envs, 7, device=sim.device)
+    for target_type in osc.cfg.target_types:
+        if target_type == "pose_abs":
+            ee_target_pose_b[:] = command[:, :7]
+        elif target_type == "wrench_abs":
+            pass  # ee_target_pose_b could stay at the root frame for force control, what matters is ee_target_b
+        else:
+            raise ValueError("Undefined target_type within update_target().")
+
+    # update the target desired pose in world frame (for marker)
+    ee_target_pos_w, ee_target_quat_w = combine_frame_transforms(
+        root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_target_pose_b[:, 0:3], ee_target_pose_b[:, 3:7]
+    )
+    ee_target_pose_w = torch.cat([ee_target_pos_w, ee_target_quat_w], dim=-1)
+
+    next_goal_idx = (current_goal_idx + 1) % len(ee_target_set)
+
+    return command, ee_target_pose_b, ee_target_pose_w, next_goal_idx
+
+
+# Convert the target commands to the task frame
+def convert_to_task_frame(osc: OperationalSpaceController, command: torch.tensor, ee_target_pose_b: torch.tensor):
+    """Converts the target commands to the task frame.
+
+    Args:
+        osc: OperationalSpaceController object.
+        command: Command to be converted.
+        ee_target_pose_b: Target pose in the body frame.
+
+    Returns:
+        command (torch.tensor): Target command in the task frame.
+        task_frame_pose_b (torch.tensor): Target pose in the task frame.
+
+    Raises:
+        ValueError: Undefined target_type.
+    """
+    command = command.clone()
+    task_frame_pose_b = ee_target_pose_b.clone()
+
+    cmd_idx = 0
+    for target_type in osc.cfg.target_types:
+        if target_type == "pose_abs":
+            command[:, :3], command[:, 3:7] = subtract_frame_transforms(
+                task_frame_pose_b[:, :3], task_frame_pose_b[:, 3:], command[:, :3], command[:, 3:7]
+            )
+            cmd_idx += 7
+        elif target_type == "wrench_abs":
+            # These are already defined in target frame for ee_goal_wrench_set_tilted_task (since it is
+            # easier), so not transforming
+            cmd_idx += 6
+        else:
+            raise ValueError("Undefined target_type within _convert_to_task_frame().")
+
+    return command, task_frame_pose_b
+
+
+def main():
+    """Main function."""
+    # Load kit helper
+    sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device)
+    sim = sim_utils.SimulationContext(sim_cfg)
+    # Set main camera
+    sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0])
+    # Design scene
+    scene_cfg = SceneCfg(num_envs=args_cli.num_envs, env_spacing=2.0)
+    scene = InteractiveScene(scene_cfg)
+    # Play the simulator
+    sim.reset()
+    # Now we are ready!
+    print("[INFO]: Setup complete...")
+    # Run the simulator
+    run_simulator(sim, scene)
+
+
+if __name__ == "__main__":
+    # run the main function
+    main()
+    # close sim app
+    simulation_app.close()

tools/per_test_timeouts.py

@@ -18,4 +18,5 @@ PER_TEST_TIMEOUTS = {
     "test_rsl_rl_wrapper.py": 200,
     "test_sb3_wrapper.py": 200,
     "test_skrl_wrapper.py": 200,
+    "test_operational_space.py": 200,
 }