Feat/add missing feature contracts

src/lerobot/robots/so100_follower/robot_kinematic_processor.py

@@ -142,6 +142,23 @@ def reset(self):
         self.reference_ee_pose = None
         self._command_when_disabled = None
 
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        features.pop("action.enabled", None)
+        features.pop("action.target_x", None)
+        features.pop("action.target_y", None)
+        features.pop("action.target_z", None)
+        features.pop("action.target_wx", None)
+        features.pop("action.target_wy", None)
+        features.pop("action.target_wz", None)
+
+        features["action.ee.x"] = float
+        features["action.ee.y"] = float
+        features["action.ee.z"] = float
+        features["action.ee.wx"] = float
+        features["action.ee.wy"] = float
+        features["action.ee.wz"] = float
+        return features
+
 
 @ProcessorStepRegistry.register("ee_bounds_and_safety")
 @dataclass
@@ -210,16 +227,6 @@ def reset(self):
         self._last_pos = None
         self._last_twist = None
 
-    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # Because this is last step we specify the dataset features of this step that we want to be stored in the dataset
-        features["action.ee.x"] = float
-        features["action.ee.y"] = float
-        features["action.ee.z"] = float
-        features["action.ee.wx"] = float
-        features["action.ee.wy"] = float
-        features["action.ee.wz"] = float
-        return features
-
 
 @ProcessorStepRegistry.register("inverse_kinematics_ee_to_joints")
 @dataclass
@@ -307,16 +314,11 @@ def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
     def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We specify the dataset features of this step that we want to be stored in the dataset
-        features["action.ee.x"] = float
-        features["action.ee.y"] = float
-        features["action.ee.z"] = float
-        features["action.ee.wx"] = float
-        features["action.ee.wy"] = float
-        features["action.ee.wz"] = float
-
         features["observation.state.gripper.pos"] = float
         features["action.gripper.pos"] = float
+        for name in self.motor_names:
+            features[f"action.{name}.pos"] = float
+
         return features
 
     def reset(self):
@@ -388,8 +390,7 @@ def __call__(self, transition: EnvTransition) -> EnvTransition:
         return transition
 
     def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
-        # We specify the dataset features of this step that we want to be stored in the dataset
-        features["observation.state.gripper.pos"] = float
+        features.pop("action.gripper", None)
         features["action.gripper.pos"] = float
         return features
 

src/lerobot/teleoperators/phone/phone_processor.py

@@ -16,6 +16,7 @@
 
 from dataclasses import dataclass, field
 
+from lerobot.configs.types import PolicyFeature
 from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 
@@ -47,7 +48,7 @@ class MapPhoneActionToRobotAction(ActionProcessor):
 
     def action(self, act: dict) -> dict:
         # Pop them from the action
-        enabled = act.pop("action.phone.enabled", 0)
+        enabled = bool(act.pop("action.phone.enabled", 0))
         pos = act.pop("action.phone.pos", None)
         rot = act.pop("action.phone.rot", None)
         inputs = act.pop("action.phone.raw_inputs", {})
@@ -81,3 +82,19 @@ def action(self, act: dict) -> dict:
             }
         )
         return act
+
+    def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]:
+        features.pop("action.phone.enabled", None)
+        features.pop("action.phone.pos", None)
+        features.pop("action.phone.rot", None)
+        features.pop("action.phone.raw_inputs", None)
+
+        features["action.enabled"] = bool
+        features["action.target_x"] = float
+        features["action.target_y"] = float
+        features["action.target_z"] = float
+        features["action.target_wx"] = float
+        features["action.target_wy"] = float
+        features["action.target_wz"] = float
+        features["action.gripper"] = float
+        return features