diff --git a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py b/src/lerobot/robots/so100_follower/robot_kinematic_processor.py index cff18c161bc..cf7b51fe60c 100644 --- a/src/lerobot/robots/so100_follower/robot_kinematic_processor.py +++ b/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 diff --git a/src/lerobot/teleoperators/phone/phone_processor.py b/src/lerobot/teleoperators/phone/phone_processor.py index 37a3d21027f..9801cf90f92 100644 --- a/src/lerobot/teleoperators/phone/phone_processor.py +++ b/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