feat(pi05): add training-time RTC

Extend trained RTC conditioning and postfix-only flow loss to plain PI05 so single-task checkpoints can use the same real-time chunking path as PI052.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn
2026-07-17 10:14:13 +00:00
committed by Pepijn
parent c395286f3c
commit 962d754658
6 changed files with 140 additions and 20 deletions
@@ -0,0 +1,45 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
import pytest
import torch
pytest.importorskip("transformers")
from lerobot.policies.pi05.configuration_pi05 import PI05Config # noqa: E402
from lerobot.policies.pi05.modeling_pi05 import ( # noqa: E402
_build_flow_matching_inputs,
_reduce_training_rtc_loss,
)
def test_pi05_training_rtc_uses_clean_prefix_and_per_token_time():
actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]])
time = torch.tensor([0.25])
prefix_mask = torch.tensor([[True, True, False, False]])
x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]]
assert torch.equal(x_t[:, :2], actions[:, :2])
assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:])
def test_pi05_training_rtc_loss_excludes_clean_prefix():
losses = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]])
prefix_mask = torch.tensor([[True, True, False, False]])
loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
assert loss.item() == pytest.approx(3.0)
@pytest.mark.parametrize("max_delay", [-1, 5])
def test_pi05_config_rejects_invalid_training_rtc_delay(max_delay):
with pytest.raises(ValueError, match="rtc_training_max_delay"):
PI05Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay)
+1 -1
View File
@@ -96,7 +96,7 @@ def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled):
def test_rtc_processor_rejects_trained_mode_when_policy_does_not_support_it():
config = RTCConfig(mode="trained")
with pytest.raises(ValueError, match="requires a Pi052 checkpoint"):
with pytest.raises(ValueError, match="requires a PI05-compatible checkpoint"):
RTCProcessor(config)
processor = RTCProcessor(config, trained_mode_supported=True)