From 949f4fcbe9d7d870d84acb13bb43b5a16a4e71b9 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Mon, 15 Jun 2026 12:25:06 +0200 Subject: [PATCH] fix(logging): batch wandb metrics - Batch all metrics into a single wandb.log() call instead of one per key, reducing API overhead. - Add support for list-valued metrics by expanding them to indexed keys (e.g. metric_0, metric_1). --- src/lerobot/common/wandb_utils.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/src/lerobot/common/wandb_utils.py b/src/lerobot/common/wandb_utils.py index b782cd751..bdac7d4da 100644 --- a/src/lerobot/common/wandb_utils.py +++ b/src/lerobot/common/wandb_utils.py @@ -180,24 +180,32 @@ class WandBLogger: self._wandb_custom_step_key.add(new_custom_key) self._wandb.define_metric(new_custom_key, hidden=True) + batch_data = {} for k, v in d.items(): + # Skip the custom step key here, it's added to the batch below. + if custom_step_key is not None and k == custom_step_key: + continue + + if isinstance(v, list): + for i, elem in enumerate(v): + if isinstance(elem, (int | float)): + batch_data[f"{mode}/{k}_{i}"] = elem + continue + if not isinstance(v, (int | float | str)): logging.warning( f'WandB logging of key "{k}" was ignored as its type "{type(v)}" is not handled by this wrapper.' ) continue - # Do not log the custom step key itself. - if self._wandb_custom_step_key is not None and k in self._wandb_custom_step_key: - continue + batch_data[f"{mode}/{k}"] = v + if batch_data: if custom_step_key is not None: - value_custom_step = d[custom_step_key] - data = {f"{mode}/{k}": v, f"{mode}/{custom_step_key}": value_custom_step} - self._wandb.log(data) - continue - - self._wandb.log(data={f"{mode}/{k}": v}, step=step) + batch_data[f"{mode}/{custom_step_key}"] = d[custom_step_key] + self._wandb.log(batch_data) + else: + self._wandb.log(data=batch_data, step=step) def log_video(self, video_path: str, step: int, mode: str = "train"): if mode not in {"train", "eval"}: