Fix replay script.

Signed-off-by: Franco Cipollone <franco.c@ekumenlabs.com>

Author: francocipollone

dora/node_hub/dora_run_policy/dora_run_policy/main.py

@@ -13,6 +13,7 @@
 import pyarrow as pa
 from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.policies.factory import make_pre_post_processors
 from lerobot.utils.control_utils import predict_action
 from lerobot.utils.utils import get_safe_torch_device
 
@@ -114,6 +115,11 @@ def main() -> None:
         policy.reset()
         device_name = policy.config.device or "auto"
         device = get_safe_torch_device(device_name)
+        # Build Policy Processors
+        preprocessor, postprocessor = make_pre_post_processors(
+            policy_cfg=policy.config,
+            pretrained_path=model_name,
+        )
     except Exception as e:
         raise RuntimeError(f"Failed to load policy '{model_name}': {e}") from None
 
@@ -171,9 +177,16 @@ def main() -> None:
                         observation_frame,
                         policy,
                         device,
+                        preprocessor,
+                        postprocessor,
                         policy.config.use_amp,
                     )
 
+                    # As of LeRobot 0.4.0, the postprocessor returns a tensor that might have batch dimension
+                    # Remove batch dimension if present
+                    if raw_action.dim() > 1:
+                        raw_action = raw_action.squeeze(0)
+
                     # Send action output
                     metadata = event["metadata"].copy()
                     metadata["primitive"] = "series"

dora/node_hub/dora_run_policy/pyproject.toml

@@ -9,7 +9,7 @@ requires-python = ">=3.10"
 
 dependencies = [
   "dora-rs >= 0.3.13",
-  "lerobot==0.3.3"
+  "lerobot==0.4.0"
 ]
 
 [build-system]

packages/lekiwi_lerobot/lekiwi_lerobot/evaluate.py

@@ -139,6 +139,7 @@ def main() -> None:
                 robot=robot,
                 events=events,
                 fps=FPS,
+                dataset=None,  # Don't record during reset phase
                 keyboard_handler=keyboard,
                 arm_keyboard_handler=arm_keyboard_handler,
                 control_time_s=RESET_TIME_SEC,

packages/lekiwi_lerobot/lekiwi_lerobot/record.py

@@ -5,13 +5,13 @@
 from lekiwi_teleoperate.teleoperate.arm import ArmTeleop
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
-from lerobot.processor import make_default_processors
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
 from lerobot.teleoperators.keyboard import (
     KeyboardTeleop,
     KeyboardTeleopConfig,
 )
+from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import (
     init_keyboard_listener,
 )
@@ -53,6 +53,12 @@ def main() -> None:
         default="Unnamed task",
         help="Task description to associate with each episode (default: 'Unnamed task').",
     )
+    parser.add_argument(
+        "--no-viz",
+        action="store_false",
+        dest="visualize",
+        help="Disable Rerun visualization during recording.",
+    )
 
     args = parser.parse_args()
     if args.repo_id is None:
@@ -74,8 +80,8 @@ def main() -> None:
     keyboard = KeyboardTeleop(keyboard_config)
     arm_keyboard_handler = ArmTeleop()
     # Configure the dataset features
-    action_features = hw_to_dataset_features(robot.action_features, "action")
-    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+    action_features = hw_to_dataset_features(robot.action_features, ACTION)
+    obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
     logging.info(f"Recording the following observation features: {list(obs_features.keys())}")
     logging.info(f"Recording the following action features: {list(action_features.keys())}")
     dataset_features = {**action_features, **obs_features}
@@ -87,7 +93,7 @@ def main() -> None:
         features=dataset_features,
         robot_type=robot.name,
         use_videos=True,
-        image_writer_threads=0,
+        image_writer_threads=4,
     )
 
     # To connect you already should have:
@@ -97,8 +103,11 @@ def main() -> None:
     robot.connect()
     keyboard.connect()
 
-    init_rerun(session_name="lekiwi_record")
-    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
+    if args.visualize:
+        logging.info("Initializing Rerun for visualization.")
+        init_rerun(session_name="lekiwi_record")
+    else:
+        logging.info("Rerun visualization is disabled.")
 
     listener, events = init_keyboard_listener()
 
@@ -119,21 +128,22 @@ def main() -> None:
             arm_keyboard_handler=arm_keyboard_handler,
             control_time_s=EPISODE_TIME_SEC,
             single_task=args.task,
-            display_data=True,
+            display_data=args.visualize,
         )
 
-        # Logic for reset env
+        # Reset the environment if not stopping or re-recording
         if not events["stop_recording"] and ((recorded_episodes < args.episodes - 1) or events["rerecord_episode"]):
             logging.info("Reset the environment")
             record_loop(
                 robot=robot,
                 events=events,
                 fps=FPS,
+                dataset=None,  # Don't record during reset phase
                 keyboard_handler=keyboard,
                 arm_keyboard_handler=arm_keyboard_handler,
                 control_time_s=RESET_TIME_SEC,
                 single_task=args.task,
-                display_data=True,
+                display_data=args.visualize,
             )
 
         if events["rerecord_episode"]:
@@ -143,10 +153,12 @@ def main() -> None:
             dataset.clear_episode_buffer()
             continue
 
+        logging.info(f"Saving episode number {recorded_episodes} to dataset.")
         dataset.save_episode()
         recorded_episodes += 1
 
     # Upload to hub and clean up
+    dataset.finalize()
     dataset.push_to_hub()
 
     robot.disconnect()

packages/lekiwi_lerobot/lekiwi_lerobot/replay.py

@@ -5,6 +5,7 @@
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
+from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 
 
@@ -84,17 +85,22 @@ def main() -> None:
     logging.info(f"Downloading dataset from {args.repo_id} into {args.directory}")
     root = args.directory + "/" + args.repo_id.split("/")[-1]
     root = root.replace("//", "/")
-    dataset = LeRobotDataset(args.repo_id, root=root, episodes=[args.episode])
+    episode_to_replay = args.episode
+    dataset = LeRobotDataset(args.repo_id, root=root, episodes=[episode_to_replay])
     logging.info(f"Dataset stored at {root}")
     actions = dataset.hf_dataset.select_columns("action")
+    # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
+    episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == episode_to_replay)
+    actions = episode_frames.select_columns(ACTION)
 
     robot.connect()
 
     if not robot.is_connected:
         raise ValueError("Robot is not connected!")
 
-    logging.info(f"Replaying episode {args.episode} with {dataset.num_frames} frames.")
-    for idx in range(dataset.num_frames):
+    len_episodes_frames = len(episode_frames)
+    logging.info(f"Replaying episode {args.episode} with {len_episodes_frames} frames.")
+    for idx in range(len_episodes_frames):
         t0 = time.perf_counter()
 
         action = {name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])}

packages/lekiwi_lerobot/lekiwi_lerobot/run_policy.py

@@ -7,6 +7,7 @@
 from lerobot.policies.factory import make_pre_post_processors
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
+from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener, predict_action
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import get_safe_torch_device
@@ -83,14 +84,14 @@ def main() -> None:
     logging.info("Robot is connected.")
 
     # Prepare for policy inference
-    action_features = hw_to_dataset_features(robot.action_features, "action")
-    obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+    action_features = hw_to_dataset_features(robot.action_features, ACTION)
+    obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
     dataset_features = {**action_features, **obs_features}
     device = get_safe_torch_device(policy.config.device)
     # Build Policy Processors
     preprocessor, postprocessor = make_pre_post_processors(
         policy_cfg=policy.config,
-        pretrained_path=None,
+        pretrained_path=args.policy,
     )
     init_rerun(session_name="lekiwi_run_policy")
 
@@ -102,7 +103,7 @@ def main() -> None:
 
         observation = robot.get_observation()
 
-        observation_frame = build_dataset_frame(dataset_features, observation, prefix="observation")
+        observation_frame = build_dataset_frame(dataset_features, observation, prefix=OBS_STR)
 
         action_values = predict_action(
             observation_frame,
@@ -114,6 +115,11 @@ def main() -> None:
             task=args.task,
             robot_type=robot.robot_type,
         )
+        # As of LeRobot 0.4.0, the postprocessor returns a tensor that might have batch dimension
+        # Remove batch dimension if present
+        if action_values.dim() > 1:
+            action_values = action_values.squeeze(0)
+
         action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
 
         logging.debug(f"Predicted action: {action}")

packages/lekiwi_lerobot/lekiwi_lerobot/utils.py

