mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
fix(rollout): reject incompatible RTC policies (#4228)
* fix(rollout): reject incompatible RTC policies * chore(policies): support rtc * chore(tests): delete compatibility test --------- Co-authored-by: ogarciarevett <ogarciarevett@gmail.com>
This commit is contained in:
@@ -42,6 +42,9 @@ class Evo1Policy(PreTrainedPolicy):
|
|||||||
config_class = Evo1Config
|
config_class = Evo1Config
|
||||||
name = "evo1"
|
name = "evo1"
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
|
def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
|
||||||
super().__init__(config)
|
super().__init__(config)
|
||||||
config.validate_features()
|
config.validate_features()
|
||||||
|
|||||||
@@ -68,6 +68,9 @@ class GrootPolicy(PreTrainedPolicy):
|
|||||||
name = "groot"
|
name = "groot"
|
||||||
config_class = GrootConfig
|
config_class = GrootConfig
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
def __init__(self, config: GrootConfig, **kwargs):
|
def __init__(self, config: GrootConfig, **kwargs):
|
||||||
"""Initialize Groot policy wrapper."""
|
"""Initialize Groot policy wrapper."""
|
||||||
require_package("transformers", extra="groot")
|
require_package("transformers", extra="groot")
|
||||||
|
|||||||
@@ -520,6 +520,9 @@ class MolmoAct2Policy(PreTrainedPolicy):
|
|||||||
config_class = MolmoAct2Config
|
config_class = MolmoAct2Config
|
||||||
name = "molmoact2"
|
name = "molmoact2"
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
return self.config.inference_action_mode == "continuous"
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config: MolmoAct2Config,
|
config: MolmoAct2Config,
|
||||||
|
|||||||
@@ -749,6 +749,9 @@ class PI0Policy(PreTrainedPolicy):
|
|||||||
config_class = PI0Config
|
config_class = PI0Config
|
||||||
name = "pi0"
|
name = "pi0"
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config: PI0Config,
|
config: PI0Config,
|
||||||
|
|||||||
@@ -714,6 +714,9 @@ class PI05Policy(PreTrainedPolicy):
|
|||||||
config_class = PI05Config
|
config_class = PI05Config
|
||||||
name = "pi05"
|
name = "pi05"
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config: PI05Config,
|
config: PI05Config,
|
||||||
|
|||||||
@@ -249,6 +249,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
"""
|
"""
|
||||||
raise NotImplementedError
|
raise NotImplementedError
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
"""Whether this policy implements Real-Time Chunking inference semantics."""
|
||||||
|
return False
|
||||||
|
|
||||||
# TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
|
# TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
|
||||||
@abc.abstractmethod
|
@abc.abstractmethod
|
||||||
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
|
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
|
||||||
|
|||||||
@@ -145,6 +145,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
|
|||||||
config_class = SmolVLAConfig
|
config_class = SmolVLAConfig
|
||||||
name = "smolvla"
|
name = "smolvla"
|
||||||
|
|
||||||
|
def supports_rtc(self) -> bool:
|
||||||
|
return True
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config: SmolVLAConfig,
|
config: SmolVLAConfig,
|
||||||
|
|||||||
@@ -57,6 +57,7 @@ from .inference import (
|
|||||||
SyncInferenceConfig,
|
SyncInferenceConfig,
|
||||||
create_inference_engine,
|
create_inference_engine,
|
||||||
)
|
)
|
||||||
|
from .inference.rtc import supports_rtc_inference
|
||||||
from .robot_wrapper import ThreadSafeRobot
|
from .robot_wrapper import ThreadSafeRobot
|
||||||
|
|
||||||
if TYPE_CHECKING or _peft_available:
|
if TYPE_CHECKING or _peft_available:
|
||||||
@@ -226,6 +227,12 @@ def build_rollout_context(
|
|||||||
policy = _load_pretrained_policy(policy_config)
|
policy = _load_pretrained_policy(policy_config)
|
||||||
|
|
||||||
if is_rtc:
|
if is_rtc:
|
||||||
|
if not supports_rtc_inference(policy):
|
||||||
|
raise ValueError(
|
||||||
|
f"RTC inference is not supported by policy type '{policy_config.type}': "
|
||||||
|
"the policy must implement RTC semantics and predict_action_chunk must accept "
|
||||||
|
"inference_delay and prev_chunk_left_over. Use '--inference.type=sync' instead."
|
||||||
|
)
|
||||||
policy.config.rtc_config = cfg.inference.rtc
|
policy.config.rtc_config = cfg.inference.rtc
|
||||||
if hasattr(policy, "init_rtc_processor"):
|
if hasattr(policy, "init_rtc_processor"):
|
||||||
policy.init_rtc_processor()
|
policy.init_rtc_processor()
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ way via ``notify_observation``.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import inspect
|
||||||
import logging
|
import logging
|
||||||
import math
|
import math
|
||||||
import time
|
import time
|
||||||
@@ -62,6 +63,23 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def supports_rtc_inference(policy: PreTrainedPolicy) -> bool:
|
||||||
|
"""Whether a policy declares RTC support and accepts the RTC call shape."""
|
||||||
|
supports_rtc = getattr(policy, "supports_rtc", None)
|
||||||
|
if not callable(supports_rtc) or not supports_rtc():
|
||||||
|
return False
|
||||||
|
|
||||||
|
try:
|
||||||
|
inspect.signature(policy.predict_action_chunk).bind(
|
||||||
|
object(),
|
||||||
|
inference_delay=0,
|
||||||
|
prev_chunk_left_over=None,
|
||||||
|
)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
|
def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
|
||||||
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
|
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
|
||||||
if prev_actions.ndim != 2:
|
if prev_actions.ndim != 2:
|
||||||
|
|||||||
Reference in New Issue
Block a user