From 1ee4d84f072c2f28dd37febb518e1878598eb368 Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Mon, 27 Oct 2025 14:28:12 +0100 Subject: [PATCH] fix(audio visualization): making audio visualization robust to multiple channel audio data --- src/lerobot/utils/visualization_utils.py | 35 +++++++++++++----------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index df5c48b27..fd83391f3 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -152,22 +152,25 @@ def log_rerun_data( for i, vi in enumerate(arr): rr.log(f"{key}_{i}", rr.Scalars(float(vi))) elif arr.ndim == 2: - rr.send_columns( - "audio/" + key, - indexes=[ - rr.TimeSecondsColumn( - "episode_time", - times=log_time - + np.linspace( - -DEFAULT_AUDIO_CHUNK_DURATION, - 0, - len(observation[key]), - endpoint=False, - ), - ) - ], - columns=rr.Scalar.columns(scalar=np.mean(observation[key], axis=1)), - ) + for i, channel_arr in enumerate(arr.T): + rr.send_columns( + "audio/" + + key + + f"_channel_{i}", # TODO(CarolinePascal): Get actual channel number/name + indexes=[ + rr.TimeSecondsColumn( + "episode_time", + times=log_time + + np.linspace( + -DEFAULT_AUDIO_CHUNK_DURATION, + 0, + len(channel_arr), + endpoint=False, + ), + ) + ], + columns=rr.Scalar.columns(scalar=channel_arr), + ) elif arr.ndim == 3: rr.log(key, rr.Image(arr), static=True) else: