From 92f96f33b3a842ead3a115cc4f701ba2ad8aa3b5 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Thu, 16 Jul 2026 12:12:37 +0200 Subject: [PATCH] Aggregate policy sub-losses through MetricsTracker (#4024) --- src/lerobot/scripts/lerobot_train.py | 10 +++++--- src/lerobot/utils/logging_utils.py | 19 ++++++++++++++++ tests/utils/test_logging_utils.py | 34 ++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 3 deletions(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 6e8458523..ec5565cf4 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -171,6 +171,9 @@ def update_policy( train_metrics.update_s = time.perf_counter() - start_time if torch.cuda.is_available(): train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3) + # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window. + if output_dict: + train_metrics.update_metrics(output_dict) return train_metrics, output_dict @@ -572,7 +575,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): batch = preprocessor(batch) train_tracker.dataloading_s = time.perf_counter() - start_time - train_tracker, output_dict = update_policy( + train_tracker, _ = update_policy( train_tracker, policy, batch, @@ -605,9 +608,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.samples_per_s = effective_batch_size / step_time logging.info(train_tracker) if wandb_logger: + # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the + # tracker by update_policy, so to_dict() already carries their windowed, + # rank-reduced averages — no per-step output_dict passthrough needed. wandb_log_dict = train_tracker.to_dict() - if output_dict: - wandb_log_dict.update(output_dict) # Log sample weighting statistics if enabled if sample_weighter is not None: weighter_stats = sample_weighter.get_stats() diff --git a/src/lerobot/utils/logging_utils.py b/src/lerobot/utils/logging_utils.py index 20673fc30..fd3021fea 100644 --- a/src/lerobot/utils/logging_utils.py +++ b/src/lerobot/utils/logging_utils.py @@ -104,6 +104,7 @@ class MetricsTracker: "episodes", "epochs", "accelerator", + "_caller_metrics", ] def __init__( @@ -129,6 +130,9 @@ class MetricsTracker: self.episodes = self.samples / self._avg_samples_per_ep self.epochs = self.samples / self._num_frames self.accelerator = accelerator + # Meter names the caller registered up front. update_metrics() leaves these untouched, so a + # policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter. + self._caller_metrics: set[str] = set(self.metrics) def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any: if name in self.__dict__: @@ -156,6 +160,21 @@ class MetricsTracker: self.episodes = self.samples / self._avg_samples_per_ep self.epochs = self.samples / self._num_frames + def update_metrics(self, values: dict[str, Any]) -> None: + """Accumulate a dict of scalar metrics, auto-registering a meter for each new key. + + Non-numeric values and bools are ignored. + Caller-registered metrics (those passed to the constructor) are never overridden. + """ + for name, value in values.items(): + if isinstance(value, bool) or not isinstance(value, (int, float)): + continue + if name in self._caller_metrics: + continue + if name not in self.metrics: + self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean") + self.metrics[name].update(float(value)) + def reduce_across_ranks(self) -> None: """ Synchronises the running averages of every metric whose ``reduction`` is not ``"none"`` diff --git a/tests/utils/test_logging_utils.py b/tests/utils/test_logging_utils.py index aa851bd2a..111bd0f01 100644 --- a/tests/utils/test_logging_utils.py +++ b/tests/utils/test_logging_utils.py @@ -233,3 +233,37 @@ def test_metrics_tracker_reduce_across_ranks_invokes_reduce(): # accumulate against the cluster view rather than the stale per-rank sum. meter = tracker.update_s assert meter.sum / meter.count == pytest.approx(meter.avg) + + +def test_metrics_tracker_update_metrics_registers_and_averages(): + tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={}) + tracker.update_metrics({"latent_loss": 0.2, "action_loss": 0.4}) + tracker.update_metrics({"latent_loss": 0.4, "action_loss": 0.6}) + + # New keys are auto-registered as mean-reduced meters and averaged over the window. + assert tracker.metrics["latent_loss"].reduction == "mean" + assert tracker.metrics["latent_loss"].avg == pytest.approx(0.3) + assert tracker.metrics["action_loss"].avg == pytest.approx(0.5) + assert tracker.to_dict()["latent_loss"] == pytest.approx(0.3) + + +def test_metrics_tracker_update_metrics_skips_non_numeric(): + tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={}) + tracker.update_metrics({"loss": 0.5, "head_mode": "sparse", "enabled": True}) + + # strings and bools ignored + assert "loss" in tracker.metrics + assert "head_mode" not in tracker.metrics + assert "enabled" not in tracker.metrics + + +def test_metrics_tracker_update_metrics_does_not_override_caller_meter(): + # A policy that echoes "loss" in its output dict must not overwrite the caller-owned, + # already-aggregated loss meter. + metrics = {"loss": AverageMeter("loss", ":.3f", reduction="mean")} + tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics) + tracker.loss = 1.0 # caller-set optimized loss + tracker.update_metrics({"loss": 99.0, "latent_loss": 0.2}) + + assert tracker.metrics["loss"].avg == pytest.approx(1.0) # snapshot ignored + assert tracker.metrics["latent_loss"].avg == pytest.approx(0.2)