diff --git a/source/geniesim_benchmark/src/geniesim_benchmark/app/controllers/api_core.py b/source/geniesim_benchmark/src/geniesim_benchmark/app/controllers/api_core.py index 6fe17a2e..f94dc308 100644 --- a/source/geniesim_benchmark/src/geniesim_benchmark/app/controllers/api_core.py +++ b/source/geniesim_benchmark/src/geniesim_benchmark/app/controllers/api_core.py @@ -2752,6 +2752,31 @@ def on_ros_tick(self, step_size): self.teleop_recording = self.server_ros_node.get_teleop_recording() def _on_recording(self): + # Multi-episode recording: autoteleop.sh publishes on /sim/stop_episode + # after the operator confirms an episode with y/n. On that signal: + # finalize the current ros2 bag gracefully (SIGINT, so metadata.yaml is + # written) and re-arm the one-shot latches so the next record-button + # press starts a fresh episode. Without this, wait_recording locks after + # the first episode and the whole simulator must be restarted per episode. + server_node = getattr(self, "server_ros_node", None) + if server_node is not None and server_node.get_stop_episode(): + server_node.clear_stop_episode() + if self.recording_started: + if self.recording_wait_num < 100: + # Episode stopped before recording_info.json was written at + # frame 100; write it now, otherwise post-processing skips + # the episode silently for lack of an info file. + self._dump_recording_info() + self._stop_recording() + logger.info("Episode finished, ready for next recording") + self.wait_recording = True + self.recording_wait_num = 0 + # Drop the cached record-button state and this frame's already-read + # copy of it, so the next episode waits for a fresh button press + # instead of auto-starting from stale signals. + server_node.reset_recording_state() + self.teleop_recording = False + if self.teleop_recording and self.wait_recording: self.set_record_topics() camera_prim_list = self.task_config["recording_setting"]["camera_list"] diff --git a/source/geniesim_benchmark/src/geniesim_benchmark/utils/ros_nodes/server_node.py b/source/geniesim_benchmark/src/geniesim_benchmark/utils/ros_nodes/server_node.py index 8bf401d5..c463c9c1 100755 --- a/source/geniesim_benchmark/src/geniesim_benchmark/utils/ros_nodes/server_node.py +++ b/source/geniesim_benchmark/src/geniesim_benchmark/utils/ros_nodes/server_node.py @@ -33,13 +33,22 @@ def __init__(self, robot_name="G1_120s", node_name="server_ros_node"): self.subscriber_teleop_recording = self.create_subscription( Bool, "/sim/is_recording", self.callback_recording, 1 ) + # Multi-episode recording: autoteleop.sh publishes on this topic when the + # operator confirms an episode with y/n. The simulator then finalizes the + # current ros2 bag gracefully (instead of being killed) and re-arms for the + # next record-button press. + self.subscriber_stop_episode = self.create_subscription( + Bool, "/sim/stop_episode", self.callback_stop_episode, 1 + ) self.pub_clock = self.create_publisher(Clock, "/clock", 1) self.playback_msg = False self.reset_msg = False self.recording_msg = False + self.stop_episode_msg = False self.playback_lock = threading.Lock() self.reset_lock = threading.Lock() self.recording_lock = threading.Lock() + self.stop_episode_lock = threading.Lock() def publish_clock(self, time_in_s): self.sec = int(time_in_s) @@ -69,3 +78,26 @@ def get_reset(self): def get_teleop_recording(self): with self.recording_lock: return self.recording_msg + + # ---- Multi-episode recording ------------------------------------------- + def callback_stop_episode(self, msg): + # Latch only; api_core clears it via clear_stop_episode() once handled, + # so the signal is processed exactly once. + if msg.data: + with self.stop_episode_lock: + self.stop_episode_msg = True + + def get_stop_episode(self): + with self.stop_episode_lock: + return self.stop_episode_msg + + def clear_stop_episode(self): + with self.stop_episode_lock: + self.stop_episode_msg = False + + def reset_recording_state(self): + # After an episode ends, drop the cached "record button pressed" state. + # Without this, re-arming wait_recording would immediately start the next + # episode from the stale True value instead of waiting for a new press. + with self.recording_lock: + self.recording_msg = False diff --git a/source/geniesim_teleop/scripts/autoteleop.sh b/source/geniesim_teleop/scripts/autoteleop.sh index 983d20d3..5d1c2d3c 100755 --- a/source/geniesim_teleop/scripts/autoteleop.sh +++ b/source/geniesim_teleop/scripts/autoteleop.sh @@ -15,7 +15,9 @@ # 3. the motion-control binary (genie_motion_control) # 4. the teleop loop (opens the VR server, waits for the Pico headset) # -# Finally it waits for y/N to decide whether the recorded episode is kept. +# Finally it enters a multi-episode loop: press the controller's record button +# to start an episode, then y/n here to keep/discard it and re-arm for the next +# episode, and q to stop everything. # Container name + repo mount point. The new geniesim CLI container is # named "geniesim3" and mounts the repo at /workspace (the legacy @@ -137,31 +139,82 @@ for i in "${!COMMANDS[@]}"; do fi done -# --- Wait for keep/discard decision ---------------------------------------- -echo -e "\nAll terminals started. Press 'y' or 'Y' = teleoperation succeeded, keep data; 'n' or 'N' = failed, do not keep data ..." +# --- Multi-episode keep/discard loop --------------------------------------- +# Legacy behaviour: a single y/N press killed every teleop process (simulator +# included) and the script exited, so collecting N episodes meant restarting +# all four terminals N times. Now y/n only ends the *current* episode: the +# simulator finalizes the ros2 bag gracefully (metadata.yaml intact) and +# re-arms, so the next record-button press starts the next episode. Press q +# to stop everything and close the terminals. + +RECORDING_BASE="${REPO_IN_CONTAINER}/output/recording_data" + +# End the current episode: broadcast on /sim/stop_episode (both the simulator +# and the teleop loop subscribe to it). +stop_current_episode() { + echo " -> asking the simulator to finalize the current episode..." + docker exec "$CONTAINER_NAME" bash -c \ + "source /opt/ros/jazzy/setup.bash && timeout 15 ros2 topic pub --times 3 /sim/stop_episode std_msgs/msg/Bool '{data: true}' >/dev/null 2>&1" \ + || echo " warning: failed to send stop signal (container or ROS not ready?)" + sleep 2 # give ros2 bag a moment to write metadata.yaml +} + +# Show whether the latest episode looks complete (.mcap = data recorded, +# metadata.yaml = bag finalized, recording_info.json = post-processing input). +show_latest_episode_status() { + docker exec "$CONTAINER_NAME" bash -c " + latest=\$(find $RECORDING_BASE -mindepth 2 -maxdepth 2 -type d -printf '%T@ %p\n' 2>/dev/null | sort -n | tail -1 | cut -d' ' -f2-) + if [ -z \"\$latest\" ]; then + echo ' status: no recording data yet (did you forget the record button?)' + exit 0 + fi + echo \" latest episode dir: \$latest\" + if ls \"\$latest\"/*.mcap >/dev/null 2>&1; then mcap=ok; else mcap=MISSING; fi + if [ -f \"\$latest/metadata.yaml\" ]; then meta=ok; else meta=MISSING; fi + if [ -f \"\$latest/recording_info.json\" ]; then info=ok; else info=MISSING; fi + echo \" .mcap: \$mcap | metadata.yaml: \$meta | recording_info.json: \$info\" + " 2>/dev/null || echo " status check failed (container not running?)" +} + +echo "" +echo "==================================================================" +echo "All terminals started. Controls:" +echo " 1. press the controller's record button to start an episode" +echo " 2. when the episode is done, come back here and press:" +echo " y = success: keep the data, then record the next episode" +echo " n = failure: mark the data discarded, then record the next episode" +echo " q = quit: stop everything and close the terminals" +echo "==================================================================" while read -n 1 -s input; do if [[ "$input" == "Y" || "$input" == "y" ]]; then - echo "Save the remote operation data.....Congratulations!" - echo -e "Sending SIGTERM to teleop processes..." - docker exec "$CONTAINER_NAME" bash -c "pkill -SIGTERM -f '$PROCESS_CLIENT' 2>/dev/null || true" - sleep 1 - echo "Patching recording_info.json: add teleop_result" + echo "" + echo "[y] keeping this episode..." + stop_current_episode docker exec "$CONTAINER_NAME" python3 ${TELEOP_PKG}/data_recording/patch_recording_info.py \ --config ${TELEOP_YAML} \ - --base ${REPO_IN_CONTAINER}/output/recording_data \ + --base ${RECORDING_BASE} \ || true - - break + show_latest_episode_status + echo " done. Press the record button for the next episode, or y/n/q here." elif [[ "$input" == "N" || "$input" == "n" ]]; then - echo -e "Sending SIGTERM to teleop processes..." - docker exec "$CONTAINER_NAME" bash -c "pkill -SIGTERM -f '$PROCESS_CLIENT' 2>/dev/null || true" - sleep 1 - echo "Patching recording_info.json: add teleop_result=false" + echo "" + echo "[n] discarding this episode (teleop_result=false; post-processing drops it)..." + stop_current_episode docker exec "$CONTAINER_NAME" python3 ${TELEOP_PKG}/data_recording/patch_recording_info.py \ --config ${TELEOP_YAML} \ - --base ${REPO_IN_CONTAINER}/output/recording_data \ + --base ${RECORDING_BASE} \ --teleop-result false \ || true + show_latest_episode_status + echo " done. Press the record button for the next episode, or y/n/q here." + elif [[ "$input" == "Q" || "$input" == "q" ]]; then + echo "" + echo "[q] stopping all teleop processes..." + # Finalize any episode still recording first (avoids a corrupt bag), + # then kill the processes. + stop_current_episode + docker exec "$CONTAINER_NAME" bash -c "pkill -SIGTERM -f '$PROCESS_CLIENT' 2>/dev/null || true" + sleep 1 break fi done diff --git a/source/geniesim_teleop/src/geniesim_teleop/teleop.py b/source/geniesim_teleop/src/geniesim_teleop/teleop.py index 30a102a2..b2af2007 100644 --- a/source/geniesim_teleop/src/geniesim_teleop/teleop.py +++ b/source/geniesim_teleop/src/geniesim_teleop/teleop.py @@ -404,6 +404,18 @@ def run(self): self.input = self.device.update() # logger.info(f"Input {self.input}") + + # Multi-episode recording: autoteleop.sh broadcasts an end-of-episode + # signal on /sim/stop_episode after the operator confirms with y/n. + # Reset the one-shot is_recording latch so the controller's record + # button can start the next episode. Kept outside the `if self.input` + # guard so the signal is handled even while the controller is asleep + # and not sending data. + if self.ros_utils.sim_ros_node.get_and_clear_stop_episode(): + self.is_recording = False + self.ros_utils.sim_ros_node.pub_recording(False) + logger.info("Episode finished; press the record button to start the next one") + if self.input: self.is_start_recording() self.ee_pub = [None, None] diff --git a/source/geniesim_teleop/src/geniesim_teleop/utils/ros_nodes.py b/source/geniesim_teleop/src/geniesim_teleop/utils/ros_nodes.py index eaef0339..7f17e9dc 100755 --- a/source/geniesim_teleop/src/geniesim_teleop/utils/ros_nodes.py +++ b/source/geniesim_teleop/src/geniesim_teleop/utils/ros_nodes.py @@ -86,6 +86,13 @@ def __init__(self, robot_name="G2", node_name="sim_ros_node"): ) self.publisher_playback = self.create_publisher(Bool, "/sim/playback_flag", 1) self.publisher_recording = self.create_publisher(Bool, "/sim/is_recording", 1) + # Multi-episode recording: listen for the end-of-episode signal (the same + # topic the simulator listens on) so the one-shot is_recording latch can be + # re-armed and the controller's record button can start the next episode. + self.subscriber_stop_episode = self.create_subscription( + Bool, "/sim/stop_episode", self.callback_stop_episode, 1 + ) + self.stop_episode_flag = False self.publisher_mc = self.create_publisher(GeniesimReactiveControl, "/wbc/retarget", 10) self.publisher_mc_debug = self.create_publisher(Pose, "/debug/retarget", 10) self.lock_js = threading.Lock() @@ -187,6 +194,17 @@ def pub_recording(self, val): msg.data = val self.publisher_recording.publish(msg) + def callback_stop_episode(self, msg): + if msg.data: + self.stop_episode_flag = True + + def get_and_clear_stop_episode(self): + # One-shot event: reading clears the flag so the main loop handles each + # end-of-episode signal exactly once. + flag = self.stop_episode_flag + self.stop_episode_flag = False + return flag + def pub_debug_pos(self, xyz, xyzw): pose = Pose() pose.position.x = xyz[0]