Merge branch 'main' into user/michel-aractingi/2025_06_30_dataset_v3

Signed-off-by: Michel Aractingi <michel.aractingi@huggingface.co>

examples/lekiwi/evaluate.py

@@ -1,32 +1,90 @@
-from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.utils.control_utils import predict_action
-from lerobot.utils.utils import get_safe_torch_device
+from lerobot.utils.control_utils import init_keyboard_listener
+from lerobot.utils.utils import log_say
+from lerobot.utils.visualization_utils import _init_rerun
 
-NB_CYCLES_CLIENT_CONNECTION = 1000
+NUM_EPISODES = 2
+FPS = 30
+EPISODE_TIME_SEC = 60
+TASK_DESCRIPTION = "My task description"
 
+# Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 robot = LeKiwiClient(robot_config)
 
+policy = ACTPolicy.from_pretrained("<hf_username>/<policy_repo_id>")
+
+# Configure the dataset features
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+dataset_features = {**action_features, **obs_features}
+
+# Create the dataset
+dataset = LeRobotDataset.create(
+    repo_id="<hf_username>/<eval_dataset_repo_id>",
+    fps=FPS,
+    features=dataset_features,
+    robot_type=robot.name,
+    use_videos=True,
+    image_writer_threads=4,
+)
+
+# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
 
-policy = ACTPolicy.from_pretrained("pepijn223/act_lekiwi_circle")
-policy.reset()
+_init_rerun(session_name="recording")
 
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+listener, events = init_keyboard_listener()
 
-print("Running inference")
-i = 0
-while i < NB_CYCLES_CLIENT_CONNECTION:
-    obs = robot.get_observation()
+if not robot.is_connected:
+    raise ValueError("Robot is not connected!")
 
-    observation_frame = build_dataset_frame(obs_features, obs, prefix="observation")
-    action_values = predict_action(
-        observation_frame, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
+recorded_episodes = 0
+while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
+    log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
+
+    # Run the policy inference loop
+    record_loop(
+        robot=robot,
+        events=events,
+        fps=FPS,
+        policy=policy,
+        dataset=dataset,
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
     )
-    action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
-    robot.send_action(action)
-    i += 1
+
+    # Logic for reset env
+    if not events["stop_recording"] and (
+        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+    ):
+        log_say("Reset the environment")
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=FPS,
+            control_time_s=EPISODE_TIME_SEC,
+            single_task=TASK_DESCRIPTION,
+            display_data=True,
+        )
+
+    if events["rerecord_episode"]:
+        log_say("Re-record episode")
+        events["rerecord_episode"] = False
+        events["exit_early"] = False
+        dataset.clear_episode_buffer()
+        continue
+
+    dataset.save_episode()
+    recorded_episodes += 1
+
+# Upload to hub and clean up
+dataset.push_to_hub()
 
 robot.disconnect()
+listener.stop()

examples/lekiwi/replay.py

@@ -4,22 +4,30 @@ 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.robot_utils import busy_wait
+from lerobot.utils.utils import log_say
+
+EPISODE_IDX = 0
 
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 robot = LeKiwiClient(robot_config)
 
-dataset = LeRobotDataset("pepijn223/lekiwi1749025613", episodes=[0])
+dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
+actions = dataset.hf_dataset.select_columns("action")
 
 robot.connect()
 
-print("Replaying episode…")
-for _, action_array in enumerate(dataset.hf_dataset["action"]):
+if not robot.is_connected:
+    raise ValueError("Robot is not connected!")
+
+log_say(f"Replaying episode {EPISODE_IDX}")
+for idx in range(dataset.num_frames):
     t0 = time.perf_counter()
 
-    action = {name: float(action_array[i]) for i, name in enumerate(dataset.features["action"]["names"])}
+    action = {
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
+    }
     robot.send_action(action)
 
     busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
 
-print("Disconnecting LeKiwi Client")
 robot.disconnect()

examples/lekiwi/teleoperate.py

@@ -1,32 +1,47 @@
+import time
+
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
+FPS = 30
+
+# Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
-
-teleop__arm_config = SO100LeaderConfig(
-    port="/dev/tty.usbmodem58760431551",
-    id="my_awesome_leader_arm",
-)
-
-teleop_keyboard_config = KeyboardTeleopConfig(
-    id="my_laptop_keyboard",
-)
+teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
+keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
 
 robot = LeKiwiClient(robot_config)
-teleop_arm = SO100Leader(teleop__arm_config)
-telep_keyboard = KeyboardTeleop(teleop_keyboard_config)
+leader_arm = SO100Leader(teleop_arm_config)
+keyboard = KeyboardTeleop(keyboard_config)
+
+# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
-teleop_arm.connect()
-telep_keyboard.connect()
+leader_arm.connect()
+keyboard.connect()
+
+_init_rerun(session_name="lekiwi_teleop")
+
+if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
+    raise ValueError("Robot, leader arm of keyboard is not connected!")
 
 while True:
+    t0 = time.perf_counter()
+
     observation = robot.get_observation()
 
-    arm_action = teleop_arm.get_action()
+    arm_action = leader_arm.get_action()
     arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
 
-    keyboard_keys = telep_keyboard.get_action()
+    keyboard_keys = keyboard.get_action()
     base_action = robot._from_keyboard_to_base_action(keyboard_keys)
 
-    robot.send_action(arm_action | base_action)
+    log_rerun_data(observation, {**arm_action, **base_action})
+
+    action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
+
+    robot.send_action(action)
+
+    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))