Fixes bug in RmpFlowController and adds demo script (#19)

* adds initialize to rmpflow controller
* adds demo script to play rmpflow
* fixes bugs in dealing with multiple command types
* adds right end-effector frame for franka lula
* adds demo script for rmpflow
* updates config and changelog

---------
Co-authored-by: Mayank Mittal <mittalma@leggedrobotics.com>

docs/source/refs/roadmap.rst

@@ -20,7 +20,7 @@ If you have any questions or suggestions, let us know on
 
   * |check_| Joint-space control
   * |check_| Differential inverse kinematics control
-  * |uncheck| Riemannian Motion Policies (RMPs)
+  * |check_| Riemannian Motion Policies (RMPs)
 
 * Supported robots
 

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

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

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

 Changelog
 ---------
 
+0.3.1 (2023-04-23)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added a modified version of ``lula_franka_gen.urdf`` which includes an end-effector frame.
+* Added a standalone script ``play_rmpflow.py`` to show RMPFlow controller.
+
+Fixed
+^^^^^
+
+* Fixed the splitting of commands in the :meth:`ActuatorGroup.compute` method. Earlier it was reshaping the
+  commands to the shape ``(num_actuators, num_commands)`` which was causing the commands to be split incorrectly.
+* Fixed the processing of actuator command in the :meth:`RobotBase._process_actuators_cfg` to deal with multiple
+  command types when using "implicit" actuator group.
+
 0.3.0 (2023-04-20)
 ~~~~~~~~~~~~~~~~~~
 
@@ -26,6 +43,7 @@ Fixed
 
 
 0.2.7 (2023-04-07)
+~~~~~~~~~~~~~~~~~~
 
 Fixed
 ^^^^^

source/extensions/omni.isaac.orbit/omni/isaac/orbit/actuators/group/actuator_group.py

@@ -307,19 +307,19 @@ class ActuatorGroup:
         # pre-processing of commands based on group.
         group_actions = self._format_command(group_actions)
         # reshape actions for sub-groups
-        group_actions = group_actions.view(self.num_articulation, self.num_actuators, -1)
+        group_actions = group_actions.split([self.num_actuators] * len(self.command_types), dim=-1)
         # pre-process relative commands
-        for index, command_type in enumerate(self.command_types):
+        for command_type, command_value in zip(self.command_types, group_actions):
             if command_type == "p_rel":
-                group_des_dof_pos = self.dof_pos_scale * group_actions[..., index] + self._dof_pos
+                group_des_dof_pos = self.dof_pos_scale * command_value + self._dof_pos
             elif command_type == "p_abs":
-                group_des_dof_pos = self.dof_pos_scale * group_actions[..., index] + self.dof_pos_offset
+                group_des_dof_pos = self.dof_pos_scale * command_value + self.dof_pos_offset
             elif command_type == "v_rel":
-                group_des_dof_vel = self.dof_vel_scale * group_actions[..., index] + self._dof_vel
+                group_des_dof_vel = self.dof_vel_scale * command_value + self._dof_vel
             elif command_type == "v_abs":
-                group_des_dof_vel = self.dof_vel_scale * group_actions[..., index]  # offset = 0
+                group_des_dof_vel = self.dof_vel_scale * command_value  # offset = 0
             elif command_type == "t_abs":
-                group_des_dof_torque = self.dof_torque_scale * group_actions[..., index]  # offset = 0
+                group_des_dof_torque = self.dof_torque_scale * command_value  # offset = 0
             else:
                 raise ValueError(f"Invalid action command type for actuators: '{command_type}'.")
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/controllers/config/rmp_flow.py

@@ -12,12 +12,14 @@ from omni.isaac.orbit.controllers.rmp_flow import RmpFlowControllerCfg
 # Note: RMP-Flow config files for supported robots are stored in the motion_generation extension
 _RMP_CONFIG_DIR = os.path.join(get_extension_path_from_name("omni.isaac.motion_generation"), "motion_policy_configs")
 
+# Path to current directory
+_CUR_DIR = os.path.dirname(os.path.realpath(__file__))
 
 FRANKA_RMPFLOW_CFG = RmpFlowControllerCfg(
     config_file=os.path.join(_RMP_CONFIG_DIR, "franka", "rmpflow", "franka_rmpflow_common.yaml"),
-    urdf_file=os.path.join(_RMP_CONFIG_DIR, "franka", "lula_franka_gen.urdf"),
+    urdf_file=os.path.join(_CUR_DIR, "data", "lula_franka_gen.urdf"),
     collision_file=os.path.join(_RMP_CONFIG_DIR, "franka", "rmpflow", "robot_descriptor.yaml"),
-    frame_name="right_gripper",
+    frame_name="panda_end_effector",
     evaluations_per_frame=5,
 )
 """Configuration of RMPFlow for Franka arm (default from `omni.isaac.motion_generation`)."""

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

@@ -146,6 +146,10 @@ class DifferentialInverseKinematics:
     Operations.
     """
 
+    def initialize(self):
+        """Initialize the internals."""
+        pass
+
     def reset_idx(self, robot_ids: torch.Tensor = None):
         """Reset the internals."""
         pass

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

@@ -134,6 +134,10 @@ class JointImpedanceController:
     Operations.
     """
 
+    def initialize(self):
+        """Initialize the internals."""
+        pass
+
     def reset_idx(self, robot_ids: torch.Tensor = None):
         """Reset the internals."""
         pass

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

@@ -177,6 +177,10 @@ class OperationSpaceController:
     Operations.
     """
 
+    def initialize(self):
+        """Initialize the internals."""
+        pass
+
     def reset_idx(self, robot_ids: torch.Tensor = None):
         """Reset the internals."""
         pass

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

@@ -11,7 +11,7 @@ import omni.isaac.core.utils.prims as prim_utils
 from omni.isaac.core.articulations import Articulation
 from omni.isaac.core.simulation_context import SimulationContext
 from omni.isaac.motion_generation import ArticulationMotionPolicy
-from omni.isaac.motion_generation.lula import RmpFlow
+from omni.isaac.motion_generation.lula.motion_policies import RmpFlow, RmpFlowSmoothed
 
 from omni.isaac.orbit.utils import configclass
 
@@ -20,6 +20,8 @@ from omni.isaac.orbit.utils import configclass
 class RmpFlowControllerCfg:
     """Configuration for RMP-Flow controller (provided through LULA library)."""
 
+    name: str = "rmp_flow"
+    """Name of the controller. Supported: "rmp_flow", "rmp_flow_smoothed". Defaults to "rmp_flow"."""
     config_file: str = MISSING
     """Path to the configuration file for the controller."""
     urdf_file: str = MISSING
@@ -28,36 +30,59 @@ class RmpFlowControllerCfg:
     """Path to collision model description of the robot."""
     frame_name: str = MISSING
     """Name of the robot frame for task space (must be present in the URDF)."""
-    evaluations_per_frame: int = MISSING
+    evaluations_per_frame: float = MISSING
     """Number of substeps during Euler integration inside LULA world model."""
     ignore_robot_state_updates: bool = False
-    """If true, then state of the world model inside controller is rolled out. (default: False)."""
+    """If true, then state of the world model inside controller is rolled out. Defaults to False."""
 
 
 class RmpFlowController:
-    """Wraps around RMP-Flow from IsaacSim for batched environments."""
+    """Wraps around RMPFlow from IsaacSim for batched environments."""
 
-    def __init__(self, cfg: RmpFlowControllerCfg, prim_paths_expr: str, device: str):
+    def __init__(self, cfg: RmpFlowControllerCfg, device: str):
         """Initialize the controller.
 
         Args:
             cfg (RmpFlowControllerCfg): The configuration for the controller.
-            prim_paths_expr (str): The expression to find the articulation prim paths.
             device (str): The device to use for computation.
-
-        Raises:
-            NotImplementedError: When the robot name is not supported.
         """
         # store input
         self.cfg = cfg
         self._device = device
+        # display info
+        print(f"[INFO]: Loading RMPFlow controller URDF from: {self.cfg.urdf_file}")
+
+    """
+    Properties.
+    """
 
-        print(f"[INFO]: Loading controller URDF from: {self.cfg.urdf_file}")
+    @property
+    def num_actions(self) -> int:
+        """Dimension of the action space of controller."""
+        return 7
+
+    """
+    Operations.
+    """
+
+    def initialize(self, prim_paths_expr: str):
+        """Initialize the controller.
+
+        Args:
+            prim_paths_expr (str): The expression to find the articulation prim paths.
+        """
         # obtain the simulation time
         physics_dt = SimulationContext.instance().get_physics_dt()
         # find all prims
         self._prim_paths = prim_utils.find_matching_prim_paths(prim_paths_expr)
         self.num_robots = len(self._prim_paths)
+        # resolve controller
+        if self.cfg.name == "rmp_flow":
+            controller_cls = RmpFlow
+        elif self.cfg.name == "rmp_flow_smoothed":
+            controller_cls = RmpFlowSmoothed
+        else:
+            raise ValueError(f"Unsupported controller in Lula library: {self.cfg.name}")
         # create all franka robots references and their controllers
         self.articulation_policies = list()
         for prim_path in self._prim_paths:
@@ -65,12 +90,12 @@ class RmpFlowController:
             robot = Articulation(prim_path)
             robot.initialize()
             # add controller
-            rmpflow = RmpFlow(
-                rmpflow_config_path=self.cfg.config_file,
-                urdf_path=self.cfg.urdf_file,
+            rmpflow = controller_cls(
                 robot_description_path=self.cfg.collision_file,
+                urdf_path=self.cfg.urdf_file,
+                rmpflow_config_path=self.cfg.config_file,
                 end_effector_frame_name=self.cfg.frame_name,
-                evaluations_per_frame=self.cfg.evaluations_per_frame,
+                maximum_substep_size=physics_dt / self.cfg.evaluations_per_frame,
                 ignore_robot_state_updates=self.cfg.ignore_robot_state_updates,
             )
             # wrap rmpflow to connect to the Franka robot articulation
@@ -86,22 +111,14 @@ class RmpFlowController:
         self.dof_pos_target = torch.zeros((self.num_robots, self.num_dof), device=self._device)
         self.dof_vel_target = torch.zeros((self.num_robots, self.num_dof), device=self._device)
 
-    """
-    Properties.
-    """
-
-    @property
-    def num_actions(self) -> int:
-        """Dimension of the action space of controller."""
-        return 7
-
-    """
-    Operations.
-    """
-
     def reset_idx(self, robot_ids: torch.Tensor = None):
         """Reset the internals."""
-        pass
+        # if no robot ids are provided, then reset all robots
+        if robot_ids is None:
+            robot_ids = torch.arange(self.num_robots, device=self._device)
+        # reset policies for specified robots
+        for index in robot_ids:
+            self.articulation_policies[index].motion_policy.reset()
 
     def set_command(self, command: torch.Tensor):
         """Set target end-effector pose command."""
@@ -127,9 +144,7 @@ class RmpFlowController:
             # apply action on the robot
             action = policy.get_next_articulation_action()
             # copy actions into buffer
-            # TODO: Make this more efficient?
-            for dof_index in range(self.num_dof):
-                self.dof_pos_target[i, dof_index] = action.joint_positions[dof_index]
-                self.dof_vel_target[i, dof_index] = action.joint_velocities[dof_index]
+            self.dof_pos_target[i, :] = torch.from_numpy(action.joint_positions[:]).to(self.dof_pos_target)
+            self.dof_vel_target[i, :] = torch.from_numpy(action.joint_velocities[:]).to(self.dof_vel_target)
 
         return self.dof_pos_target, self.dof_vel_target

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

-# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
-# All rights reserved.
-#
-# SPDX-License-Identifier: BSD-3-Clause
-
-import os
-import torch
-from typing import Tuple
-
-from omni.isaac.core.articulations import Articulation
-from omni.isaac.core.utils.extensions import get_extension_path_from_name
-from omni.isaac.core.utils.prims import find_matching_prim_paths
-from omni.isaac.motion_generation import ArticulationMotionPolicy
-from omni.isaac.motion_generation.lula import RmpFlow
-
-
-class RmpFlowController:
-    """Wraps around RMP-Flow from IsaacSim for batched environments."""
-
-    def __init__(self, prim_paths_expr: str, robot_name: str, dt: float, device: str):
-        """Initialize the controller.
-
-        Args:
-            prim_paths_expr (str): The expression to find the articulation prim paths.
-            robot_name (str): The name of the robot. Supported robots: franka, ur10.
-            dt (float): The time step of the simulation.
-            device (str): The device to use for computation.
-
-        Raises:
-            NotImplementedError: When the robot name is not supported.
-        """
-        # store input
-        self._device = device
-        # Load configuration for the controller
-        # Note: RMP-Flow config files for supported robots are stored in the motion_generation extension
-        mg_extension_path = get_extension_path_from_name("omni.isaac.motion_generation")
-        # configuration for the robot
-        if robot_name == "franka":
-            rmp_config_dir = os.path.join(mg_extension_path, "motion_policy_configs", "franka")
-            robot_description_path = os.path.join(rmp_config_dir, "rmpflow", "robot_descriptor.yaml")
-            robot_urdf_path = os.path.join(rmp_config_dir, "lula_franka_gen.urdf")
-            rmpflow_config_path = os.path.join(rmp_config_dir, "rmpflow", "franka_rmpflow_common.yaml")
-            self.arm_dof_num = 7
-            self.ee_frame_name = "right_gripper"
-        elif robot_name == "ur10":
-            rmp_config_dir = os.path.join(mg_extension_path, "motion_policy_configs", "ur10")
-            robot_description_path = os.path.join(rmp_config_dir, "rmpflow", "ur10_robot_description.yaml")
-            robot_urdf_path = os.path.join(rmp_config_dir, "ur10_robot.urdf")
-            rmpflow_config_path = os.path.join(rmp_config_dir, "rmpflow", "ur10_rmpflow_config.yaml")
-            self.arm_dof_num = 6
-            self.ee_frame_name = "ee_link"
-        else:
-            raise NotImplementedError(f"Input robot has no configuration: {robot_name}")
-
-        print(f"[INFO]: Loading controller URDF from: {robot_urdf_path}")
-        # find all prims
-        self._prim_paths = find_matching_prim_paths(prim_paths_expr)
-        self.num_robots = len(self._prim_paths)
-        # create all franka robots references and their controllers
-        self.articulation_policies = list()
-        for prim_path in self._prim_paths:
-            # add robot reference
-            robot = Articulation(prim_path)
-            robot.initialize()
-            # add controller
-            rmpflow = RmpFlow(
-                robot_description_path=robot_description_path,
-                urdf_path=robot_urdf_path,
-                rmpflow_config_path=rmpflow_config_path,
-                end_effector_frame_name=self.ee_frame_name,
-                evaluations_per_frame=5,
-            )
-            # wrap rmpflow to connect to the Franka robot articulation
-            articulation_policy = ArticulationMotionPolicy(robot, rmpflow, dt)
-            self.articulation_policies.append(articulation_policy)
-        # create buffers
-        # -- for storing command
-        self.desired_ee_pos = None
-        self.desired_ee_rot = None
-        # -- for policy output
-        self.joint_positions = torch.zeros((self.num_robots, self.arm_dof_num), device=self._device)
-        self.joint_velocities = torch.zeros((self.num_robots, self.arm_dof_num), device=self._device)
-        # update timer
-        self._update_timer = torch.zeros(self.num_robots, device=self._device)
-
-    """
-    Operations.
-    """
-
-    def reset_idx(self, robot_ids: torch.Tensor = None):
-        """Reset the internals."""
-        if robot_ids is None:
-            robot_ids = ...
-        self._update_timer[robot_ids] = 0.0
-
-    def set_command(self, desired_ee_pos: torch.Tensor, desired_ee_rot: torch.Tensor):
-        """Set target end-effector pose command."""
-        # convert pose to numpy
-        self.desired_ee_pos = desired_ee_pos.cpu().numpy()
-        self.desired_ee_rot = desired_ee_rot.cpu().numpy()
-
-    def advance(self, dt: float) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Performs inference with the controller.
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor]: The target joint positions and velocity commands.
-        """
-        # check valid
-        if self.desired_ee_pos is None or self.desired_ee_rot is None:
-            raise RuntimeError("Target command has not been set. Please call `set_command()` first.")
-        # update timer
-        self._update_timer += dt
-        # compute control actions
-        for i, policy in enumerate(self.articulation_policies):
-            # set rmpflow target to be the current position of the target cube.
-            policy.get_motion_policy().set_end_effector_target(
-                target_position=self.desired_ee_pos[i], target_orientation=self.desired_ee_rot[i]
-            )
-            # apply action on the robot
-            action = policy.get_next_articulation_action()
-            # copy actions into buffer
-            # arm-action
-            for dof_index in range(self.arm_dof_num):
-                if action.joint_positions[dof_index] is not None:
-                    self.joint_positions[i, dof_index] = torch.tensor(
-                        action.joint_positions[dof_index], device=self._device
-                    )
-                if action.joint_velocities[dof_index] is not None:
-                    self.joint_velocities[i, dof_index] = torch.tensor(
-                        action.joint_velocities[dof_index], device=self._device
-                    )
-        return self.joint_positions, self.joint_velocities

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

@@ -41,7 +41,7 @@ UR10_CFG = SingleArmManipulatorCfg(
         },
         dof_vel={".*": 0.0},
     ),
