perf(pi052): gate per-step .item() CUDA syncs to logging steps

Keep PI052Policy.forward's loss components as detached tensors and only
materialize loss/grad_norm/update_s to python floats on logging steps
(1-in-log_freq) via a new update_policy(log_metrics=...) gate. Also dedupe
the predict_actions .any().item() control-flow sync (2 -> 1 per step).

Keeps the training step fully async on non-logging steps so the next batch's
dataloading/enqueue overlaps GPU compute instead of stalling on a per-step
CUDA sync.

Co-authored-by: Cursor <cursoragent@cursor.com>

src/lerobot/policies/pi052/modeling_pi052.py

@@ -1150,9 +1150,11 @@ class PI052Policy(PreTrainedPolicy):
         ):
             return self._pi05_flow_forward(batch, reduction=reduction)
 
-        run_flow = self.config.flow_loss_weight > 0 and (
+        # Whether any sample in the batch wants actions predicted. This is a data-dependent branch, so
-            predict_actions_t is None or bool(predict_actions_t.any().item())
+        # it needs a host-side bool (one CUDA sync); compute it once and reuse for both flow and FAST
-        )
+        # instead of syncing twice.
+        predict_any = predict_actions_t is None or bool(predict_actions_t.any().item())
+        run_flow = self.config.flow_loss_weight > 0 and predict_any
         run_text = self.config.text_loss_weight > 0 and text_labels is not None
 
         loss_dict: dict[str, Any] = {}
@@ -1162,7 +1164,7 @@ class PI052Policy(PreTrainedPolicy):
         run_fast = (
             getattr(self.config, "enable_fast_action_loss", False)
             and self.config.fast_action_loss_weight > 0
-            and (predict_actions_t is None or bool(predict_actions_t.any().item()))
+            and predict_any
         )
         action_tokens = action_mask = action_code_mask = None
         if run_fast:
@@ -1200,13 +1202,13 @@ class PI052Policy(PreTrainedPolicy):
                 action_code_mask=action_code_mask if run_fast else None,
                 predict_actions_t=predict_actions_t,
             )
-            loss_dict["flow_loss"] = float(flow_loss.detach().item())
+            loss_dict["flow_loss"] = flow_loss.detach()
             total = self.config.flow_loss_weight * flow_loss
             if text_loss is not None:
-                loss_dict["text_loss"] = float(text_loss.detach().item())
+                loss_dict["text_loss"] = text_loss.detach()
                 total = total + self.config.text_loss_weight * text_loss
             if fast_loss is not None:
-                loss_dict["fast_action_loss"] = float(fast_loss.detach().item())
+                loss_dict["fast_action_loss"] = fast_loss.detach()
                 total = total + self.config.fast_action_loss_weight * fast_loss
         elif run_text or run_fast:
             text_loss, fast_loss = self._compute_text_and_fast_loss(
@@ -1218,11 +1220,11 @@ class PI052Policy(PreTrainedPolicy):
                 predict_actions_t=predict_actions_t,
             )
             if text_loss is not None:
-                loss_dict["text_loss"] = float(text_loss.detach().item())
+                loss_dict["text_loss"] = text_loss.detach()
                 weighted = self.config.text_loss_weight * text_loss
                 total = weighted if total is None else total + weighted
             if fast_loss is not None:
-                loss_dict["fast_action_loss"] = float(fast_loss.detach().item())
+                loss_dict["fast_action_loss"] = fast_loss.detach()
                 weighted = self.config.fast_action_loss_weight * fast_loss
                 total = weighted if total is None else total + weighted
 
@@ -1235,7 +1237,9 @@ class PI052Policy(PreTrainedPolicy):
                 "nothing to train."
             )
 
-        loss_dict["loss"] = float(total.detach().item()) if total.dim() == 0 else float("nan")
+        # Keep loss components as detached tensors (no CUDA sync here); the training loop converts
+        # them to python floats only on logging steps (see update_policy's log_metrics gate).
+        loss_dict["loss"] = total.detach() if total.dim() == 0 else float("nan")
         if reduction == "none":
             return total.expand(batch[OBS_LANGUAGE_TOKENS].shape[0]), loss_dict
         return total, loss_dict

src/lerobot/scripts/lerobot_train.py

@@ -82,6 +82,7 @@ def update_policy(
     lr_scheduler=None,
     lock=None,
     sample_weighter=None,
+    log_metrics: bool = True,
 ) -> tuple[MetricsTracker, dict | None]:
     """
     Performs a single training step to update the policy's weights.
@@ -99,6 +100,9 @@ def update_policy(
         lr_scheduler: An optional learning rate scheduler.
         lock: An optional lock for thread-safe optimizer updates.
         sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
+        log_metrics: When True, read loss/grad_norm/update_s off the GPU via `.item()` (a CUDA sync).
+            On non-logging steps set False so the step stays fully async — letting the next batch's
+            dataloading and enqueue overlap GPU compute instead of stalling on a per-step sync.
 
     Returns:
         A tuple containing:
@@ -166,12 +170,24 @@ def update_policy(
     if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
         accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
 
-    train_metrics.loss = loss.item()
-    train_metrics.grad_norm = grad_norm.item()
     train_metrics.lr = optimizer.param_groups[0]["lr"]
-    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)
+    # `loss.item()` / `grad_norm.item()` each block on a CUDA sync. Only pay that on logging steps;
+    # on the other steps the readouts are never consumed, so skipping them keeps the step async and
+    # lets CPU-side dataloading/enqueue overlap GPU compute. update_s is only accurate under that
+    # sync, so it too is recorded on logging steps only.
+    if log_metrics:
+        train_metrics.loss = loss.item()
+        train_metrics.grad_norm = grad_norm.item()
+        train_metrics.update_s = time.perf_counter() - start_time
+        # Policies may hand back loss components as detached tensors to keep the forward async on
+        # non-logging steps (e.g. pi052). Materialize them to python floats here, on logging steps only,
+        # so the per-step CUDA sync is paid 1-in-log_freq rather than every step.
+        if output_dict:
+            output_dict = {
+                k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
+            }
     return train_metrics, output_dict
 
 
@@ -642,6 +658,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
+        # This update produces step number `step + 1`; only sync metrics off the GPU when that step
+        # will be logged (mirrors the is_log_step gate below). Everything else stays async.
+        log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
+
         train_tracker, output_dict = update_policy(
             train_tracker,
             policy,
@@ -651,6 +671,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             accelerator=accelerator,
             lr_scheduler=lr_scheduler,
             sample_weighter=sample_weighter,
+            log_metrics=log_metrics,
         )
 
         # EMA update: pull one step of the live weights into the shadow.