Merge pull request #114 from ARISE-Initiative/collect-demos-script

Update collect demos scripts, add controller input type and reference frame options
ARISE-Initiative · Oct 25, 2024 · 3cdf2bc · 3cdf2bc
2 parents 45000dc + 90d76d7
commit 3cdf2bc
Show file tree

Hide file tree

Showing 28 changed files with 586 additions and 746 deletions.
diff --git a/robosuite/controllers/composite/composite_controller.py b/robosuite/controllers/composite/composite_controller.py
@@ -53,8 +53,16 @@ def __init__(self, sim: MjSim, robot_model: RobotModel, grippers: Dict[str, Grip
     def load_controller_config(
         self, part_controller_config, composite_controller_specific_config: Optional[Dict] = None
     ):
-        self.part_controller_config = part_controller_config
         self.composite_controller_specific_config = composite_controller_specific_config
+        body_part_ordering = self.composite_controller_specific_config.get("body_part_ordering", None)
+        if body_part_ordering is not None:
+            self.part_controller_config = OrderedDict()
+            assert len(body_part_ordering) == len(part_controller_config)
+            for part_name in body_part_ordering:
+                self.part_controller_config[part_name] = part_controller_config[part_name]
+        else:
+            self.part_controller_config = part_controller_config
+
         self.part_controllers.clear()
         self._action_split_indexes.clear()
         self._init_controllers()
@@ -165,10 +173,22 @@ def set_goal(self, all_action):
             if part_name in self.grippers.keys():
                 action = self.grippers[part_name].format_action(action)
 
-            if part_name in self.arms:
-                controller.set_goal(action, update_wrt_origin=update_wrt_origin)
-            else:
-                controller.set_goal(action)
+            if part_name in self.arms and hasattr(controller, "set_goal_update_mode"):
+                """
+                Query the last action dimension to determine if using
+                base mode (value > 0) or arm mode (value < 1).
+                If in base mode, update the new arm goal based on current desired goal,
+                to have the arm tracking with the reset of the moving base
+                as closely as possible without lagging or overshooting.
+                """
+                action_mode = all_action[-1]
+                if action_mode > 0:
+                    goal_update_mode = "desired"
+                else:
+                    goal_update_mode = "achieved"
+                controller.set_goal_update_mode(goal_update_mode)
+
+            controller.set_goal(action)
 
     @property
     def action_limits(self):
@@ -369,9 +389,6 @@ def print_action_info_dict(self, name: str = ""):
 class WholeBodyIK(WholeBody):
     name = "WHOLE_BODY_IK"
 
-    def __init__(self, sim: MjSim, robot_model: RobotModel, grippers: Dict[str, GripperModel]):
-        super().__init__(sim, robot_model, grippers)
-
     def __init__(self, sim: MjSim, robot_model: RobotModel, grippers: Dict[str, GripperModel]):
         super().__init__(sim, robot_model, grippers)
 

diff --git a/robosuite/controllers/composite/composite_controller_factory.py b/robosuite/controllers/composite/composite_controller_factory.py
@@ -85,6 +85,8 @@ def _get_robot_name(robot: str) -> str:
     """Helper function to get the standardized robot name."""
     if "GR1FloatingBody" in robot:
         return "gr1_floating_body"
+    elif "GR1FixedLowerBody" in robot:
+        return "gr1_fixed_lower_body"
     elif "GR1" in robot:
         return "gr1"
     elif "G1" in robot:

diff --git a/robosuite/controllers/config/default/composite/basic.json b/robosuite/controllers/config/default/composite/basic.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {
@@ -37,7 +38,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/default/composite/hybrid_mobile_base.json b/robosuite/controllers/config/default/composite/hybrid_mobile_base.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {
@@ -37,7 +38,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_baxter.json b/robosuite/controllers/config/robots/default_baxter.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {
@@ -37,7 +38,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_gr1_fixed_lower_body.json b/robosuite/controllers/config/robots/default_gr1_fixed_lower_body.json
@@ -0,0 +1,98 @@
+{
+    "type": "WHOLE_BODY_MINK_IK",
+    "composite_controller_specific_configs": {
+        "ref_name": ["gripper0_right_grip_site", "gripper0_left_grip_site"],
+        "interpolation": null,
+        "actuation_part_names": ["torso", "head", "right", "left"],
+        "max_dq": 4,
+        "ik_pseudo_inverse_damping": 5e-2,
+        "ik_integration_dt": 1e-1,
+        "ik_input_action_repr": "absolute",
+        "ik_input_rotation_repr": "axis_angle",
+        "verbose": false,
+        "ik_posture_weights": {
+            "robot0_torso_waist_yaw": 10.0,
+            "robot0_torso_waist_pitch": 10.0,
+            "robot0_torso_waist_roll": 200.0,
+            "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_controller_configs": {
+        "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"
+                }
+            },
+            "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"
+                }
+            }
+        },
+        "torso": {
+            "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
+        },
+        "head": {
+            "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
+        }
+    }
+}
diff --git a/robosuite/controllers/config/robots/default_gr1_floating_body.json b/robosuite/controllers/config/robots/default_gr1_floating_body.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {
@@ -38,7 +39,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_iiwa.json b/robosuite/controllers/config/robots/default_iiwa.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_kinova3.json b/robosuite/controllers/config/robots/default_kinova3.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_panda.json b/robosuite/controllers/config/robots/default_panda.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_panda_dex.json b/robosuite/controllers/config/robots/default_panda_dex.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_pandaomron.json b/robosuite/controllers/config/robots/default_pandaomron.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_sawyer.json b/robosuite/controllers/config/robots/default_sawyer.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_spotwitharm.json b/robosuite/controllers/config/robots/default_spotwitharm.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_tiago.json b/robosuite/controllers/config/robots/default_tiago.json
@@ -30,7 +30,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {
@@ -51,7 +52,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {

diff --git a/robosuite/controllers/config/robots/default_ur5e.json b/robosuite/controllers/config/robots/default_ur5e.json
@@ -16,7 +16,8 @@
                 "position_limits": null,
                 "orientation_limits": null,
                 "uncouple_pos_ori": true,
-                "control_delta": true,
+                "input_type": "delta",
+                "input_ref_frame": "base",
                 "interpolation": null,
                 "ramp_ratio": 0.2,
                 "gripper": {