-    ee_info=SingleArmManipulatorCfg.EndEffectorFrameCfg(body_name="wrist_3_link"),
+    ee_info=SingleArmManipulatorCfg.EndEffectorFrameCfg(body_name="ee_link"),
     actuator_groups={
         "arm": ActuatorGroupCfg(
             dof_names=[".*"],

source/extensions/omni.isaac.orbit/omni/isaac/orbit/robots/robot_base.py

@@ -446,7 +446,22 @@ class RobotBase:
             # store actuator group
             self.actuator_groups[group_name] = actuator_group
             # store the control mode and dof indices (optimization)
-            self.sim_dof_control_modes[actuator_group.control_mode].extend(actuator_group.dof_indices)
+            if actuator_group.model_type == "implicit":
+                for command in actuator_group.command_types:
+                    # resolve name of control mode
+                    if "p" in command:
+                        command_name = "position"
+                    elif "v" in command:
+                        command_name = "velocity"
+                    elif "t" in command:
+                        command_name = "effort"
+                    else:
+                        continue
+                    # store dof indices
+                    self.sim_dof_control_modes[command_name].extend(actuator_group.dof_indices)
+            else:
+                # in explicit mode, we always use the "effort" control mode
+                self.sim.dof_control_mode["effort"].extend(actuator_group.dof_indices)
 
         # perform some sanity checks to ensure actuators are prepared correctly
         total_act_dof = sum(group.num_actuators for group in self.actuator_groups.values())

source/standalone/demo/play_ik_control.py

@@ -140,6 +140,7 @@ def main():
     # Acquire handles
     # Initialize handles
     robot.initialize("/World/envs/env_.*/Robot")
+    ik_controller.initialize()
     # Reset states
     robot.reset_buffers()
     ik_controller.reset_idx()
@@ -214,7 +215,7 @@ def main():
         # apply actions
         robot.apply_action(robot_actions)
         # perform step
-        sim.step()
+        sim.step(render=not args_cli.headless)
         # update sim-time
         sim_time += sim_dt
         count += 1