Refactor dataset configuration in documentation and codebase

- Updated dataset configuration keys from `dataset_root` to `root` and `num_episodes` to `num_episodes_to_record` for consistency.
- Adjusted replay episode handling by renaming `episode` to `replay_episode`.
- Enhanced documentation
- added specific processor to transform from policy actions to delta actions

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
 ```
 <!-- prettier-ignore-end -->
 
@@ -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",

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"

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"`

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",

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

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)

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

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"])