From 687484a8643324abea59181b7745d687fe24bc91 Mon Sep 17 00:00:00 2001 From: Eugene Mironov Date: Fri, 7 Nov 2025 02:40:11 +0700 Subject: [PATCH] Add RTC to PI0 --- examples/rtc/eval_dataset.py | 9 ++++++++- src/lerobot/policies/pi0/configuration_pi0.py | 4 ++++ src/lerobot/policies/pi05/configuration_pi05.py | 4 ++++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/examples/rtc/eval_dataset.py b/examples/rtc/eval_dataset.py index aebf1034b..f14a00711 100644 --- a/examples/rtc/eval_dataset.py +++ b/examples/rtc/eval_dataset.py @@ -11,13 +11,20 @@ It compares action predictions with and without RTC on dataset samples, measuring consistency and ground truth alignment. Usage: - # Basic usage + # Basic usage with smolvla policy uv run python examples/rtc/eval_dataset.py \ --policy.path=helper2424/smolvla_check_rtc_last3 \ --dataset.repo_id=helper2424/check_rtc \ --rtc.execution_horizon=8 \ --device=mps + # Basic usage with pi0.5 policy + uv run python examples/rtc/eval_dataset.py \ + --policy.path=lerobot/pi05_libero_finetuned \ + --dataset.repo_id=HuggingFaceVLA/libero \ + --rtc.execution_horizon=8 \ + --device=mps + # With torch.compile for faster inference (PyTorch 2.0+) # Note: CUDA graphs disabled by default due to in-place ops in denoising loop uv run python examples/rtc/eval_dataset.py \ diff --git a/src/lerobot/policies/pi0/configuration_pi0.py b/src/lerobot/policies/pi0/configuration_pi0.py index d745f4317..9e267fb48 100644 --- a/src/lerobot/policies/pi0/configuration_pi0.py +++ b/src/lerobot/policies/pi0/configuration_pi0.py @@ -20,6 +20,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.policies.rtc.configuration_rtc import RTCConfig from lerobot.utils.constants import OBS_IMAGES @@ -47,6 +48,9 @@ class PI0Config(PreTrainedConfig): min_period: float = 4e-3 max_period: float = 4.0 + # Real-Time Chunking (RTC) configuration + rtc_config: RTCConfig | None = None + image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py` # Add empty images. Used to add empty cameras when no image features are present. diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 61346c330..2edd625af 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/src/lerobot/policies/pi05/configuration_pi05.py @@ -20,6 +20,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig +from lerobot.policies.rtc.configuration_rtc import RTCConfig @PreTrainedConfig.register_subclass("pi05") @@ -46,6 +47,9 @@ class PI05Config(PreTrainedConfig): min_period: float = 4e-3 max_period: float = 4.0 + # Real-Time Chunking (RTC) configuration + rtc_config: RTCConfig | None = None + image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py` # Add empty images. Used to add empty cameras when no image features are present.