feat(teleop): multi-episode recording without restarting the simulator

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"]

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

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

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]

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]