From 08c47eac4b3c3ed44328ff46603b340d94b6708c Mon Sep 17 00:00:00 2001
From: glannuzel <gaelle.lannuzel@gmail.com>
Date: Tue, 29 Jul 2025 10:26:03 +0200
Subject: [PATCH] Saving test scripts

---
 src/lerobot/robots/reachy2/test_reachy2.py   | 56 ++++++++++++
 src/lerobot/robots/reachy2/test_recording.py | 90 ++++++++++++++++++++
 2 files changed, 146 insertions(+)
 create mode 100644 src/lerobot/robots/reachy2/test_reachy2.py
 create mode 100644 src/lerobot/robots/reachy2/test_recording.py

diff --git a/src/lerobot/robots/reachy2/test_reachy2.py b/src/lerobot/robots/reachy2/test_reachy2.py
new file mode 100644
index 000000000..ab08216e4
--- /dev/null
+++ b/src/lerobot/robots/reachy2/test_reachy2.py
@@ -0,0 +1,56 @@
+from lerobot.robots.reachy2 import Reachy2Robot, Reachy2RobotConfig
+import time
+
+REACHY2_MOTORS = {
+    "neck_yaw.pos": "head.neck.yaw",
+    "neck_pitch.pos": "head.neck.pitch",
+    "neck_roll.pos": "head.neck.roll",
+    "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
+    "r_shoulder_roll.pos": "r_arm.shoulder.roll",
+    "r_elbow_yaw.pos": "r_arm.elbow.yaw",
+    "r_elbow_pitch.pos": "r_arm.elbow.pitch",
+    "r_wrist_roll.pos": "r_arm.wrist.roll",
+    "r_wrist_pitch.pos": "r_arm.wrist.pitch",
+    "r_wrist_yaw.pos": "r_arm.wrist.yaw",
+    "r_gripper.pos": "r_arm.gripper",
+    "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
+    "l_shoulder_roll.pos": "l_arm.shoulder.roll",
+    "l_elbow_yaw.pos": "l_arm.elbow.yaw",
+    "l_elbow_pitch.pos": "l_arm.elbow.pitch",
+    "l_wrist_roll.pos": "l_arm.wrist.roll",
+    "l_wrist_pitch.pos": "l_arm.wrist.pitch",
+    "l_wrist_yaw.pos": "l_arm.wrist.yaw",
+    "l_gripper.pos": "l_arm.gripper",
+    "l_antenna.pos": "head.l_antenna",
+    "r_antenna.pos": "head.r_antenna",
+    # "mobile_base.vx": "mobile_base.vx",
+    # "mobile_base.vy": "mobile_base.vy",
+    # "mobile_base.vtheta": "mobile_base.vtheta",
+}
+
+
+robot_config = Reachy2RobotConfig(ip_address="localhost")
+robot = Reachy2Robot(robot_config)
+
+robot.connect()
+
+print(f"is_connected(): {robot.is_connected}\n")
+
+print(f"_get_state(): {robot._get_state()}\n")
+
+obs = robot.get_observation()
+print(f"get_observation(): {obs}\n")
+print(f"observation_features: {robot.observation_features}\n")
+print(f"action_features: {robot.action_features}\n")
+
+
+def get_action(robot):
+    my_keys = REACHY2_MOTORS.keys()
+    my_values = [robot.reachy.joints[motor].present_position for motor in REACHY2_MOTORS.values()]
+    action = dict(zip(my_keys, my_values))
+    return action
+
+
+action = get_action(robot)
+time.sleep(5)
+robot.send_action(action)
diff --git a/src/lerobot/robots/reachy2/test_recording.py b/src/lerobot/robots/reachy2/test_recording.py
new file mode 100644
index 000000000..c4bd3b053
--- /dev/null
+++ b/src/lerobot/robots/reachy2/test_recording.py
@@ -0,0 +1,90 @@
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.robots.reachy2 import Reachy2Robot, Reachy2RobotConfig
+from lerobot.teleoperators.reachy2_fake_teleoperator import Reachy2FakeTeleoperator, Reachy2FakeTeleoperatorConfig
+from lerobot.utils.control_utils import init_keyboard_listener
+from lerobot.utils.utils import log_say
+from lerobot.utils.visualization_utils import _init_rerun
+from lerobot.record import record_loop
+
+NUM_EPISODES = 2
+FPS = 30
+EPISODE_TIME_SEC = 5
+RESET_TIME_SEC = 5
+TASK_DESCRIPTION = "My task description"
+
+# Create the robot and teleoperator configurations
+robot_config = Reachy2RobotConfig(
+    id="test_reachy"
+)
+teleop_config = Reachy2FakeTeleoperatorConfig()
+
+# Initialize the robot and teleoperator
+robot = Reachy2Robot(robot_config)
+teleop = Reachy2FakeTeleoperator(teleop_config)
+
+# 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="test/repo_test",
+    fps=FPS,
+    features=dataset_features,
+    robot_type=robot.name,
+    # use_videos=True,
+    # image_writer_threads=4,
+)
+
+# Initialize the keyboard listener and rerun visualization
+_, events = init_keyboard_listener()
+_init_rerun(session_name="recording")
+
+# Connect the robot and teleoperator
+robot.connect()
+teleop.connect()
+
+episode_idx = 0
+while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
+
+    record_loop(
+        robot=robot,
+        events=events,
+        fps=FPS,
+        teleop=teleop,
+        dataset=dataset,
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
+    )
+
+    # Reset the environment if not stopping or re-recording
+    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+        log_say("Reset the environment")
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=FPS,
+            teleop=teleop,
+            control_time_s=RESET_TIME_SEC,
+            single_task=TASK_DESCRIPTION,
+            display_data=True,
+        )
+
+    if events["rerecord_episode"]:
+        log_say("Re-recording episode")
+        events["rerecord_episode"] = False
+        events["exit_early"] = False
+        dataset.clear_episode_buffer()
+        continue
+
+    dataset.save_episode()
+    episode_idx += 1
+
+# Clean up
+log_say("Stop recording")
+robot.disconnect()
+# dataset.push_to_hub()
\ No newline at end of file