Add config for RTC

Author: helper2424

src/lerobot/configs/types.py

@@ -40,3 +40,10 @@ def __getitem__(self, key: Any) -> Any: ...
 class PolicyFeature:
     type: FeatureType
     shape: tuple
+
+
+class RTCAttentionSchedule(str, Enum):
+    ZEROS = "ZEROS"
+    ONES = "ONES"
+    LINEAR = "LINEAR"
+    EXP = "EXP"

src/lerobot/policies/act/modeling_act.py

@@ -108,7 +108,7 @@ def reset(self):
             self._action_queue = deque([], maxlen=self.config.n_action_steps)
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Select a single action given environment observations.
 
         This method wraps `select_actions` in order to return one action at a time for execution in the
@@ -133,7 +133,7 @@ def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         return self._action_queue.popleft()
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         self.eval()
 

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -100,7 +100,7 @@ def reset(self):
             self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         # stack n latest observations from the queue
         batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
@@ -112,7 +112,7 @@ def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         return actions
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Select a single action given environment observations.
 
         This method handles caching a history of observations and an action trajectory generated by the

src/lerobot/policies/pi0/modeling_pi0.py

@@ -360,12 +360,12 @@ def from_pretrained(cls, *args, **kwargs):
         return super().from_pretrained(*args, **kwargs)
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         raise NotImplementedError("Currently not implemented for PI0")
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs) -> Tensor:
         """Select a single action given environment observations.
 
         This method wraps `select_actions` in order to return one action at a time for execution in the

src/lerobot/policies/pi0fast/modeling_pi0fast.py

@@ -204,12 +204,12 @@ def _pi_aloha_encode_actions_inv(self, actions):
         return actions
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         raise NotImplementedError("Currently not implemented for PI0FAST")
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Select a single action given environment observations.
 
         This method wraps `select_actions` in order to return one action at a time for execution in the

src/lerobot/policies/pretrained.py

@@ -181,7 +181,7 @@ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Returns the action chunk (for action chunking policies) for a given observation, potentially in batch mode.
 
         Child classes using action chunking should use this method within `select_action` to form the action chunk
@@ -190,7 +190,7 @@ def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Return one action to run in the environment (potentially in batch mode).
 
         When the model uses a history of observations, or outputs a sequence of actions, this method deals

src/lerobot/policies/rtc/README.md

@@ -15,15 +15,6 @@ RTC can be integrated with any policy that supports flow mathicng for chunking:
 - **SmolVLA**: Vision-language-action model with RTC support
 - **Pi0**: Action prediction model with adaptive chunking
 
-## Configuration
-
-RTC behavior is configured through the `AdaptiveInferenceConfig` class located in `src/lerobot/policies/rtc_config.py`. Key parameters include:
-
-- `soft_mask_length`: Number of actions to soft mask in overlap regions
-- `beta`: Maximum guidance weight for prefix attention
-- `prefix_attention_schedule`: Attention weight scheduling strategy
-- `guidance_scale`: Scale factor for guidance correction
-
 ## Original Implementation
 
 This implementation is based on Physical Intelligence's Kinetix RTC:

src/lerobot/policies/rtc/configuration_rtc.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Real Time Chunking (RTC) and Bidirectional Decoding (BID) configuration classes.
+
+Based on:
+- Real Time Chunking: https://www.physicalintelligence.company/research/real_time_chunking
+"""
+
+from dataclasses import dataclass
+
+from lerobot.configs.types import RTCAttentionSchedule
+
+
+@dataclass
+class RTCConfig:
+    """Configuration for Real Time Chunking (RTC) inference.
+
+    RTC improves real-time inference by treating chunk generation as an inpainting problem,
+    strategically handling overlapping timesteps between action chunks using prefix attention.
+    """
+
+    # Core RTC settings
+    prefix_attention_schedule: RTCAttentionSchedule = RTCAttentionSchedule.LINEAR
+    max_guidance_weight: float = 5.0
+    execution_horizon: int = 10
+
+    def __post_init__(self):
+        """Validate RTC configuration parameters."""
+        if self.max_guidance_weight <= 0:
+            raise ValueError(f"max_guidance_weight must be positive, got {self.max_guidance_weight}")