From e4eebd0680cf6ddb9f1d88208df6e9a3f443321f Mon Sep 17 00:00:00 2001
From: CarolinePascal <caroline8.pascal@gmail.com>
Date: Fri, 28 Mar 2025 17:16:17 +0100
Subject: [PATCH] Adding microphone recording in control loop

---
 lerobot/common/datasets/lerobot_dataset.py          | 12 ++++++++++++
 lerobot/common/robot_devices/control_utils.py       | 13 +++++++++++++
 .../common/robot_devices/microphones/microphone.py  |  3 +--
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py
index 488b7696..a207845f 100644
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ b/lerobot/common/datasets/lerobot_dataset.py
@@ -79,6 +79,7 @@ from lerobot.common.datasets.video_utils import (
     get_audio_info,
 )
 from lerobot.common.robot_devices.robots.utils import Robot
+from lerobot.common.robot_devices.microphones.utils import Microphone
 
 CODEBASE_VERSION = "v2.1"
 
@@ -910,6 +911,17 @@ class LeRobotDataset(torch.utils.data.Dataset):
 
         self.episode_buffer["size"] += 1
 
+    def add_microphone_recording(self, microphone: Microphone, microphone_key: str) -> None:
+        """
+        This function will start recording audio from the microphone and save it to disk.
+        """
+
+        audio_dir = self._get_raw_audio_file_path(self.num_episodes, "observation.audio." + microphone_key).parent
+        if not audio_dir.is_dir():
+            audio_dir.mkdir(parents=True, exist_ok=True)
+        
+        microphone.start_recording(output_file = self._get_raw_audio_file_path(self.num_episodes, "observation.audio." + microphone_key))
+
     def save_episode(self, episode_data: dict | None = None) -> None:
         """
         This will save to disk the current episode in self.episode_buffer.
diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py
index 4e42a989..335b6430 100644
--- a/lerobot/common/robot_devices/control_utils.py
+++ b/lerobot/common/robot_devices/control_utils.py
@@ -77,6 +77,11 @@ def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, f
             key = f"read_camera_{name}_dt_s"
             if key in robot.logs:
                 log_dt(f"dtR{name}", robot.logs[key])
+        
+        for name in robot.microphones: 
+            key = f"read_microphone_{name}_dt_s"
+            if key in robot.logs:
+                log_dt(f"dtR{name}", robot.logs[key])
 
     info_str = " ".join(log_items)
     logging.info(info_str)
@@ -243,6 +248,11 @@ def control_loop(
 
     timestamp = 0
     start_episode_t = time.perf_counter()
+
+    if teleoperate and dataset is not None:
+        for microphone_key, microphone in robot.microphones.items():
+            dataset.add_microphone_recording(microphone, microphone_key)
+
     while timestamp < control_time_s:
         start_loop_t = time.perf_counter()
 
@@ -286,6 +296,9 @@ def control_loop(
             events["exit_early"] = False
             break
 
+    for _, microphone in robot.microphones.items():
+        microphone.stop_recording()
+
 
 def reset_environment(robot, events, reset_time_s, fps):
     # TODO(rcadene): refactor warmup_record and reset_environment
diff --git a/lerobot/common/robot_devices/microphones/microphone.py b/lerobot/common/robot_devices/microphones/microphone.py
index 5b136aab..c26d03c0 100644
--- a/lerobot/common/robot_devices/microphones/microphone.py
+++ b/lerobot/common/robot_devices/microphones/microphone.py
@@ -238,9 +238,8 @@ class Microphone:
             self.thread.daemon = True
             self.thread.start()
             
-        self.stream.start()
-
         self.logs["start_timestamp"] = capture_timestamp_utc()
+        self.stream.start()
 
     def stop_recording(self) -> None: