Merge branch 'main' into feat/unitree_g1_sonic_rebased

src/lerobot/policies/pretrained.py

@@ -23,8 +23,6 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
 
-import packaging
-import safetensors
 from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
@@ -221,26 +219,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        # Create base kwargs
+        missing_keys, unexpected_keys = load_model_as_safetensor(
-        kwargs = {"strict": strict}
+            model, model_file, strict=strict, device=map_location
-
+        )
-        # Add device parameter for newer versions that support it
-        if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
-            kwargs["device"] = map_location
-
-        # Load the model with appropriate kwargs
-        missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
         log_model_loading_keys(missing_keys, unexpected_keys)
-
-        # For older versions, manually move to device if needed
-        if "device" not in kwargs and map_location != "cpu":
-            logging.warning(
-                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
-                " This means that the model is loaded on 'cpu' first and then copied to the device."
-                " This leads to a slower loading time."
-                " Please update safetensors to version 0.4.3 or above for improved performance."
-            )
-            model.to(map_location)
         return model
 
     @abc.abstractmethod

src/lerobot/rewards/pretrained.py

@@ -21,8 +21,6 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, TypeVar
 
-import packaging
-import safetensors
 from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
@@ -129,29 +127,13 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        # Create base kwargs
+        missing_keys, unexpected_keys = load_model_as_safetensor(
-        kwargs = {"strict": strict}
+            model, model_file, strict=strict, device=map_location
-
+        )
-        # Add device parameter for newer versions that support it
-        if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
-            kwargs["device"] = map_location
-
-        # Load the model with appropriate kwargs
-        missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
         if missing_keys:
             logging.warning(f"Missing key(s) when loading model: {missing_keys}")
         if unexpected_keys:
             logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
-
-        # For older versions, manually move to device if needed
-        if "device" not in kwargs and map_location != "cpu":
-            logging.warning(
-                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
-                " This means that the model is loaded on 'cpu' first and then copied to the device."
-                " This leads to a slower loading time."
-                " Please update safetensors to version 0.4.3 or above for improved performance."
-            )
-            model.to(map_location)
         return model
 
     def get_optim_params(self):

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()

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"``

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)