From 20f0381f81488592e89c52de453c2847fcee7bdb Mon Sep 17 00:00:00 2001 From: Pepijn Date: Fri, 2 Jan 2026 10:18:59 +0100 Subject: [PATCH] wait for takeover press --- examples/rac/rac_data_collection.py | 42 ++++++++++++++++---- examples/rac/rac_data_collection_openarms.py | 29 ++++++++++---- 2 files changed, 55 insertions(+), 16 deletions(-) diff --git a/examples/rac/rac_data_collection.py b/examples/rac/rac_data_collection.py index 3e5b244c1..b4c61f576 100644 --- a/examples/rac/rac_data_collection.py +++ b/examples/rac/rac_data_collection.py @@ -317,6 +317,9 @@ def rac_rollout_loop( } last_robot_action = None + was_paused = False + was_correction_active = False + waiting_for_takeover = False timestamp = 0 start_t = time.perf_counter() @@ -329,6 +332,21 @@ def rac_rollout_loop( events["correction_active"] = False break + # Detect transition to paused state + if events["policy_paused"] and not was_paused: + obs = robot.get_observation() + robot_pos = {k: v for k, v in obs.items() if k.endswith(".pos")} + print("[RaC] Moving teleop to robot position (2s smooth transition)...") + teleop.smooth_move_to(robot_pos, duration_s=2.0, fps=50) + print("[RaC] Teleop aligned. Press 'c' to take control.") + waiting_for_takeover = True + was_paused = True + + # Detect transition to correction mode + if events["correction_active"] and not was_correction_active: + waiting_for_takeover = False + was_correction_active = True + obs = robot.get_observation() obs_frame = build_dataset_frame(dataset.features, obs, prefix=OBS_STR) @@ -344,8 +362,8 @@ def rac_rollout_loop( frame_buffer.append(frame) stats["total_frames"] += 1 - elif events["policy_paused"]: - # Paused - hold last position, don't record + elif events["policy_paused"] and not waiting_for_takeover: + # Paused and user acknowledged - hold last position, don't record if last_robot_action is not None: robot.send_action(last_robot_action) stats["paused_frames"] += 1 @@ -394,9 +412,9 @@ def reset_loop( events: dict, fps: int, ): - """Reset period where human repositions environment. Waits for user to start next episode.""" + """Reset period where human repositions environment. Two-stage: enable teleop, then start episode.""" print("\n" + "=" * 65) - print(" [RaC] RESET - Move robot to starting position") + print(" [RaC] RESET - Moving teleop to robot position...") print("=" * 65) # Enter reset mode @@ -406,13 +424,21 @@ def reset_loop( # Move teleop to match robot position to avoid sudden jumps obs = robot.get_observation() robot_pos = {k: v for k, v in obs.items() if k.endswith(".pos")} - print(" Moving teleop to robot position...") teleop.smooth_move_to(robot_pos, duration_s=2.0, fps=50) + # Stage 1: Wait for user to press start to enable teleoperation + print(" Teleop aligned. Press any key/pedal to enable teleoperation") + while not events["start_next_episode"] and not events["stop_recording"]: + precise_sleep(0.05) + + if events["stop_recording"]: + return + + # Stage 2: Enable teleop and let user move robot to starting position + events["start_next_episode"] = False teleop.disable_torque() - print(" Teleop active - move robot to starting position") - print(" Press SPACE / → / 'c' or any pedal to start next episode") - print("=" * 65) + print(" Teleop enabled - move robot to starting position") + print(" Press any key/pedal to start next episode") # Wait for user to signal ready for next episode while not events["start_next_episode"] and not events["stop_recording"]: diff --git a/examples/rac/rac_data_collection_openarms.py b/examples/rac/rac_data_collection_openarms.py index 276107648..ae19839d1 100644 --- a/examples/rac/rac_data_collection_openarms.py +++ b/examples/rac/rac_data_collection_openarms.py @@ -328,6 +328,7 @@ def rac_rollout_loop( teleop.disable_torque() was_paused = False was_correction_active = False + waiting_for_takeover = False timestamp = 0 start_t = time.perf_counter() @@ -346,14 +347,16 @@ def rac_rollout_loop( obs = robot.get_observation() obs_filtered = {k: v for k, v in obs.items() if k in robot.observation_features} robot_pos = {k: v for k, v in obs_filtered.items() if k.endswith(".pos")} - print("[RaC] Moving teleop to robot position (1s smooth transition)...") + print("[RaC] Moving teleop to robot position (2s smooth transition)...") teleop.smooth_move_to(robot_pos, duration_s=2.0, fps=50) print("[RaC] Teleop aligned. Press 'c' to take control.") + waiting_for_takeover = True was_paused = True # Detect transition to correction mode (disable torque for human control) if events["correction_active"] and not was_correction_active: teleop.disable_torque() + waiting_for_takeover = False was_correction_active = True obs = robot.get_observation() @@ -376,8 +379,8 @@ def rac_rollout_loop( frame_buffer.append(frame) stats["total_frames"] += 1 - elif events["policy_paused"]: - # Paused - teleop tracks robot position, don't record + elif events["policy_paused"] and not waiting_for_takeover: + # Paused and user acknowledged - teleop tracks robot position, don't record robot_action = {k: v for k, v in obs_filtered.items() if k.endswith(".pos")} teleop.send_feedback(robot_action) stats["paused_frames"] += 1 @@ -426,9 +429,9 @@ def reset_loop( events: dict, fps: int, ): - """Reset period where human repositions environment. Waits for user to start next episode.""" + """Reset period where human repositions environment. Two-stage: enable teleop, then start episode.""" print("\n" + "=" * 65) - print(" [RaC] RESET - Move robot to starting position") + print(" [RaC] RESET - Moving teleop to robot position...") print("=" * 65) # Enter reset mode @@ -438,11 +441,21 @@ def reset_loop( # First move teleop to match robot position to avoid sudden jumps obs = robot.get_observation() robot_pos = {k: v for k, v in obs.items() if k.endswith(".pos") and k in robot.observation_features} - print(" Moving teleop to robot position...") teleop.smooth_move_to(robot_pos, duration_s=2.0, fps=50) + + # Stage 1: Wait for user to press start to enable teleoperation + print(" Teleop aligned. Press any key/pedal to enable teleoperation") + while not events["start_next_episode"] and not events["stop_recording"]: + precise_sleep(0.05) + + if events["stop_recording"]: + return + + # Stage 2: Enable teleop and let user move robot to starting position + events["start_next_episode"] = False teleop.disable_torque() - print(" Teleop active - move robot to starting position") - print(" Press SPACE / → / 'c' or any pedal to start next episode") + print(" Teleop enabled - move robot to starting position") + print(" Press any key/pedal to start next episode") # Wait for user to signal ready for next episode while not events["start_next_episode"] and not events["stop_recording"]: