removed RobotAction2Tensor processor; imrpoved choosing observations in actor

src/lerobot/processor/__init__.py

@@ -17,6 +17,7 @@
 from .batch_processor import ToBatchProcessor
 from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict
 from .device_processor import DeviceProcessor
+from .gym_action_processor import Numpy2TorchActionProcessor, Torch2NumpyActionProcessor
 from .hil_processor import (
     AddTeleopActionAsComplimentaryData,
     AddTeleopEventsAsInfo,
@@ -26,7 +27,6 @@ from .hil_processor import (
     RewardClassifierProcessor,
     TimeLimitProcessor,
 )
-from .gym_action_processor import RobotAction2TensorProcessor, Torch2NumpyActionProcessor, Numpy2TorchActionProcessor
 from .joint_observations_processor import JointVelocityProcessor, MotorCurrentProcessor
 from .normalize_processor import NormalizerProcessor, UnnormalizerProcessor, hotswap_stats
 from .observation_processor import VanillaObservationProcessor
@@ -55,7 +55,6 @@ __all__ = [
     "DoneProcessor",
     "MapDeltaActionToRobotAction",
     "MapTensorToDeltaActionDict",
-    "RobotAction2TensorProcessor",
     "EnvTransition",
     "GripperPenaltyProcessor",
     "IdentityProcessor",

src/lerobot/processor/delta_action_processor.py

@@ -34,7 +34,7 @@ class MapTensorToDeltaActionDict(ActionProcessor):
             return action
         if action.dim() > 1:
             action = action.squeeze(0)
-        
+
         # TODO (maractingi): add rotation
         return {
             "action.delta_x": action[0],

src/lerobot/processor/gym_action_processor.py

@@ -13,25 +13,12 @@
 
 from dataclasses import dataclass
 
-import torch
 import numpy as np
+import torch
 
 from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry
 
 
-@ProcessorStepRegistry.register("robot_action_to_tensor_processor")
-@dataclass
-class RobotAction2TensorProcessor(ActionProcessor):
-    """Convert robot action to tensor."""
-    motor_names: list[str]
-
-    def action(self, action: dict | None) -> torch.Tensor | None:
-        if action is None:
-            return None
-
-        action_tensor = torch.tensor([action[f"action.{motor_name}.pos"] for motor_name in self.motor_names])
-        return action_tensor
-
 @ProcessorStepRegistry.register("torch2numpy_action_processor")
 @dataclass
 class Torch2NumpyActionProcessor(ActionProcessor):
@@ -78,4 +65,4 @@ class Numpy2TorchActionProcessor(ActionProcessor):
                 "Use appropriate processor for non-tensor actions."
             )
         torch_action = torch.from_numpy(action)
-        return torch_action
+        return torch_action

src/lerobot/processor/hil_processor.py

@@ -8,7 +8,6 @@ import torchvision.transforms.functional as F  # noqa: N812
 
 from lerobot.configs.types import PolicyFeature
 from lerobot.processor.pipeline import (
-    ActionProcessor,
     ComplementaryDataProcessor,
     EnvTransition,
     InfoProcessor,
@@ -49,8 +48,6 @@ class AddTeleopEventsAsInfo(InfoProcessor):
         return info
 
 
-
-
 @ProcessorStepRegistry.register("image_crop_resize_processor")
 @dataclass
 class ImageCropResizeProcessor(ObservationProcessor):

src/lerobot/scripts/rl/actor.py

@@ -98,7 +98,6 @@ from lerobot.utils.utils import (
 
 ACTOR_SHUTDOWN_TIMEOUT = 30
 
-
 #################################################
 # Main entry point #
 #################################################
@@ -288,7 +287,9 @@ def act_with_policy(
             logging.info("[ACTOR] Shutting down act_with_policy")
             return
 
-        observation = transition[TransitionKey.OBSERVATION]
+        observation = {
+            k: v for k, v in transition[TransitionKey.OBSERVATION].items() if k in cfg.policy.input_features
+        }
 
         # Time policy inference and check if it meets FPS requirement
         with policy_timer:
@@ -308,8 +309,16 @@ def act_with_policy(
         )
 
         # Extract values from processed transition
-        next_observation = new_transition[TransitionKey.OBSERVATION]
-        executed_action = new_transition[TransitionKey.ACTION]
+        next_observation = {
+            k: v
+            for k, v in new_transition[TransitionKey.OBSERVATION].items()
+            if k in cfg.policy.input_features
+        }
+
+        # Teleop action is the action that was executed in the environment
+        # It is either the action from the teleop device or the action from the policy
+        executed_action = new_transition[TransitionKey.COMPLEMENTARY_DATA]["teleop_action"]
+
         reward = new_transition[TransitionKey.REWARD]
         done = new_transition.get(TransitionKey.DONE, False)
         truncated = new_transition.get(TransitionKey.TRUNCATED, False)

src/lerobot/scripts/rl/gym_manipulator.py

@@ -41,7 +41,6 @@ from lerobot.processor import (
     MotorCurrentProcessor,
     Numpy2TorchActionProcessor,
     RewardClassifierProcessor,
-    RobotAction2TensorProcessor,
     RobotProcessor,
     TimeLimitProcessor,
     ToBatchProcessor,
@@ -456,7 +455,6 @@ def make_processors(
             )
         )
 
-    env_pipeline_steps.append(RobotAction2TensorProcessor(motor_names=motor_names))
     env_pipeline_steps.append(ToBatchProcessor())
     env_pipeline_steps.append(DeviceProcessor(device=device))
 
@@ -654,7 +652,6 @@ def control_loop(
             env_processor=env_processor,
             action_processor=action_processor,
         )
-        print(transition[TransitionKey.ACTION])
         terminated = transition.get(TransitionKey.DONE, False)
         truncated = transition.get(TransitionKey.TRUNCATED, False)