refine #3: add jaw semantic sensor

habitat-lab/habitat/articulated_agents/robots/spot_robot.py

@@ -165,6 +165,12 @@ def _get_spot_params(self):
                     attached_link_id=6,
                     relative_transform=mn.Matrix4.rotation_z(mn.Deg(-90)),
                 ),
+                "articulated_agent_jaw_panoptic": ArticulatedAgentCameraParams(
+                    cam_offset_pos=mn.Vector3(0.166, 0.0, -0.107),
+                    cam_orientation=mn.Vector3(0, -1.571, 0.0),
+                    attached_link_id=6,
+                    relative_transform=mn.Matrix4.rotation_z(mn.Deg(-90)),
+                ),
             },
             gripper_closed_state=[0.0],
             gripper_open_state=[-1.56],

habitat-lab/habitat/config/benchmark/rearrange/skills/place_spot_obj_ori.yaml

@@ -5,6 +5,7 @@ defaults:
     - base_to_object_distance
   - /habitat/task/lab_sensors:
     - spot_head_stereo_depth_sensor
+    - jaw_receptacle_semantic_sensor
   - override /habitat/simulator/[email protected]_agent: spot_agent_jaw_sensor
   - override /habitat/simulator/[email protected]_agent: spot
   - override /habitat/task/rearrange/actions: spot_base_arm_empty
@@ -23,6 +24,7 @@ habitat:
     obs_keys:
       - relative_initial_ee_orientation
       - relative_target_object_orientation
+      - jaw_receptacle_semantic_sensor
       - articulated_agent_arm_depth
       - articulated_agent_jaw_depth
       - joint
@@ -68,6 +70,9 @@ habitat:
         # We may want to use dimension 5 for controlling the rotation
         dimensionality: 5
         arm_joint_mask: [1,1,0,1,0,1,1]
+      jaw_receptacle_semantic_sensor:
+        height: 240
+        width: 228
     actions:
       arm_action:
         arm_controller: "ArmRelPosKinematicAction"

habitat-lab/habitat/config/default_structured_configs.py

@@ -46,6 +46,7 @@
     "HumanoidDetectorSensorConfig",
     "ArmDepthBBoxSensorConfig",
     "ArmReceptacleSemanticSensorConfig",
+    "JawReceptacleSemanticSensorConfig",
     "ArmRGBPretrainVisualFeatureSensorConfig",
     "PretrainTextualFeatureGoalSensorConfig",
     "SpotHeadStereoDepthSensorConfig",
@@ -444,6 +445,16 @@ class ArmReceptacleSemanticSensorConfig(LabSensorConfig):
     width: int = 640
 
 
+@dataclass
+class JawReceptacleSemanticSensorConfig(LabSensorConfig):
+    r"""
+    Semantic sensor for getting the target place receptacle
+    """
+    type: str = "JawReceptacleSemanticSensor"
+    height: int = 480
+    width: int = 640
+
+
 @dataclass
 class ArmRGBPretrainVisualFeatureSensorConfig(LabSensorConfig):
     r"""
@@ -1727,6 +1738,13 @@ class ArmPanopticSensorConfig(HabitatSimSemanticSensorConfig):
     height: int = 256
 
 
+@dataclass
+class JawPanopticSensorConfig(HabitatSimSemanticSensorConfig):
+    uuid: str = "articulated_agent_jaw_panoptic"
+    width: int = 256
+    height: int = 256
+
+
 @dataclass
 class ArmRGBSensorConfig(HabitatSimRGBSensorConfig):
     uuid: str = "articulated_agent_arm_rgb"
@@ -2260,6 +2278,12 @@ class HabitatConfig(HabitatBaseConfig):
     node=ArmPanopticSensorConfig,
 )
 
+cs.store(
+    group="habitat/simulator/sim_sensors",
+    name="jaw_panoptic_sensor",
+    node=JawPanopticSensorConfig,
+)
+
 cs.store(
     group="habitat/simulator/sim_sensors",
     name="third_depth_sensor",
@@ -2310,6 +2334,12 @@ class HabitatConfig(HabitatBaseConfig):
     name="arm_receptacle_semantic_sensor",
     node=ArmReceptacleSemanticSensorConfig,
 )
+cs.store(
+    package="habitat.task.lab_sensors.jaw_receptacle_semantic_sensor",
+    group="habitat/task/lab_sensors",
+    name="jaw_receptacle_semantic_sensor",
+    node=JawReceptacleSemanticSensorConfig,
+)
 cs.store(
     package="habitat.task.lab_sensors.arm_rgb_pretrain_visual_feature_sensor",
     group="habitat/task/lab_sensors",

habitat-lab/habitat/config/habitat/simulator/sensor_setups/spot_agent_jaw_sensor.yaml

@@ -7,6 +7,7 @@ defaults:
   - /habitat/simulator/sim_sensors@sim_sensors.jaw_rgb_sensor: jaw_rgb_sensor
   - /habitat/simulator/sim_sensors@sim_sensors.jaw_depth_sensor: jaw_depth_sensor
   - /habitat/simulator/sim_sensors@sim_sensors.arm_panoptic_sensor: arm_panoptic_sensor
+  - /habitat/simulator/sim_sensors@sim_sensors.jaw_panoptic_sensor: jaw_panoptic_sensor
   - /habitat/simulator/sim_sensors@sim_sensors.head_stereo_left_depth_sensor: head_stereo_left_depth_sensor
   - /habitat/simulator/sim_sensors@sim_sensors.head_stereo_right_depth_sensor: head_stereo_right_depth_sensor
 
@@ -35,6 +36,10 @@ sim_sensors:
     height: 240
     width: 228
     hfov: 60
+  jaw_panoptic_sensor:
+    height: 240
+    width: 228
+    hfov: 60
   head_stereo_right_depth_sensor:
     height: 212
     width: 120

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

@@ -1799,3 +1799,42 @@ def get_observation(self, observations, episode, task, *args, **kwargs):
         tgt_mask = (img_seg == tgt_id).astype(int)
 
         return np.float32(tgt_mask)
+
+
+@registry.register_sensor
+class JawReceptacleSemanticSensor(ArmReceptacleSemanticSensor):
+    """Semantic sensor for the target place receptacle"""
+
+    cls_uuid: str = "jaw_receptacle_semantic_sensor"
+
+    def _get_uuid(self, *args, **kwargs):
+        return JawReceptacleSemanticSensor.cls_uuid
+
+    def get_observation(self, observations, episode, task, *args, **kwargs):
+        # Get a correct observation space
+        if self.agent_id is None:
+            target_key = "articulated_agent_jaw_panoptic"
+            assert target_key in observations
+        else:
+            target_key = (
+                f"agent_{self.agent_id}_articulated_agent_jaw_panoptic"
+            )
+            assert target_key in observations
+
+        img_seg = observations[target_key]
+
+        # Check the size of the observation
+        assert (
+            img_seg.shape[0] == self._height
+            and img_seg.shape[1] == self._width
+        )
+
+        # Get the target receptacle name
+        rep_name = self._get_rep_text(task.targ_idx)
+        self._rep_name = rep_name
+
+        # Get the target mask
+        tgt_id = self._sim.receptacles_id[rep_name][0]
+        tgt_mask = (img_seg == tgt_id).astype(int)
+
+        return np.float32(tgt_mask)