[v0.2.3] [Stretch] play stretch demo

examples/interactive_play.py

@@ -72,7 +72,7 @@
 except ImportError:
     pygame = None
 
-DEFAULT_CFG = "benchmark/rearrange/play.yaml"
+DEFAULT_CFG = "benchmark/rearrange/play_stretch.yaml"
 DEFAULT_RENDER_STEPS_LIMIT = 60
 SAVE_VIDEO_DIR = "./data/vids"
 SAVE_ACTIONS_DIR = "./data/interactive_play_replays"
@@ -108,6 +108,10 @@ def get_input_vel_ctlr(
         ]
         arm_ctrlr = env.task.actions[arm_action_name].arm_ctrlr
         base_action = None
+    elif "stretch" in DEFAULT_CFG:
+        arm_action_space = np.zeros(10)
+        arm_ctrlr = None
+        base_action = [0, 0]
     else:
         arm_action_space = np.zeros(7)
         arm_ctrlr = None
@@ -182,6 +186,44 @@ def get_input_vel_ctlr(
                 arm_action[6] = 1.0
             elif keys[pygame.K_7]:
                 arm_action[6] = -1.0
+
+        elif arm_action_space.shape[0] == 10:
+            # Velocity control. A different key for each joint
+            if keys[pygame.K_q]:
+                arm_action[0] = 1.0
+            elif keys[pygame.K_1]:
+                arm_action[0] = -1.0
+
+            elif keys[pygame.K_w]:
+                arm_action[4] = 1.0
+            elif keys[pygame.K_2]:
+                arm_action[4] = -1.0
+
+            elif keys[pygame.K_e]:
+                arm_action[5] = 1.0
+            elif keys[pygame.K_3]:
+                arm_action[5] = -1.0
+
+            elif keys[pygame.K_r]:
+                arm_action[6] = 1.0
+            elif keys[pygame.K_4]:
+                arm_action[6] = -1.0
+
+            elif keys[pygame.K_t]:
+                arm_action[7] = 1.0
+            elif keys[pygame.K_5]:
+                arm_action[7] = -1.0
+
+            elif keys[pygame.K_y]:
+                arm_action[8] = 1.0
+            elif keys[pygame.K_6]:
+                arm_action[8] = -1.0
+
+            elif keys[pygame.K_u]:
+                arm_action[9] = 1.0
+            elif keys[pygame.K_7]:
+                arm_action[9] = -1.0
+
         elif isinstance(arm_ctrlr, ArmEEAction):
             EE_FACTOR = 0.5
             # End effector control

habitat-lab/habitat/config/benchmark/rearrange/play_stretch.yaml

@@ -0,0 +1,56 @@
+# @package _global_
+
+defaults:
+  - /habitat: habitat_config_base
+  - /agent@habitat.simulator.agent_0: agent_base
+  - /habitat/simulator/sim_sensors:
+    - head_rgb_sensor
+    - head_depth_sensor
+    - arm_rgb_sensor
+    - arm_depth_sensor
+  - /habitat/task/rearrange: play
+  - /habitat/dataset/rearrangement: replica_cad
+  - _self_
+
+# Config for empty task to explore the scene.
+habitat:
+  environment:
+    max_episode_steps: 0
+  task:
+    actions:
+      arm_action:
+        type: "ArmAction"
+        arm_controller: "ArmRelPosKinematicReducedActionStretch"
+        grip_controller: "MagicGraspAction"
+        arm_joint_mask: [1,0,0,0,1,1,1,1,1,1]
+        arm_joint_dimensionality: 10
+        grasp_thresh_dist: 0.15
+        disable_grip: False
+        delta_pos_limit: 0.0125
+        ee_ctrl_lim: 0.015
+  simulator:
+    type: RearrangeSim-v0
+    seed: 100
+    additional_object_paths:
+      - "data/objects/ycb/configs/"
+    agent_0:
+      radius: 0.3
+      robot_urdf: data/robots/hab_stretch/urdf/hab_stretch.urdf
+      robot_type: "StretchRobot"
+      sim_sensors:
+        head_rgb_sensor:
+          height: 212
+          width: 120
+        head_depth_sensor:
+          height: 212
+          width: 120
+        arm_depth_sensor:
+          height: 212
+          width: 120
+        arm_rgb_sensor:
+          height: 212
+          width: 120
+    habitat_sim_v0:
+      enable_physics: True
+  dataset:
+    data_path: data/datasets/replica_cad/rearrange/v1/{split}/all_receptacles_10k_1k.json.gz

habitat-lab/habitat/config/default_structured_configs.py

@@ -106,6 +106,7 @@ class ArmActionConfig(ActionConfig):
     type: str = "ArmAction"
     arm_controller: str = "ArmRelPosAction"
     grip_controller: Optional[str] = None
+    arm_joint_mask: Optional[List[int]] = None
     arm_joint_dimensionality: int = 7
     grasp_thresh_dist: float = 0.15
     disable_grip: bool = False

habitat-lab/habitat/tasks/rearrange/actions/actions.py

@@ -223,6 +223,70 @@ def step(self, set_pos, *args, **kwargs):
         self.cur_robot.arm_joint_pos = set_pos
 
 
+@registry.register_task_action
+class ArmRelPosKinematicReducedActionStretch(RobotAction):
+    """
+    The arm motor targets are offset by the delta joint values specified by the
+    action and the mask. This function is used for Stretch.
+    """
+
+    def __init__(self, *args, config, sim: RearrangeSim, **kwargs):
+        super().__init__(*args, config=config, sim=sim, **kwargs)
+        self.last_arm_action = None
+
+    def reset(self, *args, **kwargs):
+        super().reset(*args, **kwargs)
+        self.last_arm_action = None
+
+    @property
+    def action_space(self):
+        self.step_c = 0
+        return spaces.Box(
+            shape=(self._config.arm_joint_dimensionality,),
+            low=-1,
+            high=1,
+            dtype=np.float32,
+        )
+
+    def step(self, delta_pos, *args, **kwargs):
+        if self._config.get("SHOULD_CLIP", True):
+            # clip from -1 to 1
+            delta_pos = np.clip(delta_pos, -1, 1)
+        delta_pos *= self._config.delta_pos_limit
+        self._sim: RearrangeSim
+
+        # Expand delta_pos based on mask
+        expanded_delta_pos = np.zeros(len(self._config.arm_joint_mask))
+        src_idx = 0
+        tgt_idx = 0
+        for mask in self._config.arm_joint_mask:
+            if mask == 0:
+                tgt_idx += 1
+                src_idx += 1
+                continue
+            expanded_delta_pos[tgt_idx] = delta_pos[src_idx]
+            tgt_idx += 1
+            src_idx += 1
+
+        min_limit, max_limit = self.cur_robot.arm_joint_limits
+        set_arm_pos = expanded_delta_pos + self.cur_robot.arm_motor_pos
+        # Perform roll over to the joints so that the user cannot control
+        # the motor 2, 3, 4 for the arm.
+        if expanded_delta_pos[0] >= 0:
+            for i in range(3):
+                if set_arm_pos[i] > max_limit[i]:
+                    set_arm_pos[i + 1] += set_arm_pos[i] - max_limit[i]
+                    set_arm_pos[i] = max_limit[i]
+        else:
+            for i in range(3):
+                if set_arm_pos[i] < min_limit[i]:
+                    set_arm_pos[i + 1] -= min_limit[i] - set_arm_pos[i]
+                    set_arm_pos[i] = min_limit[i]
+        set_arm_pos = np.clip(set_arm_pos, min_limit, max_limit)
+
+        self.cur_robot.arm_motor_pos = set_arm_pos
+
+
 @registry.register_task_action
 class BaseVelAction(RobotAction):
     """