Update of robot controller and robot-mounted base (#679)

Author: Dingry

robosuite/controllers/parts/generic/joint_pos.py

@@ -175,6 +175,8 @@ def __init__(
         # initialize
         self.goal_qpos = None
 
+        self.use_torque_compensation = kwargs.get("use_torque_compensation", True)
+
     def set_goal(self, action, set_qpos=None):
         """
         Sets goal based on input @action. If self.impedance_mode is not "fixed", then the input will be parsed into the
@@ -260,7 +262,10 @@ def run_controller(self):
         desired_torque = np.multiply(np.array(position_error), np.array(self.kp)) + np.multiply(vel_pos_error, self.kd)
 
         # Return desired torques plus gravity compensations
-        self.torques = np.dot(self.mass_matrix, desired_torque) + self.torque_compensation
+        if self.use_torque_compensation:
+            self.torques = np.dot(self.mass_matrix, desired_torque) + self.torque_compensation
+        else:
+            self.torques = desired_torque
 
         # Always run superclass call for any cleanups at the end
         super().run_controller()

robosuite/controllers/parts/generic/joint_tor.py

@@ -106,6 +106,8 @@ def __init__(
         self.current_torque = np.zeros(self.control_dim)  # Current torques being outputted, pre-compensation
         self.torques = None  # Torques returned every time run_controller is called
 
+        self.use_torque_compensation = kwargs.get("use_torque_compensation", True)
+
     def set_goal(self, torques):
         """
         Sets goal based on input @torques.
@@ -153,7 +155,10 @@ def run_controller(self):
             self.current_torque = np.array(self.goal_torque)
 
         # Add gravity compensation
-        self.torques = self.current_torque + self.torque_compensation
+        if self.use_torque_compensation:
+            self.torques = self.current_torque + self.torque_compensation
+        else:
+            self.torques = self.current_torque
 
         # Always run superclass call for any cleanups at the end
         super().run_controller()

robosuite/controllers/parts/generic/joint_vel.py

@@ -124,6 +124,8 @@ def __init__(
         self.current_vel = np.zeros(self.joint_dim)  # Current velocity setpoint, pre-compensation
         self.torques = None  # Torques returned every time run_controller is called
 
+        self.torque_compensation = kwargs.get("use_torque_compensation", True)
+
     def set_goal(self, velocities):
         """
         Sets goal based on input @velocities.
@@ -187,7 +189,12 @@ def run_controller(self):
             self.summed_err += err
 
         # Compute command torques via PID velocity controller plus gravity compensation torques
-        torques = self.kp * err + self.ki * self.summed_err + self.kd * self.derr_buf.average + self.torque_compensation
+        if self.torque_compensation:
+            torques = (
+                self.kp * err + self.ki * self.summed_err + self.kd * self.derr_buf.average + self.torque_compensation
+            )
+        else:
+            torques = self.kp * err + self.ki * self.summed_err + self.kd * self.derr_buf.average
 
         # Clip torques
         self.torques = self.clip_torques(torques)

robosuite/devices/device.py

@@ -194,7 +194,7 @@ def get_arm_action(self, robot, arm, norm_delta, goal_update_mode="target"):
                 "delta": norm_delta,
                 "abs": abs_action,
             }
-        elif robot.composite_controller_config["type"] in ["WHOLE_BODY_MINK_IK"]:
+        elif robot.composite_controller_config["type"] in ["WHOLE_BODY_MINK_IK", "HYBRID_WHOLE_BODY_MINK_IK"]:
             ref_frame = self.env.robots[0].composite_controller.composite_controller_specific_config.get(
                 "ik_input_ref_frame", "world"
             )

robosuite/examples/third_party_controller/mink_controller.py

@@ -148,6 +148,9 @@ def __init__(
         solve_freq: float = 20.0,
         hand_pos_cost: float = 1,
         hand_ori_cost: float = 0.5,
+        com_cost: float = 0.0,
+        use_mink_posture_task: bool = False,
+        initial_qpos_as_posture_target: bool = False,
     ):
         self.full_model: mujoco.MjModel = model
         self.full_model_data: mujoco.MjData = data
@@ -157,9 +160,11 @@ def __init__(
         self.posture_weights = posture_weights
         self.hand_pos_cost = hand_pos_cost
         self.hand_ori_cost = hand_ori_cost
-
+        self.com_cost = com_cost
+        self.use_mink_posture_task = use_mink_posture_task
         self.hand_tasks: List[mink.FrameTask]
         self.posture_task: WeightedPostureTask
+        self.com_task: mink.ComTask | None = None
 
         if robot_joint_names is None:
             robot_joint_names: List[str] = [
@@ -168,15 +173,40 @@ def __init__(
                 if self.robot_model.joint(i).type != 0
             ]  # Exclude fixed joints
 
-        self.full_model_dof_ids: List[int] = np.array([self.full_model.joint(name).id for name in robot_joint_names])
+        # the order of the index is the same of model.qposadr
+        self.all_robot_qpos_indexes_in_full_model: List[int] = []
+        for i in range(self.robot_model.njnt):
+            joint_name = self.robot_model.joint(i).name
+            self.all_robot_qpos_indexes_in_full_model.extend(
+                self.full_model.joint(joint_name).qposadr[0] + np.arange(len(self.full_model.joint(joint_name).qpos0))
+            )
+        assert len(self.all_robot_qpos_indexes_in_full_model) == self.robot_model.nq
+
+        # the order of the index is determined by the order of the actuation_part_names
+        self.controlled_robot_qpos_indexes: List[int] = []
+        self.controlled_robot_qpos_indexes_in_full_model: List[int] = []
+        for name in robot_joint_names:
+            self.controlled_robot_qpos_indexes.extend(
+                self.robot_model.joint(name).qposadr[0] + np.arange(len(self.robot_model.joint(name).qpos0))
+            )
+            self.controlled_robot_qpos_indexes_in_full_model.extend(
+                self.full_model.joint(name).qposadr[0] + np.arange(len(self.full_model.joint(name).qpos0))
+            )
 
-        self.robot_model_dof_ids: List[int] = np.array([self.robot_model.joint(name).id for name in robot_joint_names])
-        self.full_model_dof_ids: List[int] = np.array([self.full_model.joint(name).id for name in robot_joint_names])
         self.site_ids = [self.robot_model.site(site_name).id for site_name in site_names]
-
         self.site_names = site_names
+
+        # update robot states
+        self.update_robot_states()
+
+        # setup tasks
         self._setup_tasks()
-        self.set_posture_target(np.zeros(self.robot_model.nq))
+        if initial_qpos_as_posture_target:
+            self.set_posture_target()
+        else:
+            self.set_posture_target(np.zeros(self.robot_model.nq))
+        if self.com_cost > 0.0:
+            self.set_com_target()
 
         self.solver = "quadprog"
 
@@ -203,12 +233,17 @@ def __repr__(self) -> str:
         return "IKSolverMink"
 
     def _setup_tasks(self):
-        weights = np.ones(self.robot_model.nq)
+        weights = np.ones(self.robot_model.nv)
+
         for joint_name, posture_weight in self.posture_weights.items():
-            joint_idx = self.robot_model.joint(joint_name).id
-            weights[joint_idx] = posture_weight
+            joint = self.robot_model.joint(joint_name)
+            joint_dof_idx = joint.dofadr[0] + np.arange(len(joint.jntid))
+            weights[joint_dof_idx] = posture_weight
 
-        self.posture_task = WeightedPostureTask(self.robot_model, cost=0.01, weights=weights, lm_damping=2)
+        if self.use_mink_posture_task:
+            self.posture_task = mink.PostureTask(self.robot_model, cost=weights * 0.1, lm_damping=1.0)
+        else:
+            self.posture_task = WeightedPostureTask(self.robot_model, cost=0.01, weights=weights, lm_damping=2)
 
         self.tasks = [self.posture_task]
 
@@ -217,6 +252,10 @@ def _setup_tasks(self):
         )
         self.tasks.extend(self.hand_tasks)
 
+        if self.com_cost > 0.0:
+            self.com_task = mink.ComTask(cost=self.com_cost)
+            self.tasks.append(self.com_task)
+
     def _create_frame_tasks(self, frame_names: List[str], position_cost: float, orientation_cost: float):
         return [
             mink.FrameTask(
@@ -229,13 +268,33 @@ def _create_frame_tasks(self, frame_names: List[str], position_cost: float, orie
             for frame in frame_names
         ]
 
+    def update_robot_states(self):
+        # update the base pose, important for mobile robots such as humanoids
+        self.configuration.model.body("robot0_base").pos = self.full_model.body("robot0_base").pos
+        self.configuration.model.body("robot0_base").quat = self.full_model.body("robot0_base").quat
+
+        # update the qpos for the robot model
+        self.configuration.update(
+            self.full_model_data.qpos[self.controlled_robot_qpos_indexes_in_full_model],
+            self.controlled_robot_qpos_indexes,
+        )
+
     def set_target_poses(self, target_poses: List[np.ndarray]):
         for task, target in zip(self.hand_tasks, target_poses):
             se3_target = mink.SE3.from_matrix(target)
             task.set_target(se3_target)
 
-    def set_posture_target(self, posture_target: np.ndarray):
-        self.posture_task.set_target(posture_target)
+    def set_posture_target(self, posture_target: np.ndarray | None = None):
+        if posture_target is None:
+            self.posture_task.set_target_from_configuration(self.configuration)
+        else:
+            self.posture_task.set_target(posture_target)
+
+    def set_com_target(self, com_target: np.ndarray | None = None):
+        assert self.com_task is not None, "COM task is not initialized"
+        if com_target is None:
+            com_target = self.configuration.data.subtree_com[1]
+        self.com_task.set_target(com_target)
 
     def action_split_indexes(self) -> Dict[str, Tuple[int, int]]:
         action_split_indexes: Dict[str, Tuple[int, int]] = {}
@@ -261,9 +320,8 @@ def transform_pose(
         if src_frame == dst_frame:
             return src_frame_pose
 
-        self.configuration.model.body("robot0_base").pos = self.full_model.body("robot0_base").pos
-        self.configuration.model.body("robot0_base").quat = self.full_model.body("robot0_base").quat
-        self.configuration.update()
+        self.robot_model.body("robot0_base").pos = self.full_model.body("robot0_base").pos
+        self.robot_model.body("robot0_base").quat = self.full_model.body("robot0_base").quat
 
         X_src_frame_pose = src_frame_pose
         # convert src frame pose to world frame pose
@@ -294,13 +352,8 @@ def solve(self, input_action: np.ndarray) -> np.ndarray:
         By updating configuration's bose to match the actual base pose (in 'world' frame),
         we're requiring our tasks' targets to be in the 'world' frame for mink.solve_ik().
         """
-        # update configuration's base to match actual base
-        self.configuration.model.body("robot0_base").pos = self.full_model.body("robot0_base").pos
-        self.configuration.model.body("robot0_base").quat = self.full_model.body("robot0_base").quat
-        # update configuration's qpos to match actual qpos
-        self.configuration.update(
-            self.full_model_data.qpos[self.full_model_dof_ids], update_idxs=self.robot_model_dof_ids
-        )
+
+        self.update_robot_states()
 
         input_action = input_action.reshape(len(self.site_names), -1)
         input_pos = input_action[:, : self.pos_dim]
@@ -409,7 +462,7 @@ def solve(self, input_action: np.ndarray) -> np.ndarray:
             if self.i % 50:
                 print(f"Task errors: {task_translation_errors}")
 
-        return self.configuration.data.qpos[self.robot_model_dof_ids]
+        return self.configuration.data.qpos[self.controlled_robot_qpos_indexes]
 
     def _get_task_translation_errors(self) -> List[float]:
         errors = []
@@ -430,6 +483,7 @@ def _get_task_errors(self) -> List[float]:
         for task in self.hand_tasks:
             error = task.compute_error(self.configuration)
             errors.append(np.linalg.norm(error[:3]))
+        errors.append(self.posture_task.compute_error(self.configuration))
         return errors
 
 
@@ -518,5 +572,9 @@ def _init_joint_action_policy(self):
             posture_weights=self.composite_controller_specific_config.get("ik_posture_weights", {}),
             hand_pos_cost=self.composite_controller_specific_config.get("ik_hand_pos_cost", 1.0),
             hand_ori_cost=self.composite_controller_specific_config.get("ik_hand_ori_cost", 0.5),
-            verbose=self.composite_controller_specific_config.get("ik_verbose", False),
+            use_mink_posture_task=self.composite_controller_specific_config.get("use_mink_posture_task", False),
+            initial_qpos_as_posture_target=self.composite_controller_specific_config.get(
+                "initial_qpos_as_posture_target", False
+            ),
+            verbose=self.composite_controller_specific_config.get("verbose", False),
         )

robosuite/models/assets/bases/null_base.xml

@@ -0,0 +1,7 @@
+<mujoco model="null_mount">
+    <worldbody>
+        <body name="support" pos="0 0 0">
+            <site name="center" type="sphere" pos="0 0 0" size="0.01" group="1" rgba="0 0 0 0"/>
+        </body>
+    </worldbody>
+</mujoco>

robosuite/models/bases/__init__.py

@@ -3,6 +3,7 @@
 from .robot_base_factory import robot_base_factory
 from .mobile_base_model import MobileBaseModel
 from .leg_base_model import LegBaseModel
+from .null_base_model import NullBaseModel
 
 from .rethink_mount import RethinkMount
 from .rethink_minimal_mount import RethinkMinimalMount
@@ -12,7 +13,7 @@
 from .null_mobile_base import NullMobileBase
 from .no_actuation_base import NoActuationBase
 from .floating_legged_base import FloatingLeggedBase
-
+from .null_base import NullBase
 from .spot_base import Spot, SpotFloating
 
 BASE_MAPPING = {
@@ -25,6 +26,7 @@
     "FloatingLeggedBase": FloatingLeggedBase,
     "Spot": Spot,
     "SpotFloating": SpotFloating,
+    "NullBase": NullBase,
 }
 
 ALL_BASES = BASE_MAPPING.keys()

robosuite/models/bases/leg_base_model.py

@@ -34,6 +34,15 @@ def _remove_joint_actuation(self, part_name):
                 parent_body = find_parent(self.actuator, motor)
                 parent_body.remove(motor)
                 self._actuators.remove(motor.get("name").replace(self.naming_prefix, ""))
+        for sensor in self.root.findall(".//jointpos"):
+            if part_name in sensor.get("joint"):
+                find_parent(self.root, sensor).remove(sensor)
+        for sensor in self.root.findall(".//jointvel"):
+            if part_name in sensor.get("joint"):
+                find_parent(self.root, sensor).remove(sensor)
+        for sensor in self.root.findall(".//jointactuatorfrc"):
+            if part_name in sensor.get("joint"):
+                find_parent(self.root, sensor).remove(sensor)
 
     def _remove_free_joint(self):
         """Remove all freejoints from the model."""

robosuite/models/bases/null_base.py

@@ -0,0 +1,27 @@
+"""
+Rethink's Generic Mount (Officially used on Sawyer).
+"""
+import numpy as np
+
+from robosuite.models.bases.null_base_model import NullBaseModel
+from robosuite.utils.mjcf_utils import xml_path_completion
+
+
+class NullBase(NullBaseModel):
+    """
+    Dummy mobile base to signify no mount.
+
+    Args:
+        idn (int or str): Number or some other unique identification string for this mount instance
+    """
+
+    def __init__(self, idn=0):
+        super().__init__(xml_path_completion("bases/null_base.xml"), idn=idn)
+
+    @property
+    def top_offset(self):
+        return np.array((0, 0, 0))
+
+    @property
+    def horizontal_radius(self):
+        return 0

robosuite/models/bases/null_base_model.py

@@ -0,0 +1,11 @@
+"""
+Defines the null base model
+"""
+
+from robosuite.models.bases.robot_base_model import RobotBaseModel
+
+
+class NullBaseModel(RobotBaseModel):
+    @property
+    def naming_prefix(self):
+        return "nullbase{}_".format(self.idn)