From dd08d4eb53502b7c4b4a554fa99550dfad823b1e Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Wed, 29 Jul 2026 18:06:01 +0200 Subject: [PATCH] fix(robot): type FK-to-EE action features as ACTION not STATE (#4213) * fix(robot): type FK-to-EE action features as ACTION not STATE ForwardKinematicsJointsToEEAction.transform_features declared its end-effector action features (ee.x/y/z/wx/wy/wz/gripper_pos) with FeatureType.STATE, copied verbatim from the sibling ForwardKinematicsJointsToEEObservation (where STATE is correct for OBSERVATION features). Every other action-producing step in this file (EEReferenceAndDelta, InverseKinematicsEEToJoints, InverseKinematicsRLStep) types its ACTION-bucket features as FeatureType.ACTION. The mismatch mis-classifies the converted EE actions as state, which propagates a wrong feature schema to downstream consumers keyed on FeatureType (e.g. normalization norm_map, policy input/output feature classification). * test(robot): FK-to-EE step feature-type contract (action vs observation) Asserts ForwardKinematicsJointsToEEAction emits EE features in the ACTION bucket typed FeatureType.ACTION, and ForwardKinematicsJointsToEEObservation emits them in the OBSERVATION bucket typed FeatureType.STATE. * chore: delete user file * chore(processor): reduce verbosity --------- Co-authored-by: Jaagat-P --- .../so_follower/robot_kinematic_processor.py | 4 +- .../robots/test_robot_kinematic_processor.py | 45 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 tests/robots/test_robot_kinematic_processor.py diff --git a/src/lerobot/robots/so_follower/robot_kinematic_processor.py b/src/lerobot/robots/so_follower/robot_kinematic_processor.py index ac2ec1f8b..50519d5e6 100644 --- a/src/lerobot/robots/so_follower/robot_kinematic_processor.py +++ b/src/lerobot/robots/so_follower/robot_kinematic_processor.py @@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep): # We only use the ee pose in the dataset, so we don't need the joint positions for n in self.motor_names: features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None) - # We specify the dataset features of this step that we want to be stored in the dataset + # Store end-effector features as actions in the dataset schema for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]: features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature( - type=FeatureType.STATE, shape=(1,) + type=FeatureType.ACTION, shape=(1,) ) return features diff --git a/tests/robots/test_robot_kinematic_processor.py b/tests/robots/test_robot_kinematic_processor.py new file mode 100644 index 000000000..c8dd4c77a --- /dev/null +++ b/tests/robots/test_robot_kinematic_processor.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python + +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature +from lerobot.robots.so_follower.robot_kinematic_processor import ( + ForwardKinematicsJointsToEEAction, + ForwardKinematicsJointsToEEObservation, +) + +MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"] +EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]} + + +def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]: + return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES} + + +@pytest.mark.parametrize( + ("step_cls", "bucket", "feature_type"), + [ + (ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION), + (ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE), + ], +) +def test_fk_feature_schema(step_cls, bucket, feature_type): + features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}} + features[bucket] = _joint_bucket(feature_type) + out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket] + assert set(out) == EE_KEYS + assert {feature.type for feature in out.values()} == {feature_type}