diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx index 7c63c41a8..c70745d50 100644 --- a/docs/source/hilserl.mdx +++ b/docs/source/hilserl.mdx @@ -127,11 +127,11 @@ class RewardClassifierConfig: # Dataset configuration class DatasetConfig: repo_id: str # LeRobot dataset repository ID - dataset_root: str # Local dataset root directory task: str # Task identifier - num_episodes: int # Number of episodes for recording - episode: int # Episode index for replay - push_to_hub: bool # Whether to push datasets to Hub + root: str | None = None # Local dataset root directory + num_episodes_to_record: int = 5 # Number of episodes for recording + replay_episode: int | None = None # Episode index for replay + push_to_hub: bool = False # Whether to push datasets to Hub ``` @@ -351,7 +351,7 @@ Create a configuration file for recording demonstrations (or edit an existing on 1. Set `mode` to `"record"` at the root level 2. Specify a unique `repo_id` for your dataset in the `dataset` section (e.g., "username/task_name") -3. Set `num_episodes` in the `dataset` section to the number of demonstrations you want to collect +3. Set `num_episodes_to_record` in the `dataset` section to the number of demonstrations you want to collect 4. Set `env.processor.image_preprocessing.crop_params_dict` to `{}` initially (we'll determine crops later) 5. Configure `env.robot`, `env.teleop`, and other hardware settings in the `env` section @@ -390,10 +390,10 @@ Example configuration section: }, "dataset": { "repo_id": "username/pick_lift_cube", - "dataset_root": null, + "root": null, "task": "pick_and_lift", - "num_episodes": 15, - "episode": 0, + "num_episodes_to_record": 15, + "replay_episode": 0, "push_to_hub": true }, "mode": "record", @@ -626,7 +626,7 @@ python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/r - **mode**: set it to `"record"` to collect a dataset (at root level) - **dataset.repo_id**: `"hf_username/dataset_name"`, name of the dataset and repo on the hub -- **dataset.num_episodes**: Number of episodes to record +- **dataset.num_episodes_to_record**: Number of episodes to record - **env.processor.reset.terminate_on_success**: Whether to automatically terminate episodes when success is detected (default: `true`) - **env.fps**: Number of frames per second to record - **dataset.push_to_hub**: Whether to push the dataset to the hub @@ -664,8 +664,8 @@ Example configuration section for data collection: "repo_id": "hf_username/dataset_name", "dataset_root": "data/your_dataset", "task": "reward_classifier_task", - "num_episodes": 20, - "episode": 0, + "num_episodes_to_record": 20, + "replay_episode": null, "push_to_hub": true }, "mode": "record", diff --git a/docs/source/hilserl_sim.mdx b/docs/source/hilserl_sim.mdx index bbb0cc6f9..656e650a0 100644 --- a/docs/source/hilserl_sim.mdx +++ b/docs/source/hilserl_sim.mdx @@ -107,10 +107,10 @@ To collect a dataset, set the mode to `record` whilst defining the repo_id and n }, "dataset": { "repo_id": "username/sim_dataset", - "dataset_root": null, + "root": null, "task": "pick_cube", - "num_episodes": 10, - "episode": 0, + "num_episodes_to_record": 10, + "replay_episode": null, "push_to_hub": true }, "mode": "record" diff --git a/docs/source/il_sim.mdx b/docs/source/il_sim.mdx index 6047bf884..70de48f88 100644 --- a/docs/source/il_sim.mdx +++ b/docs/source/il_sim.mdx @@ -36,10 +36,10 @@ To teleoperate and collect a dataset, we need to modify this config file. Here's }, "dataset": { "repo_id": "your_username/il_gym", - "dataset_root": null, + "root": null, "task": "pick_cube", - "num_episodes": 30, - "episode": 0, + "num_episodes_to_record": 30, + "replay_episode": null, "push_to_hub": true }, "mode": "record", @@ -50,7 +50,7 @@ To teleoperate and collect a dataset, we need to modify this config file. Here's Key configuration points: - Set your `repo_id` in the `dataset` section: `"repo_id": "your_username/il_gym"` -- Set `num_episodes: 30` to collect 30 demonstration episodes +- Set `num_episodes_to_record: 30` to collect 30 demonstration episodes - Ensure `mode` is set to `"record"` - If you don't have an NVIDIA GPU, change `"device": "cuda"` to `"mps"` for macOS or `"cpu"` - To use keyboard instead of gamepad, change `"task"` to `"PandaPickCubeKeyboard-v0"` diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 979f7ebc4..6718b723a 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -15,7 +15,7 @@ # limitations under the License. from .batch_processor import ToBatchProcessor -from .delta_action_processor import MapDeltaActionToRobotAction +from .delta_action_processor import MapDeltaActionToRobotAction, MapTensorToDeltaActionDict from .device_processor import DeviceProcessor from .hil_processor import ( AddTeleopActionAsComplimentaryData, @@ -55,6 +55,7 @@ __all__ = [ "DeviceProcessor", "DoneProcessor", "MapDeltaActionToRobotAction", + "MapTensorToDeltaActionDict", "EnvTransition", "GripperPenaltyProcessor", "IdentityProcessor", diff --git a/src/lerobot/processor/delta_action_processor.py b/src/lerobot/processor/delta_action_processor.py index a575bb07a..4499aa860 100644 --- a/src/lerobot/processor/delta_action_processor.py +++ b/src/lerobot/processor/delta_action_processor.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dataclasses import dataclass, field +from dataclasses import dataclass from torch import Tensor @@ -22,6 +22,28 @@ from lerobot.configs.types import FeatureType, PolicyFeature from lerobot.processor.pipeline import ActionProcessor, ProcessorStepRegistry +@ProcessorStepRegistry.register("map_tensor_to_delta_action_dict") +@dataclass +class MapTensorToDeltaActionDict(ActionProcessor): + """ + Map a tensor to a delta action dictionary. + """ + + def action(self, action: Tensor) -> dict: + if isinstance(action, dict): + return action + if action.dim() > 1: + action = action.squeeze(0) + + # TODO (maractingi): add rotation + return { + "action.delta_x": action[0], + "action.delta_y": action[1], + "action.delta_z": action[2], + "action.gripper": action[3], + } + + @ProcessorStepRegistry.register("map_delta_action_to_robot_action") @dataclass class MapDeltaActionToRobotAction(ActionProcessor): @@ -53,25 +75,17 @@ class MapDeltaActionToRobotAction(ActionProcessor): # Scale factors for delta movements position_scale: float = 1.0 rotation_scale: float = 0.0 # No rotation deltas for gamepad/keyboard - gripper_deadzone: float = 0.1 # Threshold for gripper activation - _prev_enabled: bool = field(default=False, init=False, repr=False) - def action(self, action: dict | Tensor | None) -> dict: + def action(self, action: dict | None) -> dict: if action is None: return {} # NOTE (maractingi): Action can be a dict from the teleop_devices or a tensor from the policy # TODO (maractingi): changing this target_xyz naming convention from the teleop_devices - if isinstance(action, dict): - delta_x = action.pop("action.delta_x", 0.0) - delta_y = action.pop("action.delta_y", 0.0) - delta_z = action.pop("action.delta_z", 0.0) - gripper = action.pop("action.gripper", 1.0) # Default to "stay" (1.0) - else: - delta_x = action[0].item() - delta_y = action[1].item() - delta_z = action[2].item() - gripper = action[3].item() + delta_x = action.pop("action.delta_x", 0.0) + delta_y = action.pop("action.delta_y", 0.0) + delta_z = action.pop("action.delta_z", 0.0) + gripper = action.pop("action.gripper", 1.0) # Default to "stay" (1.0) # Determine if the teleoperator is actively providing input # Consider enabled if any significant movement delta is detected @@ -101,7 +115,6 @@ class MapDeltaActionToRobotAction(ActionProcessor): "action.gripper": float(gripper), } - self._prev_enabled = enabled return action def transform_features(self, features: dict[str, PolicyFeature]) -> dict[str, PolicyFeature]: @@ -120,6 +133,3 @@ class MapDeltaActionToRobotAction(ActionProcessor): } ) return features - - def reset(self): - self._prev_enabled = False diff --git a/src/lerobot/processor/hil_processor.py b/src/lerobot/processor/hil_processor.py index 9e31548b2..4afc92a86 100644 --- a/src/lerobot/processor/hil_processor.py +++ b/src/lerobot/processor/hil_processor.py @@ -271,7 +271,8 @@ class InterventionActionProcessor: # Get intervention signals from complementary data info = transition.get(TransitionKey.INFO, {}) - teleop_action = info.get("teleop_action", {}) + complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) + teleop_action = complementary_data.get("teleop_action", {}) is_intervention = info.get(TeleopEvents.IS_INTERVENTION, False) terminate_episode = info.get(TeleopEvents.TERMINATE_EPISODE, False) success = info.get(TeleopEvents.SUCCESS, False) diff --git a/src/lerobot/processor/joint_observations_processor.py b/src/lerobot/processor/joint_observations_processor.py index 185db4048..012e6180a 100644 --- a/src/lerobot/processor/joint_observations_processor.py +++ b/src/lerobot/processor/joint_observations_processor.py @@ -13,7 +13,7 @@ from lerobot.robots import Robot @dataclass @ProcessorStepRegistry.register("joint_velocity_processor") -class JointVelocityProcessor: +class JointVelocityProcessor(ObservationProcessor): """Add joint velocity information to observations.""" joint_velocity_limits: float = 100.0 diff --git a/src/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py index 37ff1cc7e..aac9c3d29 100644 --- a/src/lerobot/scripts/rl/gym_manipulator.py +++ b/src/lerobot/scripts/rl/gym_manipulator.py @@ -37,6 +37,7 @@ from lerobot.processor import ( InterventionActionProcessor, JointVelocityProcessor, MapDeltaActionToRobotAction, + MapTensorToDeltaActionDict, MotorCurrentProcessor, Numpy2TorchActionProcessor, RewardClassifierProcessor, @@ -80,11 +81,11 @@ class DatasetConfig: """Configuration for dataset creation and management.""" repo_id: str - dataset_root: str task: str - num_episodes: int - episode: int - push_to_hub: bool + root: str | None = None + num_episodes_to_record: int = 5 + replay_episode: int | None = None + push_to_hub: bool = False @dataclass @@ -473,6 +474,7 @@ def make_processors( if cfg.processor.inverse_kinematics is not None and kinematics_solver is not None: # Add EE bounds and safety processor inverse_kinematics_steps = [ + MapTensorToDeltaActionDict(), MapDeltaActionToRobotAction(), EEReferenceAndDelta( kinematics=kinematics_solver, @@ -625,7 +627,7 @@ def control_loop( dataset = LeRobotDataset.create( cfg.dataset.repo_id, cfg.env.fps, - root=cfg.dataset.dataset_root, + root=cfg.dataset.root, use_videos=True, image_writer_threads=4, image_writer_processes=0, @@ -636,7 +638,7 @@ def control_loop( episode_step = 0 episode_start_time = time.perf_counter() - while episode_idx < cfg.dataset.num_episodes: + while episode_idx < cfg.dataset.num_episodes_to_record: step_start_time = time.perf_counter() # Create a neutral action (no movement) @@ -711,10 +713,12 @@ def control_loop( def replay_trajectory(env: gym.Env, action_processor: RobotProcessor, cfg: GymManipulatorConfig) -> None: """Replay recorded trajectory on robot environment.""" + assert cfg.dataset.replay_episode is not None, "Replay episode must be provided for replay" + dataset = LeRobotDataset( cfg.dataset.repo_id, - root=cfg.dataset.dataset_root, - episodes=[cfg.dataset.episode], + root=cfg.dataset.root, + episodes=[cfg.dataset.replay_episode], download_videos=False, ) dataset_actions = dataset.hf_dataset.select_columns(["action"])