Support Train-Time RTC

pyproject.toml

@@ -72,7 +72,7 @@ serving-web = [
     "uvicorn==0.44.0",
 ]
 server = [
-    "av==17.0.0",
+    "av==16.0.0",
     "chex==0.1.90",
     "dash==4.1.0",
     "fastapi==0.135.3",
@@ -87,7 +87,7 @@ groot = [
     "albumentations==2.0.8",
     "diffusers==0.33.1",
     "dm-tree==0.1.8",
-    "huggingface-hub==0.34.0",
+    "huggingface-hub==0.36.2",
     "opencv-python-headless==4.11.0.86",
     "pillow==12.2.0",
     "scipy==1.17.1",
@@ -96,6 +96,7 @@ groot = [
     "transformers==4.57.0",
 ]
 dev = [
+    "httpx==0.28.1",
     "modal==1.5.2",
     "pre-commit==4.6.0",
     "pytest==9.0.3",
@@ -104,6 +105,7 @@ dev = [
 
 [dependency-groups]
 dev = [
+    "httpx==0.28.1",
     "modal==1.5.2",
     "pre-commit==4.6.0",
     "pytest==9.0.3",
@@ -130,7 +132,7 @@ known-first-party = ["armory", "armory_client", "evaluation", "openpi_adapter",
 [tool.uv.sources]
 armory-client = { path = "armory-client", editable = true }
 gr00t = { path = "third_party/Isaac-GR00T", editable = true }
-lerobot = { git = "https://github.com/huggingface/lerobot", rev = "0cf864870cf29f4738d3ade893e6fd13fbd7cdb5" }
+lerobot = { git = "https://github.com/huggingface/lerobot", rev = "33cad37054c2b594ceba57463e8f11ee374fa93c" }
 libero = { path = "third_party/libero", editable = true }
 openpi = { path = "third_party/openpi", editable = true }
 openpi-client = { path = "third_party/openpi/packages/openpi-client", editable = true }

src/armory/backends/mock.py

@@ -10,6 +10,7 @@
 import numpy as np
 
 from armory.backends.types import PolicyResult, warmup_request
+from armory.serving.rtc import InferType
 from armory.serving.schemas import SlotData
 
 with Path(__file__).with_name("inference_profiles.json").open() as f:
@@ -36,8 +37,8 @@ def make_infer_request(self) -> SlotData:
         # Only ever fed to infer_batch, which ignores every field but the count.
         return warmup_request({})
 
-    def warmup(self, max_batch_size: int) -> None:
-        del max_batch_size
+    def warmup(self, max_batch_size: int, infer_type: InferType) -> None:
+        del max_batch_size, infer_type
 
     def infer_batch(self, requests: Sequence[SlotData]) -> list[PolicyResult]:
         inference_latency = self._inference_latency[len(requests)]

src/armory/backends/types.py

@@ -84,7 +84,7 @@ class PolicyResult(TypedDict):
 class ServingPolicy(Protocol):
     """Synchronous batched-policy interface consumed by ``GpuWorker``."""
 
-    def warmup(self, max_batch_size: int) -> None: ...
+    def warmup(self, max_batch_size: int, infer_type: InferType) -> None: ...
 
     def make_infer_request(self) -> SlotData: ...
 

src/armory/serving/config.py

@@ -3,10 +3,12 @@
 from pydantic import BaseModel
 
 from armory.serving.protocol import SchedulerConfig
+from armory.serving.rtc import InferType
 
 
 class EngineConfig(BaseModel):
     num_steps: int = 10
+    infer_type: InferType = InferType.SYNC
 
 
 class ServerConfig(BaseModel):

src/armory/serving/engine.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import math
 import multiprocessing as mp
 import pathlib
 import signal
@@ -79,7 +80,7 @@ def run(self) -> None:
         logger.info("GPU worker starting")
 
         policy = self.policy_factory()
-        policy.warmup(self.config.max_batch_size)
+        policy.warmup(self.config.max_batch_size, self.config.engine.infer_type)
 
         ctx = zmq.Context()
 
@@ -290,17 +291,19 @@ def _process_server_messages(self, req_sock: zmq.Socket) -> None:
 
     def _make_params(self, slot_data: SlotData, batch_size: int) -> RTCParams | None:
         if (
-            slot_data.infer_type == InferType.INFERENCE_TIME_RTC
+            slot_data.infer_type in (InferType.INFERENCE_TIME_RTC, InferType.TRAIN_TIME_RTC)
             and slot_data.robot_id in self._last_served_action_index
         ):
+            prev_action = self._prev_actions[slot_data.robot_id]
             s = slot_data.action_index_start - self._last_served_action_index[slot_data.robot_id]
-            d = (
+            if s >= len(prev_action):
+                return None
+            d = math.ceil(
                 self._latency_tracker.total_latency(slot_data.robot_id, batch_size)
                 * slot_data.control_hz
             )
-            return RTCParams(
-                prev_action=self._prev_actions[slot_data.robot_id], s_param=s, d_param=d
-            )
+            d = min(d, len(prev_action) - s)
+            return RTCParams(prev_action=prev_action, s_param=s, d_param=d)
         return None
 
     def _update_state(

src/armory/serving/rtc.py

@@ -1,9 +1,3 @@
-"""Dormant server-side RTC protocol types.
-
-Clients no longer select an inference mode. These stay server-local so the
-engine and policy RTC implementation can be retained without a wire feature.
-"""
-
 from dataclasses import dataclass
 from enum import Enum
 
@@ -14,7 +8,6 @@ class InferType(Enum):
     SYNC = "sync"
     INFERENCE_TIME_RTC = "inference_time_rtc"
     TRAIN_TIME_RTC = "train_time_rtc"
-    VLASH = "vlash"
 
 
 @dataclass

src/armory/serving/session.py

@@ -15,7 +15,6 @@
 from fastapi import WebSocket
 from starlette.websockets import WebSocketDisconnect
 
-from armory.serving.rtc import InferType
 from armory.serving.schemas import AckNotification, RobotID, SlotData, WarmupSeed
 from armory.serving.server_runtime import ServerState
 from armory_client import msgpack_numpy
@@ -176,7 +175,7 @@ async def _receive_loop(
                 deadline=req.deadline,
                 min_execution_horizon=req.min_execution_horizon,
                 max_execution_horizon=req.max_execution_horizon,
-                infer_type=InferType.SYNC,
+                infer_type=state.config.engine.infer_type,
                 params=None,
                 noise=req.noise,
                 control_hz=state.robot_metadata[robot_id].control_hz,

src/backends/gr00t_adapter/policy_adapter.py

@@ -27,6 +27,7 @@
 import numpy as np
 
 from armory.backends.types import PolicyResult, warmup_request
+from armory.serving.rtc import InferType
 from armory.serving.schemas import SlotData
 
 logger = logging.getLogger(__name__)
@@ -146,7 +147,8 @@ def infer_batch(self, requests: Sequence[SlotData]) -> list[PolicyResult]:
     def make_infer_request(self) -> SlotData:
         return warmup_request(_make_example_obs())
 
-    def warmup(self, max_batch_size: int) -> None:
+    def warmup(self, max_batch_size: int, infer_type: InferType) -> None:
+        del infer_type
         request = self.make_infer_request()
         for batch_size in range(1, max_batch_size + 1):
             logger.info("Warming up GR00T batch_size=%d", batch_size)

src/backends/openpi_adapter/policy_adapter.py

@@ -140,7 +140,7 @@ def create_batch_obs(self, observations: list[dict]):
                     [obs["observation/wrist_image"] for obs in observations]
                 ),
                 "prompt": np.stack([obs["prompt"] for obs in observations]),
-            }
+            }, None
 
         # Apply transform to each observation independently (tokenization, image parsing etc.)
         transformed = [
@@ -166,10 +166,11 @@ def create_batch_obs(self, observations: list[dict]):
                 batched,
             )
 
-        return _model.Observation.from_dict(batched)
+        prev_actions = batched.pop("actions", None)
+        return _model.Observation.from_dict(batched), prev_actions
 
     def _infer_batch_group(
-        self, requests: Sequence[SlotData], *, use_rtc: bool
+        self, requests: Sequence[SlotData], *, infer_type: InferType
     ) -> list[PolicyResult]:
         """Run a homogeneous sub-batch (all RTC or all non-RTC) in a single GPU call."""
         batch_size = len(requests)
@@ -190,7 +191,13 @@ def _infer_batch_group(
                     if req.noise is not None:
                         noise_to_use = noise_to_use.at[i].set(req.noise)
 
-        observation = self.create_batch_obs([_rename_keys(req.observation) for req in requests])
+        obs_dicts = [_rename_keys(req.observation) for req in requests]
+        if infer_type != InferType.SYNC:
+            obs_dicts = [
+                {**obs, "actions": np.array(req.params.prev_action)}
+                for obs, req in zip(obs_dicts, requests, strict=True)
+            ]
+        observation, prev_actions = self.create_batch_obs(obs_dicts)
 
         if self._is_triton_optimized:
             sample_kwargs = dict(self._sample_kwargs)
@@ -206,17 +213,16 @@ def _infer_batch_group(
                 result["noise"] = (
                     noise_np[i] if (noise_np is not None and noise_np.ndim == 3) else noise_np
                 )
-                result["rtc_prev_actions"] = raw_actions[i]
+                result["rtc_prev_actions"] = result["actions"]
                 results.append(result)
             return results
 
         sample_kwargs = dict(self._sample_kwargs)
         if noise_to_use is not None:
             sample_kwargs["noise"] = noise_to_use
 
-        if use_rtc:
+        if infer_type != InferType.SYNC:
             rtc_params = [req.params for req in requests]
-            prev_actions = np.stack([np.asarray(p.prev_action) for p in rtc_params], axis=0)
             s_values = np.asarray([p.s_param for p in rtc_params], dtype=np.int32)
             d_values = np.asarray([p.d_param for p in rtc_params], dtype=np.int32)
             eh_values = np.asarray([req.max_execution_horizon for req in requests], dtype=np.int32)
@@ -227,8 +233,11 @@ def _infer_batch_group(
                 d_values.tolist(),
                 eh_values.tolist(),
             )
-            sample_kwargs["use_rtc"] = True
-            sample_kwargs["prev_action"] = jnp.asarray(prev_actions)
+            if infer_type == InferType.TRAIN_TIME_RTC:
+                sample_kwargs["use_train_rtc"] = True
+            else:
+                sample_kwargs["use_rtc"] = True
+            sample_kwargs["prev_action"] = prev_actions
             sample_kwargs["s"] = jnp.asarray(s_values)
             sample_kwargs["d"] = jnp.asarray(d_values)
             sample_kwargs["execution_horizon"] = jnp.asarray(eh_values)
@@ -251,7 +260,7 @@ def _infer_batch_group(
             result["noise"] = (
                 noise_np[i] if (noise_np is not None and noise_np.ndim == 3) else noise_np
             )
-            result["rtc_prev_actions"] = raw_actions[i]
+            result["rtc_prev_actions"] = result["actions"]
             results.append(result)
 
         return results
@@ -266,25 +275,23 @@ def infer_batch(self, requests: Sequence[SlotData]) -> list[PolicyResult]:
             return []
 
         results: list[PolicyResult | None] = [None] * len(requests)
-        grouped: dict[bool, list[int]] = {False: [], True: []}
+        grouped: dict[InferType, list[int]] = {}
 
         for i, req in enumerate(requests):
             can_rtc = (
                 not self._is_pytorch_model
                 and not self._is_triton_optimized
-                and req.infer_type == InferType.INFERENCE_TIME_RTC
+                and req.infer_type in (InferType.INFERENCE_TIME_RTC, InferType.TRAIN_TIME_RTC)
                 and isinstance(req.params, RTCParams)
             )
             logger.debug(
                 f"can_rtc: {can_rtc}, pytorch_model: {self._is_pytorch_model}, triton_optimized: {self._is_triton_optimized}, infer_type: {req.infer_type}, params: {req.params}"
             )
-            grouped[can_rtc].append(i)
+            grouped.setdefault(req.infer_type if can_rtc else InferType.SYNC, []).append(i)
 
-        for use_rtc, indices in grouped.items():
-            if not indices:
-                continue
+        for infer_type, indices in grouped.items():
             sub_requests = [requests[i] for i in indices]
-            sub_results = self._infer_batch_group(sub_requests, use_rtc=use_rtc)
+            sub_results = self._infer_batch_group(sub_requests, infer_type=infer_type)
             for i, result in zip(indices, sub_results, strict=True):
                 results[i] = result
 
@@ -294,30 +301,28 @@ def infer_batch(self, requests: Sequence[SlotData]) -> list[PolicyResult]:
     def make_infer_request(self) -> SlotData:
         return warmup_request(self._make_example_fn())
 
-    def warmup(self, max_batch_size: int) -> None:
-        """Warm up both SYNC and RTC paths to trigger JAX JIT compilation."""
+    def warmup(self, max_batch_size: int, infer_type: InferType) -> None:
         example_obs = self._make_example_fn()
-        warmup_requests = [warmup_request(example_obs)]
-
-        # Add RTC warmup if the model supports it
-        if not self._is_pytorch_model and not self._is_triton_optimized:
-            example_actions = (
-                np.asarray(self._model.make_example_actions())
-                if hasattr(self._model, "make_example_actions")
-                else np.zeros((8, 7), dtype=np.float32)
-            )
-            warmup_requests.append(
-                warmup_request(
-                    example_obs,
-                    infer_type=InferType.INFERENCE_TIME_RTC,
-                    params=RTCParams(prev_action=example_actions, s_param=5, d_param=3),
-                )
+        for batch_size in range(1, max_batch_size + 1):
+            logger.info("Warming up %s batch_size=%d", InferType.SYNC, batch_size)
+            result = self.infer_batch([warmup_request(example_obs)] * batch_size)
+
+        if (
+            infer_type != InferType.SYNC
+            and not self._is_pytorch_model
+            and not self._is_triton_optimized
+        ):
+            req = warmup_request(
+                example_obs,
+                infer_type=infer_type,
+                params=RTCParams(
+                    prev_action=np.zeros_like(result[0]["actions"]), s_param=5, d_param=3
+                ),
             )
-
-        for req in warmup_requests:
             for batch_size in range(1, max_batch_size + 1):
-                logger.info("Warming up %s batch_size=%d", req.infer_type, batch_size)
+                logger.info("Warming up %s batch_size=%d", infer_type, batch_size)
                 result = self.infer_batch([req] * batch_size)
+
         logger.info("Warmup complete; output shape: %s", result[0]["actions"].shape)
 
     @property

tests/backends/openpi_adapter_test.py

@@ -96,28 +96,34 @@ def test_infer_batch_splits_rtc_requests_and_restores_original_order(
         # Existing behavior: RTC mode without RTCParams follows the non-RTC path.
         _request("rtc-without-params-2", InferType.INFERENCE_TIME_RTC),
         _request("rtc-3", InferType.INFERENCE_TIME_RTC, rtc_params),
+        _request("train-rtc-4", InferType.TRAIN_TIME_RTC, rtc_params),
     ]
-    calls: list[tuple[bool, list[str]]] = []
+    calls: list[tuple[InferType, list[str]]] = []
 
     def fake_infer_batch_group(
         grouped_requests: list[SlotData],
         *,
-        use_rtc: bool,
+        infer_type: InferType,
     ) -> list[dict]:
-        calls.append((use_rtc, [request.robot_id for request in grouped_requests]))
-        return [{"robot_id": request.robot_id, "used_rtc": use_rtc} for request in grouped_requests]
+        calls.append((infer_type, [request.robot_id for request in grouped_requests]))
+        return [
+            {"robot_id": request.robot_id, "infer_type": infer_type}
+            for request in grouped_requests
+        ]
 
     monkeypatch.setattr(adapter, "_infer_batch_group", fake_infer_batch_group)
 
     results = adapter.infer_batch(requests)
 
-    assert calls == [
-        (False, ["sync-1", "rtc-without-params-2"]),
-        (True, ["rtc-0", "rtc-3"]),
+    assert sorted(calls, key=lambda c: c[0].value) == [
+        (InferType.INFERENCE_TIME_RTC, ["rtc-0", "rtc-3"]),
+        (InferType.SYNC, ["sync-1", "rtc-without-params-2"]),
+        (InferType.TRAIN_TIME_RTC, ["train-rtc-4"]),
     ]
     assert results == [
-        {"robot_id": "rtc-0", "used_rtc": True},
-        {"robot_id": "sync-1", "used_rtc": False},
-        {"robot_id": "rtc-without-params-2", "used_rtc": False},
-        {"robot_id": "rtc-3", "used_rtc": True},
+        {"robot_id": "rtc-0", "infer_type": InferType.INFERENCE_TIME_RTC},
+        {"robot_id": "sync-1", "infer_type": InferType.SYNC},
+        {"robot_id": "rtc-without-params-2", "infer_type": InferType.SYNC},
+        {"robot_id": "rtc-3", "infer_type": InferType.INFERENCE_TIME_RTC},
+        {"robot_id": "train-rtc-4", "infer_type": InferType.TRAIN_TIME_RTC},
     ]

tests/serving/server_smoke_test.py

@@ -52,8 +52,8 @@ def get_next_batches(
 class _SmokePolicy:
     """Small CPU policy exercising the real engine process and shared slots."""
 
-    def warmup(self, max_batch_size: int) -> None:
-        del max_batch_size
+    def warmup(self, max_batch_size: int, infer_type: object) -> None:
+        del max_batch_size, infer_type
 
     def make_infer_request(self) -> object:
         return object()