diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index e4974561b..95a87504e 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -22,6 +22,7 @@ and :class:`DatasetContext` — assembled into :class:`RolloutContext`. from __future__ import annotations import logging +from copy import copy from dataclasses import dataclass, field from threading import Event from typing import TYPE_CHECKING @@ -69,6 +70,35 @@ else: logger = logging.getLogger(__name__) +def _wrap_predict_action_chunk_with_torch_compile( + policy: PreTrainedPolicy, + *, + backend: str, + mode: str, +) -> bool: + """Install the JIT wrapper and report whether it was configured successfully. + + ``torch.compile`` compiles lazily on the first invocation, so success here + does not guarantee that backend compilation will succeed during warm-up. + """ + if not hasattr(torch, "compile"): + logger.warning("torch.compile is not available in this PyTorch build") + return False + + try: + policy.predict_action_chunk = torch.compile( + policy.predict_action_chunk, + backend=backend, + mode=mode, + ) + except Exception as exc: + logger.warning("Failed to configure torch.compile: %s", exc) + return False + + logger.info("torch.compile configured for predict_action_chunk") + return True + + def _resolve_action_key_order( policy_action_names: list[str] | None, dataset_action_names: list[str] ) -> list[str]: @@ -241,18 +271,19 @@ def build_rollout_context( policy.eval() logger.info("Policy loaded: type=%s, device=%s", policy_config.type, cfg.device) + torch_compile_active = cfg.use_torch_compile if cfg.use_torch_compile and policy.type not in ("pi0", "pi05"): - try: - if hasattr(torch, "compile"): - compile_kwargs = { - "backend": cfg.torch_compile_backend, - "mode": cfg.torch_compile_mode, - "options": {"triton.cudagraphs": False}, - } - policy.predict_action_chunk = torch.compile(policy.predict_action_chunk, **compile_kwargs) - logger.info("torch.compile applied to predict_action_chunk") - except Exception as e: - logger.warning("Failed to apply torch.compile: %s", e) + torch_compile_active = _wrap_predict_action_chunk_with_torch_compile( + policy, + backend=cfg.torch_compile_backend, + mode=cfg.torch_compile_mode, + ) + + if cfg.use_torch_compile and not torch_compile_active: + # RolloutConfig.__post_init__ reloads the policy configuration, so avoid + # dataclasses.replace when carrying the effective state downstream. + cfg = copy(cfg) + cfg.use_torch_compile = False # --- 2. Robot-side processors (user-supplied or defaults) -------- if ( @@ -470,7 +501,7 @@ def build_rollout_context( task=task_str, fps=cfg.fps, device=cfg.device, - use_torch_compile=cfg.use_torch_compile, + use_torch_compile=torch_compile_active, compile_warmup_inferences=cfg.compile_warmup_inferences, shutdown_event=shutdown_event, )