From 5f114c1d743422651b162d84cadad60e67eb5957 Mon Sep 17 00:00:00 2001
From: CarolinePascal <caroline8.pascal@gmail.com>
Date: Mon, 19 May 2025 19:53:11 +0200
Subject: [PATCH] feat(init audio buffers): adding correct audio buffer
 initialization with actually recorded background noise instead of pure
 silence

---
 src/lerobot/datasets/audio_utils.py     | 13 +++++++++++--
 src/lerobot/datasets/lerobot_dataset.py |  6 +++++-
 src/lerobot/scripts/lerobot_record.py   |  4 ++++
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/lerobot/datasets/audio_utils.py b/src/lerobot/datasets/audio_utils.py
index ef4f39c13..694449877 100644
--- a/src/lerobot/datasets/audio_utils.py
+++ b/src/lerobot/datasets/audio_utils.py
@@ -41,6 +41,7 @@ def decode_audio(
     audio_path: Path | str,
     timestamps: list[float],
     duration: float,
+    start_time_s: float | None = 0.0,
     backend: str | None = "torchcodec",
 ) -> torch.Tensor:
     """
@@ -57,9 +58,9 @@ def decode_audio(
     Currently supports torchaudio.
     """
     if backend == "torchcodec":
-        return decode_audio_torchcodec(audio_path, timestamps, duration)
+        return decode_audio_torchcodec(audio_path, timestamps, duration, start_time_s)
     elif backend == "torchaudio":
-        return decode_audio_torchaudio(audio_path, timestamps, duration)
+        return decode_audio_torchaudio(audio_path, timestamps, duration, start_time_s)
     else:
         raise ValueError(f"Unsupported video backend: {backend}")
 
@@ -68,6 +69,7 @@ def decode_audio_torchcodec(
     audio_path: Path | str,
     timestamps: list[float],
     duration: float,
+    start_time_s: float | None = 0.0,
     log_loaded_timestamps: bool = False,
 ) -> torch.Tensor:
     # TODO(CarolinePascal) : add channels selection
@@ -77,6 +79,9 @@ def decode_audio_torchcodec(
     # TODO(CarolinePascal) : assert ts < total record duration
 
     audio_chunks = []
+    timestamps = [
+        timestamp + start_time_s for timestamp in timestamps
+    ]  # Add an offset of start_time_s to each timestamp
     for ts in timestamps:
         current_audio_chunk = audio_decoder.get_samples_played_in_range(
             start_seconds=max(0.0, ts - duration), stop_seconds=ts
@@ -118,6 +123,7 @@ def decode_audio_torchaudio(
     audio_path: Path | str,
     timestamps: list[float],
     duration: float,
+    start_time_s: float | None = 0.0,
     log_loaded_timestamps: bool = False,
 ) -> torch.Tensor:
     # TODO(CarolinePascal) : add channels selection
@@ -137,6 +143,9 @@ def decode_audio_torchaudio(
     )
 
     audio_chunks = []
+    timestamps = [
+        timestamp + start_time_s for timestamp in timestamps
+    ]  # Add an offset of start_time_s to each timestamp
     for ts in timestamps:
         reader.seek(max(0.0, ts - duration))  # Default to closest audio sample. Needs to be non-negative !
         status = reader.fill_buffer()
diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py
index c26bb3c4c..b68ab6ffd 100644
--- a/src/lerobot/datasets/lerobot_dataset.py
+++ b/src/lerobot/datasets/lerobot_dataset.py
@@ -482,6 +482,7 @@ class LeRobotDatasetMetadata:
             if not self.features[key].get("info", None):
                 audio_path = self.root / self.audio_path.format(audio_key=key, chunk_index=0, file_index=0)
                 self.info["features"][key]["info"] = get_audio_info(audio_path)
+                self.info["features"][key]["info"]["start_time_s"] = DEFAULT_AUDIO_CHUNK_DURATION
 
     def update_chunk_settings(
         self,
@@ -1154,7 +1155,10 @@ class LeRobotDataset(torch.utils.data.Dataset):
             shifted_query_ts = [from_timestamp + ts for ts in query_ts]
 
             audio_path = self.root / self.meta.get_audio_file_path(ep_idx, audio_key)
-            audio_chunk = decode_audio(audio_path, shifted_query_ts, query_duration, self.audio_backend)
+            start_time_s = self.meta.features[audio_key]["info"].get("start_time_s", 0.0)
+            audio_chunk = decode_audio(
+                audio_path, shifted_query_ts, query_duration, start_time_s, self.audio_backend
+            )
             item[audio_key] = audio_chunk.squeeze(0)
 
         return item
diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py
index ae62ffe6f..2e121c9b0 100644
--- a/src/lerobot/scripts/lerobot_record.py
+++ b/src/lerobot/scripts/lerobot_record.py
@@ -343,6 +343,10 @@ def record_loop(
     else:
         async_microphones_start_recording(robot.microphones)
 
+    # Fill audio buffers if needed
+    if robot.microphones:
+        busy_wait(DEFAULT_AUDIO_CHUNK_DURATION)
+
     timestamp = 0
     start_episode_t = time.perf_counter()
     while timestamp < control_time_s: