diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 4279a3e27..813756490 100644 --- a/src/lerobot/processor/__init__.py +++ b/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", diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index 4499aa860..63a8b717f 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/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], diff --git a/src/lerobot/processor/gym_action_processor.py b/src/lerobot/processor/gym_action_processor.py index febf51cd2..7b3816750 100644 --- a/src/lerobot/processor/gym_action_processor.py +++ b/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 \ No newline at end of file + return torch_action diff --git a/src/lerobot/processor/hil_processor.py b/src/lerobot/processor/hil_processor.py index 7fa5267e8..074976e45 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/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): diff --git a/src/lerobot/scripts/rl/actor.py b/src/lerobot/scripts/rl/actor.py index 5f44d3c5f..997bc620b 100644 --- a/src/lerobot/scripts/rl/actor.py +++ b/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) diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 098f97565..a5f2c304c 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/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)