diff --git a/src/lerobot/robots/reachy2/test_policy.py b/src/lerobot/robots/reachy2/test_policy.py index 3005592ad..d285ea088 100644 --- a/src/lerobot/robots/reachy2/test_policy.py +++ b/src/lerobot/robots/reachy2/test_policy.py @@ -11,32 +11,55 @@ from reachy2_sdk import ReachySDK import numpy as np import time -NUM_EPISODES = 5 -FPS = 20 -EPISODE_TIME_SEC = 10 -TASK_DESCRIPTION = "Grab a cube in Mujoco simulation" + +NUM_EPISODES = 2 +FPS = 15 +EPISODE_TIME_SEC = 4 +TASK_DESCRIPTION = "Grab a cube with Reachy 2" + + +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", +} # Create the robot configuration robot_config = Reachy2RobotConfig( - ip_address="localhost", - # ip_address="172.18.131.66", + ip_address="192.168.0.199", id="test_reachy", ) # Initialize the robot robot = Reachy2Robot(robot_config) - -reachy = ReachySDK("localhost") -reachy.turn_on() -reachy.mobile_base.goto(-0.2, -0.3, 0, wait=True) -time.sleep(2) -reachy.r_arm.goto_posture("elbow_90", wait=True) -reachy.r_arm.gripper.open() -reachy.mobile_base.goto(0, -0.3, 0) +# Instantiate a client for starting and intermediate positions +reachy = ReachySDK(robot_config.ip_address) # Initialize the policy -policy = ACTPolicy.from_pretrained("pepijn223/grab_cube_simulation_2") +policy = ACTPolicy.from_pretrained("pepijn223/grab_cube_2") + +# Get initial dataset first episode +initial_dataset = LeRobotDataset(repo_id="glannuzel/grab_cube_2", episodes=[0]) +actions = initial_dataset.hf_dataset.select_columns("action") # Configure the dataset features action_features = hw_to_dataset_features(robot.action_features, "action") @@ -45,7 +68,7 @@ dataset_features = {**action_features, **obs_features} # Create the dataset dataset = LeRobotDataset.create( - repo_id="glannuzel/eval_grab_cube_simulation_2", + repo_id="glannuzel/eval_grab_cube_2.1", fps=FPS, features=dataset_features, robot_type=robot.name, @@ -60,11 +83,33 @@ _, events = init_keyboard_listener() # Connect the robot robot.connect() -M = reachy.r_arm.get_default_posture_matrix("elbow_90") -np.round(M, 3) -first_pose = M.copy() -first_pose[0, 3] += 0.05 -first_pose[1, 3] += 0.1 +# Go to the initial pose +action_array = actions[0]["action"] +action = {} +for i, name in enumerate(dataset.features["action"]["names"]): + action[name] = action_array[i].item() + +neck_goal = [action["neck_roll.pos"], action["neck_pitch.pos"], action["neck_yaw.pos"]] +r_arm_goal = [action["r_shoulder_pitch.pos"], + action["r_shoulder_roll.pos"], + action["r_elbow_yaw.pos"], + action["r_elbow_pitch.pos"], + action["r_wrist_roll.pos"], + action["r_wrist_pitch.pos"], + action["r_wrist_yaw.pos"]] +l_arm_goal = [action["l_shoulder_pitch.pos"], + action["l_shoulder_roll.pos"], + action["l_elbow_yaw.pos"], + action["l_elbow_pitch.pos"], + action["l_wrist_roll.pos"], + action["l_wrist_pitch.pos"], + action["l_wrist_yaw.pos"]] + +reachy.head.goto(neck_goal) +reachy.r_arm.goto(r_arm_goal) +reachy.r_arm.gripper.goto(100.0) +reachy.l_arm.goto(l_arm_goal, wait=True) +time.sleep(1.0) for episode_idx in range(NUM_EPISODES): log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") @@ -81,13 +126,15 @@ for episode_idx in range(NUM_EPISODES): display_data=False, ) - reachy.r_arm.gripper.goto(100, percentage=True, wait=True) - reachy.head.goto_posture() - reachy.r_arm.goto(first_pose) - reachy.r_arm.goto(M, wait=True) + # Set the robot back to the initial pose + reachy.head.goto(neck_goal) + reachy.r_arm.goto(r_arm_goal) + reachy.r_arm.gripper.goto(100.0) + reachy.l_arm.goto(l_arm_goal, wait=True) + time.sleep(1.0) dataset.save_episode() # Clean up robot.disconnect() -# dataset.push_to_hub() \ No newline at end of file +dataset.push_to_hub() diff --git a/src/lerobot/robots/reachy2/test_replay.py b/src/lerobot/robots/reachy2/test_replay.py new file mode 100644 index 000000000..20d930990 --- /dev/null +++ b/src/lerobot/robots/reachy2/test_replay.py @@ -0,0 +1,101 @@ +from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig +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.robots.reachy2 import Reachy2Robot, Reachy2RobotConfig +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 +from reachy2_sdk import ReachySDK +from lerobot.utils.robot_utils import busy_wait +import numpy as np +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", +} + +# Create the robot configuration +robot_config = Reachy2RobotConfig( + # ip_address="localhost", + # ip_address="172.18.131.66", + ip_address="192.168.0.199", + id="reachy2-pvt02", +) + +# Initialize the robot +robot = Reachy2Robot(robot_config) + + +reachy = ReachySDK(robot_config.ip_address) + + +# Create the dataset +dataset = LeRobotDataset(repo_id="glannuzel/grab_cube_2", episodes=[0]) +actions = dataset.hf_dataset.select_columns("action") + +# Connect the robot +robot.connect() + +action_array = actions[0]["action"] +action = {} +for i, name in enumerate(dataset.features["action"]["names"]): + action[name] = action_array[i].item() + +neck_goal = [action["neck_roll.pos"], action["neck_pitch.pos"], action["neck_yaw.pos"]] +r_arm_goal = [action["r_shoulder_pitch.pos"], + action["r_shoulder_roll.pos"], + action["r_elbow_yaw.pos"], + action["r_elbow_pitch.pos"], + action["r_wrist_roll.pos"], + action["r_wrist_pitch.pos"], + action["r_wrist_yaw.pos"]] +l_arm_goal = [action["l_shoulder_pitch.pos"], + action["l_shoulder_roll.pos"], + action["l_elbow_yaw.pos"], + action["l_elbow_pitch.pos"], + action["l_wrist_roll.pos"], + action["l_wrist_pitch.pos"], + action["l_wrist_yaw.pos"]] + +reachy.head.goto(neck_goal) +reachy.r_arm.goto(r_arm_goal) +reachy.l_arm.goto(l_arm_goal, wait=True) + +for idx in range(dataset.num_frames): + start_episode_t = time.perf_counter() + + action_array = actions[idx]["action"] + + action = {} + for i, name in enumerate(dataset.features["action"]["names"]): + action[name] = action_array[i].item() + + robot.send_action(action) + dt_s = time.perf_counter() - start_episode_t + busy_wait(1 / dataset.fps - dt_s) + +# Clean up +# robot.disconnect()