add joint bypassing flag

ARISE-Initiative · Nov 21, 2024 · bc8e9ad · bc8e9ad
1 parent c9f15a4
commit bc8e9ad
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 3 deletions.
diff --git a/robosuite/controllers/composite/composite_controller.py b/robosuite/controllers/composite/composite_controller.py
@@ -229,6 +229,7 @@ def __init__(self, sim: MjSim, robot_model: RobotModel, grippers: Dict[str, Grip
 
         self._whole_body_controller_action_split_indexes: OrderedDict = OrderedDict()
         self._latest_all_joint_angle_action: Optional[np.ndarray] = None
+        self._use_joint_angle_action_input: bool = False
 
     def _init_controllers(self):
         for part_name in self.part_controller_config.keys():
@@ -307,9 +308,10 @@ def setup_whole_body_controller_action_split_idx(self):
             previous_idx = last_idx
 
     def set_goal(self, all_action):
-        target_qpos = self.joint_action_policy.solve(all_action[: self.joint_action_policy.control_dim])
-        # create new all_action vector with the IK solver's actions first
-        all_action = np.concatenate([target_qpos, all_action[self.joint_action_policy.control_dim :]])
+        if not self._use_joint_angle_action_input:
+            target_qpos = self.joint_action_policy.solve(all_action[: self.joint_action_policy.control_dim])
+            # create new all_action vector with the IK solver's actions first
+            all_action = np.concatenate([target_qpos, all_action[self.joint_action_policy.control_dim :]])
         self._latest_all_joint_angle_action = all_action
         for part_name, controller in self.part_controllers.items():
             start_idx, end_idx = self._action_split_indexes[part_name]
@@ -328,6 +330,9 @@ def action_limits(self):
         Returns the action limits for the whole body controller.
         Corresponds to each term in the action vector passed to env.step().
         """
+        if self._use_joint_angle_action_input:
+            return super().action_limits
+
         low, high = [], []
         # assumption: IK solver's actions come first
         low_c, high_c = self.joint_action_policy.control_limits
@@ -430,6 +435,11 @@ def _validate_composite_controller_specific_config(self) -> None:
         # Update the configuration with only the valid reference names
         self.composite_controller_specific_config["ref_name"] = valid_ref_names
 
+        # Set the use joint angle action input flag
+        self._use_joint_angle_action_input = self.composite_controller_specific_config.get(
+            "use_joint_angle_action_input", False
+        )
+
     def _init_joint_action_policy(self):
         joint_names: str = []
         for part_name in self.composite_controller_specific_config["actuation_part_names"]:

diff --git a/robosuite/examples/third_party_controller/default_mink_ik_gr1_arms_only.json b/robosuite/examples/third_party_controller/default_mink_ik_gr1_arms_only.json
@@ -0,0 +1,68 @@
+{
+    "type": "WHOLE_BODY_MINK_IK",
+    "composite_controller_specific_configs": {
+        "ref_name": ["gripper0_right_grip_site", "gripper0_left_grip_site"],
+        "interpolation": null,
+        "actuation_part_names": ["right", "left"],
+        "max_dq": 4,
+        "ik_pseudo_inverse_damping": 5e-2,
+        "ik_integration_dt": 1e-1,
+        "ik_input_type": "absolute",
+        "ik_input_ref_frame": "world", 
+        "ik_input_rotation_repr": "axis_angle",
+        "verbose": false,
+        "ik_posture_weights": {
+            "robot0_l_shoulder_pitch": 4.0,
+            "robot0_r_shoulder_pitch": 4.0,
+            "robot0_l_shoulder_roll": 3.0,
+            "robot0_r_shoulder_roll": 3.0,
+            "robot0_l_shoulder_yaw": 2.0,
+            "robot0_r_shoulder_yaw": 2.0
+        },
+        "ik_hand_pos_cost": 1.0,
+        "ik_hand_ori_cost": 0.5,
+        "use_joint_angle_action_input": false
+    },
+    "body_parts": {
+        "arms": {
+            "right": {
+                "type" : "JOINT_POSITION",
+                "input_max": 1,
+                "input_min": -1,
+                "input_type": "absolute",
+                "output_max": 0.5,
+                "output_min": -0.5,
+                "kd": 200,
+                "kv": 200,
+                "kp": 1000,
+                "velocity_limits": [-1,1],
+                "kp_limits": [0, 1000],
+                "interpolation": null,
+                "ramp_ratio": 0.2,
+                "gripper": {
+                    "type": "GRIP",
+                    "use_action_scaling": false
+                }
+            },
+            "left": {
+                "type" : "JOINT_POSITION",
+                "input_max": 1,
+                "input_min": -1,
+                "input_type": "absolute",
+                "output_max": 0.5,
+                "output_min": -0.5,
+                "kd": 200,
+                "kv": 200,
+                "kp": 1000,
+                "velocity_limits": [-1,1],
+                "kp_limits": [0, 1000],
+                "interpolation": null,
+                "ramp_ratio": 0.2,
+                "gripper": {
+                    "type": "GRIP",
+                    "use_action_scaling": false
+                }
+            }
+        }
+    }
+}
diff --git a/robosuite/examples/third_party_controller/mink_controller.py b/robosuite/examples/third_party_controller/mink_controller.py
@@ -475,6 +475,11 @@ def _validate_composite_controller_specific_config(self) -> None:
         # Update the configuration with only the valid posture weights
         self.composite_controller_specific_config["ik_posture_weights"] = valid_posture_weights
 
+        # Set the use joint angle action input flag
+        self._use_joint_angle_action_input = self.composite_controller_specific_config.get(
+            "use_joint_angle_action_input", False
+        )
+
     def _init_joint_action_policy(self):
         joint_names: str = []
         for part_name in self.composite_controller_specific_config["actuation_part_names"]: