diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 8017de28f..ebdf9c05e 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/src/lerobot/policies/pi05/modeling_pi05.py @@ -1066,7 +1066,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` training_max_delay = int(getattr(self.config, "rtc_training_max_delay", 0)) if training_max_delay <= 0: raise ValueError( - "RTC mode='trained' requires a Pi052 checkpoint trained with " + "RTC mode='trained' requires a checkpoint trained with " "policy.rtc_training_max_delay > 0." ) trained_prefix, trained_prefix_mask = _prepare_trained_rtc_prefix( diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index ad6f22c84..b3ac33982 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -121,9 +121,7 @@ class PI052Config(PI05Config): # Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow. flow_num_repeats: int = 5 - # Training-time RTC (arXiv:2512.05964). Zero preserves standard flow matching. - rtc_training_max_delay: int = 0 - """Maximum clean-prefix delay sampled for training-time RTC; zero disables it.""" + # Training-time RTC configuration is inherited from PI05Config. # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it. text_ce_z_loss_weight: float = 1e-4 @@ -155,11 +153,6 @@ class PI052Config(PI05Config): self.train_expert_only = False if self.flow_num_repeats < 1: raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}") - if not 0 <= self.rtc_training_max_delay < self.chunk_size: - raise ValueError( - "rtc_training_max_delay must satisfy " - f"0 <= delay < chunk_size ({self.chunk_size}), got {self.rtc_training_max_delay}" - ) if self.manual_attention_scope not in {"all", "action"}: raise ValueError( f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}" diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index cd26998fe..3a240e544 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -36,6 +36,8 @@ from ..pi05.modeling_pi05 import ( ActionSelectKwargs, PI05Policy, PI05Pytorch as PI05PytorchBase, + _build_flow_matching_inputs, + _sample_training_rtc_prefix_mask, make_att_2d_masks, ) from .configuration_pi052 import PI052Config @@ -177,43 +179,6 @@ def _enable_hf_kernels() -> None: logger.info("PI052: HF kernels (Liger) enabled — rope, geglu fused.") -def _sample_training_rtc_prefix_mask( - batch_size: int, - action_horizon: int, - max_delay: int, - device: torch.device, -) -> Tensor | None: - """Sample per-draw clean prefixes for training-time RTC.""" - if max_delay <= 0: - return None - delays = torch.randint(0, max_delay + 1, (batch_size,), device=device) - positions = torch.arange(action_horizon, device=device) - return positions.unsqueeze(0) < delays.unsqueeze(1) - - -def _build_flow_matching_inputs( - actions: Tensor, - noise: Tensor, - time: Tensor, - prefix_mask: Tensor | None, -) -> tuple[Tensor, Tensor]: - """Build noisy actions and scalar/per-token flow times. - - LeRobot's PI0.5 flow uses ``t=0`` for clean data and ``t=1`` for - noise, the reverse of the notation in arXiv:2512.05964. Consequently, - clean RTC prefix tokens receive ``t=0`` here. - """ - if prefix_mask is None: - model_time = time - expanded_time = time[:, None, None] - else: - model_time = time[:, None].expand_as(prefix_mask) - model_time = torch.where(prefix_mask, torch.zeros_like(model_time), model_time) - expanded_time = model_time.unsqueeze(-1) - x_t = expanded_time * noise + (1 - expanded_time) * actions - return x_t, model_time - - def _flow_loss_components(flow_per_dim: Tensor, prefix_mask: Tensor | None) -> tuple[Tensor, Tensor]: """Return each sample's postfix loss sum and valid-element count.""" if prefix_mask is None: diff --git a/tests/policies/pi052/test_pi052_training_time_rtc.py b/tests/policies/pi052/test_pi052_training_time_rtc.py index 398203d64..6993f433c 100644 --- a/tests/policies/pi052/test_pi052_training_time_rtc.py +++ b/tests/policies/pi052/test_pi052_training_time_rtc.py @@ -16,13 +16,13 @@ from torch import nn pytest.importorskip("transformers") from lerobot.policies.pi05.modeling_pi05 import ( # noqa: E402 + _build_flow_matching_inputs, _prepare_trained_rtc_prefix, create_sinusoidal_pos_embedding, ) from lerobot.policies.pi052.configuration_pi052 import PI052Config # noqa: E402 from lerobot.policies.pi052.modeling_pi052 import ( # noqa: E402 PI05Pytorch as PI052Pytorch, - _build_flow_matching_inputs, _flow_loss_per_sample, _reduce_flow_loss, )