@@ -8,9 +8,8 @@
 from lerobot.datasets.utils import build_dataset_frame
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.processor import (
-    RobotAction,
-    RobotObservation,
-    RobotProcessorPipeline,
+    PolicyAction,
+    PolicyProcessorPipeline,
 )
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
 from lerobot.teleoperators.keyboard import (
@@ -31,17 +30,12 @@ def record_loop(
     robot: LeKiwiClient,
     events: dict[Any, Any],
     fps: int,
-    teleop_action_processor: RobotProcessorPipeline[
-        tuple[RobotAction, RobotObservation], RobotAction
-    ],  # runs after teleop
-    robot_action_processor: RobotProcessorPipeline[
-        tuple[RobotAction, RobotObservation], RobotAction
-    ],  # runs before robot
-    robot_observation_processor: RobotProcessorPipeline[RobotObservation, RobotObservation],  # runs after robot
     dataset: LeRobotDataset | None = None,
     keyboard_handler: KeyboardTeleop | None = None,
     arm_keyboard_handler: ArmTeleop | None = None,
     policy: PreTrainedPolicy | None = None,
+    preprocessor: PolicyProcessorPipeline[dict[str, Any], dict[str, Any]] | None = None,
+    postprocessor: PolicyProcessorPipeline[PolicyAction, PolicyAction] | None = None,
     control_time_s: int | None = None,
     single_task: str | None = None,
     display_data: bool = False,
@@ -64,6 +58,11 @@ def record_loop(
     if control_time_s is None:
         raise ValueError("A control time must be provided.")
 
+    if display_data:
+        logging.info("Visualizing data with Rerun.")
+    else:
+        logging.info("Not visualizing data.")
+
     # if policy is given it needs cleaning up
     if policy is not None:
         policy.reset()
@@ -84,17 +83,19 @@ def record_loop(
                 raise ValueError("Dataset features must be defined if using a dataset or a policy.")
             observation_frame = build_dataset_frame(dataset.features, observation, prefix="observation")  # type: ignore[union-attr]
 
-        if policy is not None:
+        if policy is not None and preprocessor is not None and postprocessor is not None:
             action_values = predict_action(
-                observation_frame,
-                policy,
-                get_safe_torch_device(policy.config.device),
-                policy.config.use_amp,
+                observation=observation_frame,
+                policy=policy,
+                device=get_safe_torch_device(policy.config.device),
                 preprocessor=preprocessor,
                 postprocessor=postprocessor,
+                use_amp=policy.config.use_amp,
                 task=single_task,
                 robot_type=robot.robot_type,
             )
+            if action_values.dim() > 1:
+                action_values = action_values.squeeze(0)
             action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
             print(f"Predicted action: {action}")
         elif policy is None and keyboard_handler is not None and arm_keyboard_handler is not None:
@@ -103,6 +104,8 @@ def record_loop(
             arm_action = arm_keyboard_handler.from_keyboard_to_arm_action(pressed_keys)
 
             action = {**base_action, **arm_action}  # Merge base and arm actions
+            # TODO(francocipollone): We would probably want to use the teleop_action_processor here.
+            # action = teleop_action_processor((action, observation))
             logging.debug("Sending action: %s", action)
 
         else:
@@ -113,14 +116,13 @@ def record_loop(
             )
             continue
 
-        # Action can eventually be clipped using `max_relative_target`,
-        # so action actually sent is saved in the dataset.
+        # TODO(francocipollone): We would probably want to use the robot_action_processor here before sending the action
         sent_action = robot.send_action(action)
 
         if dataset is not None:
             action_frame = build_dataset_frame(dataset.features, sent_action, prefix="action")
-            frame = {**observation_frame, **action_frame}
-            dataset.add_frame(frame, task=single_task)
+            frame = {**observation_frame, **action_frame, "task": single_task}
+            dataset.add_frame(frame)
 
         if display_data:
             log_rerun_data(observation, action)

packages/lekiwi_sim/lekiwi_sim/lekiwi_sim_host.py

@@ -106,7 +106,7 @@ def main() -> None:
                 watchdog_active = False
             except zmq.Again:
                 if not watchdog_active:
-                    logging.warning("No command available")
+                    logging.debug("No command available")
             except Exception as e:
                 logging.error("Message fetching failed: %s", e)