fix: use _get_autocast_context for fp32 action head (MPS compat)

VLA-JEPA hardcodes torch.autocast with dtype=torch.bfloat16, which
crashes on MPS (no AMP support) and silently misbehaves on pre-Ampere
CUDA GPUs (no bf16). Add a _get_autocast_context() helper that reuses
the existing is_amp_available() utility to pick a safe strategy per
device, matching the pattern used by pi05 and molmoact2.

Fixes #3744

docs/source/inference.mdx

@@ -149,14 +149,13 @@ lerobot-rollout \
 
 Foot pedal input is also supported via `--strategy.input_device=pedal`. Configure pedal codes with `--strategy.pedal.*` flags.
 
-| Flag                                 | Description                                                                                                                                                    |
-| ------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `--strategy.num_episodes`            | Number of correction episodes to record (default: 10)                                                                                                          |
-| `--strategy.record_autonomous`       | Record autonomous frames too (default: false)                                                                                                                  |
-| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5)                                                                                                                      |
-| `--strategy.input_device`            | Input device: `keyboard` or `pedal` (default: keyboard)                                                                                                        |
-| `--strategy.smooth_handover`         | Smoothly hand control over at pause / correction start (default: true). Disable for clutch-style teleops that re-reference at the current robot pose on engage |
-| `--teleop.type`                      | **Required.** Teleoperator type                                                                                                                                |
+| Flag                                 | Description                                             |
+| ------------------------------------ | ------------------------------------------------------- |
+| `--strategy.num_episodes`            | Number of correction episodes to record (default: 10)   |
+| `--strategy.record_autonomous`       | Record autonomous frames too (default: false)           |
+| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5)               |
+| `--strategy.input_device`            | Input device: `keyboard` or `pedal` (default: keyboard) |
+| `--teleop.type`                      | **Required.** Teleoperator type                         |
 
 ### Episodic (`--strategy.type=episodic`)
 

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 
 import logging
 from collections import deque
+from contextlib import nullcontext
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -26,6 +27,7 @@ from torch import Tensor, nn
 from lerobot.policies.pretrained import PreTrainedPolicy, T
 from lerobot.policies.utils import populate_queues
 from lerobot.utils.constants import ACTION, OBS_STATE
+from lerobot.utils.device_utils import is_amp_available
 from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
@@ -39,6 +41,21 @@ from .configuration_vla_jepa import VLAJEPAConfig
 from .qwen_interface import Qwen3VLInterface
 from .world_model import ActionConditionedVideoPredictor
 
+
+def _get_autocast_context(device_type: str, dtype: torch.dtype = torch.bfloat16):
+    """Return an autocast context appropriate for the device.
+
+    MPS does not support ``torch.autocast`` at all.  On CUDA devices
+    without bfloat16 support (compute capability < 8.0) we fall back to
+    float16.
+    """
+    if not is_amp_available(device_type):
+        return nullcontext()
+    if device_type == "cuda" and dtype == torch.bfloat16 and not torch.cuda.is_bf16_supported():
+        dtype = torch.float16
+    return torch.autocast(device_type=device_type, dtype=dtype)
+
+
 # ============================================================================
 # Native VLA-JEPA Model - follows original starVLA VLA_JEPA.py implementation
 # ============================================================================
@@ -183,7 +200,7 @@ class VLAJEPAModel(nn.Module):
             action_idx = action_mask.nonzero(as_tuple=True)
 
         device_type = next(self.parameters()).device.type
-        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+        with _get_autocast_context(device_type, torch.bfloat16):
             last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
             b, _, h = last_hidden.shape
             embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h)
@@ -250,7 +267,7 @@ class VLAJEPAModel(nn.Module):
     ) -> Tensor:
         """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
         device_type = next(self.parameters()).device.type
-        with torch.autocast(device_type=device_type, dtype=torch.float32):
+        with _get_autocast_context(device_type, torch.float32):
             r = self.config.repeated_diffusion_steps
             horizon = self.config.chunk_size
             actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)

src/lerobot/rollout/configs.py

@@ -180,14 +180,6 @@ class DAggerStrategyConfig(RolloutStrategyConfig):
     # Target video file size in MB for episode rotation (record_autonomous
     # mode only).  Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when None.
     target_video_file_size_mb: int | None = None
-    # Whether to turn on or off the smooth handover behavior at phase transitions:
-    # the leader is driven to the follower position on pause (teleops with
-    # `send_feedback` capability), and the follower is slid to the teleop pose when
-    # a correction starts (non-actuated teleops). Disable for clutch-style
-    # teleoperators (e.g. VR controllers) that re-reference at the current robot
-    # pose on engage: the handover is already continuous there, and the blocking
-    # interpolation only delays the start of the correction.
-    smooth_handover: bool = True
     input_device: str = "keyboard"
     keyboard: DAggerKeyboardConfig = field(default_factory=DAggerKeyboardConfig)
     pedal: DAggerPedalConfig = field(default_factory=DAggerPedalConfig)

src/lerobot/rollout/strategies/dagger.py

@@ -623,8 +623,8 @@ class DAggerStrategy(RolloutStrategy):
     # State-machine transition side-effects
     # ------------------------------------------------------------------
 
+    @staticmethod
     def _apply_transition(
-        self,
         old_phase: DAggerPhase,
         new_phase: DAggerPhase,
         engine,
@@ -634,10 +634,6 @@ class DAggerStrategy(RolloutStrategy):
     ) -> None:
         """Execute side-effects for a validated phase transition, including smooth handovers.
 
-        The smooth handovers below can be disabled with
-        ``--strategy.smooth_handover=false`` (useful for clutch-style teleops
-        that re-reference at the current robot pose on engage).
-
         AUTONOMOUS -> PAUSED (actuated teleop):
             Pause the engine, then drive the leader arm to the follower's last
             commanded position so the operator takes over without a jerk.
@@ -661,7 +657,7 @@ class DAggerStrategy(RolloutStrategy):
             logger.info("Pausing engine - robot holds position")
             engine.pause()
 
-            if self.config.smooth_handover and teleop_supports_feedback(teleop) and prev_action is not None:
+            if teleop_supports_feedback(teleop) and prev_action is not None:
                 # TODO(Maxime): prev_action is in robot action key space (output of robot_action_processor).
                 # send_feedback expects teleop feedback key space. For homogeneous setups (e.g. SO-101
                 # leader + SO-101 follower) the keys are identical so this works. If the processor pipeline
@@ -672,11 +668,7 @@ class DAggerStrategy(RolloutStrategy):
 
         elif old_phase == DAggerPhase.PAUSED and new_phase == DAggerPhase.CORRECTING:
             logger.info("Entering correction mode - human teleop control")
-            if (
-                self.config.smooth_handover
-                and not teleop_supports_feedback(teleop)
-                and prev_action is not None
-            ):
+            if not teleop_supports_feedback(teleop) and prev_action is not None:
                 logger.info("Smooth handover: sliding follower to teleop position")
                 obs = robot.get_observation()
                 teleop_action = teleop.get_action()