From 925367d98af50372505bd9f5ce66ad7e4bfdca44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8E=E4=B9=90?= Date: Wed, 14 Oct 2026 00:58:21 +0800 Subject: [PATCH] feat: add RTC asynchronous execution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 于乐 --- docs/en/configuration.md | 20 + docs/en/quickstart.md | 2 +- docs/zh/configuration.md | 18 + docs/zh/quickstart.md | 2 +- integrations/sglang_pi05/README.md | 7 + .../src/embodirun_sglang_pi05/pi05.py | 545 ++++++++++++- src/embodirun/application/__init__.py | 10 + .../application/async_inference/__init__.py | 50 ++ .../application/async_inference/contracts.py | 241 ++++++ .../application/async_inference/rtc.py | 425 ++++++++++ .../application/async_inference/runtime.py | 761 ++++++++++++++++++ src/embodirun/application/contracts.py | 18 +- src/embodirun/application/control_service.py | 148 +++- src/embodirun/application/model_loop.py | 29 +- .../application/simulation/__init__.py | 28 +- .../application/simulation/async_runtime.py | 728 +++++++++++++++++ .../application/simulation/contracts.py | 6 +- .../application/simulation/service.py | 58 +- src/embodirun/deployment/plan.py | 1 + src/embodirun/robots/lerobot/so101/adapter.py | 43 +- .../robots/sensors/cameras/v4l2/__init__.py | 3 +- .../services/host/cli/command/run.py | 18 +- .../inference/adapters/sglang/pi05.py | 3 +- src/embodirun/simulators/libero/adapter.py | 78 +- src/embodirun/simulators/libero/config.py | 5 + tests/fixtures/async_inference_golden.json | 215 +++++ tests/test_async_algorithms.py | 669 +++++++++++++++ tests/test_async_simulation_runtime.py | 607 ++++++++++++++ tests/test_control_services.py | 6 + tests/test_libero_action_semantics.py | 56 ++ tests/test_sglang_async_wire.py | 86 ++ tests/test_sglang_http.py | 195 +++++ tests/test_so101_pi05_runtime.py | 44 + 33 files changed, 5040 insertions(+), 85 deletions(-) create mode 100644 src/embodirun/application/async_inference/__init__.py create mode 100644 src/embodirun/application/async_inference/contracts.py create mode 100644 src/embodirun/application/async_inference/rtc.py create mode 100644 src/embodirun/application/async_inference/runtime.py create mode 100644 src/embodirun/application/simulation/async_runtime.py create mode 100644 tests/fixtures/async_inference_golden.json create mode 100644 tests/test_async_algorithms.py create mode 100644 tests/test_async_simulation_runtime.py create mode 100644 tests/test_libero_action_semantics.py create mode 100644 tests/test_sglang_async_wire.py diff --git a/docs/en/configuration.md b/docs/en/configuration.md index d993037..e795b8b 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -149,6 +149,26 @@ Host. Device-only runtimes can omit the model; see Wireless runtimes additionally declare `inference_client.bind` and `inference_client.port`, and each listener on a node must use a distinct port. +Asynchronous execution is opt-in through `options.async`. Use `mode: rtc` for +Real-Time Chunking. The `overlap_steps` and `inference_delay_steps` values +describe the modeled control window in simulation, while a physical runtime +measures the request latency. `mode: naive` keeps the overlap baseline without +algorithm-specific conditioning. RTC through SGLang requires its +`/v1/actions/metadata` response to advertise the requested +`async_inference` capability. + +```yaml +runtimes: + libero-pi05-runtime: + # ... robot/simulator, model, binding, and server fields ... + options: + async: + mode: rtc + overlap_steps: 10 + inference_delay_steps: 5 # simulation Δ; omit for a physical runtime + execution_horizon: 50 +``` + ## Providers Inference backends are registered providers. A provider declares which diff --git a/docs/en/quickstart.md b/docs/en/quickstart.md index 6cd21d7..aa9a149 100644 --- a/docs/en/quickstart.md +++ b/docs/en/quickstart.md @@ -121,7 +121,7 @@ uv run embodirun --config my-deployment.yaml run \ | `--task`, `--seed` | Simulator task and episode seed. | | `--chunk-steps` | Actions executed from each inference chunk (capped by the binding). | | `--max-steps` | Maximum number of inference/action chunks (default 1). | -| `--control-hz` | Action playback rate (default 5). | +| `--control-hz` | Action playback rate (default 5); set to `0` to disable software pacing. | | `--request-timeout` | Timeout for one inference request (default 60 s). | A π0.5 response has no task-complete signal. The chunk bound limits the run; diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 022477e..01f7558 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -146,6 +146,24 @@ runtimes: 无线运行时还需声明 `inference_client.bind` 和 `inference_client.port`，并且节点上的每个监听器必须使用不同的端口。 +异步执行通过 `options.async` 显式开启。`mode: rtc` 使用 Real-Time Chunking。 +仿真器用 `overlap_steps` 和 `inference_delay_steps` 描述模拟的控制窗口，真机运行时会测量实际请求延迟。 +`mode: naive` 可用作不携带算法条件的异步基线。通过 SGLang +使用 RTC 时，其 `/v1/actions/metadata` 必须声明对应的 +`async_inference` 能力。 + +```yaml +runtimes: + libero-pi05-runtime: + # ... simulator/robot、model、binding 和 server 字段 ... + options: + async: + mode: rtc + overlap_steps: 10 + inference_delay_steps: 5 # 仿真 Δ；真机可以省略 + execution_horizon: 50 +``` + ## 扩展推理后端：provider {#provider} 每个推理后端通过 provider 注册，声明支持的传输协议、启动命令、依赖组和模型参数。 diff --git a/docs/zh/quickstart.md b/docs/zh/quickstart.md index 80180f2..b2f4a7c 100644 --- a/docs/zh/quickstart.md +++ b/docs/zh/quickstart.md @@ -111,7 +111,7 @@ uv run embodirun --config my-deployment.yaml run \ | `--task`、`--seed` | 仿真器任务与 episode 种子。 | | `--chunk-steps` | 每个动作块中实际执行的步数，不能超过绑定上限。 | | `--max-steps` | 推理与动作执行的最大轮数（默认 1）。 | -| `--control-hz` | 动作回放频率（默认 5）。 | +| `--control-hz` | 动作回放频率（默认 5）；设为 `0` 可关闭软件节拍。 | | `--request-timeout` | 单次推理请求的超时时间（默认 60 s）。 | π0.5 每次返回一个动作块（chunk），不返回任务完成信号。运行在达到设定的动作块数量后结束， diff --git a/integrations/sglang_pi05/README.md b/integrations/sglang_pi05/README.md index effa430..154ad65 100644 --- a/integrations/sglang_pi05/README.md +++ b/integrations/sglang_pi05/README.md @@ -4,3 +4,10 @@ This optional package owns the LeRobot Pi05 checkpoint compatibility pipeline. Install it only in the managed SGLang inference environment. The managed entry point is `embodirun-sglang-pi05-serve`. EmbodiRun's own `embodirun-sglang-serve` launcher stays a thin wrapper in the root package. + +The Pi05 pipeline exposes the shared `embodiinfer.async.v1` RTC metadata +contract. It keeps a per-session model-space action chunk and applies the +requested hard-prefix or soft-VJP denoising correction. RTC is currently +supported for SGLang's +monolithic worker (the normal single-4090 deployment); a split prefix/action +worker rejects an RTC request instead of silently running it synchronously. diff --git a/integrations/sglang_pi05/src/embodirun_sglang_pi05/pi05.py b/integrations/sglang_pi05/src/embodirun_sglang_pi05/pi05.py index 6f6cc1c..ef6409a 100644 --- a/integrations/sglang_pi05/src/embodirun_sglang_pi05/pi05.py +++ b/integrations/sglang_pi05/src/embodirun_sglang_pi05/pi05.py @@ -9,6 +9,10 @@ from __future__ import annotations import json +import logging +import math +import threading +import time from dataclasses import dataclass, replace from pathlib import Path from typing import Any @@ -26,12 +30,237 @@ VLAActionPostprocessStage, VLAObservationPreprocessStage, VLAPrefixEncodingStage, + _cuda_graph_enabled, + synchronize_vla_action_tensor, + vla_state, + vla_timings, ) from embodirun.model_services.backends.sglang import ( prepare_sglang_environment, ) +_ASYNC_SCHEMA = "embodiinfer.async.v1" +_RTC_ACTIVATION_LOCK = threading.Lock() +_RTC_TELEMETRY_LOCK = threading.Lock() +_RTC_TELEMETRY_REQUESTS = 0 +logger = logging.getLogger(__name__) + + +def _async_block(observation: dict[str, Any]) -> dict[str, Any]: + """Read and validate the shared async metadata from an action request.""" + + metadata = observation.get("metadata") + if metadata is None: + return {} + if not isinstance(metadata, dict): + raise ValueError("observation.metadata must be an object") + block = metadata.get("async") + if block is None: + return {} + if not isinstance(block, dict): + raise ValueError("observation.metadata.async must be an object") + if block.get("schema") != _ASYNC_SCHEMA: + raise ValueError(f"unsupported async metadata schema {block.get('schema')!r}") + return block + + +def _prefix_weights(start: int, end: int, total: int, schedule: str) -> torch.Tensor: + start = min(start, end) + if schedule not in {"zeros", "ones", "linear", "exp"}: + raise ValueError(f"unsupported RTC prefix_attention_schedule {schedule!r}") + if schedule == "zeros": + out = torch.zeros(total) + out[:start] = 1 + return out + if schedule == "ones": + out = torch.ones(total) + out[end:] = 0 + return out + interior = torch.linspace(1, 0, max(total - max(total - end, 0) - start, 0) + 2)[1:-1] + if schedule == "exp": + interior = interior * torch.expm1(interior) / (math.e - 1.0) + return torch.cat([torch.ones(min(start, total)), interior, torch.zeros(max(total - end, 0))]) + + +def _rtc_strength(action_ness: torch.Tensor, maximum: float) -> torch.Tensor: + one_minus = (1 - action_ness) ** 2 + inverse_r2 = (one_minus + action_ness**2) / one_minus + c = torch.nan_to_num((1 - action_ness) / action_ness, posinf=maximum) + value = torch.nan_to_num(c * inverse_r2, posinf=maximum) + return torch.minimum(value, torch.as_tensor(maximum, device=value.device, dtype=value.dtype)) + + +class _AutogradSafeGeluAndMul(torch.nn.Module): + """Use SGLang's pure-Torch GEGLU whenever a VJP is being recorded. + + The CUDA ``GeluAndMul`` path writes into a preallocated output tensor and + does not expose an autograd formula. Plain SGLang sampling still delegates + to that optimized implementation; only soft RTC's grad-enabled action + expert pass uses ``forward_native``. + """ + + def __init__(self, operation: torch.nn.Module) -> None: + super().__init__() + self.operation = operation + + def forward(self, value: torch.Tensor) -> torch.Tensor: + if torch.is_grad_enabled() and value.requires_grad: + return self.operation.forward_native(value) + return self.operation(value) + + +def _install_rtc_autograd_activations(model: Any) -> None: + """Make the Pi0.5 action expert's fused activations VJP-safe once.""" + + if bool(getattr(model, "_embodirun_rtc_autograd_activations", False)): + return + with _RTC_ACTIVATION_LOCK: + if bool(getattr(model, "_embodirun_rtc_autograd_activations", False)): + return + core = getattr(getattr(model, "action_expert", None), "core_model", None) + expert = getattr(getattr(core, "paligemma_with_expert", None), "gemma_expert", None) + expert_model = getattr(expert, "model", None) + layers = getattr(expert_model, "layers", None) + if layers is None: + raise RuntimeError("soft RTC could not locate the Pi0.5 action-expert layers") + wrapped = 0 + for layer in layers: + mlp = getattr(layer, "mlp", None) + activation = getattr(mlp, "act_fn", None) + if isinstance(activation, _AutogradSafeGeluAndMul): + wrapped += 1 + continue + if activation is None or not callable(getattr(activation, "forward_native", None)): + raise RuntimeError("soft RTC requires an autograd-safe action-expert activation") + mlp.act_fn = _AutogradSafeGeluAndMul(activation).eval() + wrapped += 1 + if wrapped == 0: + raise RuntimeError("soft RTC found no Pi0.5 action-expert activations") + model._embodirun_rtc_autograd_activations = True + + +def _grad_enabled_denoise( + model: Any, + prefix_context: Any, + x: torch.Tensor, + timestep: torch.Tensor, +) -> torch.Tensor: + """Run the upstream Pi0.5 core without its inference-only ``no_grad``. + + SGLang decorates ``Pi05CoreModel.denoise_step`` with ``@torch.no_grad``. + Calling ``Pi05PolicyModel.denoise_step`` from an outer ``enable_grad`` + therefore produces a detached velocity and silently reduces RTC's VJP to + the identity. ``torch.no_grad`` preserves the wrapped implementation, so + soft RTC can call that exact body while retaining all model semantics. + """ + + _install_rtc_autograd_activations(model) + core = getattr(getattr(model, "action_expert", None), "core_model", None) + decorated = getattr(type(core), "denoise_step", None) if core is not None else None + implementation = getattr(decorated, "__wrapped__", None) + if implementation is None: + raise RuntimeError("soft RTC requires access to the grad-enabled Pi0.5 denoise implementation") + # The enclosing SGLang action stage runs in inference_mode, so its + # per-step timestep is an inference tensor. Pi0.5's suffix embedder saves + # that scalar while differentiating with respect to the action tensor; + # PyTorch rejects inference tensors saved for backward. Materialize the + # tiny batch as a normal tensor before entering the grad-enabled core. + with torch.inference_mode(False): + timestep = timestep.detach().clone() + return implementation( + core, + prefix_context.prefix_pad_masks, + prefix_context.past_key_values, + x, + timestep, + bool(prefix_context.layout.get("full_attention", False)), + action_position_offset=0, + ) + + +def _rtc_autograd_prefix_context(prefix_context: Any) -> Any: + """Copy request prefix tensors out of the server's inference-mode arena. + + SDPA saves the fixed prefix K/V tensors while differentiating with respect + to the suffix query. PyTorch refuses to save inference tensors for + backward, so soft RTC materializes one normal-tensor cache per request and + reuses it across all denoise steps. + """ + + with torch.inference_mode(False): + cache = prefix_context.past_key_values + normal_cache = type(cache)( + ( + keys.detach().clone(), + values.detach().clone(), + sliding_window, + ) + for keys, values, sliding_window in cache + ) + normal_masks = prefix_context.prefix_pad_masks.detach().clone() + return replace( + prefix_context, + past_key_values=normal_cache, + prefix_pad_masks=normal_masks, + ) + + +def _rtc_velocity( + denoise: Any, + x: torch.Tensor, + time_value: torch.Tensor, + previous: torch.Tensor, + weights: torch.Tensor, + maximum: float, + *, + telemetry: bool = False, +) -> torch.Tensor: + """Compute the soft RTC VJP used by EmbodiInfer's flow decoder.""" + + global _RTC_TELEMETRY_REQUESTS + + # SGLang may invoke the action pipeline under inference_mode. Unlike + # no_grad, inference_mode cannot be overridden by enable_grad alone, so + # leave it explicitly for this one VJP and return to the caller's mode. + with torch.inference_mode(False), torch.enable_grad(): + leaf = x.detach().clone().requires_grad_(True) + # ``time_value`` is a scalar view from an inference-mode linspace in + # the server loop. The clean-action product also saves it for VJP. + time_value = time_value.detach().clone().to(device=leaf.device, dtype=leaf.dtype) + previous = previous.detach().clone().to(device=leaf.device, dtype=leaf.dtype) + weights = weights.detach().clone().to(device=leaf.device, dtype=leaf.dtype) + velocity = denoise(leaf) + if not velocity.requires_grad or velocity.grad_fn is None: + raise RuntimeError("soft RTC denoiser returned a detached velocity; the VJP would degenerate to identity") + if velocity.shape != leaf.shape: + raise RuntimeError("soft RTC denoiser returned an unexpected action shape") + action_ness = 1.0 - time_value + clean = leaf - (1.0 - action_ness) * velocity + error = (previous - clean) * weights.reshape(1, -1, 1) + correction = torch.autograd.grad(clean, leaf, error, retain_graph=False)[0] + telemetry_request = None + if telemetry and _RTC_TELEMETRY_REQUESTS < 2: + with _RTC_TELEMETRY_LOCK: + if _RTC_TELEMETRY_REQUESTS < 2: + _RTC_TELEMETRY_REQUESTS += 1 + telemetry_request = _RTC_TELEMETRY_REQUESTS + if telemetry_request is not None: + if not torch.isfinite(velocity).all() or not torch.isfinite(correction).all(): + raise RuntimeError("soft RTC produced non-finite velocity or VJP correction") + identity_delta = (correction - error).abs().max().detach() + if not torch.isfinite(identity_delta) or float(identity_delta) <= 0.0: + raise RuntimeError("soft RTC VJP correction is identical to the identity error") + logger.info( + "soft RTC VJP verified request=%d velocity_requires_grad=%s " + "velocity_grad_fn=%s correction_identity_max_abs=%.9g", + telemetry_request, + velocity.requires_grad, + type(velocity.grad_fn).__name__, + float(identity_delta), + ) + return (velocity - _rtc_strength(action_ness, maximum) * correction).detach() + @dataclass(frozen=True) class _Statistics: @@ -64,7 +293,10 @@ class _Statistics: @classmethod def load(cls, root: Path, pipeline: str, registry: str, feature: str, kind: str) -> _Statistics: - payload = json.loads((root / f"policy_{pipeline}.json").read_text()) + processor_path = root / f"policy_{pipeline}.json" + if not processor_path.is_file(): + raise FileNotFoundError(f"checkpoint is missing policy_{pipeline}.json") + payload = json.loads(processor_path.read_text()) steps = payload["steps"] allowed = { registry, @@ -151,6 +383,35 @@ def denormalize(self, values: Any) -> torch.Tensor: class _LeRobotPolicyModel(Pi05PolicyModel): + @staticmethod + def _apply_checkpoint_config(model_path: str, config: Any) -> None: + """Apply local tokenizer and LeRobot feature metadata.""" + + Pi05PolicyModel._apply_checkpoint_config(model_path, config) + root = Path(model_path) + local_tokenizer = root / "tokenizer_local" + if (local_tokenizer / "tokenizer_config.json").is_file(): + config.tokenizer_name = str(local_tokenizer) + config_path = root / "config.json" + if not config_path.is_file(): + return + payload = json.loads(config_path.read_text()) + input_features = payload.get("input_features") or {} + image_prefix = "observation.images." + image_keys = tuple( + key[len(image_prefix) :] + for key, feature in input_features.items() + if key.startswith(image_prefix) and isinstance(feature, dict) and feature.get("type") == "VISUAL" + ) + if image_keys: + # SGLang's Pi0.5 defaults name the three OpenPI cameras. Use the + # checkpoint's actual camera contract when it is available. + config.image_keys = image_keys + state = input_features.get("observation.state") or {} + shape = state.get("shape") or [] + if shape: + config.state_dim = int(shape[0]) + @staticmethod def _inspect_checkpoint(model_path: str, config: Any) -> Any: manifest = Pi05PolicyModel._inspect_checkpoint(model_path, config) @@ -164,14 +425,32 @@ def _inspect_checkpoint(model_path: str, config: Any) -> Any: class _LeRobotPreprocessor: - def __init__(self, processor: Pi05Preprocessor, statistics: _Statistics) -> None: + def __init__( + self, + processor: Pi05Preprocessor, + statistics: _Statistics, + ) -> None: self.processor = processor self.statistics = statistics def __call__(self, raw_observation: dict[str, Any]) -> Any: observation = dict(raw_observation) + async_block = _async_block(observation) observation["state"] = self.statistics.normalize(observation["state"]) - return self.processor(observation) + batch = self.processor(observation) + # ``VLAObservationBatch.metadata`` is the only object carried from the + # request adapter into the denoising stage. Keep the validated async + # block there; SGLang's generic action endpoint otherwise drops the + # request metadata after building ActionSamplingParams. + batch.metadata = { + **dict(getattr(batch, "metadata", {}) or {}), + "async": async_block, + "session_id": observation.get("session_id"), + "session_revision": observation.get("session_revision", 0), + "step_id": observation.get("step_id"), + "reset": bool(observation.get("reset", False)), + } + return batch class _LeRobotPostprocess(VLAActionPostprocessStage): @@ -188,6 +467,210 @@ def forward(self, batch: Any, server_args: Any) -> Any: return output +class _AsyncVLAActionDenoisingStage(VLAActionDenoisingStage): + """Pi05 action stage with the shared RTC contract and session cache. + + SGLang's stock ``sample_actions`` loop is intentionally synchronous and + has no hook for a committed prefix. This stage mirrors that loop for the + monolithic action path, adding the hard-prefix and soft-VJP RTC variants. + """ + + def __init__(self, policy_model: Any) -> None: + super().__init__(policy_model) + self._model_chunks: dict[tuple[str, int], torch.Tensor] = {} + self._model_chunks_lock = threading.RLock() + + def run_grouped_requests(self, batches: list[Any], server_args: Any) -> list[Any]: + """Keep async requests isolated from SGLang's ordinary action grouping. + + A grouped action call shares one denoising trajectory and therefore + cannot carry a different session prefix per row. The + stock stage may group by shape when prefix caching is disabled, so the + async marker must explicitly force the one-request path. + """ + + for batch in batches: + observation = vla_state(batch).get("observation_batch") + metadata = getattr(observation, "metadata", {}) if observation is not None else {} + async_block = metadata.get("async") if isinstance(metadata, dict) else None + if isinstance(async_block, dict) and "rtc" in async_block: + return [self(batch, server_args) for batch in batches] + return super().run_grouped_requests(batches, server_args) + + def forward(self, batch: Any, server_args: Any) -> Any: + state = vla_state(batch) + observation = state.get("observation_batch") + metadata = getattr(observation, "metadata", {}) if observation is not None else {} + async_block = metadata.get("async") if isinstance(metadata, dict) else None + rtc = async_block.get("rtc") if isinstance(async_block, dict) else None + # Warmup must retain SGLang's exact behavior and never populate a + # user-session action cache. + if batch.is_warmup or not isinstance(rtc, dict): + return super().forward(batch, server_args) + + split = __import__( + "sglang.multimodal_gen.runtime.vla.parallel", fromlist=["get_vla_split_group"] + ).get_vla_split_group() + if split is not None: + raise RuntimeError("SGLang Pi05 RTC currently requires monolithic execution (no split group)") + prefix_context = state.get("prefix_context") + if prefix_context is None: + raise RuntimeError("SGLang Pi05 RTC requires a prefix context") + + session = metadata.get("session_id") + revision = metadata.get("session_revision", 0) + key = (str(session), int(revision)) if session is not None else None + if metadata.get("reset") and key is not None: + with self._model_chunks_lock: + for cached_key in tuple(self._model_chunks): + if cached_key[0] == key[0]: + self._model_chunks.pop(cached_key, None) + previous = self._rtc_prefix(rtc, key, observation) + + start = time.perf_counter() + actions = self._sample_actions( + observation, + prefix_context, + noise=observation.noise, + num_steps=batch.num_inference_steps, + use_cuda_graph=_cuda_graph_enabled(batch), + generator=batch.generator, + rtc=rtc if previous is not None else None, + previous=previous, + ) + synchronize_vla_action_tensor(actions) + vla_timings(batch)["action_denoise_ms"] = (time.perf_counter() - start) * 1000 + state["parallel"] = self.policy_model.action_parallel_info(prefix_context) + state["actions"] = actions + if key is not None: + with self._model_chunks_lock: + self._model_chunks[key] = actions.detach().float().cpu() + return batch + + def _rtc_prefix(self, rtc: dict[str, Any], key: tuple[str, int] | None, observation: Any) -> torch.Tensor | None: + payload = rtc.get("prev_chunk_left_over") + if payload is not None: + if rtc.get("model_space") is not True: + raise ValueError("explicit RTC prefix requires model_space=true") + value = torch.as_tensor(payload, dtype=torch.float32) + return value if value.numel() else None + # The prefix starts at the action index already executed from the + # previous full model chunk. That index is the replan stride (for the + # paper recipe, 10), not inference latency Δ. Keep the delay fallback so + # older wire clients remain readable. + cache_start_index = rtc.get("cache_start_index", rtc.get("inference_delay", 0)) + if isinstance(cache_start_index, bool) or not isinstance(cache_start_index, int) or cache_start_index < 0: + raise ValueError("async.rtc.cache_start_index must be a non-negative integer") + explicit_cache_offset = "cache_start_index" in rtc + if key is None: + if explicit_cache_offset and cache_start_index > 0: + raise ValueError("RTC cache_start_index requires session metadata and a previous model chunk") + return None + with self._model_chunks_lock: + cached = self._model_chunks.get(key) + if cached is None: + if explicit_cache_offset and cache_start_index > 0: + raise ValueError("RTC cache_start_index requires a previous model chunk for this session") + return None + if cached.ndim == 3: + horizon = int(cached.shape[1]) + remaining = cached[:, cache_start_index:, :] + elif cached.ndim == 2: + # Kept for compatibility with caches populated by an older stage + # that squeezed the batch dimension. + horizon = int(cached.shape[0]) + remaining = cached[cache_start_index:, :] + else: + raise ValueError("SGLang RTC session cache must have shape [B, H, A] or [H, A]") + if explicit_cache_offset and cache_start_index > 0 and cache_start_index >= horizon: + raise ValueError("RTC cache_start_index must be smaller than the cached model horizon") + return remaining if remaining.numel() else None + + def _sample_actions( + self, + observation: Any, + prefix_context: Any, + *, + noise: Any, + num_steps: int, + use_cuda_graph: bool, + generator: Any, + rtc: dict[str, Any] | None, + previous: torch.Tensor | None, + ) -> torch.Tensor: + model = self.policy_model + if rtc is None or previous is None: + return model.sample_actions( + observation, + prefix_context, + noise=noise, + num_steps=num_steps, + use_cuda_graph=use_cuda_graph, + generator=generator, + ) + if previous.ndim == 2: + previous = previous.unsqueeze(0) + if previous.ndim != 3: + raise ValueError("async.rtc.prev_chunk_left_over must have shape [T, D]") + offload = model._offload_action_expert_between_requests() + if offload: + model._move_action_modules_to_device(model.device) + use_cuda_graph = False + x = noise + if x is None: + x = model.sample_noise(observation.batch_size, generator=generator) + else: + x = x.to(device=model.device, dtype=torch.float32).clone() + if tuple(previous.shape[0:1]) != tuple(x.shape[0:1]): + raise ValueError("RTC prefix batch size does not match the action batch") + aligned = torch.zeros_like(x) + rows = min(previous.shape[1], x.shape[1]) + cols = min(previous.shape[2], x.shape[2]) + aligned[:, :rows, :cols] = previous[:, :rows, :cols].to(device=x.device, dtype=x.dtype) + delay = rtc.get("inference_delay", 0) + horizon = rtc.get("execution_horizon", x.shape[1]) + maximum = float(rtc.get("max_guidance_weight", 10.0)) + schedule = str(rtc.get("prefix_attention_schedule", "linear")) + if isinstance(delay, bool) or not isinstance(delay, int) or delay < 0: + raise ValueError("async.rtc.inference_delay must be a non-negative integer") + if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 0: + raise ValueError("async.rtc.execution_horizon must be a non-negative integer") + weights = _prefix_weights(delay, min(horizon, rows, x.shape[1]), x.shape[1], schedule).to( + device=x.device, dtype=x.dtype + ) + hard = bool(rtc.get("hard_prefix", False)) + rtc_prefix_context = prefix_context if hard else _rtc_autograd_prefix_context(prefix_context) + mask = torch.arange(x.shape[1], device=x.device) < min(delay, x.shape[1]) + dt = -1.0 / num_steps + for step_index, timestep_value in enumerate(torch.linspace(1.0, 1.0 / num_steps, num_steps, device=x.device)): + timestep = timestep_value.expand(observation.batch_size) + if hard: + x[:, mask, :] = aligned[:, mask, :] + velocity = model.denoise_step( + prefix_context, + x, + timestep, + use_cuda_graph=use_cuda_graph, + ) + else: + velocity = _rtc_velocity( + lambda x_in, timestep=timestep: _grad_enabled_denoise(model, rtc_prefix_context, x_in, timestep), + x, + timestep_value, + aligned, + weights, + maximum, + telemetry=step_index == 0, + ) + x.add_(velocity, alpha=dt) + if hard: + x[:, mask, :] = aligned[:, mask, :] + if offload: + model._move_action_modules_to_device(torch.device("cpu")) + torch.cuda.empty_cache() + return x + + class LeRobotPi05Pipeline(Pi05Pipeline): """Native SGLang Pi05 execution with LeRobot MEAN_STD / QUANTILES processors.""" @@ -213,7 +696,10 @@ def initialize_pipeline(self, server_args: Any) -> None: root = Path(self.get_module("policy_model").model_path) state = _Statistics.load(root, "preprocessor", "normalizer_processor", "observation.state", "STATE") self.action_statistics = _Statistics.load(root, "postprocessor", "unnormalizer_processor", "action", "ACTION") - self.preprocessor = _LeRobotPreprocessor(self.preprocessor, state) + self.preprocessor = _LeRobotPreprocessor( + self.preprocessor, + state, + ) def create_pipeline_stages(self, server_args: Any) -> None: """Retain SGLang's prefix/denoising stages and restore action units last.""" @@ -223,15 +709,35 @@ def create_pipeline_stages(self, server_args: Any) -> None: "pi05_prefix", ) self.add_stage( - VLAActionDenoisingStage(self.get_module("policy_model")), + _AsyncVLAActionDenoisingStage(self.get_module("policy_model")), "pi05_action_denoise", ) self.add_stage(_LeRobotPostprocess(self.action_statistics), "pi05_postprocess") +def _sync_api_checkpoint_config(server_args: Any) -> bool: + """Apply checkpoint metadata in SGLang's HTTP process too. + + Model loading happens in a scheduler process, so mutations made by + ``Pi05PolicyModel.from_pretrained`` do not reach the API process's separate + ``server_args.pipeline_config``. The action protocol uses that copy for + metadata and image routing; synchronize it before either operation. + """ + + selected = getattr(server_args, "pipeline_class_name", None) + if selected not in (None, LeRobotPi05Pipeline.pipeline_name): + return False + model_path = getattr(server_args, "model_path", None) + config = getattr(server_args, "pipeline_config", None) + if model_path is not None and config is not None: + _LeRobotPolicyModel._apply_checkpoint_config(model_path, config) + return True + + def _register_pipeline() -> None: # The console entry point imports this module in spawned workers too. from sglang.multimodal_gen import registry + from sglang.multimodal_gen.runtime.entrypoints.action import api as action_api prepare_sglang_environment() registry._discover_and_register_pipelines() @@ -240,6 +746,35 @@ def _register_pipeline() -> None: LeRobotPi05Pipeline.pipeline_config_cls, LeRobotPi05Pipeline.sampling_params_cls, ) + # EmbodiRun probes this endpoint before enabling async execution. SGLang's + # generic action metadata has no knowledge of deployment-side RTC, so + # annotate the response only for the pipeline that implements it. + if not getattr(action_api, "_embodirun_async_metadata", False): + original_metadata = action_api.action_metadata + original_infer_action = action_api.infer_action + + def _metadata_with_async(server_args: Any) -> dict[str, Any]: + selected = _sync_api_checkpoint_config(server_args) + metadata = dict(original_metadata(server_args)) + # This entry point may still be used with another registered + # SGLang pipeline. Advertise the deployment contract only when + # the server actually selected the pipeline that implements it. + if not selected: + return metadata + async_capabilities = {"schema": _ASYNC_SCHEMA, "rtc": True} + metadata["async_inference"] = async_capabilities + capabilities = dict(metadata.get("capabilities") or {}) + capabilities["async_inference"] = async_capabilities + metadata["capabilities"] = capabilities + return metadata + + async def _infer_action_with_checkpoint_config(payload: dict[str, Any], server_args: Any) -> dict[str, Any]: + _sync_api_checkpoint_config(server_args) + return await original_infer_action(payload, server_args) + + action_api.action_metadata = _metadata_with_async + action_api.infer_action = _infer_action_with_checkpoint_config + action_api._embodirun_async_metadata = True def main() -> None: diff --git a/src/embodirun/application/__init__.py b/src/embodirun/application/__init__.py index 3d1bcc6..1a16812 100644 --- a/src/embodirun/application/__init__.py +++ b/src/embodirun/application/__init__.py @@ -17,9 +17,17 @@ ) _LAZY_CONTROL = {"ControlService", "ControlServiceError", "ControlTaskRejected"} +_LAZY_ASYNC = {"AsyncControlRuntime", "AsyncRuntimeConfig"} def __getattr__(name: str): + if name in _LAZY_ASYNC: + from .async_inference import AsyncControlRuntime, AsyncRuntimeConfig + + return { + "AsyncControlRuntime": AsyncControlRuntime, + "AsyncRuntimeConfig": AsyncRuntimeConfig, + }[name] if name in _LAZY_CONTROL: from .control_service import ( ControlService, @@ -48,4 +56,6 @@ def __getattr__(name: str): "Role", "TaskRequest", "TaskResult", + "AsyncControlRuntime", + "AsyncRuntimeConfig", ] diff --git a/src/embodirun/application/async_inference/__init__.py b/src/embodirun/application/async_inference/__init__.py new file mode 100644 index 0000000..3029581 --- /dev/null +++ b/src/embodirun/application/async_inference/__init__.py @@ -0,0 +1,50 @@ +"""Execution-side asynchronous VLA inference with RTC. + +This subpackage owns the execution half of RTC action chunking: + +* :mod:`~embodirun.application.async_inference.rtc` — Real-Time Chunking action + queue, latency tracker, and prefix-weight table. +* :mod:`~embodirun.application.async_inference.contracts` — the versioned + ``embodiinfer.async.v1`` metadata block shared with the inference engine. +* :mod:`~embodirun.application.async_inference.runtime` — the + :class:`AsyncControlRuntime` schedule. + +The package is pure Python plus numpy and never imports the inference engine; it +talks to it only through the versioned HTTP/WirelessComm API. +""" + +from __future__ import annotations + +from .contracts import ( + ASYNC_SCHEMA, + RTCAsyncRequest, + build_async_metadata, +) +from .rtc import ( + PREFIX_ATTENTION_SCHEDULES, + ActionQueue, + LatencyTracker, + prefix_weights, +) +from .runtime import ( + AsyncControlRuntime, + AsyncRuntimeConfig, + InferenceEvent, + default_action_encoder, + default_chunk_extractor, +) + +__all__ = [ + "ASYNC_SCHEMA", + "PREFIX_ATTENTION_SCHEDULES", + "ActionQueue", + "AsyncControlRuntime", + "AsyncRuntimeConfig", + "InferenceEvent", + "LatencyTracker", + "RTCAsyncRequest", + "build_async_metadata", + "default_action_encoder", + "default_chunk_extractor", + "prefix_weights", +] diff --git a/src/embodirun/application/async_inference/contracts.py b/src/embodirun/application/async_inference/contracts.py new file mode 100644 index 0000000..6bab7e5 --- /dev/null +++ b/src/embodirun/application/async_inference/contracts.py @@ -0,0 +1,241 @@ +"""Wire contract for asynchronous-inference hints. + +EmbodiRun and EmbodiInfer speak a versioned HTTP/WirelessComm API, so the async +algorithms cannot share Python objects across the boundary — they share a +*payload*. This module owns that payload: one optional ``async`` block inside +the existing step ``metadata``, carrying everything the server needs to +reproduce the execution side's timing decisions. + +Keeping it in ``metadata`` rather than widening the step schema is deliberate: +the step schema is frozen and every existing client keeps working unchanged, +because a request without an ``async`` block is exactly the synchronous request +it always was. RTC is opt-in per request. + +Schema (``embodiinfer.async.v1``):: + + { + "schema": "embodiinfer.async.v1", + "rtc": { + "prev_chunk_left_over": [[...], ...], + "inference_delay": 3, + "cache_start_index": 10, + "execution_horizon": 8, + "prefix_attention_schedule": "linear", + "max_guidance_weight": 10.0, + "hard_prefix": false + } + } + +Validation mirrors the server's ``embodiinfer.engine.async_inference.contracts`` +rules so a malformed block is rejected on the execution side too, instead of +silently downgrading to synchronous behaviour. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .rtc import PREFIX_ATTENTION_SCHEDULES + +__all__ = [ + "ASYNC_SCHEMA", + "RTCAsyncRequest", + "build_async_metadata", +] + +ASYNC_SCHEMA = "embodiinfer.async.v1" + +# Bounds mirroring the server so a malformed or hostile request cannot make a +# peer allocate without limit. +MAX_GRID_ROWS = 4096 +MAX_GRID_COLUMNS = 4096 + + +def _is_sequence(value: object) -> bool: + if isinstance(value, (str, bytes)): + return False + return isinstance(value, (Sequence, np.ndarray)) + + +def _finite_float(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)): + raise TypeError(f"{name} must be a number") + number = float(value) + if not math.isfinite(number): + raise ValueError(f"{name} must be finite") + return number + + +def _non_negative_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, (int, np.integer)): + raise TypeError(f"{name} must be an integer") + number = int(value) + if number < 0: + raise ValueError(f"{name} must be non-negative") + return number + + +def _positive_float(value: object, name: str) -> float: + number = _finite_float(value, name) + if number <= 0: + raise ValueError(f"{name} must be a finite positive number") + return number + + +def _numeric_grid(value: object, name: str) -> tuple[tuple[float, ...], ...]: + """Validate a nested sequence of finite numbers into a nested tuple.""" + + if not _is_sequence(value): + raise TypeError(f"{name} must be a list of lists") + rows_value = list(value) # type: ignore[arg-type] + if len(rows_value) > MAX_GRID_ROWS: + raise ValueError(f"{name} has too many rows (limit {MAX_GRID_ROWS})") + rows: list[tuple[float, ...]] = [] + width: int | None = None + for index, row in enumerate(rows_value): + if not _is_sequence(row): + raise TypeError(f"{name}[{index}] must be a list of numbers") + columns = [_finite_float(item, f"{name}[{index}]") for item in row] + if not columns: + raise ValueError(f"{name}[{index}] must not be empty") + if len(columns) > MAX_GRID_COLUMNS: + raise ValueError(f"{name}[{index}] has too many columns (limit {MAX_GRID_COLUMNS})") + if width is None: + width = len(columns) + elif len(columns) != width: + raise ValueError(f"{name} rows must all have the same width") + rows.append(tuple(columns)) + return tuple(rows) + + +def _grid_payload(grid: tuple[tuple[float, ...], ...]) -> list[list[float]]: + return [list(row) for row in grid] + + +@dataclass(frozen=True, slots=True) +class RTCAsyncRequest: + """One validated RTC hint for a single step request. + + Args: + prev_chunk_left_over: Unconsumed raw model-space actions from the + previous chunk, or ``None`` on the first inference of an episode. + inference_delay: Latency in control steps (``Δ``). + cache_start_index: Index in the server-cached full model chunk at which + the unexecuted prefix begins. This is the replan stride, not the + inference delay. Omit it when an explicit prefix is supplied. + execution_horizon: Number of leading steps eligible to be held. + prefix_attention_schedule: Weight shape across the horizon. + max_guidance_weight: Upper clamp on the guidance strength. + hard_prefix: Use the fixed-prefix clamp instead of soft guidance. + """ + + prev_chunk_left_over: tuple[tuple[float, ...], ...] | None + inference_delay: int + execution_horizon: int + prefix_attention_schedule: str = "linear" + max_guidance_weight: float = 10.0 + hard_prefix: bool = False + model_space: bool = False + cache_start_index: int | None = None + + def __post_init__(self) -> None: + if self.prev_chunk_left_over is not None: + grid = _numeric_grid(self.prev_chunk_left_over, "async.rtc.prev_chunk_left_over") + else: + grid = None + object.__setattr__(self, "prev_chunk_left_over", grid) + object.__setattr__( + self, + "inference_delay", + _non_negative_int(self.inference_delay, "async.rtc.inference_delay"), + ) + object.__setattr__( + self, + "execution_horizon", + _non_negative_int(self.execution_horizon, "async.rtc.execution_horizon"), + ) + if self.cache_start_index is not None: + object.__setattr__( + self, + "cache_start_index", + _non_negative_int(self.cache_start_index, "async.rtc.cache_start_index"), + ) + if not isinstance(self.prefix_attention_schedule, str) or ( + self.prefix_attention_schedule not in PREFIX_ATTENTION_SCHEDULES + ): + raise ValueError( + "async.rtc.prefix_attention_schedule must be one of " + f"{', '.join(PREFIX_ATTENTION_SCHEDULES)}, got {self.prefix_attention_schedule!r}" + ) + object.__setattr__( + self, + "max_guidance_weight", + _positive_float(self.max_guidance_weight, "async.rtc.max_guidance_weight"), + ) + if not isinstance(self.hard_prefix, bool): + raise TypeError("async.rtc.hard_prefix must be a boolean") + if not isinstance(self.model_space, bool): + raise TypeError("async.rtc.model_space must be a boolean") + if self.prev_chunk_left_over is not None and not self.model_space: + raise ValueError( + "async.rtc.prev_chunk_left_over requires model_space=true; " + "otherwise omit the prefix and use the session cache" + ) + + @property + def enabled(self) -> bool: + """Whether a previous chunk is present and worth conditioning on.""" + + return self.prev_chunk_left_over is not None and bool(self.prev_chunk_left_over) + + def to_payload(self) -> dict[str, Any]: + """Return the JSON-serializable ``rtc`` block.""" + + payload = { + "prev_chunk_left_over": ( + None if self.prev_chunk_left_over is None else _grid_payload(self.prev_chunk_left_over) + ), + "inference_delay": self.inference_delay, + "execution_horizon": self.execution_horizon, + "prefix_attention_schedule": self.prefix_attention_schedule, + "max_guidance_weight": self.max_guidance_weight, + "hard_prefix": self.hard_prefix, + } + if self.cache_start_index is not None: + payload["cache_start_index"] = self.cache_start_index + if self.prev_chunk_left_over is not None: + payload["model_space"] = True + return payload + + +def build_async_metadata( + *, + rtc: RTCAsyncRequest | None = None, + base: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Build step metadata carrying the optional ``async`` block. + + Args: + rtc: RTC hint, or ``None`` when RTC is disabled for this request. + base: Existing observation metadata to extend. It is copied, never + mutated. + + Returns: + A metadata mapping. When RTC is ``None`` the result equals + ``base`` (an empty mapping by default) with no ``async`` key, so the + request is byte-identical to the synchronous one. + """ + + metadata: dict[str, Any] = dict(base) if base is not None else {} + + block: dict[str, Any] = {"schema": ASYNC_SCHEMA} + if rtc is not None: + block["rtc"] = rtc.to_payload() + if len(block) > 1: + metadata["async"] = block + return metadata diff --git a/src/embodirun/application/async_inference/rtc.py b/src/embodirun/application/async_inference/rtc.py new file mode 100644 index 0000000..7064d7b --- /dev/null +++ b/src/embodirun/application/async_inference/rtc.py @@ -0,0 +1,425 @@ +"""Real-Time Chunking (RTC) execution-side bookkeeping. + +This module is the execution-side port of LeRobot's +``src/lerobot/policies/rtc/action_queue.py`` and ``latency_tracker.py``. It is +pure Python plus numpy: EmbodiRun never imports the inference engine, so the +numeric pieces the server also computes (the prefix-weight table) are duplicated +here and pinned against the shared golden vectors. + +The central object is :class:`ActionQueue`, a thread-safe queue holding *two* +parallel sequences: + +``original_queue`` + The policy's raw (model-space) chunk. Only the leftover prefix + ``original_queue[last_index:]`` is sent back to the server as + ``prev_chunk_left_over``. +``queue`` + The post-processed (robot-space) chunk that is actually executed. + +Both are numpy arrays of shape ``[T, D]``. Clones are returned so a caller +cannot mutate queue state through a returned action. +""" + +from __future__ import annotations + +import math +import threading +import time +from collections import deque +from collections.abc import Callable + +import numpy as np + +__all__ = [ + "PREFIX_ATTENTION_SCHEDULES", + "ActionQueue", + "LatencyTracker", + "prefix_weights", +] + +PREFIX_ATTENTION_SCHEDULES = ("zeros", "ones", "linear", "exp") + + +def _as_action_array(value: object, name: str) -> np.ndarray: + """Coerce one action chunk into a numeric numpy array, preserving dtype.""" + + array = np.asarray(value) + if array.ndim != 2: + raise ValueError(f"{name} must be a 2D array of actions, got shape {array.shape}") + if array.size and array.dtype.kind not in "fiu": + raise ValueError(f"{name} must be numeric") + if array.dtype.kind == "f" and not np.all(np.isfinite(array)): + raise ValueError(f"{name} must contain finite values") + return array + + +def prefix_weights( + start: int, + end: int, + total: int, + schedule: str = "linear", +) -> np.ndarray: + """Return the ``[total]`` RTC prefix-firmness weights for one request. + + ``start`` is the inference delay (pushed down to ``end`` when larger, since + ``end`` takes precedence), ``end`` is the execution horizon (exclusive), and + ``total`` is the action horizon ``H``. This is a numpy port of the same + piecewise leading-ones / ramp / trailing-zeros construction used by the + inference engine, so both sides of the boundary assert the identical golden + table and a drift on either side fails loudly. + + Args: + start: Inference delay in control steps (``Δ``). + end: Execution horizon, exclusive. + total: Action horizon ``H``. + schedule: One of :data:`PREFIX_ATTENTION_SCHEDULES`. + + Returns: + A ``float64`` array of length ``total``. + + Raises: + ValueError: for a non-positive ``total``, negative ``start``/``end``, or + an unknown schedule. + """ + + if isinstance(total, bool) or not isinstance(total, int) or total <= 0: + raise ValueError("total must be a positive integer") + for name, value in (("start", start), ("end", end)): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value < 0: + raise ValueError(f"{name} must be non-negative") + if schedule not in PREFIX_ATTENTION_SCHEDULES: + raise ValueError(f"schedule must be one of {', '.join(PREFIX_ATTENTION_SCHEDULES)}, got {schedule!r}") + + start = min(start, end) + + if schedule == "zeros": + weights = np.zeros(total, dtype=np.float64) + weights[:start] = 1.0 + return weights + if schedule == "ones": + weights = np.ones(total, dtype=np.float64) + weights[end:] = 0.0 + return weights + + ramp = _linear_ramp(start, end, total) + if schedule == "exp": + ramp = ramp * np.expm1(ramp) / (math.e - 1.0) + return _add_leading_ones(_add_trailing_zeros(ramp, total, end), start, total) + + +def _linear_ramp(start: int, end: int, total: int) -> np.ndarray: + """The interior ``linspace(1, 0)`` segment of the linear/exp schedules.""" + + skip_steps_at_end = max(total - end, 0) + linspace_steps = total - skip_steps_at_end - start + if end <= start or linspace_steps <= 0: + return np.array([], dtype=np.float64) + return np.linspace(1.0, 0.0, linspace_steps + 2)[1:-1] + + +def _add_trailing_zeros(weights: np.ndarray, total: int, end: int) -> np.ndarray: + zeros_len = total - end + if zeros_len <= 0: + return weights + return np.concatenate([weights, np.zeros(zeros_len, dtype=np.float64)]) + + +def _add_leading_ones(weights: np.ndarray, start: int, total: int) -> np.ndarray: + ones_len = min(start, total) + if ones_len <= 0: + return weights + return np.concatenate([np.ones(ones_len, dtype=np.float64), weights]) + + +class ActionQueue: + """Thread-safe double queue of raw and post-processed action chunks. + + Args: + rtc_enabled: When true, :meth:`merge` *replaces* the queue (dropping the + actions consumed during inference); when false it *appends*, + preserving continuity without replacement. + """ + + def __init__(self, *, rtc_enabled: bool = True) -> None: + self.rtc_enabled = bool(rtc_enabled) + self.queue: np.ndarray | None = None + self.original_queue: np.ndarray | None = None + self._task_queue: list[str | None] | None = None + self.lock = threading.Lock() + self.last_index = 0 + + # -- consumption ----------------------------------------------------- + def get(self) -> np.ndarray | None: + """Return the next processed action (a clone), or ``None`` if empty.""" + + queued = self.get_with_task() + return None if queued is None else queued[0] + + def get_with_task(self) -> tuple[np.ndarray, str | None] | None: + """Return the next processed action together with its chunk's task.""" + + with self.lock: + if self.queue is None or self.last_index >= len(self.queue): + return None + self._assert_lockstep() + action = np.array(self.queue[self.last_index], copy=True) + task = None if self._task_queue is None else self._task_queue[self.last_index] + self.last_index += 1 + return action, task + + def qsize(self) -> int: + """Number of unconsumed processed actions.""" + + with self.lock: + if self.queue is None: + return 0 + return len(self.queue) - self.last_index + + def empty(self) -> bool: + """Whether no processed actions remain.""" + + with self.lock: + if self.queue is None: + return True + return len(self.queue) - self.last_index <= 0 + + def get_action_index(self) -> int: + """Index of the next processed action to consume.""" + + with self.lock: + return self.last_index + + # -- leftovers ------------------------------------------------------- + def get_left_over(self) -> np.ndarray | None: + """Unconsumed *original* (model-space) actions, the RTC prefix.""" + + with self.lock: + self._assert_lockstep() + if self.original_queue is None: + return None + return np.array(self.original_queue[self.last_index :], copy=True) + + def get_processed_left_over(self) -> np.ndarray | None: + """Unconsumed *processed* (robot-space) actions.""" + + with self.lock: + self._assert_lockstep() + if self.queue is None: + return None + return np.array(self.queue[self.last_index :], copy=True) + + def clear(self) -> None: + """Drop all queued actions and reset consumption.""" + + with self.lock: + self.queue = None + self.original_queue = None + self._task_queue = None + self.last_index = 0 + + # -- merging --------------------------------------------------------- + def merge( + self, + original_actions: object, + processed_actions: object, + real_delay: int, + action_index_before_inference: int | None = None, + *, + task: str | None = None, + ) -> int: + """Merge a fresh chunk into the queue and return the resolved delay. + + RTC enabled replaces the queue, dropping the first ``resolved_delay`` + actions; RTC disabled appends, dropping already-consumed entries so the + new chunk continues where the old one left off. + + Args: + original_actions: Raw model-space chunk ``[T, D_model]``. + processed_actions: Robot-space chunk ``[T, D_robot]``. + real_delay: Delay in control steps, typically latency / period. + action_index_before_inference: ``last_index`` captured when the + request was sent; enables the never-over-discard rule. + task: Instruction that produced the chunk; labels every new action. + + Returns: + The resolved delay actually applied. + """ + + original = _as_action_array(original_actions, "original_actions") + processed = _as_action_array(processed_actions, "processed_actions") + with self.lock: + delay = self._check_and_resolve_delays(real_delay, action_index_before_inference) + if self.rtc_enabled: + return self._replace_actions_queue(original, processed, delay, task) + self._append_actions_queue(original, processed, task) + return delay + + def _replace_actions_queue( + self, + original_actions: np.ndarray, + processed_actions: np.ndarray, + real_delay: int, + task: str | None, + ) -> int: + clamped_delay = max(0, min(real_delay, len(original_actions), len(processed_actions))) + self.original_queue = np.array(original_actions[clamped_delay:], copy=True) + self.queue = np.array(processed_actions[clamped_delay:], copy=True) + self._task_queue = [task] * len(self.queue) + self.last_index = 0 + return clamped_delay + + def _append_actions_queue( + self, + original_actions: np.ndarray, + processed_actions: np.ndarray, + task: str | None, + ) -> None: + if self.queue is None: + self.original_queue = np.array(original_actions, copy=True) + self.queue = np.array(processed_actions, copy=True) + self._task_queue = [task] * len(self.queue) + self.last_index = 0 + return + + self._assert_lockstep() + existing_tasks = self._task_queue if self._task_queue is not None else [None] * len(self.queue) + original_queue = self.original_queue if self.original_queue is not None else np.empty((0,)) + self.original_queue = np.concatenate([original_queue, np.array(original_actions, copy=True)])[self.last_index :] + self.queue = np.concatenate([self.queue, np.array(processed_actions, copy=True)])[self.last_index :] + self._task_queue = existing_tasks[self.last_index :] + [task] * len(processed_actions) + self.last_index = 0 + + def _check_and_resolve_delays( + self, + real_delay: int, + action_index_before_inference: int | None = None, + ) -> int: + """Resolve the delay to apply. + + When the caller captured ``action_index_before_inference``, the delay can + never exceed the number of actions actually consumed. Discarding more + would splice the queue ahead of the physical pose. + """ + + if isinstance(real_delay, bool) or not isinstance(real_delay, int): + raise TypeError("real_delay must be an integer") + effective_delay = max(0, real_delay) + + if action_index_before_inference is not None: + if isinstance(action_index_before_inference, bool) or not isinstance(action_index_before_inference, int): + raise TypeError("action_index_before_inference must be an integer") + indexes_diff = max(0, self.last_index - action_index_before_inference) + if indexes_diff != real_delay: + # Never discard more than was actually consumed. + return min(real_delay, indexes_diff) + + return effective_delay + + def _assert_lockstep(self) -> None: + """Raise if the task-label list ever disagrees with the action length.""" + + action_count = 0 if self.queue is None else len(self.queue) + if self._task_queue is not None and len(self._task_queue) != action_count: + raise RuntimeError( + "ActionQueue task labels out of sync with actions " + f"({len(self._task_queue)} labels for {action_count} actions) — " + "a queue mutation broke the action/task lockstep invariant" + ) + + +class LatencyTracker: + """Inference-latency window with an injectable clock. + + Args: + window: Maximum number of samples kept. ``None`` keeps all. + monotonic: Injectable monotonic clock, for deterministic tests. + """ + + def __init__( + self, + window: int | None = 100, + *, + monotonic: Callable[[], float] = time.monotonic, + ) -> None: + if window is not None and (isinstance(window, bool) or not isinstance(window, int) or window <= 0): + raise ValueError("window must be a positive integer or None") + self._monotonic = monotonic + self._values: deque[float] = deque(maxlen=window) + self._started_at: float | None = None + self.max_latency = 0.0 + + def reset(self) -> None: + """Clear all recorded latencies.""" + + self._values.clear() + self._started_at = None + self.max_latency = 0.0 + + def start(self) -> None: + """Mark the start of one inference call.""" + + self._started_at = self._monotonic() + + def stop(self) -> float: + """Mark the end of one inference call and record its latency.""" + + if self._started_at is None: + raise RuntimeError("LatencyTracker.stop() called before start()") + latency = max(0.0, self._monotonic() - self._started_at) + self._started_at = None + self.record(latency) + return latency + + def record(self, latency: float) -> None: + """Add one latency sample in seconds. Negative values are ignored.""" + + value = float(latency) + if not math.isfinite(value) or value < 0: + return + self._values.append(value) + self.max_latency = max(self.max_latency, value) + + def __len__(self) -> int: + return len(self._values) + + def average(self) -> float: + """Mean latency over the window, or ``0.0`` when empty.""" + + if not self._values: + return 0.0 + return float(sum(self._values) / len(self._values)) + + def median(self) -> float: + """Median latency over the window, or ``0.0`` when empty.""" + + if not self._values: + return 0.0 + values = np.asarray(self._values, dtype=np.float64) + return float(np.median(values)) + + def max(self) -> float: + """Maximum latency recorded so far.""" + + return self.max_latency + + def delay_in_steps(self, latency_s: float, period_s: float) -> int: + """Convert a latency in seconds into a non-negative integer step count.""" + + period = float(period_s) + if not math.isfinite(period) or period <= 0: + raise ValueError("period_s must be a finite positive number") + latency = float(latency_s) + if not math.isfinite(latency) or latency < 0: + return 0 + # A tiny epsilon keeps an exact ``n * period`` product from flooring to + # ``n - 1`` because of binary floating-point representation. + return max(0, int(math.floor(latency / period + 1e-9))) + + def latest_delay(self, period_s: float) -> int: + """Delay implied by the most recent sample, or ``0`` when empty.""" + + if not self._values: + return 0 + return self.delay_in_steps(self._values[-1], period_s) diff --git a/src/embodirun/application/async_inference/runtime.py b/src/embodirun/application/async_inference/runtime.py new file mode 100644 index 0000000..be6ffb3 --- /dev/null +++ b/src/embodirun/application/async_inference/runtime.py @@ -0,0 +1,761 @@ +"""Asynchronous control loop executing naive overlap and RTC schedules. + +:class:`AsyncControlRuntime` mirrors :class:`embodirun.application.model_loop.ControlRuntime` +and reuses its session, cancellation, observation-source, and command-sink +handling. It adds the async schedule: while a chunk is being executed, the next +inference is launched before that chunk is exhausted, the number of actions +consumed during inference is tracked, and RTC prefix information is attached to +request metadata when enabled. + +Inference runs on one per-runtime worker while the control thread dispatches +the current action chunk. Only the control thread mutates queue and safety +state, so cancellation and takeover cannot apply a late result. + +* The inference client remains a blocking request-response API, so a single + worker is used per session. The request is never issued concurrently with a + second request; actions continue on the control thread until the worker + result is ready. +""" + +from __future__ import annotations + +import math +import threading +import time +import uuid +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, TimeoutError +from contextlib import suppress +from dataclasses import dataclass, field, replace +from typing import Any + +import numpy as np + +from embodirun.bindings import BindingMapper +from embodirun.model_services import InferenceClient, PolicyObservation, PolicyResult +from embodirun.robots import RobotAction, RobotAdapter, RobotObservation +from embodirun.robots.sensors.cameras import CameraFrame + +from ..model_loop import CommandSink, ControlRuntime +from .contracts import RTCAsyncRequest, build_async_metadata +from .rtc import ActionQueue, LatencyTracker + +__all__ = [ + "AsyncControlRuntime", + "AsyncRuntimeConfig", + "InferenceEvent", +] + + +def _flatten_numbers(value: object, out: list[float]) -> None: + """Flatten numeric leaves of a nested mapping/sequence in insertion order.""" + + if value is None or isinstance(value, (bool, str, bytes)): + return + if isinstance(value, (int, float, np.integer, np.floating)): + number = float(value) + if math.isfinite(number): + out.append(number) + return + if isinstance(value, Mapping): + for key, item in value.items(): + if key == "type": + continue + _flatten_numbers(item, out) + return + if isinstance(value, (list, tuple, np.ndarray)): + for item in value: + _flatten_numbers(item, out) + return + + +def default_action_encoder(action: RobotAction) -> np.ndarray: + """Encode one robot action as a flat numeric vector. + + The default flattens ``action.values`` in insertion order, skipping the + ``type`` discriminator. A binding whose robot-space vector order does not + match that convention should inject its own ``action_encoder``. + """ + + numbers: list[float] = [] + _flatten_numbers(action.values, numbers) + if not numbers: + raise ValueError("robot action does not expose any numeric values to encode") + return np.asarray(numbers, dtype=np.float64) + + +def default_chunk_extractor(result: PolicyResult) -> np.ndarray: + """Extract the raw model-space chunk from a policy result. + + The default looks for the first ``action_chunk`` action and returns its + ``values.data`` grid, the π0.5 wire shape. A different policy payload should + inject its own ``chunk_extractor``. + """ + + for action in result.actions: + if action.kind == "action_chunk": + data = action.values.get("data") + if data is None: + continue + chunk = np.asarray(data, dtype=np.float64) + if chunk.ndim != 2 or len(chunk) == 0: + raise ValueError("action_chunk values.data must be a non-empty 2D grid") + return chunk + raise ValueError("policy result does not contain an action_chunk action with values.data") + + +@dataclass(frozen=True, slots=True) +class AsyncRuntimeConfig: + """Static async-inference settings for :class:`AsyncControlRuntime`. + + Args: + enabled: Master switch. When false the runtime defers entirely to the + synchronous :meth:`ControlRuntime.step`. + rtc_enabled: Attach the RTC block and use RTC queue replacement. + overlap_steps: Control steps before chunk end at which to launch the next + inference. + execution_horizon: RTC execution horizon; defaults to ``chunk_steps``. + prefix_attention_schedule: RTC weight schedule. + max_guidance_weight: RTC guidance clamp. + hard_prefix: Use the RTC fixed-prefix clamp. + initial_inference_delay: Delay to advertise before any latency sample + exists. + artificial_inference_delay_s: Extra wall-clock delay added after each + inference response before it becomes available to the controller. + """ + + enabled: bool = False + rtc_enabled: bool = False + overlap_steps: int = 0 + execution_horizon: int | None = None + prefix_attention_schedule: str = "linear" + max_guidance_weight: float = 10.0 + hard_prefix: bool = False + initial_inference_delay: int = 0 + artificial_inference_delay_s: float = 0.0 + # Explicit prefixes are only safe when the extractor is guaranteed to + # return model-space actions. The default uses the server/session cache. + send_model_space_prefix: bool = False + + def __post_init__(self) -> None: + for name in ("enabled", "rtc_enabled", "hard_prefix", "send_model_space_prefix"): + if not isinstance(getattr(self, name), bool): + raise TypeError(f"{name} must be a boolean") + if isinstance(self.overlap_steps, bool) or not isinstance(self.overlap_steps, int) or self.overlap_steps < 0: + raise ValueError("overlap_steps must be a non-negative integer") + if self.execution_horizon is not None and ( + isinstance(self.execution_horizon, bool) + or not isinstance(self.execution_horizon, int) + or self.execution_horizon < 0 + ): + raise ValueError("execution_horizon must be a non-negative integer or None") + if not isinstance(self.max_guidance_weight, (int, float)) or isinstance(self.max_guidance_weight, bool): + raise TypeError("max_guidance_weight must be a number") + weight = float(self.max_guidance_weight) + if not math.isfinite(weight) or weight <= 0: + raise ValueError("max_guidance_weight must be a finite positive number") + object.__setattr__(self, "max_guidance_weight", weight) + if ( + isinstance(self.initial_inference_delay, bool) + or not isinstance(self.initial_inference_delay, int) + or self.initial_inference_delay < 0 + ): + raise ValueError("initial_inference_delay must be a non-negative integer") + if ( + isinstance(self.artificial_inference_delay_s, bool) + or not isinstance(self.artificial_inference_delay_s, (int, float)) + or not math.isfinite(self.artificial_inference_delay_s) + or self.artificial_inference_delay_s < 0 + ): + raise ValueError("artificial_inference_delay_s must be a finite non-negative number") + object.__setattr__(self, "artificial_inference_delay_s", float(self.artificial_inference_delay_s)) + + @classmethod + def from_options(cls, options: Mapping[str, Any] | None) -> AsyncRuntimeConfig: + """Parse a runtime ``async`` option block without silent downgrade.""" + + raw = dict(options or {}) + if not raw: + return cls() + if "async_mode" in raw: + raw.setdefault("mode", raw.pop("async_mode")) + if "async" in raw: + block = raw.pop("async") + if not isinstance(block, Mapping): + raise ValueError("runtime_options.async must be an object") + if raw: + raise ValueError("runtime_options contains keys outside async: " + ", ".join(sorted(raw))) + raw = dict(block) + allowed = { + "enabled", + "mode", + "rtc_enabled", + "overlap_steps", + "execution_horizon", + "prefix_attention_schedule", + "max_guidance_weight", + "hard_prefix", + "initial_inference_delay", + "artificial_inference_delay_s", + "send_model_space_prefix", + } + unknown = sorted(set(raw) - allowed) + if unknown: + raise ValueError("unsupported async runtime option(s): " + ", ".join(unknown)) + mode = raw.pop("mode", None) + if mode is not None: + if mode not in {"sync", "naive", "rtc"}: + raise ValueError("async mode must be one of 'sync', 'naive', or 'rtc'") + raw.setdefault("enabled", mode not in {"sync"}) + raw.setdefault("rtc_enabled", mode == "rtc") + if "enabled" not in raw: + raw["enabled"] = bool(raw.get("rtc_enabled", False) or mode == "naive") + enabled = bool(raw.get("enabled", False)) + if enabled and not raw.get("rtc_enabled", False) and mode != "naive": + raise ValueError("async runtime enabled but rtc_enabled is false and mode is not naive") + return cls(**raw) + + +@dataclass(frozen=True, slots=True) +class InferenceEvent: + """Bookkeeping for one inference launched by the async schedule.""" + + kind: str + step_id: int + index_before_inference: int + inference_delay: int + measured_latency_s: float + latency_delay: int + resolved_delay: int + actions_consumed: int + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class _InferencePayload: + """Immutable worker output adopted only by the control thread.""" + + result: PolicyResult + actions: tuple[RobotAction, ...] + vectors: np.ndarray + originals: np.ndarray + latency_s: float + index_before: int + delay_for_request: int + generation: int + + +@dataclass(frozen=True, slots=True) +class _InferenceJob: + future: Future[_InferencePayload] + generation: int + index_before: int + delay_for_request: int + + +class AsyncControlRuntime(ControlRuntime): + """Execute one naive or RTC asynchronous action-chunk schedule on one robot. + + Accepts the same robot/client/mapper/clock surface as :class:`ControlRuntime` + plus :class:`AsyncRuntimeConfig`. When ``config.enabled`` is false, + :meth:`step` delegates to the untouched synchronous implementation. + """ + + def __init__( + self, + robot: RobotAdapter, + client: InferenceClient, + *, + instruction: str, + mapper: BindingMapper, + chunk_steps: int, + control_hz: float = 5.0, + config: AsyncRuntimeConfig | None = None, + command_sink: CommandSink | None = None, + observation_source: Callable[[], RobotObservation] | None = None, + cancel_event: threading.Event | None = None, + monotonic: Callable[[], float] | None = None, + sleep: Callable[[float], None] | None = None, + latency_tracker: LatencyTracker | None = None, + action_encoder: Callable[[RobotAction], np.ndarray] | None = None, + chunk_extractor: Callable[[PolicyResult], np.ndarray] | None = None, + ) -> None: + kwargs: dict[str, Any] = {} + if monotonic is not None: + kwargs["monotonic"] = monotonic + if sleep is not None: + kwargs["sleep"] = sleep + super().__init__( + robot, + client, + instruction=instruction, + mapper=mapper, + chunk_steps=chunk_steps, + control_hz=control_hz, + inference_delay_s=(config or AsyncRuntimeConfig()).artificial_inference_delay_s, + command_sink=command_sink, + observation_source=observation_source, + cancel_event=cancel_event, + **kwargs, + ) + self.config = config or AsyncRuntimeConfig() + self.queue = ActionQueue(rtc_enabled=self.config.rtc_enabled) + self.latency = latency_tracker if latency_tracker is not None else LatencyTracker(monotonic=self.monotonic) + self.action_encoder = action_encoder or default_action_encoder + self.chunk_extractor = chunk_extractor or default_chunk_extractor + + self.events: list[InferenceEvent] = [] + self._chunk: tuple[RobotAction, ...] | None = None + self._chunk_vectors: np.ndarray | None = None + self._chunk_index = 0 + # Offset of ``_chunk[0]`` within the server's cached full model chunk. + # RTC adopts ``actions[delay:]``, so the local index alone no longer + # identifies the server-cache prefix on later replans. + self._chunk_model_start_index = 0 + self._pending: _InferencePayload | None = None + # An RTC response may arrive before the actions advertised in its + # ``inference_delay`` prefix have actually been sent to the robot. Keep + # it ready, but do not splice it in until that committed prefix has run. + self._rtc_pending: _InferencePayload | None = None + # One worker per runtime is sufficient: requests on a session must stay + # ordered, while this worker lets the control thread dispatch the old + # chunk during the network/model wait. + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="embodirun-infer") + self._job: _InferenceJob | None = None + self._generation = 0 + self._closed_async = False + self._join_timeout_s = 1.0 + self._last_result: PolicyResult | None = None + self._next_action_deadline_s: float | None = None + # Without a fixed playback period, wall time cannot be converted into + # steps. Keep the most recent number of actions actually dispatched + # while inference was in flight instead. + self._latest_unpaced_delay: int | None = None + + # -- public API ------------------------------------------------------ + def step( + self, + frames: Sequence[CameraFrame], + *, + reset: bool = False, + ) -> PolicyResult | None: + """Advance one control step of the async schedule. + + Executes exactly one action and, at the trigger index, launches the next + inference. Returns the :class:`PolicyResult` produced this step, or + ``None`` when no inference ran. + """ + + if not self.config.enabled: + return super().step(frames, reset=reset) + + if reset: + self.reset() + self._raise_if_cancelled() + frames = tuple(frames) + self._last_result = None + + # A completed worker is always adopted before a new action is sent. + # This is the stale-result boundary used by cancellation and takeover. + self._collect_job(block=False) + + self._ensure_chunk(frames) + assert self._chunk is not None + + if self._job is None and self._pending is None and self._rtc_pending is None and self._should_launch(): + self._start_job(frames) + self._collect_job(block=False) + if self._chunk is None or self._chunk_index >= len(self._chunk): + self._ensure_chunk(frames) + + action = self._chunk[self._chunk_index] + self._pace_action() + self._raise_if_cancelled() + self._execute(action) + self.queue.get() + self._chunk_index += 1 + return self._last_result + + def reset(self) -> None: + """Reset the session and all async schedule state.""" + + self._cancel_job(wait=True) + self._generation += 1 + super().reset() + self.queue.clear() + self._chunk = None + self._chunk_vectors = None + self._chunk_index = 0 + self._chunk_model_start_index = 0 + self._pending = None + self._rtc_pending = None + self._last_result = None + self._next_action_deadline_s = None + self._latest_unpaced_delay = None + self.events.clear() + self.latency.reset() + + def close(self) -> None: + if self._closed_async: + return + self._closed_async = True + self._generation += 1 + self._pending = None + self._rtc_pending = None + self._cancel_job() + # A running HTTP call may not be cancellable. Wait briefly for a + # cooperative client, then detach the executor so close never hangs a + # robot shutdown or an emergency-stop path. + self._shutdown_executor() + super().close() + + # -- schedule helpers ------------------------------------------------ + @property + def effective_overlap_steps(self) -> int: + """Configured overlap in control steps.""" + + return self.config.overlap_steps + + def _trigger_index(self) -> int: + if self._chunk is None: + return self.chunk_steps + length = len(self._chunk) + return max(0, length - self.effective_overlap_steps) + + def _should_launch(self) -> bool: + if self._chunk is None: + return False + return self._chunk_index == self._trigger_index() + + def _ensure_chunk(self, frames: Sequence[CameraFrame]) -> None: + if self._chunk is not None and self._chunk_index < len(self._chunk): + return + if self._job is not None: + self._collect_job(block=True) + if self._rtc_pending is not None: + if not self._adopt_rtc_pending(): + raise RuntimeError( + "RTC action chunk was exhausted before the advertised inference-delay prefix was consumed" + ) + return + if self._pending is not None: + self._adopt_pending() + return + self._bootstrap(frames) + + def _adopt(self, actions: Sequence[RobotAction], vectors: np.ndarray, *, skip: int) -> None: + self._chunk = tuple(actions)[skip:] + self._chunk_vectors = np.asarray(vectors, dtype=np.float64)[skip:] + self._chunk_index = 0 + self._chunk_model_start_index = skip + + def _bootstrap(self, frames: Sequence[CameraFrame]) -> None: + self._start_job(frames, delay_for_request=self.config.initial_inference_delay) + self._collect_job(block=True) + if self._pending is not None: + self._adopt_pending() + if self._chunk is None: + raise RuntimeError("bootstrap inference did not produce an action chunk") + + def _start_job(self, frames: Sequence[CameraFrame], *, delay_for_request: int | None = None) -> None: + if self._job is not None: + return + if self._closed_async: + raise RuntimeError("async runtime is closed") + index_before = self.queue.get_action_index() + delay = self._predict_delay() if delay_for_request is None else delay_for_request + observation = self.robot.observe() if self.observation_source is None else self.observation_source() + cache_start_index = self._chunk_model_start_index + self._chunk_index + request = self._build_request(observation, frames, delay, cache_start_index=cache_start_index) + generation = self._generation + future = self._executor.submit(self._infer_request, request, index_before, delay, generation) + self._job = _InferenceJob(future, generation, index_before, delay) + self.step_id += 1 + # Give the worker a scheduling opportunity before the control thread + # executes the first overlapped action. This matters for very fast + # local clients and deterministic test clocks: dispatching the request + # is part of the overlap boundary even when the request itself returns + # immediately. + time.sleep(0) + + def _collect_job(self, *, block: bool) -> None: + if self._rtc_pending is not None: + self._adopt_rtc_pending() + if self._rtc_pending is not None: + return + job = self._job + if job is None or not (block or job.future.done()): + return + try: + payload = job.future.result(timeout=self._join_timeout_s if block else 0) + except TimeoutError as error: + if block: + raise RuntimeError("inference worker did not finish before the action chunk was exhausted") from error + return + except CancelledError: + self._job = None + return + except BaseException: + self._job = None + raise + self._job = None + if payload.generation != self._generation or self._closed_async: + return + self.latency.record(payload.latency_s) + if self.config.rtc_enabled: + consumed = max(0, self.queue.get_action_index() - payload.index_before) + # Bootstrap has no old action chunk to honor. Every subsequent RTC + # request promises that ``delay_for_request`` old actions will run; + # a faster-than-predicted response must wait at this handover. + if self._chunk is not None and consumed < payload.delay_for_request: + self._rtc_pending = payload + else: + self._adopt_rtc_payload(payload) + else: + # Naive async adopts at the old-chunk handover. + self._pending = payload + + def _adopt_rtc_pending(self) -> bool: + """Adopt a ready RTC result after its promised old-action prefix ran.""" + + payload = self._rtc_pending + if payload is None: + return False + consumed = max(0, self.queue.get_action_index() - payload.index_before) + if consumed < payload.delay_for_request: + return False + self._rtc_pending = None + self._adopt_rtc_payload(payload) + return True + + def _adopt_rtc_payload(self, payload: _InferencePayload) -> None: + consumed = max(0, self.queue.get_action_index() - payload.index_before) + latency_delay = self._latency_delay(payload.latency_s, consumed) + # Queue alignment follows physical actions actually sent since capture. + # Wall-time rounding remains a reported/predicted latency signal, but it + # must not make us replay a prefix action after the robot already ran it. + resolved = self.queue.merge(payload.originals, payload.vectors, consumed, payload.index_before) + self._adopt(payload.actions, payload.vectors, skip=resolved) + self._last_result = payload.result + self._record_event( + kind="inference", + index_before=payload.index_before, + inference_delay=payload.delay_for_request, + latency=payload.latency_s, + latency_delay=latency_delay, + resolved=resolved, + consumed=consumed, + result=payload.result, + ) + + def _adopt_pending(self) -> None: + assert self._pending is not None + payload = self._pending + consumed = max(0, self.queue.get_action_index() - payload.index_before) + latency_delay = self._latency_delay(payload.latency_s, consumed) + resolved = self.queue.merge(payload.originals, payload.vectors, latency_delay, payload.index_before) + self._pending = None + self._adopt(payload.actions, payload.vectors, skip=0) + self._last_result = payload.result + self._record_event( + kind="inference", + index_before=payload.index_before, + inference_delay=payload.delay_for_request, + latency=payload.latency_s, + latency_delay=latency_delay, + resolved=resolved, + consumed=consumed, + result=payload.result, + ) + + def _predict_delay(self) -> int: + if self.action_period_s == 0: + if not self._latest_unpaced_delay: + return self.config.initial_inference_delay + return self._latest_unpaced_delay + latest = self.latency.latest_delay(self.action_period_s) + if latest == 0: + return self.config.initial_inference_delay + return latest + + def _pace_action(self) -> None: + """Dispatch asynchronous actions at the configured control rate.""" + + if self.action_period_s == 0: + return + + now = self.monotonic() + if self._next_action_deadline_s is None: + self._next_action_deadline_s = now + return + deadline = self._next_action_deadline_s + self.action_period_s + remaining = deadline - now + if remaining > 0: + if self.cancel_event is None: + self.sleep(remaining) + else: + self.cancel_event.wait(remaining) + self._raise_if_cancelled() + now = self.monotonic() + self._next_action_deadline_s = max(deadline, now) + + def _latency_delay(self, latency_s: float, consumed: int) -> int: + """Express latency in steps using the active playback semantics.""" + + if self.action_period_s == 0: + self._latest_unpaced_delay = consumed + return consumed + return self.latency.delay_in_steps(latency_s, self.action_period_s) + + # -- inference ------------------------------------------------------- + def _infer_request( + self, + request: PolicyObservation, + index_before: int, + delay_for_request: int, + generation: int, + ) -> _InferencePayload: + """Run one network request in the worker without mutating schedule state.""" + + started = self.monotonic() + result = self.client.step(request) + self._wait_for_inference_delay() + latency = max(0.0, self.monotonic() - started) + + actions = tuple(self.mapper.map_result(result)) + if not actions: + raise RuntimeError("binding returned an empty action chunk") + if any(not isinstance(action, RobotAction) for action in actions): + raise TypeError("binding action chunk must contain RobotAction values") + if len(actions) < self.chunk_steps: + raise RuntimeError( + f"binding returned {len(actions)} action(s), fewer than requested chunk_steps={self.chunk_steps}" + ) + actions = actions[: self.chunk_steps] + vectors = np.asarray([self.action_encoder(action) for action in actions], dtype=np.float64) + originals = np.asarray(self.chunk_extractor(result), dtype=np.float64) + return _InferencePayload( + result=result, + actions=actions, + vectors=vectors, + originals=originals, + latency_s=latency, + index_before=index_before, + delay_for_request=delay_for_request, + generation=generation, + ) + + def _cancel_job(self, *, wait: bool = False) -> None: + job = self._job + self._job = None + if job is not None: + job.future.cancel() + if wait: + with suppress(TimeoutError, CancelledError, BaseException): + job.future.result(timeout=self._join_timeout_s) + + def _shutdown_executor(self) -> None: + # ``Future.cancel`` cannot stop an in-flight socket read. Bounded join + # keeps emergency-stop and manual takeover responsive; the worker has no + # access to mutable schedule state, so detaching it is safe. + job = self._job + if job is not None: + with suppress(TimeoutError, CancelledError, BaseException): + job.future.result(timeout=self._join_timeout_s) + self._job = None + try: + self._executor.shutdown(wait=False, cancel_futures=True) + except TypeError: # pragma: no cover - Python 3.8 compatibility + self._executor.shutdown(wait=False) + + def _build_request( + self, + observation: RobotObservation, + frames: Sequence[CameraFrame], + delay_for_request: int, + *, + cache_start_index: int, + ) -> PolicyObservation: + request = self.mapper.map_observation( + observation, + session_id=self.session.session_id, + request_id=f"step-{self.step_id}-{uuid.uuid4().hex}", + step_id=self.step_id, + instruction=self.instruction, + frames=tuple(frames), + ) + observation_metadata = getattr(observation, "metadata", {}) + observation_id = ( + observation_metadata.get("observation_id") if isinstance(observation_metadata, Mapping) else None + ) + metadata: dict[str, Any] = {} + request_metadata = getattr(request, "metadata", {}) + if isinstance(request_metadata, Mapping): + metadata.update(request_metadata) + if isinstance(observation_id, str) and observation_id: + metadata.setdefault("observation_id", observation_id) + metadata.setdefault("snapshot_id", observation_id) + + metadata = build_async_metadata( + rtc=self._rtc_request(delay_for_request, cache_start_index=cache_start_index), + base=metadata, + ) + if hasattr(request, "__dataclass_fields__"): + return replace(request, metadata=metadata) + try: + request.metadata = metadata + except AttributeError as error: # pragma: no cover - defensive parity with ControlRuntime + raise TypeError("mapper request must support metadata propagation") from error + return request + + def _rtc_request( + self, + delay_for_request: int, + *, + cache_start_index: int, + ) -> RTCAsyncRequest | None: + if not self.config.rtc_enabled: + return None + left_over = self.queue.get_left_over() if self.config.send_model_space_prefix else None + return RTCAsyncRequest( + prev_chunk_left_over=None if left_over is None else left_over.tolist(), + inference_delay=delay_for_request, + execution_horizon=( + self.chunk_steps if self.config.execution_horizon is None else self.config.execution_horizon + ), + prefix_attention_schedule=self.config.prefix_attention_schedule, + max_guidance_weight=self.config.max_guidance_weight, + hard_prefix=self.config.hard_prefix, + model_space=self.config.send_model_space_prefix, + # Explicit prefixes are already sliced locally. Session-cache mode + # instead needs the executed offset in the server's full chunk; it + # cannot be reconstructed from network/model latency. + cache_start_index=None if left_over is not None else cache_start_index, + ) + + def _record_event( + self, + *, + kind: str, + index_before: int, + inference_delay: int, + latency: float, + latency_delay: int, + resolved: int, + consumed: int, + result: PolicyResult | None, + ) -> None: + metadata = getattr(result, "request_id", None) + self.events.append( + InferenceEvent( + kind=kind, + step_id=self.step_id, + index_before_inference=index_before, + inference_delay=inference_delay, + measured_latency_s=latency, + latency_delay=latency_delay, + resolved_delay=resolved, + actions_consumed=consumed, + metadata={"request_id": metadata} if metadata is not None else {}, + ) + ) diff --git a/src/embodirun/application/contracts.py b/src/embodirun/application/contracts.py index cd3cf63..dc7e823 100644 --- a/src/embodirun/application/contracts.py +++ b/src/embodirun/application/contracts.py @@ -455,10 +455,20 @@ def __post_init__(self) -> None: value = getattr(self, name) if isinstance(value, bool) or not isinstance(value, int) or value <= 0: raise ControlContractError(f"{name} must be a positive integer") - for name in ("control_hz", "inference_timeout_s"): - value = getattr(self, name) - if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value <= 0: - raise ControlContractError(f"{name} must be a finite positive number") + if ( + isinstance(self.control_hz, bool) + or not isinstance(self.control_hz, (int, float)) + or not math.isfinite(self.control_hz) + or self.control_hz < 0 + ): + raise ControlContractError("control_hz must be a finite non-negative number") + if ( + isinstance(self.inference_timeout_s, bool) + or not isinstance(self.inference_timeout_s, (int, float)) + or not math.isfinite(self.inference_timeout_s) + or self.inference_timeout_s <= 0 + ): + raise ControlContractError("inference_timeout_s must be a finite positive number") def to_payload(self) -> dict[str, Any]: """Return the transport-neutral JSON object sent to Control.""" diff --git a/src/embodirun/application/control_service.py b/src/embodirun/application/control_service.py index 6c112cd..a6eb683 100644 --- a/src/embodirun/application/control_service.py +++ b/src/embodirun/application/control_service.py @@ -293,9 +293,16 @@ def _execute_locked( f"chunk_steps {request.chunk_steps} exceeds binding " f"{binding.kind!r} maximum {binding.maximum_chunk_steps}" ) + async_runtime_type: type[Any] | None = None + async_config: Any | None = None if profile_config.runtime_options: - names = ", ".join(sorted(profile_config.runtime_options)) - raise ControlServiceError(f"unsupported runtime options: {names}") + from .async_inference import AsyncControlRuntime, AsyncRuntimeConfig + + async_runtime_type = AsyncControlRuntime + try: + async_config = AsyncRuntimeConfig.from_options(profile_config.runtime_options) + except (TypeError, ValueError) as error: + raise ControlServiceError(str(error)) from error client = self._inference_client(profile_config, request.inference_timeout_s) cameras: CameraSource | None = None controller: Any | None = None @@ -306,6 +313,7 @@ def _execute_locked( health = client.health() if health.get("status") != "ok": raise ControlServiceError(f"inference service at {self.config.inference_endpoint} is not healthy") + _check_async_capability(client, profile_config, async_config) cameras, _ = self._acquire_cameras(profile_config) robot, arbiter = self._ensure_robot_arbiter( robot_definition_value, @@ -323,32 +331,77 @@ def observation_source() -> RobotObservation: except ObservationViewError as error: raise ControlServiceError(str(error)) from error - controller = self.runtime_factory( - robot, - client, - instruction=request.prompt, - mapper=binding.mapper_factory(), - chunk_steps=request.chunk_steps, - control_hz=request.control_hz, - command_sink=ArbiterCommandSink(arbiter, task_cancel), - observation_source=observation_source, - cancel_event=task_cancel, - monotonic=self.monotonic, - sleep=self.sleep, + runtime_factory = self.runtime_factory + if async_config is not None and async_config.enabled: + assert async_runtime_type is not None + if runtime_factory is ControlRuntime: + runtime_factory = async_runtime_type + elif runtime_factory is not async_runtime_type: + raise ControlServiceError( + "async runtime_options require a runtime_factory that accepts AsyncRuntimeConfig" + ) + runtime_kwargs: dict[str, Any] = { + "robot": robot, + "client": client, + "instruction": request.prompt, + "mapper": binding.mapper_factory(), + "chunk_steps": request.chunk_steps, + "control_hz": request.control_hz, + "inference_delay_s": (0.0 if async_config is None else async_config.artificial_inference_delay_s), + "command_sink": ArbiterCommandSink(arbiter, task_cancel), + "observation_source": observation_source, + "cancel_event": task_cancel, + "monotonic": self.monotonic, + "sleep": self.sleep, + } + if async_runtime_type is not None and runtime_factory is async_runtime_type: + runtime_kwargs.pop("inference_delay_s") + runtime_kwargs["config"] = async_config + controller = runtime_factory( + runtime_kwargs.pop("robot"), + runtime_kwargs.pop("client"), + **runtime_kwargs, + ) + # A synchronous runtime advances one complete action chunk per + # call. The asynchronous runtime advances one physical action so + # inference can overlap the remaining chunk. Preserve the public + # ``max_steps`` meaning (number of chunks) by expanding only the + # async loop to its action count. + runtime_steps = ( + request.max_steps * request.chunk_steps + if async_config is not None and async_config.enabled + else request.max_steps ) - for _ in range(request.max_steps): + for runtime_step in range(runtime_steps): requested = _runtime_camera_requests(profile_config) - current_snapshot = self._shared_observations.snapshot( - require_state=True, - required_source_ids=[source_id for source_id, _ in requested], - force_publish=True, - ) - frames = self._shared_observations.frames_for( - current_snapshot, - requested, - ) + for capture_attempt in range(3): + current_snapshot = self._shared_observations.snapshot( + require_state=True, + required_source_ids=[source_id for source_id, _ in requested], + # Async execution advances once per physical action. The + # camera workers already publish at their configured FPS; + # forcing a fresh USB capture on every action both stalls + # the control loop and can turn one transient V4L2 timeout + # into a failed task. Synchronous execution still forces + # one fresh observation for each policy chunk. + force_publish=async_config is None or not async_config.enabled, + ) + try: + frames = self._shared_observations.frames_for( + current_snapshot, + requested, + ) + break + except SharedSensorError: + if capture_attempt == 2: + raise + # USB cameras can miss one publish while the other + # source succeeds. Retry the complete snapshot rather + # than terminating a physical task for one absent frame. + self.sleep(0.01) controller.step(frames) - completed += 1 + if async_config is None or not async_config.enabled or (runtime_step + 1) % request.chunk_steps == 0: + completed += 1 finally: with self._control_lock: if self._current_task_cancel is task_cancel: @@ -851,6 +904,12 @@ def _make_robot_arbiter( RobotAdapterCommandPort( robot, action_resolver=resolve_teleop_action, + # SO-101 serial recovery may briefly fall back from one group + # transaction to per-motor reads. Keep the caller wait budget + # above that bounded recovery path so the bus is not + # quarantined while the driver is still completing it. + io_wait_timeout_s=3.0 if profile_config.robot_kind == "lerobot.so101" else 1.0, + observation_wait_timeout_s=(3.0 if profile_config.robot_kind == "lerobot.so101" else 1.0), emergency_stop_policy=( EmergencyStopPolicy.PREEMPTIVE if profile_config.robot_kind == "franka.fr3" @@ -1660,6 +1719,45 @@ def _runtime_camera_requests( return tuple((_sensor_identity(config, item).key, item.name) for item in config.inputs) +def _check_async_capability( + client: Any, + profile_config: ControlServiceConfig, + async_config: Any | None, +) -> None: + """Reject algorithmic async on SGLang unless its endpoint advertises it.""" + + if async_config is None or not async_config.enabled or profile_config.inference_backend != "sglang": + return + if not async_config.rtc_enabled: + return # ``mode=naive`` is the provider-neutral overlap baseline. + capabilities = getattr(client, "capabilities", None) + if not callable(capabilities): + raise ControlServiceError( + "SGLang async RTC requires endpoint capability metadata; use async mode=naive for the baseline" + ) + try: + payload = capabilities() + except Exception as error: + raise ControlServiceError(f"SGLang async capability probe failed: {error}") from error + if not isinstance(payload, Mapping): + raise ControlServiceError("SGLang capabilities response must be an object") + advertised = payload.get("async_inference") + if not isinstance(advertised, Mapping): + nested = payload.get("capabilities") + advertised = nested.get("async_inference") if isinstance(nested, Mapping) else None + if not isinstance(advertised, Mapping): + raise ControlServiceError( + "SGLang endpoint does not advertise async_inference.rtc; use async mode=naive for the baseline" + ) + missing = ["rtc"] if advertised.get("rtc") is not True else [] + if missing: + raise ControlServiceError( + "SGLang endpoint does not support requested async algorithm(s): " + + ", ".join(missing) + + "; use async mode=naive for the baseline" + ) + + def _observation_payload(value: Any) -> Any: if isinstance(value, Mapping): return dict(value) diff --git a/src/embodirun/application/model_loop.py b/src/embodirun/application/model_loop.py index dfdd5e0..ca65f72 100644 --- a/src/embodirun/application/model_loop.py +++ b/src/embodirun/application/model_loop.py @@ -40,6 +40,7 @@ def __init__( mapper: BindingMapper, chunk_steps: int, control_hz: float = 5.0, + inference_delay_s: float = 0.0, command_sink: CommandSink | None = None, observation_source: Callable[[], RobotObservation] | None = None, cancel_event: threading.Event | None = None, @@ -54,9 +55,16 @@ def __init__( isinstance(control_hz, bool) or not isinstance(control_hz, (int, float)) or not math.isfinite(control_hz) - or control_hz <= 0 + or control_hz < 0 ): - raise ValueError("control_hz must be a finite positive number") + raise ValueError("control_hz must be a finite non-negative number") + if ( + isinstance(inference_delay_s, bool) + or not isinstance(inference_delay_s, (int, float)) + or not math.isfinite(inference_delay_s) + or inference_delay_s < 0 + ): + raise ValueError("inference_delay_s must be a finite non-negative number") self.robot = robot self.client = client self.instruction = instruction @@ -77,7 +85,10 @@ def __init__( if command_sink is not None and command_sink.robot_id != robot.robot_id: raise ValueError("command_sink robot_id must match robot") self.cancel_event = cancel_event - self.action_period_s = 1.0 / control_hz + # Zero explicitly disables software pacing. The adapter/transport then + # determines how quickly actions can be dispatched. + self.action_period_s = 0.0 if control_hz == 0 else 1.0 / control_hz + self.inference_delay_s = float(inference_delay_s) self.monotonic = monotonic self.sleep = sleep self.session = client.open_session( @@ -123,6 +134,7 @@ def step( except AttributeError as error: raise TypeError("mapper request must support metadata propagation") from error result = self.client.step(request) + self._wait_for_inference_delay() self._raise_if_cancelled() actions = tuple(self.mapper.map_result(result)) if not actions: @@ -163,6 +175,17 @@ def step( self.step_id += 1 return result + def _wait_for_inference_delay(self) -> None: + """Delay result availability to emulate extra end-to-end latency.""" + + if self.inference_delay_s <= 0: + return + if self.cancel_event is None: + self.sleep(self.inference_delay_s) + else: + self.cancel_event.wait(self.inference_delay_s) + self._raise_if_cancelled() + def reset(self) -> None: self._stop() self.session = self.client.reset( diff --git a/src/embodirun/application/simulation/__init__.py b/src/embodirun/application/simulation/__init__.py index f7e3666..5018022 100644 --- a/src/embodirun/application/simulation/__init__.py +++ b/src/embodirun/application/simulation/__init__.py @@ -7,7 +7,30 @@ SimulationServiceConfig, ) from .runtime import EpisodeOutcome, SimulationRuntime -from .service import SimulationService, SimulationServiceError, create_inference_client + +_LAZY_ASYNC = {"AsyncEpisodeOutcome", "AsyncSimulationConfig", "AsyncSimulationRuntime"} +_LAZY_SERVICE = {"SimulationService", "SimulationServiceError", "create_inference_client"} + + +def __getattr__(name: str): + if name in _LAZY_ASYNC: + from .async_runtime import AsyncEpisodeOutcome, AsyncSimulationConfig, AsyncSimulationRuntime + + return { + "AsyncEpisodeOutcome": AsyncEpisodeOutcome, + "AsyncSimulationConfig": AsyncSimulationConfig, + "AsyncSimulationRuntime": AsyncSimulationRuntime, + }[name] + if name in _LAZY_SERVICE: + from .service import SimulationService, SimulationServiceError, create_inference_client + + return { + "SimulationService": SimulationService, + "SimulationServiceError": SimulationServiceError, + "create_inference_client": create_inference_client, + }[name] + raise AttributeError(name) + __all__ = [ "EpisodeOutcome", @@ -15,6 +38,9 @@ "EpisodeResult", "SimulationContractError", "SimulationRuntime", + "AsyncEpisodeOutcome", + "AsyncSimulationConfig", + "AsyncSimulationRuntime", "SimulationServiceConfig", "SimulationService", "SimulationServiceError", diff --git a/src/embodirun/application/simulation/async_runtime.py b/src/embodirun/application/simulation/async_runtime.py new file mode 100644 index 0000000..f85f0fc --- /dev/null +++ b/src/embodirun/application/simulation/async_runtime.py @@ -0,0 +1,728 @@ +"""Asynchronous chunk execution for one simulator episode. + +:class:`embodirun.application.simulation.runtime.SimulationRuntime` runs a +simulator episode *synchronously*: it blocks on inference, then executes the whole +chunk, then blocks again. A real robot does not have that luxury — inference takes +wall-clock time during which the body keeps moving — so RTC exists to keep +the policy aligned with a robot whose state has advanced since the observation it +was conditioned on. + +A simulator has no wall clock either: its ``step`` returns immediately. This +runtime therefore uses an explicit, **validated configuration field**, +:attr:`AsyncSimulationConfig.inference_delay_steps` (``Δ``), for the number of +simulator steps in the overlap window while a worker computes the next chunk. + + ``Δ`` is a *modeled* delay, not a measured one. It is the standard way the RTC + paper expresses the execution/inference overlap in simulation + (arXiv:2506.07339). The runtime additionally measures the + real wall-clock time each blocking client call takes and reports it as a + diagnostic (:class:`AsyncInferenceRecord`), but it deliberately never + substitutes that measurement for ``Δ`` — the loop is driven by ``Δ`` alone. + +The loop reproduces the physical ordering of an asynchronous deployment: + +1. Execute the current chunk step by step. +2. At the trigger point ``chunk_index == chunk_steps - overlap_steps`` capture the + observation and record the current chunk index. The captured observation is + what the server will see, so it is intentionally stale by ``Δ`` steps. +3. Issue the request on the inference worker before advancing those ``Δ`` steps. +4. Keep executing the current chunk for ``Δ`` more simulator steps (the inference + window). Those are the actions the server must reason about. +5. Join the worker and adopt the returned chunk with the mode's handover rule. + +Modes +----- + +``sync`` + No overlap and no modeled delay: the next request happens only after the + current chunk is fully consumed. No ``async`` block is attached, so the + request is byte-identical to the synchronous runtime's. +``naive`` + The unstable asynchronous baseline: the request is issued at the trigger from + the stale observation with *no* conditioning block, and the returned chunk is + adopted at handover. +``rtc`` + Real-Time Chunking (Black et al., arXiv:2506.07339). The request carries + ``inference_delay = Δ`` plus the configured execution horizon and attention + schedule. The committed prefix is the model-space leftover of the previously + issued chunk; this client does not hold model space, so it sends + ``prev_chunk_left_over = None`` plus ``cache_start_index``, the exact offset + already executed from the server's cached full chunk. This offset is the + replan stride and is independent of ``Δ``. The new chunk is adopted + immediately, skipping its first ``Δ`` actions (they correspond to time already + executed during the window). +Bookkeeping mirrors :class:`~embodirun.application.async_inference.runtime.AsyncControlRuntime`: +an :class:`~embodirun.application.async_inference.rtc.ActionQueue` is kept in +lockstep with the executing chunk (``get`` per executed action, ``merge`` per +adopted chunk) and a :class:`~embodirun.application.async_inference.rtc.LatencyTracker` +collects the measured call latencies. + +The module is pure Python plus numpy, imports no simulator/inference package, and +keeps simulator mutation on one thread while inference runs on a worker. +""" + +from __future__ import annotations + +import math +import time +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace +from typing import Any + +import numpy as np + +from embodirun.bindings import BindingMapper +from embodirun.model_services import InferenceClient, PolicyObservation, PolicyResult +from embodirun.robots import RobotAction +from embodirun.simulators import SimulatorAdapter, SimulatorObservation + +from ..async_inference.contracts import RTCAsyncRequest, build_async_metadata +from ..async_inference.rtc import PREFIX_ATTENTION_SCHEDULES, ActionQueue, LatencyTracker +from ..async_inference.runtime import default_action_encoder, default_chunk_extractor +from .runtime import EpisodeOutcome + +__all__ = [ + "ASYNC_SIMULATION_MODES", + "AsyncEpisodeOutcome", + "AsyncInferenceRecord", + "AsyncSimulationConfig", + "AsyncSimulationRuntime", + "latency_steps", +] + +ASYNC_SIMULATION_MODES = ("sync", "naive", "rtc") +"""Every execution mode the runtime understands.""" + + +def latency_steps(latency_s: float, *, control_hz: float) -> int: + """Convert a measured latency in seconds to whole control steps. + + This is the reported conversion only; it is never used as ``Δ``. It mirrors + :meth:`embodirun.application.async_inference.rtc.LatencyTracker.delay_in_steps` + (including the epsilon that keeps an exact ``n * period`` product from flooring + to ``n - 1``). + """ + + if isinstance(control_hz, bool) or not isinstance(control_hz, (int, float)): + raise TypeError("control_hz must be a number") + hertz = float(control_hz) + if not math.isfinite(hertz) or hertz <= 0: + raise ValueError("control_hz must be a finite positive number") + return LatencyTracker().delay_in_steps(latency_s, 1.0 / hertz) + + +@dataclass(frozen=True, slots=True) +class AsyncSimulationConfig: + """Static configuration for one asynchronous simulator episode. + + Args: + mode: One of :data:`ASYNC_SIMULATION_MODES`. + chunk_steps: Actions per returned chunk (``H``). + overlap_steps: Steps before chunk end at which the next inference is + triggered, i.e. the number of old-chunk steps remaining at the trigger. + Must be ``0`` for ``sync``. + inference_delay_steps: ``Δ``, the **modeled** number of simulator steps that + elapse while the next chunk is computed. It is a simulation parameter, + not a measured wall-clock latency, and it must not exceed + ``overlap_steps`` (the chunk remaining at the trigger). + control_hz: Control rate used to express the *measured* latency in steps. + execution_horizon: RTC execution horizon; defaults to ``chunk_steps``. + prefix_attention_schedule: RTC weight schedule. + max_guidance_weight: RTC guidance clamp. + hard_prefix: Use the RTC fixed-prefix clamp. + """ + + mode: str = "sync" + chunk_steps: int = 50 + overlap_steps: int = 0 + inference_delay_steps: int = 0 + control_hz: float = 20.0 + execution_horizon: int | None = None + prefix_attention_schedule: str = "linear" + max_guidance_weight: float = 10.0 + hard_prefix: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.mode, str) or self.mode not in ASYNC_SIMULATION_MODES: + raise ValueError(f"mode must be one of {', '.join(ASYNC_SIMULATION_MODES)}, got {self.mode!r}") + + if isinstance(self.chunk_steps, bool) or not isinstance(self.chunk_steps, int) or self.chunk_steps <= 0: + raise ValueError("chunk_steps must be a positive integer") + if isinstance(self.overlap_steps, bool) or not isinstance(self.overlap_steps, int) or self.overlap_steps < 0: + raise ValueError("overlap_steps must be a non-negative integer") + if self.overlap_steps > self.chunk_steps: + raise ValueError(f"overlap_steps {self.overlap_steps} exceeds chunk_steps {self.chunk_steps}") + if ( + isinstance(self.inference_delay_steps, bool) + or not isinstance(self.inference_delay_steps, int) + or self.inference_delay_steps < 0 + ): + raise ValueError("inference_delay_steps must be a non-negative integer") + + if self.mode == "sync": + if self.overlap_steps != 0 or self.inference_delay_steps != 0: + raise ValueError( + "sync mode executes the whole chunk before the next request; " + "overlap_steps and inference_delay_steps must both be 0" + ) + elif self.inference_delay_steps > self.overlap_steps: + raise ValueError( + f"inference_delay_steps {self.inference_delay_steps} exceeds the {self.overlap_steps} step(s) " + "of the current chunk still remaining at the trigger" + ) + if self.mode == "rtc" and 2 * self.inference_delay_steps > self.chunk_steps: + # RTC adopts each chunk with its first Δ actions skipped, so every chunk + # after the bootstrap is only ``chunk_steps - Δ`` long. The window must + # still fit before that shortened chunk runs out. + raise ValueError( + f"rtc inference_delay_steps {self.inference_delay_steps} leaves only " + f"{self.chunk_steps - self.inference_delay_steps} action(s) in each adopted chunk; " + "the inference window must still fit, so 2 * inference_delay_steps must not exceed chunk_steps" + ) + + if ( + isinstance(self.control_hz, bool) + or not isinstance(self.control_hz, (int, float)) + or not math.isfinite(float(self.control_hz)) + or float(self.control_hz) <= 0 + ): + raise ValueError("control_hz must be a finite positive number") + object.__setattr__(self, "control_hz", float(self.control_hz)) + + if self.execution_horizon is not None and ( + isinstance(self.execution_horizon, bool) + or not isinstance(self.execution_horizon, int) + or self.execution_horizon < 0 + ): + raise ValueError("execution_horizon must be a non-negative integer or None") + if self.prefix_attention_schedule not in PREFIX_ATTENTION_SCHEDULES: + raise ValueError( + "prefix_attention_schedule must be one of " + f"{', '.join(PREFIX_ATTENTION_SCHEDULES)}, got {self.prefix_attention_schedule!r}" + ) + if ( + isinstance(self.max_guidance_weight, bool) + or not isinstance(self.max_guidance_weight, (int, float)) + or not math.isfinite(float(self.max_guidance_weight)) + or float(self.max_guidance_weight) <= 0 + ): + raise ValueError("max_guidance_weight must be a finite positive number") + object.__setattr__(self, "max_guidance_weight", float(self.max_guidance_weight)) + if not isinstance(self.hard_prefix, bool): + raise TypeError("hard_prefix must be a boolean") + + @classmethod + def from_options( + cls, + options: Mapping[str, Any] | None, + *, + chunk_steps: int, + control_hz: float = 20.0, + ) -> AsyncSimulationConfig: + """Build simulation async settings from static service options.""" + + raw = dict(options or {}) + if "async" in raw: + block = raw.pop("async") + if not isinstance(block, Mapping): + raise ValueError("runtime_options.async must be an object") + if raw: + raise ValueError("runtime_options contains keys outside async: " + ", ".join(sorted(raw))) + raw = dict(block) + if "async_mode" in raw: + raw.setdefault("mode", raw.pop("async_mode")) + allowed = { + "mode", + "overlap_steps", + "inference_delay_steps", + "control_hz", + "execution_horizon", + "prefix_attention_schedule", + "max_guidance_weight", + "hard_prefix", + } + unknown = sorted(set(raw) - allowed) + if unknown: + raise ValueError("unsupported async simulation option(s): " + ", ".join(unknown)) + raw.setdefault("chunk_steps", chunk_steps) + raw.setdefault("control_hz", control_hz) + return cls(**raw) + + def trigger_index_for(self, chunk_length: int) -> int | None: + """Chunk index that triggers the next inference, or ``None`` when there is none. + + The trigger is ``chunk_length - overlap_steps`` (clamped at zero), exactly + for the full chunk. Using the *actual* chunk length matters for RTC, + whose adopted chunks are ``chunk_steps - Δ`` long, so that ``overlap_steps`` + always means "actions still remaining when the next inference starts". + + ``sync`` and every zero-overlap configuration only request a new chunk once + the current one is exhausted, which is what makes ``Δ = 0`` degenerate to + the synchronous behaviour. + """ + + if self.mode == "sync" or self.overlap_steps == 0: + return None + return max(0, chunk_length - self.overlap_steps) + + @property + def resolved_execution_horizon(self) -> int: + """RTC execution horizon, defaulting to ``chunk_steps``.""" + + return self.chunk_steps if self.execution_horizon is None else self.execution_horizon + + +@dataclass(frozen=True, slots=True) +class AsyncInferenceRecord: + """Diagnostics for one inference request issued by the runtime. + + ``inference_delay_steps`` is the modeled ``Δ`` actually attached to the request, + while ``measured_latency_s``/``latency_steps`` observe how long the blocking + client call really took. They are reported side by side and the measured value + never drives the loop. + """ + + request_index: int + mode: str + trigger: bool + trigger_chunk_index: int | None + captured_observation_id: str | None + environment_steps_at_capture: int + environment_steps_at_send: int + inference_delay_steps: int + measured_latency_s: float + latency_steps: int + pending_actions: tuple[tuple[float, ...], ...] + async_block: Mapping[str, Any] | None + + +@dataclass(frozen=True, slots=True) +class AsyncEpisodeOutcome: + """An :class:`EpisodeOutcome` plus the asynchronous diagnostics. + + The five episode fields are duplicated at the top level so the value is a + drop-in replacement for :class:`EpisodeOutcome`; :meth:`to_episode_outcome` + returns the plain summary. + """ + + policy_steps: int + environment_steps: int + total_reward: float + terminated: bool + truncated: bool + requests: int + inference_delay_steps: tuple[int, ...] + measured_latency_s: tuple[float, ...] + latency_steps: tuple[int, ...] + records: tuple[AsyncInferenceRecord, ...] + + def to_episode_outcome(self) -> EpisodeOutcome: + """Return the plain synchronous-compatible summary.""" + + return EpisodeOutcome( + policy_steps=self.policy_steps, + environment_steps=self.environment_steps, + total_reward=self.total_reward, + terminated=self.terminated, + truncated=self.truncated, + ) + + +@dataclass(frozen=True, slots=True) +class _Inference: + """One mapped, validated inference result before it is adopted.""" + + actions: tuple[RobotAction, ...] + vectors: np.ndarray + originals: np.ndarray + record: AsyncInferenceRecord + + +class AsyncSimulationRuntime: + """Run one simulator episode under asynchronous chunk execution. + + The runtime is a sibling of + :class:`~embodirun.application.simulation.runtime.SimulationRuntime`, not a + replacement, and never touches the wall clock to advance the simulator. + """ + + def __init__( + self, + simulator: SimulatorAdapter, + client: InferenceClient, + *, + mapper: BindingMapper, + config: AsyncSimulationConfig | None = None, + monotonic: Callable[[], float] | None = None, + latency_tracker: LatencyTracker | None = None, + action_encoder: Callable[[RobotAction], np.ndarray] | None = None, + chunk_extractor: Callable[[PolicyResult], np.ndarray] | None = None, + ) -> None: + self.simulator = simulator + self.client = client + self.mapper = mapper + self.config = config or AsyncSimulationConfig() + self._monotonic = monotonic or time.monotonic + self.latency = latency_tracker if latency_tracker is not None else LatencyTracker(monotonic=self._monotonic) + self.action_encoder = action_encoder or default_action_encoder + self.chunk_extractor = chunk_extractor or default_chunk_extractor + self.session = None + + self.queue = ActionQueue(rtc_enabled=self.config.mode == "rtc") + self.records: list[AsyncInferenceRecord] = [] + self._prompt = "" + self._control_period_s = 1.0 / self.config.control_hz + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="embodirun-sim-infer") + self._closed = False + + @property + def control_period_s(self) -> float: + """Seconds per control step implied by ``control_hz``.""" + + return self._control_period_s + + def run( + self, + *, + instruction: str | None, + max_policy_steps: int, + task: str | None = None, + seed: int | None = None, + ) -> AsyncEpisodeOutcome: + """Run one bounded episode and return its outcome plus async diagnostics.""" + + if instruction is not None and not instruction.strip(): + raise ValueError("instruction must be non-empty when provided") + if isinstance(max_policy_steps, bool) or not isinstance(max_policy_steps, int) or max_policy_steps <= 0: + raise ValueError("max_policy_steps must be a positive integer") + + config = self.config + current = self.simulator.reset(task=task, seed=seed) + self._prompt = instruction.strip() if instruction is not None else _instruction(current) + self.session = self.client.open_session( + robot_id=self.simulator.simulator_id, + action_space=self.mapper.policy_action_space, + metadata={"source": "simulator", "async_mode": config.mode}, + ) + self.queue = ActionQueue(rtc_enabled=config.mode == "rtc") + self.records = [] + self.latency.reset() + + total_reward = 0.0 + environment_steps = 0 + policy_steps = 0 + terminated = False + truncated = False + + chunk: tuple[RobotAction, ...] = () + vectors = np.empty((0, 0), dtype=np.float64) + chunk_index = 0 + # Offset of ``chunk[0]`` in the server-cached full model chunk. RTC drops + # the first Δ actions at handover, so later trigger indices are relative + # to a shortened local chunk and must be translated back before asking the + # server for its model-space prefix. + model_chunk_start_index = 0 + pending_adopt: _Inference | None = None + + def execute(action: RobotAction) -> None: + nonlocal current, total_reward, environment_steps, terminated, truncated + transition = self.simulator.step(action) + current = transition.observation + total_reward += transition.reward + environment_steps += 1 + terminated = transition.terminated + truncated = transition.truncated + self.queue.get() + + # Bootstrap: the first chunk is computed from the fresh reset observation. + # There is no stale observation and no committed prefix yet, so it is not a + # trigger request and carries no conditioning payload. + bootstrap = self._infer( + current, + request_index=0, + delay=0, + pending=(), + trigger=False, + environment_steps_at_capture=0, + environment_steps_at_send=0, + cache_start_index=0, + ) + policy_steps += 1 + self.queue.merge(bootstrap.originals, bootstrap.vectors, 0, None) + chunk, vectors = bootstrap.actions, bootstrap.vectors + chunk_index = 0 + model_chunk_start_index = 0 + + while not (terminated or truncated): + if chunk_index >= len(chunk): + if pending_adopt is not None: + # Naive handover: the old chunk is now fully consumed, so the + # append merge collapses to the new chunk from its first action. + adopted = pending_adopt + pending_adopt = None + self.queue.merge(adopted.originals, adopted.vectors, 0, None) + chunk, vectors = adopted.actions, adopted.vectors + chunk_index = 0 + model_chunk_start_index = 0 + continue + if policy_steps >= max_policy_steps: + break + # Sync-style handover: no overlap remains to start inference early, + # so the next chunk is requested from the observation reached here. + fresh = self._infer( + current, + request_index=policy_steps, + delay=0, + pending=(), + trigger=False, + environment_steps_at_capture=environment_steps, + environment_steps_at_send=environment_steps, + cache_start_index=model_chunk_start_index + chunk_index, + ) + policy_steps += 1 + self.queue.merge(fresh.originals, fresh.vectors, 0, None) + chunk, vectors = fresh.actions, fresh.vectors + chunk_index = 0 + model_chunk_start_index = 0 + continue + + trigger_index = config.trigger_index_for(len(chunk)) + if trigger_index is not None and chunk_index == trigger_index: + if policy_steps >= max_policy_steps: + break + remaining = len(chunk) - chunk_index + if remaining < config.inference_delay_steps: + raise ValueError( + f"inference window Δ={config.inference_delay_steps} does not fit in the " + f"{remaining} remaining action(s) of the current chunk" + ) + captured = current + captured_index = chunk_index + captured_environment_steps = environment_steps + pending = tuple( + tuple(float(value) for value in vectors[index]) + for index in range(chunk_index, min(len(chunk), chunk_index + config.inference_delay_steps)) + ) + + # Keep the blocking model request in flight while the simulator + # advances the configured overlap window. Only the inference + # client runs on the worker; simulator state stays single-threaded. + future = self._executor.submit( + self._infer, + captured, + request_index=policy_steps, + delay=config.inference_delay_steps, + pending=pending, + trigger=True, + trigger_chunk_index=captured_index, + environment_steps_at_capture=captured_environment_steps, + # The worker is dispatched at the stale observation and + # completes after the modeled overlap window. Record the + # logical send boundary at that completion point even + # though the simulator stays on the control thread. + environment_steps_at_send=environment_steps + config.inference_delay_steps, + cache_start_index=model_chunk_start_index + captured_index, + ) + for _ in range(config.inference_delay_steps): + if terminated or truncated: + break + execute(chunk[chunk_index]) + chunk_index += 1 + if terminated or truncated: + future.cancel() + break + policy_steps += 1 + outcome = future.result() + + if config.mode == "rtc": + resolved = self.queue.merge( + outcome.originals, + outcome.vectors, + config.inference_delay_steps, + captured_index, + ) + chunk = outcome.actions[resolved:] + vectors = outcome.vectors[resolved:] + chunk_index = 0 + model_chunk_start_index = resolved + else: + pending_adopt = outcome + continue + + execute(chunk[chunk_index]) + chunk_index += 1 + + return self._outcome( + policy_steps=policy_steps, + environment_steps=environment_steps, + total_reward=total_reward, + terminated=terminated, + truncated=truncated, + ) + + def close(self) -> None: + """Close the policy session, mirroring :class:`SimulationRuntime`.""" + + if self._closed: + return + self._closed = True + self._executor.shutdown(wait=True, cancel_futures=True) + if self.session is not None: + self.client.close(self.session.session_id) + self.session = None + + # -- inference ------------------------------------------------------- + def _infer( + self, + captured: SimulatorObservation, + *, + request_index: int, + delay: int, + pending: tuple[tuple[float, ...], ...], + trigger: bool, + environment_steps_at_capture: int, + environment_steps_at_send: int, + trigger_chunk_index: int | None = None, + cache_start_index: int | None = None, + ) -> _Inference: + if self.session is None: + raise RuntimeError("async simulation runtime must open a session before inference") + request = self.mapper.map_observation( + captured.robot, + session_id=self.session.session_id, + request_id=f"async-{request_index}-{uuid.uuid4().hex}", + step_id=request_index, + instruction=self._prompt, + frames=captured.frames, + ) + metadata: dict[str, Any] = {} + request_metadata = getattr(request, "metadata", {}) + if isinstance(request_metadata, Mapping): + metadata.update(request_metadata) + observation_id = _observation_identity(captured) + if observation_id is not None: + metadata.setdefault("observation_id", observation_id) + metadata.setdefault("snapshot_id", observation_id) + + metadata = build_async_metadata( + rtc=( + self._rtc_request(delay, cache_start_index=cache_start_index) + if self.config.mode == "rtc" and self.config.overlap_steps > 0 + else None + ), + base=metadata, + ) + request = _replace_metadata(request, metadata) + + started = self._monotonic() + result = self.client.step(request) + measured = max(0.0, self._monotonic() - started) + self.latency.record(measured) + measured_steps = self.latency.delay_in_steps(measured, self._control_period_s) + + actions = tuple(self.mapper.map_result(result)) + if not actions: + raise RuntimeError("binding returned an empty action chunk") + if any(not isinstance(action, RobotAction) for action in actions): + raise TypeError("binding action chunk must contain RobotAction values") + if len(actions) < self.config.chunk_steps: + raise RuntimeError( + f"binding returned {len(actions)} action(s), fewer than requested chunk_steps={self.config.chunk_steps}" + ) + actions = actions[: self.config.chunk_steps] + vectors = np.asarray([self.action_encoder(action) for action in actions], dtype=np.float64) + originals = np.asarray(self.chunk_extractor(result), dtype=np.float64) + + async_block = metadata.get("async") + self.records.append( + AsyncInferenceRecord( + request_index=request_index, + mode=self.config.mode, + trigger=trigger, + trigger_chunk_index=trigger_chunk_index, + captured_observation_id=observation_id, + environment_steps_at_capture=environment_steps_at_capture, + environment_steps_at_send=environment_steps_at_send, + inference_delay_steps=delay, + measured_latency_s=measured, + latency_steps=measured_steps, + pending_actions=tuple(pending), + async_block=dict(async_block) if isinstance(async_block, Mapping) else None, + ) + ) + return _Inference(actions=actions, vectors=vectors, originals=originals, record=self.records[-1]) + + def _rtc_request(self, delay: int, *, cache_start_index: int | None) -> RTCAsyncRequest: + config = self.config + # The committed prefix lives in the server's cache: this client holds + # robot-space actions, not model space. ``cache_start_index`` identifies + # the executed prefix of the cached full chunk; using ``delay`` here would + # be wrong whenever the replan stride differs from inference latency. + return RTCAsyncRequest( + prev_chunk_left_over=None, + inference_delay=delay, + execution_horizon=config.resolved_execution_horizon, + prefix_attention_schedule=config.prefix_attention_schedule, + max_guidance_weight=config.max_guidance_weight, + hard_prefix=config.hard_prefix, + cache_start_index=cache_start_index, + ) + + def _outcome( + self, + *, + policy_steps: int, + environment_steps: int, + total_reward: float, + terminated: bool, + truncated: bool, + ) -> AsyncEpisodeOutcome: + records = tuple(self.records) + return AsyncEpisodeOutcome( + policy_steps=policy_steps, + environment_steps=environment_steps, + total_reward=total_reward, + terminated=terminated, + truncated=truncated, + requests=policy_steps, + inference_delay_steps=tuple(record.inference_delay_steps for record in records), + measured_latency_s=tuple(record.measured_latency_s for record in records), + latency_steps=tuple(record.latency_steps for record in records), + records=records, + ) + + +def _replace_metadata(request: PolicyObservation, metadata: Mapping[str, Any]) -> PolicyObservation: + """Attach async metadata, using ``replace`` when the request is a dataclass.""" + + if hasattr(request, "__dataclass_fields__"): + return replace(request, metadata=metadata) + try: + request.metadata = metadata # type: ignore[misc] + except AttributeError as error: # pragma: no cover - defensive parity with ControlRuntime + raise TypeError("mapper request must support metadata propagation") from error + return request + + +def _observation_identity(observation: SimulatorObservation) -> str | None: + """Best-effort identity of a captured observation for diagnostics and tests.""" + + metadata = observation.robot.metadata + if not isinstance(metadata, Mapping): + return None + for key in ("observation_id", "snapshot_id", "step_index"): + value = metadata.get(key) + if isinstance(value, str) and value: + return value + if isinstance(value, int) and not isinstance(value, bool): + return str(value) + return None + + +def _instruction(observation: SimulatorObservation) -> str: + instruction = observation.robot.metadata.get("instruction") + if not isinstance(instruction, str) or not instruction.strip(): + raise RuntimeError("simulator reset did not provide an instruction; pass one explicitly") + return instruction.strip() diff --git a/src/embodirun/application/simulation/contracts.py b/src/embodirun/application/simulation/contracts.py index 838b4d5..6a88a59 100644 --- a/src/embodirun/application/simulation/contracts.py +++ b/src/embodirun/application/simulation/contracts.py @@ -5,7 +5,7 @@ import json import math from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any SIMULATION_CONFIG_SCHEMA = "embodirun.simulation.config.v1" @@ -31,6 +31,7 @@ class SimulationServiceConfig: simulator_kind: str simulator_options: Mapping[str, Any] inference_backend: str = "embodiinfer" + runtime_options: Mapping[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: for name in ( @@ -58,6 +59,7 @@ def __post_init__(self) -> None: raise SimulationContractError("inference backend must be embodiinfer or sglang") for name in ("inference_options", "simulator_options"): object.__setattr__(self, name, _mapping(getattr(self, name), name)) + object.__setattr__(self, "runtime_options", _mapping(self.runtime_options, "runtime_options")) def to_json(self) -> str: return json.dumps( @@ -77,6 +79,7 @@ def to_json(self) -> str: "type": self.simulator_kind, "options": dict(self.simulator_options), }, + "runtime_options": dict(self.runtime_options), }, allow_nan=False, ensure_ascii=False, @@ -110,6 +113,7 @@ def from_json(cls, value: str) -> SimulationServiceConfig: simulator_id=_string(simulator, "id", "simulation config.simulator"), simulator_kind=_string(simulator, "type", "simulation config.simulator"), simulator_options=_mapping(simulator.get("options", {}), "simulation config.simulator.options"), + runtime_options=_mapping(root.get("runtime_options", {}), "simulation config.runtime_options"), ) diff --git a/src/embodirun/application/simulation/service.py b/src/embodirun/application/simulation/service.py index d03bc76..346da1d 100644 --- a/src/embodirun/application/simulation/service.py +++ b/src/embodirun/application/simulation/service.py @@ -8,13 +8,14 @@ from __future__ import annotations import threading -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Any from embodirun.bindings import binding_definition from embodirun.model_services import build_inference_client from embodirun.simulators import simulator_definition +from .async_runtime import AsyncSimulationConfig, AsyncSimulationRuntime from .contracts import EpisodeRequest, EpisodeResult, SimulationServiceConfig from .runtime import SimulationRuntime @@ -83,21 +84,41 @@ def _execute_locked(self, request: EpisodeRequest) -> EpisodeResult: f"simulator {self.config.simulator_id!r} configuration is invalid: {error}" ) from error + try: + async_config = AsyncSimulationConfig.from_options( + self.config.runtime_options, + chunk_steps=request.chunk_steps, + ) + except (TypeError, ValueError) as error: + raise SimulationServiceError(str(error)) from error + client: Any | None = None simulator: Any | None = None runtime: Any | None = None try: client = self._inference_client(request.inference_timeout_s) simulator = definition.adapter_type(simulator_config) - runtime = self.runtime_factory( - simulator, - client, - mapper=binding.mapper_factory(), - chunk_steps=request.chunk_steps, - ) health = client.health() if health.get("status") != "ok": raise SimulationServiceError("inference service is not healthy") + if async_config.mode == "rtc" and self.config.inference_backend == "sglang": + _check_sglang_async_capability(client, async_config) + if self.runtime_factory is SimulationRuntime and async_config.mode != "sync": + runtime = AsyncSimulationRuntime( + simulator, + client, + mapper=binding.mapper_factory(), + config=async_config, + ) + else: + if async_config.mode != "sync": + raise SimulationServiceError("async runtime_options require the default SimulationRuntime factory") + runtime = self.runtime_factory( + simulator, + client, + mapper=binding.mapper_factory(), + chunk_steps=request.chunk_steps, + ) outcome = runtime.run( instruction=request.prompt, max_policy_steps=request.max_policy_steps, @@ -168,4 +189,27 @@ def _shutdown_client(client: Any) -> None: shutdown() +def _check_sglang_async_capability(client: Any, config: AsyncSimulationConfig) -> None: + capabilities = getattr(client, "capabilities", None) + if not callable(capabilities): + raise SimulationServiceError("SGLang async RTC requires endpoint capability metadata; use async mode=naive") + try: + payload = capabilities() + except Exception as error: + raise SimulationServiceError(f"SGLang async capability probe failed: {error}") from error + if not isinstance(payload, Mapping): + raise SimulationServiceError("SGLang capabilities response must be an object") + advertised = payload.get("async_inference") + if not isinstance(advertised, Mapping): + nested = payload.get("capabilities") + advertised = nested.get("async_inference") if isinstance(nested, Mapping) else None + if not isinstance(advertised, Mapping): + raise SimulationServiceError("SGLang endpoint does not advertise async_inference.rtc; use async mode=naive") + required = [config.mode] + if any(advertised.get(name) is not True for name in required): + raise SimulationServiceError( + f"SGLang endpoint does not support async algorithm {config.mode!r}; use async mode=naive" + ) + + __all__ = ["SimulationService", "SimulationServiceError", "create_inference_client"] diff --git a/src/embodirun/deployment/plan.py b/src/embodirun/deployment/plan.py index 5bb1a50..309f008 100644 --- a/src/embodirun/deployment/plan.py +++ b/src/embodirun/deployment/plan.py @@ -479,6 +479,7 @@ def _simulation_service( simulator_id=simulator.simulator_id, simulator_kind=simulator.kind, simulator_options=simulator.options, + runtime_options=runtime.options, ) except (TypeError, ValueError) as error: raise ServiceError(f"runtime {runtime.runtime_id!r} simulation configuration is invalid: {error}") from error diff --git a/src/embodirun/robots/lerobot/so101/adapter.py b/src/embodirun/robots/lerobot/so101/adapter.py index f12bd4f..59f6683 100644 --- a/src/embodirun/robots/lerobot/so101/adapter.py +++ b/src/embodirun/robots/lerobot/so101/adapter.py @@ -29,7 +29,9 @@ _MOTOR_IDS = dict(zip(SO101_MOTORS, range(1, 7))) _STS3215_MODEL_NUMBER = 777 _ENCODER_RESOLUTION = 4096 -_READ_RETRIES = 2 +_READ_RETRIES = 5 +_READ_RETRY_BACKOFF_S = 0.005 +_MAX_STALE_POSITION_READS = 5 _RETURN_DELAY = (7, 1) _MIN_POSITION = (9, 2) @@ -186,6 +188,8 @@ def __init__(self, config: SO101Config, *, read_only: bool = False) -> None: self._connected = False self._calibrated = False self._prepared = False + self._last_positions: dict[str, float] | None = None + self._stale_position_reads = 0 # Startup may fail after partially enabling the bus. Keep this # separate from ``_prepared`` so close can still attempt torque off, # while a passive connection never writes torque registers. @@ -413,20 +417,42 @@ def _read_positions(self) -> dict[str, float]: if result == self.sdk.COMM_SUCCESS: break if attempt == _READ_RETRIES: - self._communication_error("read SO-101 positions", result, 0) + return self._read_positions_individually() + # A busy USB/serial adapter can occasionally miss one status + # packet. Immediate retries tend to collide with the same device + # timing window, so give the bus a short increasing recovery gap. + time.sleep(_READ_RETRY_BACKOFF_S * (attempt + 1)) positions: dict[str, float] = {} for motor, motor_id in _MOTOR_IDS.items(): if not reader.isAvailable(motor_id, address, size): - raise SO101AdapterError(f"position response for motor {motor!r} is incomplete") + return self._read_positions_individually() raw = _decode_sign_magnitude(reader.getData(motor_id, address, size), 15) positions[f"{motor}.pos"] = self._normalized_position(motor, raw) return positions + def _read_positions_individually(self) -> dict[str, float]: + """Recover from an occasional failed multi-servo status packet.""" + + positions: dict[str, float] = {} + for motor, motor_id in _MOTOR_IDS.items(): + raw = self._read_register(motor_id, _PRESENT_POSITION, sign_bit=15) + positions[f"{motor}.pos"] = self._normalized_position(motor, raw) + return positions + def get_observation(self) -> Mapping[str, object]: if not self.is_connected: raise SO101AdapterError("SO-101 is not connected") - return self._read_positions() + try: + positions = self._read_positions() + except SO101AdapterError: + if self._last_positions is None or self._stale_position_reads >= _MAX_STALE_POSITION_READS: + raise + self._stale_position_reads += 1 + return dict(self._last_positions) + self._last_positions = dict(positions) + self._stale_position_reads = 0 + return positions def send_action(self, action: Mapping[str, float]) -> None: if self.read_only: @@ -441,8 +467,13 @@ def send_action(self, action: Mapping[str, float]) -> None: data = [self.sdk.SCS_LOBYTE(encoded), self.sdk.SCS_HIBYTE(encoded)] if not writer.addParam(motor_id, data): raise SO101AdapterError(f"failed to prepare position command for motor {motor!r}") - result = writer.txPacket() - self._communication_error("write SO-101 positions", result, 0) + for attempt in range(_READ_RETRIES + 1): + result = writer.txPacket() + if result == self.sdk.COMM_SUCCESS: + return + if attempt == _READ_RETRIES: + self._communication_error("write SO-101 positions", result, 0) + time.sleep(_READ_RETRY_BACKOFF_S * (attempt + 1)) def _close_port(self) -> None: if getattr(self.port_handler, "is_open", False): diff --git a/src/embodirun/robots/sensors/cameras/v4l2/__init__.py b/src/embodirun/robots/sensors/cameras/v4l2/__init__.py index bb1f9e9..579e6eb 100644 --- a/src/embodirun/robots/sensors/cameras/v4l2/__init__.py +++ b/src/embodirun/robots/sensors/cameras/v4l2/__init__.py @@ -12,7 +12,7 @@ def create_source(inputs: Sequence[SensorInput]) -> V4L2CameraSource: cameras: list[V4L2CameraConfig] = [] for item in inputs: options = dict(item.options) - unknown = sorted(set(options) - {"device", "width", "height", "fps"}) + unknown = sorted(set(options) - {"device", "width", "height", "fps", "input_format"}) if unknown: raise CameraError(f"sensor {item.sensor_id!r} contains unknown V4L2 fields: {', '.join(unknown)}") missing = [name for name in ("device", "width", "height", "fps") if name not in options] @@ -26,6 +26,7 @@ def create_source(inputs: Sequence[SensorInput]) -> V4L2CameraSource: width=options["width"], height=options["height"], fps=options["fps"], + input_format=options.get("input_format"), ) ) except (TypeError, ValueError) as error: diff --git a/src/embodirun/services/host/cli/command/run.py b/src/embodirun/services/host/cli/command/run.py index f7b46b6..b2c8aeb 100644 --- a/src/embodirun/services/host/cli/command/run.py +++ b/src/embodirun/services/host/cli/command/run.py @@ -59,10 +59,10 @@ def register(commands: Any) -> None: ) parser.add_argument( "--control-hz", - type=_positive_number, + type=_nonnegative_number, default=DEFAULT_CONTROL_HZ, metavar="HZ", - help=f"action playback rate (default: {DEFAULT_CONTROL_HZ:g})", + help=f"action playback rate; 0 disables software pacing (default: {DEFAULT_CONTROL_HZ:g})", ) parser.add_argument( "--request-timeout", @@ -132,7 +132,9 @@ def run(args: argparse.Namespace, context: CommandContext) -> int: progress.advance(runtime.node) playback_s = ( - args.max_steps * max(0, chunk_steps - 1) / args.control_hz if runtime.target_kind == "robot" else 0.0 + args.max_steps * max(0, chunk_steps - 1) / args.control_hz + if runtime.target_kind == "robot" and args.control_hz > 0 + else 0.0 ) task_timeout_s = max(60.0, args.max_steps * args.request_timeout + playback_s + 30.0) progress.update( @@ -200,4 +202,14 @@ def _positive_number(value: str) -> float: return result +def _nonnegative_number(value: str) -> float: + try: + result = float(value) + except ValueError as error: + raise argparse.ArgumentTypeError("must be a number") from error + if not math.isfinite(result) or result < 0: + raise argparse.ArgumentTypeError("must be a finite non-negative number") + return result + + __all__ = ["RunError", "register", "run"] diff --git a/src/embodirun/services/inference/adapters/sglang/pi05.py b/src/embodirun/services/inference/adapters/sglang/pi05.py index 05b19e6..0815cc9 100644 --- a/src/embodirun/services/inference/adapters/sglang/pi05.py +++ b/src/embodirun/services/inference/adapters/sglang/pi05.py @@ -11,6 +11,7 @@ from embodirun_sglang_pi05.pi05 import ( # type: ignore[import-not-found] LeRobotPi05Pipeline, _LeRobotPolicyModel, + _LeRobotPreprocessor, _Statistics, main, ) @@ -19,4 +20,4 @@ "SGLang Pi05 integration is optional; install integrations/sglang_pi05 in the managed inference environment" ) from error -__all__ = ["LeRobotPi05Pipeline", "_LeRobotPolicyModel", "_Statistics", "main"] +__all__ = ["LeRobotPi05Pipeline", "_LeRobotPolicyModel", "_LeRobotPreprocessor", "_Statistics", "main"] diff --git a/src/embodirun/simulators/libero/adapter.py b/src/embodirun/simulators/libero/adapter.py index 98d34ae..6830f4a 100644 --- a/src/embodirun/simulators/libero/adapter.py +++ b/src/embodirun/simulators/libero/adapter.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect import json import math import os @@ -171,7 +172,14 @@ def _normalize_observation( values={"observation.state": state}, metadata=dict(metadata), ), - frames=tuple(_encode_frame(f"observation.images.{name}", value) for name, value in pixels.items()), + frames=tuple( + _encode_frame( + f"observation.images.{name}", + value, + encoding=self.config.image_encoding, + ) + for name, value in pixels.items() + ), ) def _ensure_open(self) -> None: @@ -196,25 +204,35 @@ def make_libero_environment(config: LiberoConfig) -> Any: task_count = len(suite.tasks) if config.task_id >= task_count: raise ValueError(f"LIBERO task_id {config.task_id} is outside suite {config.suite!r} with {task_count} tasks") - return LiberoEnv( - task_suite=suite, - task_id=config.task_id, - task_suite_name=config.suite, - episode_length=config.max_episode_steps, - camera_name=tuple(_CAMERA_MAPPING), - camera_name_mapping=dict(_CAMERA_MAPPING), - obs_type="pixels_agent_pos", - render_mode="rgb_array", - observation_width=config.width, - observation_height=config.height, - init_states=True, - episode_index=0, - n_envs=1, - num_steps_wait=10, - control_freq=20, - control_mode="relative", - hard_reset=True, - ) + # LeRobot's LIBERO wrapper changed its constructor between the 0.4 and + # 0.6 releases. Keep the adapter usable with either optional environment + # while retaining the common settings. ``max_episode_steps`` is enforced + # by :meth:`step`, so older wrappers that do not expose ``episode_length`` + # remain bounded by the runtime contract. + options = { + "task_suite": suite, + "task_id": config.task_id, + "task_suite_name": config.suite, + "episode_length": config.max_episode_steps, + "camera_name": tuple(_CAMERA_MAPPING), + "camera_name_mapping": dict(_CAMERA_MAPPING), + "obs_type": "pixels_agent_pos", + "render_mode": "rgb_array", + "observation_width": config.width, + "observation_height": config.height, + "init_states": True, + "episode_index": 0, + "n_envs": 1, + "num_steps_wait": 10, + "control_freq": 20, + "control_mode": "relative", + "hard_reset": True, + } + try: + accepted = set(inspect.signature(LiberoEnv).parameters) + except (TypeError, ValueError): # pragma: no cover - unusual extension wrapper + accepted = set(options) + return LiberoEnv(**{key: value for key, value in options.items() if key in accepted}) def _prepare_libero_import() -> None: @@ -271,7 +289,10 @@ def _action_array(values: Any) -> Any: raise ValueError(f"LIBERO action must have shape ({LIBERO_ACTION_DIM},), got {action.shape}") if not np.isfinite(action).all(): raise ValueError("LIBERO action must contain only finite values") - return np.clip(action, -1.0, 1.0) + # LIBERO's OSC controller consumes a binary open/close gripper command. + clipped = np.clip(action, -1.0, 1.0) + clipped[-1] = np.sign(clipped[-1]) + return clipped def _mapping(value: Any, name: str) -> dict[str, Any]: @@ -323,19 +344,24 @@ def _policy_image(value: Any, *, name: str) -> Any: return np.ascontiguousarray(image[::-1, ::-1], dtype=np.uint8) -def _encode_frame(name: str, value: Any) -> CameraFrame: +def _encode_frame(name: str, value: Any, *, encoding: str = "jpeg") -> CameraFrame: try: import cv2 except ImportError as error: raise RuntimeError("LIBERO image encoding requires OpenCV") from error + if encoding not in {"jpeg", "png"}: + raise ValueError("LIBERO image encoding must be 'jpeg' or 'png'") + extension = ".jpg" if encoding == "jpeg" else ".png" + options = [int(cv2.IMWRITE_JPEG_QUALITY), 90] if encoding == "jpeg" else [] encoded, payload = cv2.imencode( - ".jpg", + extension, cv2.cvtColor(value, cv2.COLOR_RGB2BGR), - [int(cv2.IMWRITE_JPEG_QUALITY), 90], + options, ) if not encoded: - raise RuntimeError(f"LIBERO camera {name!r} could not be JPEG encoded") - return CameraFrame(name=name, mime_type="image/jpeg", data=payload.tobytes()) + raise RuntimeError(f"LIBERO camera {name!r} could not be {encoding.upper()} encoded") + mime_type = "image/jpeg" if encoding == "jpeg" else "image/png" + return CameraFrame(name=name, mime_type=mime_type, data=payload.tobytes()) def _unpack_reset(result: Any) -> tuple[Any, Mapping[str, Any]]: diff --git a/src/embodirun/simulators/libero/config.py b/src/embodirun/simulators/libero/config.py index e895d36..f9d692f 100644 --- a/src/embodirun/simulators/libero/config.py +++ b/src/embodirun/simulators/libero/config.py @@ -26,6 +26,7 @@ class LiberoConfig: width: int = 256 height: int = 256 viewer: bool = False + image_encoding: str = "jpeg" @classmethod def from_mapping( @@ -41,6 +42,7 @@ def from_mapping( "width", "height", "viewer", + "image_encoding", } unknown = sorted(set(options) - allowed) if unknown: @@ -53,6 +55,7 @@ def from_mapping( width=options.get("width", 256), height=options.get("height", 256), viewer=options.get("viewer", False), + image_encoding=options.get("image_encoding", "jpeg"), ) def __post_init__(self) -> None: @@ -70,6 +73,8 @@ def __post_init__(self) -> None: raise ValueError(f"LIBERO {name} must be a positive integer") if not isinstance(self.viewer, bool): raise TypeError("LIBERO viewer must be a boolean") + if self.image_encoding not in {"jpeg", "png"}: + raise ValueError("LIBERO image_encoding must be 'jpeg' or 'png'") @property def task(self) -> str: diff --git a/tests/fixtures/async_inference_golden.json b/tests/fixtures/async_inference_golden.json new file mode 100644 index 0000000..14e3783 --- /dev/null +++ b/tests/fixtures/async_inference_golden.json @@ -0,0 +1,215 @@ +{ + "description": "Shared RTC golden vectors asserted by EmbodiInfer and EmbodiRun.", + "rtc_prefix_weights": [ + { + "end": 6, + "expected": [ + 1.0, + 1.0, + 0.800000011920929, + 0.6000000238418579, + 0.4000000059604645, + 0.20000000298023224, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "name": "linear_2_6_10", + "schedule": "linear", + "start": 2, + "total": 10 + }, + { + "end": 10, + "expected": [ + 0.9090908765792847, + 0.8181818127632141, + 0.7272727489471436, + 0.6363636255264282, + 0.5454545021057129, + 0.4545454680919647, + 0.3636363744735718, + 0.27272728085517883, + 0.1818181872367859, + 0.09090909361839294 + ], + "name": "linear_0_10_10", + "schedule": "linear", + "start": 0, + "total": 10 + }, + { + "end": 0, + "expected": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "name": "linear_0_0_10_all_held", + "schedule": "linear", + "start": 0, + "total": 10 + }, + { + "end": 3, + "expected": [ + 1.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "name": "linear_3_3_10_degenerate", + "schedule": "linear", + "start": 3, + "total": 10 + }, + { + "end": 5, + "expected": [ + 0.8333333134651184, + 0.6666666269302368, + 0.5, + 0.3333333432674408, + 0.1666666716337204 + ], + "name": "linear_0_5_5", + "schedule": "linear", + "start": 0, + "total": 5 + }, + { + "end": 8, + "expected": [ + 1.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "name": "zeros_3_8_10", + "schedule": "zeros", + "start": 3, + "total": 10 + }, + { + "end": 5, + "expected": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "name": "zeros_0_5_5", + "schedule": "zeros", + "start": 0, + "total": 5 + }, + { + "end": 4, + "expected": [ + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + 0.0 + ], + "name": "ones_0_4_6", + "schedule": "ones", + "start": 0, + "total": 6 + }, + { + "end": 7, + "expected": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0 + ], + "name": "ones_2_7_10", + "schedule": "ones", + "start": 2, + "total": 10 + }, + { + "end": 6, + "expected": [ + 1.0, + 1.0, + 0.570589005947113, + 0.28707242012023926, + 0.11449221521615982, + 0.02577025070786476, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "name": "exp_2_6_10", + "schedule": "exp", + "start": 2, + "total": 10 + }, + { + "end": 4, + "expected": [ + 0.570589005947113, + 0.28707242012023926, + 0.11449221521615982, + 0.02577025070786476, + 0.0, + 0.0 + ], + "name": "exp_0_4_6", + "schedule": "exp", + "start": 0, + "total": 6 + }, + { + "end": 9, + "expected": [ + 1.0, + 0.7410117983818054, + 0.5325998067855835, + 0.3677061200141907, + 0.24019765853881836, + 0.14475014805793762, + 0.07674574851989746, + 0.03218316286802292, + 0.007599262986332178, + 0.0 + ], + "name": "exp_1_9_10", + "schedule": "exp", + "start": 1, + "total": 10 + } + ] +} diff --git a/tests/test_async_algorithms.py b/tests/test_async_algorithms.py new file mode 100644 index 0000000..8fd6f71 --- /dev/null +++ b/tests/test_async_algorithms.py @@ -0,0 +1,669 @@ +"""Execution-side RTC algorithms. + +The numeric core is pinned against the *shared* golden vectors in +``tests/fixtures/async_inference_golden.json``, which EmbodiInfer asserts with its +own inference-side implementation. Both halves of the async loop therefore agree +on the same arithmetic across a process and a repository boundary. + +Everything here runs on CPU with fakes and an injected clock; no robot hardware +and no inference engine are involved. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +from embodirun.application.async_inference import ( + ASYNC_SCHEMA, + ActionQueue, + AsyncControlRuntime, + AsyncRuntimeConfig, + LatencyTracker, + RTCAsyncRequest, + build_async_metadata, + prefix_weights, +) +from embodirun.model_services import ( + ImagePayload, + PolicyAction, + PolicyObservation, + PolicyResult, + Session, +) +from embodirun.robots import RobotAction, RobotObservation +from embodirun.robots.sensors.cameras import CameraFrame + +GOLDEN_PATH = Path(__file__).parent / "fixtures" / "async_inference_golden.json" + + +def _golden() -> dict: + return json.loads(GOLDEN_PATH.read_text(encoding="utf-8")) + + +# ============================================================================ +# Shared golden vectors +# ============================================================================ +@pytest.mark.parametrize("case", _golden()["rtc_prefix_weights"], ids=lambda case: case["name"]) +def test_prefix_weights_match_golden_vectors(case: dict) -> None: + weights = prefix_weights(case["start"], case["end"], case["total"], case["schedule"]) + expected = np.asarray(case["expected"], dtype=np.float64) + assert weights.shape == (case["total"],), case["name"] + assert np.allclose(weights, expected, atol=1e-6), (case["name"], weights.tolist()) + + +def test_runtime_metadata_round_trips_as_json() -> None: + """The async block must be JSON-serializable as the HTTP client sends it.""" + + payload = build_async_metadata( + rtc=RTCAsyncRequest( + prev_chunk_left_over=((1.0, 2.0),), + inference_delay=1, + execution_horizon=2, + model_space=True, + ), + ) + assert json.loads(json.dumps(payload)) == payload + + +# ============================================================================ +# RTC prefix weights +# ============================================================================ +def test_prefix_weights_documented_example() -> None: + weights = prefix_weights(2, 6, 10, "linear") + assert [round(float(value), 6) for value in weights] == [1, 1, 0.8, 0.6, 0.4, 0.2, 0, 0, 0, 0] + + +def test_prefix_weights_start_is_pushed_down_to_end() -> None: + assert np.array_equal(prefix_weights(9, 4, 10, "zeros"), prefix_weights(4, 4, 10, "zeros")) + + +def test_prefix_weights_rejects_bad_arguments() -> None: + with pytest.raises(ValueError): + prefix_weights(0, 1, 0) + with pytest.raises(ValueError): + prefix_weights(-1, 1, 4) + with pytest.raises(ValueError): + prefix_weights(0, 1, 4, "cosine") + with pytest.raises(ValueError): + prefix_weights(0, 1, 4.0) # type: ignore[arg-type] + + +# ============================================================================ +# ActionQueue — replace (RTC) vs append (non-RTC) +# ============================================================================ +def _chunks(rows: int, width: int = 2, offset: float = 0.0) -> tuple[np.ndarray, np.ndarray]: + original = np.arange(rows * width, dtype=np.float64).reshape(rows, width) + offset + processed = original + 100.0 + return original, processed + + +def test_action_queue_replace_drops_the_delay_prefix() -> None: + original, processed = _chunks(5) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 2) + + assert queue.get_action_index() == 0 + assert queue.qsize() == 3 + assert not queue.empty() + assert np.array_equal(queue.get_left_over(), original[2:]) + assert np.array_equal(queue.get_processed_left_over(), processed[2:]) + + +def test_action_queue_get_returns_a_clone_and_advances_the_index() -> None: + original, processed = _chunks(3) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 0, task="pick") + + action, task = queue.get_with_task() + assert task == "pick" + assert np.array_equal(action, processed[0]) + action[0] = -999.0 + assert queue.get_action_index() == 1 + assert np.array_equal(queue.get_left_over(), original[1:]) + assert np.array_equal(queue.get_processed_left_over(), processed[1:]) + assert np.array_equal(queue.get(), processed[1]) + assert np.array_equal(queue.get(), processed[2]) + assert queue.get() is None + assert queue.empty() + + +def test_action_queue_leftover_after_partial_consumption() -> None: + original, processed = _chunks(6) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 1) + queue.get() + queue.get() + assert np.array_equal(queue.get_left_over(), original[3:]) + assert np.array_equal(queue.get_processed_left_over(), processed[3:]) + + +def test_action_queue_append_preserves_continuity() -> None: + original_a, processed_a = _chunks(5) + original_b, processed_b = _chunks(3, offset=1000.0) + queue = ActionQueue(rtc_enabled=False) + queue.merge(original_a, processed_a, 0) + queue.get() + queue.get() + queue.merge(original_b, processed_b, 0) + + assert queue.get_action_index() == 0 + assert queue.qsize() == 6 + assert np.array_equal(queue.get_processed_left_over(), np.concatenate([processed_a[2:], processed_b])) + + +def test_action_queue_append_from_empty_starts_the_queue() -> None: + original, processed = _chunks(4) + queue = ActionQueue(rtc_enabled=False) + queue.merge(original, processed, 3) + assert np.array_equal(queue.get_left_over(), original) + assert np.array_equal(queue.get_processed_left_over(), processed) + + +def test_action_queue_clear_resets_everything() -> None: + original, processed = _chunks(3) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 0) + queue.get() + queue.clear() + assert queue.empty() + assert queue.qsize() == 0 + assert queue.get_action_index() == 0 + assert queue.get_left_over() is None + assert queue.get_processed_left_over() is None + + +def test_action_queue_delay_resolution_never_over_discards() -> None: + original, processed = _chunks(5) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 0, task="pick") + queue.get() + queue.get() + queue.get() + + # Inference claims a delay of 5 but only three actions were consumed. + original_new, processed_new = _chunks(4, offset=2000.0) + resolved = queue.merge(original_new, processed_new, 5, action_index_before_inference=0) + assert resolved == 3 + assert np.array_equal(queue.get_left_over(), original_new[3:]) + assert np.array_equal(queue.get_processed_left_over(), processed_new[3:]) + + +def test_action_queue_delay_resolution_keeps_the_smaller_delay() -> None: + original, processed = _chunks(5) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 0) + for _ in range(4): + queue.get() + + original_new, processed_new = _chunks(4, offset=3000.0) + # real_delay 2 is smaller than the 4 consumed and is what gets applied. + resolved = queue.merge(original_new, processed_new, 2, action_index_before_inference=0) + assert resolved == 2 + assert np.array_equal(queue.get_left_over(), original_new[2:]) + + +def test_action_queue_delay_resolution_without_index_is_clamped_non_negative() -> None: + original, processed = _chunks(3) + queue = ActionQueue(rtc_enabled=True) + assert queue.merge(original, processed, -4) == 0 + assert np.array_equal(queue.get_left_over(), original) + + +def test_action_queue_delay_clamps_to_the_shortest_chunk() -> None: + original, processed = _chunks(3) + queue = ActionQueue(rtc_enabled=True) + assert queue.merge(original, processed, 99) == 3 + assert queue.qsize() == 0 + + +def test_action_queue_task_lockstep_invariant_raises() -> None: + original, processed = _chunks(3) + queue = ActionQueue(rtc_enabled=True) + queue.merge(original, processed, 0, task="pick") + queue._task_queue.append("extra") # deliberately break the invariant + + with pytest.raises(RuntimeError, match="lockstep"): + queue.get() + with pytest.raises(RuntimeError, match="lockstep"): + queue.get_left_over() + + +def test_action_queue_rejects_non_numeric_and_non_2d_chunks() -> None: + queue = ActionQueue(rtc_enabled=True) + with pytest.raises(ValueError): + queue.merge("not-a-grid", np.zeros((2, 2)), 0) + with pytest.raises(ValueError): + queue.merge([[1.0, 2.0]], np.zeros((2, 2, 2)), 0) + with pytest.raises(TypeError): + queue.merge(np.zeros((2, 2)), np.zeros((2, 2)), 1.5) # type: ignore[arg-type] + + +# ============================================================================ +# LatencyTracker +# ============================================================================ +class FakeClock: + """Deterministic monotonic clock with recorded sleeps.""" + + def __init__(self) -> None: + self.now = 0.0 + self.sleeps: list[float] = [] + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += seconds + + +def test_latency_tracker_measures_with_an_injected_clock() -> None: + clock = FakeClock() + tracker = LatencyTracker(window=3, monotonic=clock.monotonic) + assert tracker.average() == 0.0 + assert tracker.median() == 0.0 + assert tracker.latest_delay(0.1) == 0 + + tracker.start() + clock.now += 0.3 + assert tracker.stop() == pytest.approx(0.3) + tracker.record(0.1) + tracker.record(0.2) + + assert len(tracker) == 3 + assert tracker.average() == pytest.approx(0.2) + assert tracker.median() == pytest.approx(0.2) + assert tracker.max() == pytest.approx(0.3) + assert tracker.delay_in_steps(0.3, 0.1) == 3 + assert tracker.delay_in_steps(0.25, 0.1) == 2 + assert tracker.delay_in_steps(-1.0, 0.1) == 0 + assert tracker.latest_delay(0.1) == 2 + + tracker.reset() + assert len(tracker) == 0 + assert tracker.max() == 0.0 + + +def test_latency_tracker_rejects_invalid_period_and_stop_before_start() -> None: + tracker = LatencyTracker(monotonic=FakeClock().monotonic) + with pytest.raises(RuntimeError): + tracker.stop() + with pytest.raises(ValueError): + tracker.delay_in_steps(0.1, 0.0) + with pytest.raises(ValueError): + LatencyTracker(window=0) + + +# ============================================================================ +# build_async_metadata +# ============================================================================ +def test_build_async_metadata_matches_the_documented_schema() -> None: + metadata = build_async_metadata( + rtc=RTCAsyncRequest( + prev_chunk_left_over=((1.0, 2.0), (3.0, 4.0)), + inference_delay=3, + execution_horizon=8, + prefix_attention_schedule="linear", + max_guidance_weight=10.0, + hard_prefix=False, + model_space=True, + cache_start_index=6, + ), + base={"observation_id": "obs-1"}, + ) + + assert metadata == { + "observation_id": "obs-1", + "async": { + "schema": ASYNC_SCHEMA, + "rtc": { + "prev_chunk_left_over": [[1.0, 2.0], [3.0, 4.0]], + "inference_delay": 3, + "cache_start_index": 6, + "execution_horizon": 8, + "prefix_attention_schedule": "linear", + "max_guidance_weight": 10.0, + "hard_prefix": False, + "model_space": True, + }, + }, + } + + +def test_build_async_metadata_omits_disabled_algorithms() -> None: + assert build_async_metadata() == {} + assert build_async_metadata(base={"a": 1}) == {"a": 1} + + rtc_only = build_async_metadata( + rtc=RTCAsyncRequest(prev_chunk_left_over=None, inference_delay=0, execution_horizon=2) + ) + assert set(rtc_only["async"]) == {"schema", "rtc"} + + +def test_build_async_metadata_does_not_mutate_the_base() -> None: + base = {"observation_id": "obs-1"} + build_async_metadata(rtc=RTCAsyncRequest(None, 0, 1), base=base) + assert base == {"observation_id": "obs-1"} + + +def test_rtc_request_rejects_malformed_input() -> None: + with pytest.raises(TypeError): + RTCAsyncRequest(None, 1.5, 4) # type: ignore[arg-type] + with pytest.raises(ValueError): + RTCAsyncRequest(None, -1, 4) + with pytest.raises(ValueError): + RTCAsyncRequest(None, 1, -4) + with pytest.raises(ValueError): + RTCAsyncRequest(((1.0, float("inf")),), 1, 4) + with pytest.raises(ValueError): + RTCAsyncRequest(((1.0, 2.0), (3.0,)), 1, 4) + with pytest.raises(ValueError): + RTCAsyncRequest(None, 1, 4, "cosine") + with pytest.raises(ValueError): + RTCAsyncRequest(None, 1, 4, "linear", 0.0) + with pytest.raises(TypeError): + RTCAsyncRequest(None, 1, 4, "linear", 10.0, 1) # type: ignore[arg-type] + with pytest.raises(ValueError, match="cache_start_index"): + RTCAsyncRequest(None, 1, 4, cache_start_index=-1) + with pytest.raises(TypeError, match="cache_start_index"): + RTCAsyncRequest(None, 1, 4, cache_start_index=True) # type: ignore[arg-type] + + +# ============================================================================ +# AsyncControlRuntime — fakes only, injected clock +# ============================================================================ +class FakeRobot: + robot_id = "fake-1" + + def __init__(self) -> None: + self.actions: list[RobotAction] = [] + self.observations = 0 + + def observe(self) -> RobotObservation: + self.observations += 1 + return RobotObservation(timestamp_s=float(self.observations), values={"joint_0": 0.0, "joint_1": 0.0}) + + def execute(self, action: RobotAction) -> None: + self.actions.append(action) + + def stop(self) -> None: + return None + + def close(self) -> None: + return None + + +class FakeInferenceClient: + """Blocking client that consumes fake-clock time while "inferring".""" + + policy_action_space = "pi05.action_chunk.v1" + + def __init__( + self, + *, + chunk: list[list[float]], + clock: FakeClock, + period_s: float, + latency_steps: int, + robot: FakeRobot, + ) -> None: + self.requests: list[PolicyObservation] = [] + self.executed_at_request: list[int] = [] + self._base = np.asarray(chunk, dtype=np.float64) + self._clock = clock + self._period_s = period_s + self._latency_steps = latency_steps + self._robot = robot + self._counter = 0 + + def open_session(self, *, robot_id: str, action_space: str, metadata=None) -> Session: + return Session("session-1", 0) + + def step(self, observation: PolicyObservation) -> PolicyResult: + self.requests.append(observation) + self.executed_at_request.append(len(self._robot.actions)) + self._counter += 1 + self._clock.now += self._latency_steps * self._period_s + data = (self._base + 100.0 * (self._counter - 1)).tolist() + return PolicyResult( + request_id=observation.request_id, + session_id=observation.session_id, + step_id=observation.step_id, + session_revision=self._counter, + action_space=self.policy_action_space, + actions=(PolicyAction("action_chunk", {"data": data}),), + ) + + def reset(self, session_id: str, *, request_id: str) -> Session: + return Session(session_id, 0) + + def close(self, session_id: str) -> None: + return None + + +class FakeMapper: + policy_action_space = "pi05.action_chunk.v1" + + def map_observation( + self, + observation: RobotObservation, + *, + session_id: str, + request_id: str, + step_id: int, + instruction: str, + frames, + ) -> PolicyObservation: + return PolicyObservation( + session_id=session_id, + request_id=request_id, + step_id=step_id, + instruction=instruction, + state=dict(observation.values), + images=tuple(ImagePayload(frame.name, frame.mime_type, frame.data) for frame in frames), + metadata={"robot_timestamp_s": observation.timestamp_s}, + ) + + def map_result(self, result: PolicyResult) -> tuple[RobotAction, ...]: + data = result.actions[0].values["data"] + return tuple( + RobotAction(timestamp_s=0.0, values={"type": "joint_position", "vector": [float(v) for v in row]}) + for row in data + ) + + +FRAME = CameraFrame("observation.images.front", "image/jpeg", b"jpeg") + + +def _build_runtime( + *, + chunk: list[list[float]], + config: AsyncRuntimeConfig, + control_hz: float = 10.0, + latency_steps: int = 2, + chunk_steps: int | None = None, +): + clock = FakeClock() + robot = FakeRobot() + client = FakeInferenceClient( + chunk=chunk, + clock=clock, + period_s=0.0 if control_hz == 0 else 1.0 / control_hz, + latency_steps=latency_steps, + robot=robot, + ) + runtime = AsyncControlRuntime( + robot, + client, + instruction="pick up the block", + mapper=FakeMapper(), + chunk_steps=chunk_steps if chunk_steps is not None else len(chunk), + control_hz=control_hz, + config=config, + monotonic=clock.monotonic, + sleep=clock.sleep, + ) + return runtime, robot, client, clock + + +def test_async_rtc_sends_the_leftover_and_the_consumed_delay() -> None: + chunk = [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [4.0, 4.0]] + config = AsyncRuntimeConfig( + enabled=True, + rtc_enabled=True, + overlap_steps=3, + initial_inference_delay=2, + execution_horizon=4, + max_guidance_weight=7.5, + send_model_space_prefix=True, + ) + runtime, _robot, client, _clock = _build_runtime(chunk=chunk, config=config, latency_steps=2) + + for _ in range(4): + runtime.step((FRAME,)) + runtime._collect_job(block=True) + + assert len(client.requests) == 2 + block = client.requests[1].metadata["async"]["rtc"] + trigger = len(chunk) - config.overlap_steps + assert block["prev_chunk_left_over"] == chunk[trigger:] + assert block["inference_delay"] == 2 + assert block["execution_horizon"] == 4 + assert block["max_guidance_weight"] == 7.5 + assert block["hard_prefix"] is False + + event = runtime.events[-1] + assert event.actions_consumed == 2 + assert event.inference_delay == 2 + assert event.resolved_delay == 2 + + +def test_async_rtc_fast_response_waits_for_and_skips_the_promised_prefix() -> None: + chunk = [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [4.0, 4.0]] + config = AsyncRuntimeConfig( + enabled=True, + rtc_enabled=True, + overlap_steps=3, + initial_inference_delay=2, + ) + runtime, _robot, _client, _clock = _build_runtime(chunk=chunk, config=config, latency_steps=0) + + for _ in range(4): + runtime.step((FRAME,)) + runtime._collect_job(block=True) + + event = runtime.events[-1] + assert event.latency_delay == 0 + assert event.actions_consumed == 2 + assert event.resolved_delay == 2 + + +def test_async_rtc_session_cache_sends_chunk_offset_not_latency() -> None: + chunk = [[float(index), float(index)] for index in range(8)] + config = AsyncRuntimeConfig( + enabled=True, + rtc_enabled=True, + overlap_steps=5, + initial_inference_delay=2, + execution_horizon=3, + prefix_attention_schedule="exp", + max_guidance_weight=5.0, + send_model_space_prefix=False, + ) + runtime, _robot, client, _clock = _build_runtime( + chunk=chunk, + config=config, + latency_steps=2, + ) + + for _ in range(4): + runtime.step((FRAME,)) + runtime._collect_job(block=True) + + assert len(client.requests) == 2 + bootstrap = client.requests[0].metadata["async"]["rtc"] + replan = client.requests[1].metadata["async"]["rtc"] + assert bootstrap["cache_start_index"] == 0 + assert replan["cache_start_index"] == 3 # H - overlap, not Δ=2 + assert replan["inference_delay"] == 2 + assert replan["prev_chunk_left_over"] is None + + +def test_async_disabled_is_identical_to_the_synchronous_runtime() -> None: + chunk = [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]] + config = AsyncRuntimeConfig(enabled=False, rtc_enabled=True, overlap_steps=1) + runtime, robot, client, clock = _build_runtime( + chunk=chunk, config=config, control_hz=20.0, latency_steps=0, chunk_steps=3 + ) + + result = runtime.step((FRAME,)) + + assert result is not None + assert [action.values["vector"] for action in robot.actions] == chunk + assert clock.sleeps == [0.05, 0.05] + assert len(client.requests) == 1 + assert "async" not in client.requests[0].metadata + assert runtime.events == [] + + +def test_async_runtime_paces_one_action_per_control_period() -> None: + chunk = [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]] + config = AsyncRuntimeConfig( + enabled=True, + overlap_steps=1, + ) + runtime, robot, _client, clock = _build_runtime( + chunk=chunk, + config=config, + control_hz=20.0, + latency_steps=0, + ) + + runtime.step((FRAME,)) + runtime.step((FRAME,)) + runtime.step((FRAME,)) + + assert len(robot.actions) == 3 + assert clock.sleeps == [0.05, 0.05] + + +def test_async_runtime_zero_control_hz_disables_software_pacing() -> None: + chunk = [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]] + config = AsyncRuntimeConfig( + enabled=True, + overlap_steps=1, + ) + runtime, robot, _client, clock = _build_runtime( + chunk=chunk, + config=config, + control_hz=0.0, + latency_steps=0, + ) + + runtime.step((FRAME,)) + runtime.step((FRAME,)) + runtime.step((FRAME,)) + + assert len(robot.actions) == 3 + assert clock.sleeps == [] + assert runtime.action_period_s == 0.0 + + +def test_async_runtime_config_accepts_artificial_inference_delay() -> None: + config = AsyncRuntimeConfig.from_options({"async": {"mode": "naive", "artificial_inference_delay_s": 0.2}}) + + assert config.enabled is True + assert config.artificial_inference_delay_s == pytest.approx(0.2) + + with pytest.raises(ValueError, match="artificial_inference_delay_s"): + AsyncRuntimeConfig(artificial_inference_delay_s=-0.1) + + +def test_async_runtime_rejects_invalid_overlap() -> None: + with pytest.raises(ValueError): + AsyncRuntimeConfig(enabled=True, overlap_steps=-1) + with pytest.raises(ValueError): + AsyncRuntimeConfig(enabled=True, max_guidance_weight=float("inf")) diff --git a/tests/test_async_simulation_runtime.py b/tests/test_async_simulation_runtime.py new file mode 100644 index 0000000..b9d613b --- /dev/null +++ b/tests/test_async_simulation_runtime.py @@ -0,0 +1,607 @@ +"""Asynchronous simulator-episode runtime. + +Everything here runs on CPU with fakes and an injected clock: a fake simulator, a +fake inference client, and a fake monotonic clock. No LIBERO, no network, no GPU, +no inference engine. + +The ``Δ`` (``inference_delay_steps``) exercised throughout is the *modeled* delay +the runtime advances the simulator by while the next chunk is computed, exactly as +the RTC paper expresses it in simulation. +""" + +from __future__ import annotations + +import pytest + +from embodirun.application.simulation.async_runtime import ( + ASYNC_SIMULATION_MODES, + AsyncSimulationConfig, + AsyncSimulationRuntime, + latency_steps, +) +from embodirun.application.simulation.runtime import EpisodeOutcome, SimulationRuntime +from embodirun.model_services import ( + ImagePayload, + PolicyAction, + PolicyObservation, + PolicyResult, + Session, +) +from embodirun.robots import RobotAction, RobotObservation +from embodirun.robots.sensors.cameras import CameraFrame +from embodirun.simulators import SimulationStep, SimulatorAdapter, SimulatorObservation + +FRAME_NAME = "observation.images.front" +INSTRUCTION = "pick up the black bowl" + + +class FakeClock: + """Deterministic monotonic clock advanced by explicit, injected amounts.""" + + def __init__(self) -> None: + self.now = 0.0 + + def monotonic(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +class FakeSimulator(SimulatorAdapter): + """Deterministic LIBERO-shaped simulator: state is the executed step count.""" + + simulator_id = "fake-libero-sim" + + def __init__(self, *, terminate_at: int | None = None, truncate_at: int | None = None, reward: float = 1.0): + self.terminate_at = terminate_at + self.truncate_at = truncate_at + self.reward = reward + self.step_index = 0 + self.executed: list[float] = [] + self.closed = False + self.reset_count = 0 + + def reset(self, *, task=None, seed=None, options=None) -> SimulatorObservation: + self.reset_count += 1 + self.step_index = 0 + self.executed = [] + return self._observation() + + def step(self, action: RobotAction) -> SimulationStep: + self.step_index += 1 + self.executed.append(float(action.values["vector"][0])) + terminated = self.terminate_at is not None and self.step_index >= self.terminate_at + truncated = self.truncate_at is not None and self.step_index >= self.truncate_at + return SimulationStep( + observation=self._observation(), + reward=self.reward, + terminated=terminated, + truncated=truncated, + ) + + def close(self) -> None: + self.closed = True + + def _observation(self) -> SimulatorObservation: + step = self.step_index + return SimulatorObservation( + robot=RobotObservation( + timestamp_s=float(step), + values={"observation.state": [float(step)]}, + metadata={ + "instruction": INSTRUCTION, + "step_index": step, + "observation_id": f"obs-{step}", + }, + ), + frames=(CameraFrame(FRAME_NAME, "image/jpeg", f"frame-{step}".encode()),), + ) + + +class FakeInferenceClient: + """Blocking fake client that advances the injected clock while 'inferring'.""" + + policy_action_space = "fake.action_chunk.v1" + + def __init__(self, *, chunk_steps: int, clock: FakeClock, latency_s: float) -> None: + self.requests: list[PolicyObservation] = [] + self._chunk_steps = chunk_steps + self._clock = clock + self._latency_s = latency_s + self._counter = 0 + self.closed = False + + def open_session(self, *, robot_id: str, action_space: str, metadata=None) -> Session: + self.session_metadata = dict(metadata or {}) + return Session("session-1", 0) + + def step(self, observation: PolicyObservation) -> PolicyResult: + self.requests.append(observation) + self._clock.advance(self._latency_s) + base = 100.0 * self._counter + self._counter += 1 + data = [[base + float(row)] for row in range(self._chunk_steps)] + return PolicyResult( + request_id=observation.request_id, + session_id=observation.session_id, + step_id=observation.step_id, + session_revision=self._counter, + action_space=self.policy_action_space, + actions=(PolicyAction("action_chunk", {"data": data}),), + ) + + def reset(self, session_id: str, *, request_id: str) -> Session: + return Session(session_id, 0) + + def close(self, session_id: str) -> None: + self.closed = True + + +class FakeMapper: + """Maps the fake observation and the fake action chunk without any conversion.""" + + policy_action_space = "fake.action_chunk.v1" + + def map_observation( + self, + observation: RobotObservation, + *, + session_id: str, + request_id: str, + step_id: int, + instruction: str, + frames, + ) -> PolicyObservation: + return PolicyObservation( + session_id=session_id, + request_id=request_id, + step_id=step_id, + instruction=instruction, + state=dict(observation.values), + images=tuple(ImagePayload(frame.name, frame.mime_type, frame.data) for frame in frames), + metadata={ + "robot_timestamp_s": observation.timestamp_s, + "observation_id": observation.metadata.get("observation_id"), + "step_index": observation.metadata.get("step_index"), + }, + ) + + def map_result(self, result: PolicyResult) -> tuple[RobotAction, ...]: + data = result.actions[0].values["data"] + return tuple( + RobotAction(timestamp_s=0.0, values={"type": "joint_position", "vector": [float(value) for value in row]}) + for row in data + ) + + +def _harness( + *, + mode: str, + chunk_steps: int, + overlap_steps: int = 0, + inference_delay_steps: int = 0, + latency_s: float = 0.0, + control_hz: float = 10.0, + execution_horizon: int | None = None, + prefix_attention_schedule: str = "linear", + max_guidance_weight: float = 10.0, + hard_prefix: bool = False, + terminate_at: int | None = None, + truncate_at: int | None = None, +): + config = AsyncSimulationConfig( + mode=mode, + chunk_steps=chunk_steps, + overlap_steps=overlap_steps, + inference_delay_steps=inference_delay_steps, + control_hz=control_hz, + execution_horizon=execution_horizon, + prefix_attention_schedule=prefix_attention_schedule, + max_guidance_weight=max_guidance_weight, + hard_prefix=hard_prefix, + ) + clock = FakeClock() + simulator = FakeSimulator(terminate_at=terminate_at, truncate_at=truncate_at) + client = FakeInferenceClient(chunk_steps=chunk_steps, clock=clock, latency_s=latency_s) + runtime = AsyncSimulationRuntime( + simulator, + client, + mapper=FakeMapper(), + config=config, + monotonic=clock.monotonic, + ) + return runtime, simulator, client, clock + + +def _run(runtime, *, max_policy_steps: int, instruction: str | None = INSTRUCTION): + return runtime.run(instruction=instruction, max_policy_steps=max_policy_steps, seed=7) + + +def _async_block(request: PolicyObservation): + return request.metadata.get("async") + + +# ============================================================================ +# sync +# ============================================================================ +def test_sync_mode_has_no_overlap_no_delay_and_no_async_block() -> None: + runtime, simulator, client, _clock = _harness(mode="sync", chunk_steps=4) + + outcome = _run(runtime, max_policy_steps=3) + runtime.close() + + assert outcome.policy_steps == 3 + assert outcome.requests == 3 + assert outcome.environment_steps == 12 + assert outcome.terminated is False and outcome.truncated is False + assert outcome.inference_delay_steps == (0, 0, 0) + + # A request is issued only once the previous chunk is fully consumed: the + # simulator state at request time is 0, 4, 8. + assert [request.state["observation.state"] for request in client.requests] == [[0.0], [4.0], [8.0]] + assert all("async" not in request.metadata for request in client.requests) + assert client.closed is True + assert simulator.executed == [ + 0.0, + 1.0, + 2.0, + 3.0, + 100.0, + 101.0, + 102.0, + 103.0, + 200.0, + 201.0, + 202.0, + 203.0, + ] + + +def test_sync_mode_outcome_is_episode_outcome_compatible() -> None: + runtime, _simulator, _client, _clock = _harness(mode="sync", chunk_steps=3) + + outcome = _run(runtime, max_policy_steps=2) + + assert outcome.to_episode_outcome() == EpisodeOutcome(2, 6, 6.0, False, False) + assert isinstance(outcome.to_episode_outcome(), EpisodeOutcome) + + +# ============================================================================ +# naive +# ============================================================================ +def test_naive_mode_sends_no_conditioning_block_and_adopts_at_handover() -> None: + runtime, _simulator, client, _clock = _harness( + mode="naive", chunk_steps=5, overlap_steps=2, inference_delay_steps=2 + ) + + outcome = _run(runtime, max_policy_steps=2) + + assert outcome.requests == 2 + assert all("async" not in request.metadata for request in client.requests) + + trigger = client.requests[1] + assert trigger.metadata["observation_id"] == "obs-3" + assert trigger.state["observation.state"] == [3.0] + assert client.requests[1].images[0].data == b"frame-3" + + record = outcome.records[1] + assert record.trigger is True + assert record.trigger_chunk_index == 3 + assert record.environment_steps_at_capture == 3 + assert record.environment_steps_at_send == 5 + assert record.inference_delay_steps == 2 + + # The first chunk was executed through to its end (handover) before the new + # chunk's first action ran, so the naive baseline never replaces the tail. + assert _simulator_actions(runtime) == [ + 0.0, + 1.0, + 2.0, + 3.0, + 4.0, # bootstrap chunk, fully executed + 100.0, + 101.0, + 102.0, # new chunk taken over at handover + ] + + +def _simulator_actions(runtime: AsyncSimulationRuntime) -> list[float]: + return list(runtime.simulator.executed) + + +# ============================================================================ +# rtc +# ============================================================================ +def test_rtc_mode_emits_inference_delay_horizon_and_schedule() -> None: + runtime, _simulator, client, _clock = _harness( + mode="rtc", + chunk_steps=5, + overlap_steps=2, + inference_delay_steps=2, + execution_horizon=4, + prefix_attention_schedule="exp", + max_guidance_weight=7.5, + hard_prefix=True, + ) + + outcome = _run(runtime, max_policy_steps=2) + + block = _async_block(client.requests[1]) + assert set(block) == {"schema", "rtc"} + assert block["schema"] == "embodiinfer.async.v1" + assert set(block["rtc"]) == { + "prev_chunk_left_over", + "inference_delay", + "cache_start_index", + "execution_horizon", + "prefix_attention_schedule", + "max_guidance_weight", + "hard_prefix", + } + assert block["rtc"]["inference_delay"] == 2 + assert block["rtc"]["cache_start_index"] == 3 + assert block["rtc"]["execution_horizon"] == 4 + assert block["rtc"]["prefix_attention_schedule"] == "exp" + assert block["rtc"]["max_guidance_weight"] == 7.5 + assert block["rtc"]["hard_prefix"] is True + # The client holds no model space, so it sends the delay and lets the server + # slice its own cache. + assert block["rtc"]["prev_chunk_left_over"] is None + assert outcome.records[1].async_block == block + + # RTC adopts immediately, skipping the Δ actions that elapsed during the window. + assert _simulator_actions(runtime) == [0.0, 1.0, 2.0, 3.0, 4.0, 102.0] + + +def test_rtc_bootstrap_request_carries_a_zero_delay_block() -> None: + runtime, _simulator, client, _clock = _harness( + mode="rtc", chunk_steps=5, overlap_steps=2, inference_delay_steps=1, execution_horizon=3 + ) + + _run(runtime, max_policy_steps=1) + + bootstrap = _async_block(client.requests[0])["rtc"] + assert bootstrap["inference_delay"] == 0 + assert bootstrap["cache_start_index"] == 0 + assert bootstrap["prev_chunk_left_over"] is None + assert bootstrap["execution_horizon"] == 3 + + +def test_rtc_cache_offset_is_replan_stride_not_inference_delay() -> None: + """Every replan points at model row ``s`` even after dropping Δ rows. + + Here H=8, overlap=5, so the replan stride is s=3 while Δ=2. After the + first handover the local chunk is ``new[2:]`` and its trigger is local row + 1; the server-cache offset must still be 2 + 1 = 3. + """ + + runtime, _simulator, client, _clock = _harness( + mode="rtc", + chunk_steps=8, + overlap_steps=5, + inference_delay_steps=2, + execution_horizon=3, + prefix_attention_schedule="exp", + max_guidance_weight=5.0, + ) + + _run(runtime, max_policy_steps=3) + + rtc_blocks = [_async_block(request)["rtc"] for request in client.requests] + assert [block["cache_start_index"] for block in rtc_blocks] == [0, 3, 3] + assert [block["inference_delay"] for block in rtc_blocks] == [0, 2, 2] + + +# ============================================================================ +# stale observation capture +# ============================================================================ +def test_request_uses_the_observation_captured_before_the_window() -> None: + runtime, _simulator, client, _clock = _harness(mode="rtc", chunk_steps=6, overlap_steps=3, inference_delay_steps=3) + + outcome = _run(runtime, max_policy_steps=2) + + captured = client.requests[1] + assert captured.metadata["observation_id"] == "obs-3" + assert captured.metadata["step_index"] == 3 + assert captured.state["observation.state"] == [3.0] + assert captured.images[0].data == b"frame-3" + + record = outcome.records[1] + assert record.captured_observation_id == "obs-3" + assert record.environment_steps_at_capture == 3 + assert record.environment_steps_at_send == 6 + assert record.trigger_chunk_index == 3 + + +# ============================================================================ +# Δ = 0 degenerates to the synchronous result +# ============================================================================ +def test_zero_delay_degenerates_to_the_synchronous_result() -> None: + sync_runtime, _sim, sync_client, _clock = _harness(mode="sync", chunk_steps=5) + sync_outcome = _run(sync_runtime, max_policy_steps=2) + + naive_runtime, _sim2, naive_client, _clock2 = _harness( + mode="naive", chunk_steps=5, overlap_steps=0, inference_delay_steps=0 + ) + naive_outcome = _run(naive_runtime, max_policy_steps=2) + + assert naive_outcome.to_episode_outcome() == sync_outcome.to_episode_outcome() + assert [request.state for request in naive_client.requests] == [ + dict(request.state) for request in sync_client.requests + ] + assert all("async" not in request.metadata for request in naive_client.requests) + + # RTC has the same interaction as sync when there is no overlap and no delay. + runtime, _simulator, client, _clock = _harness(mode="rtc", chunk_steps=5, overlap_steps=0, inference_delay_steps=0) + outcome = _run(runtime, max_policy_steps=2) + assert outcome.to_episode_outcome() == sync_outcome.to_episode_outcome() + assert [request.state for request in client.requests] == [dict(request.state) for request in sync_client.requests] + assert all("async" not in request.metadata for request in client.requests) + + +def test_sync_runtime_and_async_sync_mode_execute_identically() -> None: + runtime, _simulator, client, _clock = _harness(mode="sync", chunk_steps=4) + async_outcome = _run(runtime, max_policy_steps=3) + async_actions = list(runtime.simulator.executed) + + clock = FakeClock() + simulator = FakeSimulator() + reference_client = FakeInferenceClient(chunk_steps=4, clock=clock, latency_s=0.0) + reference = SimulationRuntime(simulator, reference_client, mapper=FakeMapper(), chunk_steps=4) + reference_outcome = reference.run(instruction=INSTRUCTION, max_policy_steps=3, seed=7) + reference.close() + + assert async_outcome.to_episode_outcome() == reference_outcome + assert async_actions == list(simulator.executed) + assert [request.state for request in client.requests] == [ + dict(request.state) for request in reference_client.requests + ] + + +# ============================================================================ +# measured latency is reported, never substituted for Δ +# ============================================================================ +def test_measured_latency_is_reported_but_not_used_as_the_modeled_delay() -> None: + runtime, _simulator, client, _clock = _harness( + mode="rtc", + chunk_steps=6, + overlap_steps=3, + inference_delay_steps=2, + latency_s=0.5, + control_hz=10.0, + ) + + outcome = _run(runtime, max_policy_steps=2) + + record = outcome.records[1] + assert record.measured_latency_s == pytest.approx(0.5) + assert record.latency_steps == 5 # 0.5 s at 10 Hz + assert record.inference_delay_steps == 2 # Δ, not the measurement + assert record.environment_steps_at_send - record.environment_steps_at_capture == 2 + assert _async_block(client.requests[1])["rtc"]["inference_delay"] == 2 + assert outcome.latency_steps[1] == 5 + assert outcome.inference_delay_steps == (0, 2) + + assert runtime.latency.average() == pytest.approx(0.5) + + +def test_latency_steps_helper_and_validation() -> None: + assert latency_steps(0.25, control_hz=10.0) == 2 + assert latency_steps(0.05, control_hz=10.0) == 0 + assert latency_steps(-1.0, control_hz=10.0) == 0 + with pytest.raises(ValueError): + latency_steps(0.1, control_hz=0.0) + with pytest.raises(TypeError): + latency_steps(0.1, control_hz="fast") # type: ignore[arg-type] + + +# ============================================================================ +# termination +# ============================================================================ +def test_episode_that_terminates_inside_the_inference_window() -> None: + runtime, _simulator, client, _clock = _harness( + mode="rtc", + chunk_steps=5, + overlap_steps=2, + inference_delay_steps=2, + terminate_at=4, + ) + + outcome = _run(runtime, max_policy_steps=5) + runtime.close() + + assert outcome.terminated is True + assert outcome.truncated is False + assert outcome.policy_steps == 1 + assert outcome.requests == 1 + assert outcome.environment_steps == 4 + # Termination keeps the trigger out of the logical episode count. The + # physical request is submitted before the delay window starts, though, so + # cancellation can race with an already-running inference worker. + assert 1 <= len(client.requests) == len(runtime.records) <= 2 + assert 1 <= len(outcome.records) <= len(runtime.records) + assert outcome.records == tuple(runtime.records[: len(outcome.records)]) + if len(client.requests) == 2: + assert client.requests[1].step_id == 1 + assert client.requests[1].metadata["async"]["rtc"]["inference_delay"] == 2 + assert runtime.records[1].request_index == 1 + assert runtime.records[1].trigger is True + + +def test_episode_that_hits_max_policy_steps() -> None: + runtime, _simulator, _client, _clock = _harness(mode="rtc", chunk_steps=5, overlap_steps=2, inference_delay_steps=2) + + outcome = _run(runtime, max_policy_steps=3) + + assert outcome.terminated is False and outcome.truncated is False + assert outcome.policy_steps == 3 + assert outcome.requests == 3 + assert outcome.inference_delay_steps == (0, 2, 2) + # Bootstrap (5 steps), then a request every H - O = 3 steps at the trigger. + assert outcome.environment_steps == 9 + + +def test_truncated_step_stops_the_episode() -> None: + runtime, _simulator, _client, _clock = _harness( + mode="naive", chunk_steps=4, overlap_steps=2, inference_delay_steps=1, truncate_at=2 + ) + + outcome = _run(runtime, max_policy_steps=4) + + assert outcome.truncated is True + assert outcome.environment_steps == 2 + assert outcome.policy_steps == 1 + + +# ============================================================================ +# validation +# ============================================================================ +def test_config_rejects_invalid_values() -> None: + with pytest.raises(ValueError, match="chunk_steps"): + AsyncSimulationConfig(mode="sync", chunk_steps=0) + with pytest.raises(ValueError, match="mode"): + AsyncSimulationConfig(mode="teleport") + with pytest.raises(ValueError, match="inference_delay_steps"): + AsyncSimulationConfig(mode="rtc", chunk_steps=5, overlap_steps=2, inference_delay_steps=-1) + with pytest.raises(ValueError, match="remaining at the trigger"): + AsyncSimulationConfig(mode="rtc", chunk_steps=5, overlap_steps=2, inference_delay_steps=3) + with pytest.raises(ValueError, match="adopted chunk"): + # Δ <= overlap passes, but each RTC-adopted chunk is only H - Δ long. + AsyncSimulationConfig(mode="rtc", chunk_steps=5, overlap_steps=3, inference_delay_steps=3) + with pytest.raises(ValueError, match="overlap_steps"): + AsyncSimulationConfig(mode="sync", chunk_steps=5, overlap_steps=1) + with pytest.raises(ValueError, match="overlap_steps"): + AsyncSimulationConfig(mode="rtc", chunk_steps=5, overlap_steps=6, inference_delay_steps=0) + with pytest.raises(ValueError, match="execution_horizon"): + AsyncSimulationConfig(mode="rtc", execution_horizon=-1) + with pytest.raises(ValueError, match="prefix_attention_schedule"): + AsyncSimulationConfig(mode="rtc", prefix_attention_schedule="cosine") + with pytest.raises(ValueError, match="control_hz"): + AsyncSimulationConfig(control_hz=0.0) + + +def test_modes_are_the_documented_three() -> None: + assert ASYNC_SIMULATION_MODES == ("sync", "naive", "rtc") + + +def test_run_rejects_invalid_arguments() -> None: + runtime, _simulator, _client, _clock = _harness(mode="sync", chunk_steps=3) + with pytest.raises(ValueError, match="instruction"): + runtime.run(instruction=" ", max_policy_steps=1) + with pytest.raises(ValueError, match="max_policy_steps"): + runtime.run(instruction=INSTRUCTION, max_policy_steps=0) + + +def test_delay_is_a_modeled_parameter_not_a_wall_clock_sleep() -> None: + """The episode advances by simulator steps only; the clock is never slept on.""" + + runtime, simulator, _client, clock = _harness( + mode="naive", chunk_steps=4, overlap_steps=2, inference_delay_steps=2, latency_s=0.0 + ) + + outcome = _run(runtime, max_policy_steps=2) + + assert outcome.environment_steps == 6 + assert clock.now == 0.0 # no latency configured, so no measured time elapsed + assert simulator.step_index == 6 diff --git a/tests/test_control_services.py b/tests/test_control_services.py index 1d9ebae..04ad52d 100644 --- a/tests/test_control_services.py +++ b/tests/test_control_services.py @@ -85,6 +85,12 @@ def test_host_control_contract_contains_only_dynamic_task_fields() -> None: assert "inference" not in payload +def test_host_control_contract_allows_zero_hz_to_disable_pacing() -> None: + request = replace(task_request(), control_hz=0.0) + + assert TaskRequest.from_payload(request.to_payload()).control_hz == 0.0 + + def test_host_control_contract_rejects_unknown_fields() -> None: payload = task_request().to_payload() payload["robot"] = {"id": "must-not-cross-this-boundary"} diff --git a/tests/test_libero_action_semantics.py b/tests/test_libero_action_semantics.py new file mode 100644 index 0000000..cd108e4 --- /dev/null +++ b/tests/test_libero_action_semantics.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from embodirun.simulators.libero.adapter import _action_array, _encode_frame +from embodirun.simulators.libero.config import LiberoConfig + + +@pytest.mark.parametrize( + ("gripper", "expected"), + [ + (-0.001, -1.0), + (0.0, 0.0), + (0.001, 1.0), + (-4.0, -1.0), + (4.0, 1.0), + ], +) +def test_libero_action_binarizes_gripper_like_official_eval( + gripper: float, + expected: float, +) -> None: + source = np.asarray([2.0, -2.0, 0.25, -0.25, 0.5, -0.5, gripper], dtype=np.float32) + + action = _action_array({"action": source}) + + np.testing.assert_array_equal( + action, + np.asarray([1.0, -1.0, 0.25, -0.25, 0.5, -0.5, expected], dtype=np.float32), + ) + # Formatting an action must not mutate the policy result kept for RTC + # prefix bookkeeping or diagnostics. + assert source[-1] == pytest.approx(gripper) + + +def test_libero_png_transport_is_lossless() -> None: + cv2 = pytest.importorskip("cv2") + image = np.asarray( + [ + [[0, 17, 255], [64, 91, 128]], + [[33, 127, 231], [255, 3, 79]], + ], + dtype=np.uint8, + ) + + frame = _encode_frame("observation.images.image", image, encoding="png") + decoded_bgr = cv2.imdecode(np.frombuffer(frame.data, dtype=np.uint8), cv2.IMREAD_COLOR) + + assert frame.mime_type == "image/png" + np.testing.assert_array_equal(decoded_bgr[..., ::-1], image) + + +def test_libero_config_rejects_unknown_image_encoding() -> None: + with pytest.raises(ValueError, match="image_encoding"): + LiberoConfig("libero", "libero_spatial", 0, image_encoding="webp") diff --git a/tests/test_sglang_async_wire.py b/tests/test_sglang_async_wire.py new file mode 100644 index 0000000..660014a --- /dev/null +++ b/tests/test_sglang_async_wire.py @@ -0,0 +1,86 @@ +"""SGLang action payloads preserve the shared async-inference contract.""" + +from __future__ import annotations + +import json + +from embodirun.application.async_inference.contracts import ( + RTCAsyncRequest, + build_async_metadata, +) +from embodirun.model_services import ImagePayload, PolicyObservation +from embodirun.model_services.backends.sglang.http import SglangHttpClient +from embodirun.model_services.protocols.http import HttpResponse + + +class _Transport: + def __init__(self) -> None: + self.posts: list[dict] = [] + + def request(self, method, url, *, headers, body, timeout_s, maximum_bytes): + del headers, timeout_s, maximum_bytes + payload = json.loads(body) if body is not None else None + if method == "POST": + self.posts.append(payload) + response = { + "id": "request-1", + "data": [{"action": {"values": [[0.1, 0.2]]}}], + } + else: + response = {"async_inference": {"rtc": True}} + return HttpResponse(200, {"content-type": "application/json"}, json.dumps(response).encode()) + + +def test_sglang_client_forwards_rtc_metadata() -> None: + transport = _Transport() + client = SglangHttpClient("http://sglang", transport=transport, output_action_dim=2) + session = client.open_session(robot_id="robot-1", action_space="pi05") + observation = PolicyObservation( + session_id=session.session_id, + request_id="request-1", + step_id=0, + instruction="pick up the block", + state={"observation.state": [0.0, 1.0]}, + images=(ImagePayload("front", "image/png", b"png"),), + metadata=build_async_metadata( + rtc=RTCAsyncRequest(None, 3, 8, cache_start_index=10), + ), + ) + + client.step(observation) + + sent = transport.posts[0]["input"]["observation"] + assert sent["metadata"]["async"]["schema"] == "embodiinfer.async.v1" + assert sent["metadata"]["async"]["rtc"]["inference_delay"] == 3 + assert sent["metadata"]["async"]["rtc"]["cache_start_index"] == 10 + + +def test_sglang_client_maps_both_libero_cameras_to_checkpoint_keys() -> None: + transport = _Transport() + client = SglangHttpClient( + "http://sglang", + transport=transport, + output_action_dim=2, + image_keys={ + "observation.images.image": "image", + "observation.images.wrist_image": "image2", + }, + ) + session = client.open_session(robot_id="robot-1", action_space="pi05") + observation = PolicyObservation( + session_id=session.session_id, + request_id="request-1", + step_id=0, + instruction="pick up the block", + state={"observation.state": [0.0, 1.0]}, + images=( + ImagePayload("observation.images.image", "image/png", b"front"), + ImagePayload("observation.images.wrist_image", "image/png", b"wrist"), + ), + ) + + client.step(observation) + + images = transport.posts[0]["input"]["observation"]["images"] + assert set(images) == {"image", "image2"} + assert images["image"]["base64"] != images["image2"]["base64"] diff --git a/tests/test_sglang_http.py b/tests/test_sglang_http.py index fbfd651..2f4f862 100644 --- a/tests/test_sglang_http.py +++ b/tests/test_sglang_http.py @@ -50,6 +50,201 @@ def test_sglang_environment_preserves_or_discovers_toolkit(tmp_path, monkeypatch assert sglang.os.environ.get("CUDA_PATH") == ("/selected/cuda" if selection == "CUDA_PATH" else None) +def _sglang_pi05_integration_module(): + pytest.importorskip("sglang.multimodal_gen") + import importlib + + from embodirun.services.inference.adapters.sglang.pi05 import LeRobotPi05Pipeline + + return importlib.import_module(LeRobotPi05Pipeline.__module__) + + +def test_sglang_soft_rtc_uses_the_denoiser_jacobian(): + torch = pytest.importorskip("torch") + + integration = _sglang_pi05_integration_module() + x = torch.ones(1, 2, 1) + previous = torch.zeros_like(x) + weights = torch.ones(2) + + guided = integration._rtc_velocity( + lambda value: value * 0.5, + x, + torch.tensor(0.5), + previous, + weights, + 10.0, + ) + + # clean=0.75*x, error=-0.75 and d(clean)/d(x)=0.75. The RTC + # strength at t=0.5 is 2, so a real VJP yields 0.5 + 2*0.5625. + torch.testing.assert_close(guided, torch.full_like(x, 1.625), rtol=0, atol=0) + # A detached denoiser would incorrectly use correction=error and yield 2. + assert not torch.equal(guided, torch.full_like(x, 2.0)) + + +def test_sglang_soft_rtc_rejects_a_detached_denoiser(): + torch = pytest.importorskip("torch") + + integration = _sglang_pi05_integration_module() + x = torch.ones(1, 2, 1) + with pytest.raises(RuntimeError, match="detached velocity"): + integration._rtc_velocity( + lambda value: (value * 0.5).detach(), + x, + torch.tensor(0.5), + torch.zeros_like(x), + torch.ones(2), + 10.0, + ) + + +def test_sglang_soft_rtc_materializes_inference_mode_timestep(): + torch = pytest.importorskip("torch") + + integration = _sglang_pi05_integration_module() + value = torch.ones(1, 2, 1) + with torch.inference_mode(): + inference_time = torch.tensor(0.5) + assert inference_time.is_inference() + + guided = integration._rtc_velocity( + lambda current: current * 0.5, + value, + inference_time, + torch.zeros_like(value), + torch.ones(2), + 10.0, + ) + + torch.testing.assert_close(guided, torch.full_like(value, 1.625), rtol=0, atol=0) + + +def test_sglang_soft_rtc_stops_device_diagnostics_after_two_requests(monkeypatch): + """Steady-state RTC must not synchronize CUDA just for diagnostics.""" + + torch = pytest.importorskip("torch") + + integration = _sglang_pi05_integration_module() + integration._RTC_TELEMETRY_REQUESTS = 2 + + def unexpected_isfinite(*args, **kwargs): + del args, kwargs + raise AssertionError("steady-state RTC diagnostics touched torch.isfinite") + + monkeypatch.setattr(integration.torch, "isfinite", unexpected_isfinite) + try: + value = torch.ones(1, 2, 1) + guided = integration._rtc_velocity( + lambda current: current * 0.5, + value, + torch.tensor(0.5), + torch.zeros_like(value), + torch.ones(2), + 10.0, + telemetry=True, + ) + finally: + integration._RTC_TELEMETRY_REQUESTS = 0 + + torch.testing.assert_close(guided, torch.full_like(value, 1.625), rtol=0, atol=0) + + +def test_sglang_soft_rtc_copies_prefix_cache_out_of_inference_mode(): + torch = pytest.importorskip("torch") + + pytest.importorskip("sglang.multimodal_gen") + from sglang.multimodal_gen.runtime.vla.prefix_cache import ( + PrefixContext, + VLADensePrefixCache, + ) + + integration = _sglang_pi05_integration_module() + with torch.inference_mode(): + keys = torch.ones(1, 1, 2, 2) + values = torch.ones(1, 1, 2, 2) + masks = torch.ones(1, 2, dtype=torch.bool) + context = PrefixContext( + past_key_values=VLADensePrefixCache([(keys, values, None)]), + prefix_pad_masks=masks, + prefix_len=2, + ) + + normal = integration._rtc_autograd_prefix_context(context) + + normal_keys, normal_values, _ = normal.past_key_values[0] + assert keys.is_inference() and values.is_inference() and masks.is_inference() + assert not normal_keys.is_inference() + assert not normal_values.is_inference() + assert not normal.prefix_pad_masks.is_inference() + torch.testing.assert_close(normal_keys, keys) + torch.testing.assert_close(normal_values, values) + + +def test_sglang_soft_rtc_unwraps_core_no_grad_and_uses_native_geglu(): + torch = pytest.importorskip("torch") + + integration = _sglang_pi05_integration_module() + + class FakeFusedActivation(torch.nn.Module): + def forward(self, value): + return value.detach() * 3.0 + + def forward_native(self, value): + return value * 3.0 + + activation = FakeFusedActivation() + layer = SimpleNamespace(mlp=SimpleNamespace(act_fn=activation)) + + class FakeCore: + def __init__(self): + expert = SimpleNamespace(model=SimpleNamespace(layers=[layer])) + self.paligemma_with_expert = SimpleNamespace(gemma_expert=expert) + + @torch.no_grad() + def denoise_step( + self, + prefix_pad_masks, + past_key_values, + x, + timestep, + prefix_full_attention, + *, + action_position_offset=0, + ): + del prefix_pad_masks, past_key_values, timestep, prefix_full_attention + assert action_position_offset == 0 + return layer.mlp.act_fn(x) + + core = FakeCore() + model = SimpleNamespace(action_expert=SimpleNamespace(core_model=core)) + prefix = SimpleNamespace( + prefix_pad_masks=torch.ones(1, 1, dtype=torch.bool), + past_key_values=object(), + layout={"full_attention": True}, + ) + value = torch.ones(1, 2, 1, requires_grad=True) + + with torch.inference_mode(): + inference_timestep = torch.tensor([0.5]) + assert inference_timestep.is_inference() + + original_denoise_step = FakeCore.denoise_step.__wrapped__ + + def checked_denoise_step(self, *args, **kwargs): + assert not args[3].is_inference() + return original_denoise_step(self, *args, **kwargs) + + FakeCore.denoise_step.__wrapped__ = checked_denoise_step + with torch.enable_grad(): + velocity = integration._grad_enabled_denoise(model, prefix, value, inference_timestep) + + assert velocity.requires_grad + assert velocity.grad_fn is not None + torch.testing.assert_close(torch.autograd.grad(velocity.sum(), value)[0], torch.full_like(value, 3.0)) + assert isinstance(layer.mlp.act_fn, integration._AutogradSafeGeluAndMul) + + def test_sglang_lerobot_statistics_preserve_mean_std_math(): pytest.importorskip("sglang.multimodal_gen") import torch diff --git a/tests/test_so101_pi05_runtime.py b/tests/test_so101_pi05_runtime.py index 252cf6b..acc3497 100644 --- a/tests/test_so101_pi05_runtime.py +++ b/tests/test_so101_pi05_runtime.py @@ -146,6 +146,50 @@ def test_control_runtime_plays_action_chunk_at_control_rate() -> None: assert clock.sleeps == [0.05, 0.05] +def test_control_runtime_adds_inference_result_delay_before_actions() -> None: + robot = FakeRobot() + client = FakeClient() + clock = FakeClock() + runtime = ControlRuntime( + robot, + client, + instruction="pick up the block", + mapper=ThreeActionMapper(), + chunk_steps=3, + control_hz=50.0, + inference_delay_s=0.2, + monotonic=clock.monotonic, + sleep=clock.sleep, + ) + + runtime.step((CameraFrame("observation.images.front", "image/jpeg", b"jpeg"),)) + + assert [action.values["index"] for action in robot.actions] == [0, 1, 2] + assert clock.sleeps == pytest.approx([0.2, 0.02, 0.02]) + + +def test_control_runtime_zero_control_hz_disables_software_pacing() -> None: + robot = FakeRobot() + client = FakeClient() + clock = FakeClock() + runtime = ControlRuntime( + robot, + client, + instruction="pick up the block", + mapper=ThreeActionMapper(), + chunk_steps=3, + control_hz=0.0, + monotonic=clock.monotonic, + sleep=clock.sleep, + ) + + runtime.step((CameraFrame("observation.images.front", "image/jpeg", b"jpeg"),)) + + assert [action.values["index"] for action in robot.actions] == [0, 1, 2] + assert clock.sleeps == [] + assert runtime.action_period_s == 0.0 + + def test_control_runtime_executes_only_requested_chunk_steps() -> None: robot = FakeRobot() client = FakeClient()