diff --git a/README.md b/README.md
index 722b80f6..b6663964 100644
--- a/README.md
+++ b/README.md
@@ -284,6 +284,23 @@ See [examples/README.md](examples/README.md) for the example runner and baseline
The published site provides the same English and Chinese information architecture. Model-specific instructions are
generated from registered example READMEs under **Models & Cookbook**. The build also publishes `/TeleFuser/llms.txt`.
+- [docs/en/blog/index.md](docs/en/blog/index.md): optimization design, profiling evidence, results, and related work
+- [docs/en/service.md](docs/en/service.md): REST serving, task APIs, OpenAI-compatible APIs
+- [docs/en/stream_server.md](docs/en/stream_server.md): LiveKit streaming, session APIs, data topics, and deployment
+- [docs/en/stream_scheduler.md](docs/en/stream_scheduler.md): actor-based stage scheduling, backpressure, lifecycle, metrics, and LingBot placement
+- [docs/en/vla.md](docs/en/vla.md): semantic VLA action contracts, embodiment mapping, sessions, runtime safety, and simulator adapters
+- [docs/en/parallel.md](docs/en/parallel.md): distributed inference architecture
+- [docs/en/communication.md](docs/en/communication.md): collectives, CUDA IPC, synchronization, and transport lifecycle
+- [docs/en/latent_cache.md](docs/en/latent_cache.md): CacheSeek latent cache integration
+- [docs/en/feature_cache.md](docs/en/feature_cache.md): `AdaTaylorCache`
+- [docs/en/model_loading.md](docs/en/model_loading.md): model loading patterns
+- [docs/en/attention.md](docs/en/attention.md): attention backends and configuration
+- [docs/en/torch_compile_compatibility.md](docs/en/torch_compile_compatibility.md): compile-related constraints
+- [docs/en/adding_new_model.md](docs/en/adding_new_model.md): integrating new models
+- [docs/en/adding_new_example.md](docs/en/adding_new_example.md): authoring examples and pipeline contracts
+- [docs/en/abot_world.md](docs/en/abot_world.md): ABot-World single-GPU interactive pipeline, controls, and tests
+- [examples/lingbot_vla_v2/README.md](examples/lingbot_vla_v2/README.md): LingBot-VLA v2 inference, structured service, parity, and validation boundaries
+- [examples/swiftvr/README.md](examples/swiftvr/README.md): SwiftVR checkpoint loading, streaming usage, performance, and acceleration options
## Known Limitations
diff --git a/docs/en/vla.md b/docs/en/vla.md
new file mode 100644
index 00000000..4801f03f
--- /dev/null
+++ b/docs/en/vla.md
@@ -0,0 +1,173 @@
+# VLA Action Integration
+
+TeleFuser keeps model lifecycle, GPU placement, request cancellation, and monitoring in the existing pipeline and
+service layers. VLA integration adds semantic action boundaries around those facilities; it does not create a second
+model server.
+
+## Data Flow
+
+```text
+robot observation
+ -> EmbodimentAdapter.encode_observation
+ -> VLAPolicy.predict
+ -> ModelActionChunk
+ -> EmbodimentAdapter.decode_actions
+ -> RobotActionChunk
+ -> ChunkExecutor (horizon, age, safety)
+ -> SimulatorAdapter.execute
+```
+
+Tensor width is never used to infer action meaning. Every chunk carries an `ActionSpaceSpec` with representation,
+ordered dimension names, units, frame, control rate, and normalization identity. Session opening rejects a policy and
+embodiment pair whose action spaces are semantically incompatible.
+
+Each embodiment also declares an `ObservationSpaceSpec` with ordered state dimensions, named image dtype, layout,
+and channel requirements, plus the timestamp unit and clock domain. Image height and width remain dynamic so
+simulators can choose their native camera resolution.
+
+The public modules are:
+
+- `telefuser.vla.contracts`: observations, requests, capabilities, action spaces, and model/robot chunks.
+- `telefuser.vla.policy`: the model-facing `VLAPolicy` protocol.
+- `telefuser.vla.embodiment`: observation encoding and model-to-robot action mapping.
+- `telefuser.vla.registry`: explicit registration of already-loaded policies and embodiments.
+- `telefuser.vla.session`: transport-neutral OPEN, PREDICT, RESET, and CLOSE lifecycle.
+- `telefuser.vla.serialization`: versioned JSON/Base64 wire formats for action and observation spaces, observations,
+ and chunks.
+- `telefuser.vla.runtime`: scheduling, deterministic chunk state, action trimming, age checks, and safety policies.
+- `telefuser.integrations.sim`: simulator protocol and the dependency-free RoboTwin callback adapter.
+- `telefuser.service.vla_session`: the additive generic VLA WebSocket application factory.
+- `telefuser.service.vla_replica`: worker-local OPEN, PREDICT, RESET, and CLOSE dispatch for pipeline replicas.
+- `telefuser.client.AsyncVLAClient`: remote session client with concurrent request correlation.
+
+## LingBot-VLA v2 Compatibility
+
+`LingBotVlaV2VLAPolicy` wraps `LingBotVlaV2Pipeline`; the pipeline's existing tensor input and return types are not
+changed. It labels the normalized canonical `[T,55]` result as a `ModelActionChunk`.
+
+`RobotWinProfile` is the first `EmbodimentAdapter`. It preserves the existing normalization code and maps a canonical
+chunk to an absolute-position `[H,14]` `RobotActionChunk` in the declared dual-arm joint order. Its model and robot
+action spaces are exported as `LINGBOT_VLA_V2_ACTION_SPACE` and `ROBOTWIN_ACTION_SPACE`.
+
+The generic VLA WebSocket is the primary online integration path. The existing LingBot RoboTwin WebSocket endpoint
+is a compatibility adapter for unmodified upstream `WebsocketClientPolicy` clients and calls the same policy,
+embodiment, session, and runtime path internally. Its URL, MessagePack request fields, metadata frame, response fields,
+reset behavior, and latest-wins
+scheduler behavior remain compatible. The old model-specific scheduler import is retained as an alias to the common
+runtime scheduler.
+
+The standalone endpoint owns one resident pipeline, so each connection session is inherently pinned to that policy
+instance. The shared HTTP structured-task route remains unchanged and continues to return the existing canonical
+JSON result.
+
+## Session Lifecycle
+
+Create and register loaded components, then open a typed session:
+
+```python
+from telefuser.pipelines.lingbot_vla_v2 import LingBotVlaV2VLAPolicy, RobotWinProfile
+from telefuser.vla import VLARegistry, VLASessionManager
+
+registry = VLARegistry()
+registry.register_policy("lingbot-vla-v2", LingBotVlaV2VLAPolicy(pipeline))
+registry.register_embodiment("robotwin", RobotWinProfile.default())
+
+sessions = VLASessionManager(registry)
+session = sessions.open(
+ "connection-1",
+ model_id="lingbot-vla-v2",
+ embodiment_id="robotwin",
+ episode_id="episode-1",
+ execute_horizon=8,
+)
+robot_chunk = session.predict(robot_observation, "pick up the block", sequence_id=0, seed=7)
+sessions.reset("connection-1")
+sessions.close("connection-1")
+```
+
+Sequence IDs must strictly increase between resets. A policy result must echo the request episode, sequence, and
+observation timestamp. `ChunkExecutor` can reject stale observations when `max_observation_age_ns` and a same-clock
+`now_ns` are supplied. Network round-trip deadlines must still be enforced by the simulator client because monotonic
+clocks on separate machines are not comparable.
+
+## Generic WebSocket Protocol
+
+`create_vla_session_app` exposes `/v1/vla/session` for an in-process registry, while
+`create_pipeline_pool_vla_session_app` exposes the same protocol through worker-local sessions pinned by
+`PipelinePool`. Neither factory changes an existing TeleFuser service route. The server
+first sends `HELLO` with protocol version `1.0`, wire encoding, supported operations, registered component IDs, and
+payload limits. A client then uses `OPEN`, ordered `PREDICT`, `RESET`, and `CLOSE`. `OPEN` can declare the expected
+robot action and observation spaces; semantic incompatibility is rejected before inference.
+
+The protocol returns machine-readable error codes for malformed messages, unsupported versions, unknown components
+or sessions, contract mismatches, superseded work, out-of-order or expired observations, request timeout, unavailable
+sessions, and replica failure. Tensor payloads are dense, typed, shape-checked Base64 data inside bounded JSON
+messages. This is the stable interoperability format; the LingBot-RoboTwin compatibility endpoint continues to use
+its existing MessagePack format. New model and simulator integrations must target the generic protocol rather than
+add behavior to the model-specific compatibility endpoint.
+
+`request_ttl_ms` is measured with server monotonic time and bounds inference delivery. Observation age is checked
+separately using `observation_timestamp_ns` and `observation_clock_now_ns`, which must come from the same clock domain.
+The server permits inference and observation delivery to overlap. Per session it runs one replica request and retains
+only the newest waiting observation; replaced requests return `superseded`. `RESET` and `CLOSE` are barriers: they
+invalidate earlier tickets and wait for any non-cancellable replica call before acknowledging. A discarded stateful
+inference triggers policy reset before the retained observation runs. After a timed-out inference finishes, the
+session is reset before it accepts more prediction work.
+
+## Replica Session Leases
+
+`PipelinePool.open_session` removes one healthy replica from the ordinary request pool and binds it to a
+`session_id`. Calls through `acquire_session` always reach that replica and are serialized per session.
+`close_session` returns a live replica to ordinary capacity. Replica exit removes the binding and raises
+`ReplicaDeadError`; pool shutdown marks sessions as closing, waits for active session calls, clears every binding, and
+then stops replicas.
+
+This API is additive. Existing HTTP work continues to use `PipelinePool.acquire()`, whose allocation behavior and
+callers are unchanged. When `PipelinePool.start_all(..., vla_provider_factory="get_vla_provider")` is requested, each
+worker constructs its own provider around the already-loaded pipeline and handles VLA lifecycle RPC. The provider is
+optional; pipeline files and callers that do not enable it keep the original task and shutdown protocols.
+
+## Chunk State Machine
+
+`ActionChunkStateMachine` makes action lifecycle explicit: pending inference, ready buffer, execution, and terminal
+executed, superseded, expired, or rejected states. It supports latest-sequence admission, separate observation-age and
+network-TTL checks, configurable discard/retain handling after `execute_horizon`, reset, and deterministic hold/stop
+states after disconnect. Stateful policies must provide a recovery callback; it is invoked when an inference that may
+have mutated history is discarded.
+
+The generic server tracks inference admission through PENDING, READY, SUPERSEDED, EXPIRED, and REJECTED. It does not
+claim that a returned action was executed. `SimulatorChunkRuntime` runs on the simulator client and owns READY,
+EXECUTING, EXECUTED, HOLD, and STOP transitions while calling a `SimulatorAdapter` one action at a time.
+
+The bundled RoboTwin statistics do not declare an authoritative simulator control frequency, so both exported
+LingBot/RoboTwin specs use `control_hz=None`. The remote simulator must resolve its actual control period rather than
+guessing one in the inference process.
+
+## Simulator Boundary
+
+`RoboTwinSimulatorAdapter` has no RoboTwin, SAPIEN, Vulkan, or ROS dependency. The RTX-side process provides observe,
+execute, and reset callbacks; the adapter validates semantics and converts each action to contiguous CPU `float32`
+before invoking RoboTwin. This keeps simulator dependencies outside the H100 inference environment.
+
+Adding another simulator requires a new `SimulatorAdapter`; it must not add model-specific decoding. Adding another
+robot requires an `EmbodimentAdapter`. Adding another VLA requires a `VLAPolicy` wrapper around its maintained
+pipeline.
+
+## Service Boundary
+
+The generic WebSocket application remains an explicit factory rather than an automatically mounted `telefuser serve`
+route. The LingBot example supplies a standalone server that starts `PipelinePool` with the optional VLA provider.
+This keeps existing HTTP routing and every non-VLA pipeline unchanged. A real simulator still owns control timing,
+actuator feedback, and verification that each returned action was actually applied.
+
+For LingBot-VLA v2, the standalone generic WebSocket is the primary continuous-control entrypoint. The direct Python
+entrypoint remains the reference/offline baseline, and the native HTTP structured service remains for existing
+TeleFuser callers. The RoboTwin MessagePack server is a legacy compatibility entrypoint only; it can be removed after
+all upstream clients migrate to `/v1/vla/session`.
+
+The direct inference CLI and structured HTTP task remain separate because they provide reference and batch workflows,
+not simulator session transports. The legacy RoboTwin MessagePack endpoint should be removed only after the RTX client
+passes generic-protocol action delivery, reset, timeout, reconnect, and latest-wins parity checks.
+
+No new dependencies, environment variables, shared model configuration fields, CLI options, HTTP schemas, or
+existing service routes are introduced by this semantic and transport layer.
diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md
index e29c99db..bea2d696 100644
--- a/examples/lingbot_vla_v2/README.md
+++ b/examples/lingbot_vla_v2/README.md
@@ -25,10 +25,12 @@ The parity reference uses [Robbyant/lingbot-vla-v2](https://github.com/Robbyant/
| CUDA Graph | Supported | Dynamic eager prefix with an opt-in fixed-shape action-denoising graph |
| Quantization | Partial | Profile-specific release status; see Configuration and Performance |
| Native server API | Supported | Asynchronous structured task API and `TFClient` |
-| RoboTwin policy protocol | Supported | Standalone persistent MessagePack WebSocket service |
+| RoboTwin policy protocol | Compatibility | Legacy MessagePack endpoint for unmodified upstream clients |
| Request replicas | Supported | One complete policy copy per GPU |
| Single-policy FSDP, TP, or PP | Unsupported | The integration does not split one policy across GPUs |
| RoboTwin action mapping | Supported | Unnormalizes canonical output to absolute-position `50 x 14` chunks |
+| Semantic VLA contract | Supported | Model and robot action spaces are explicit; see [VLA Action Integration](../../docs/en/vla.md) |
+| Generic VLA session server | Preferred | Versioned JSON WebSocket with replica-affine sessions |
## Requirements
@@ -190,6 +192,19 @@ Compare a deterministic quantized capture with the corresponding TeleFuser BF16
## Serving
+The generic VLA session server is the only recommended online path for new simulator integrations. The four current
+entrypoints share one `LingBotVlaV2Pipeline`; they are access modes, not separate model implementations:
+
+| Entry point | Role | Recommendation |
+| --- | --- | --- |
+| `lingbot_vla_v2_inference.py` | Direct Python reference and offline baseline | Keep for regression/debugging |
+| `lingbot_vla_v2_native_service.py` via `telefuser serve` | Native HTTP structured requests | Keep for TeleFuser compatibility |
+| `lingbot_vla_v2_vla_server.py` | Stateful generic VLA WebSocket | **Primary simulator path** |
+| `lingbot_vla_v2_robotwin_server.py` | Upstream RoboTwin MessagePack compatibility | Legacy; remove after client migration |
+
+For a production or simulation deployment, start only the generic WebSocket server. The other entries remain for
+baseline comparison and backward compatibility and do not change the model or session implementation.
+
Start the native structured service:
```bash
@@ -228,11 +243,37 @@ CUDA_VISIBLE_DEVICES=0,1 TF_MODEL_ZOO_PATH=/path/to/model_zoo \
This creates one complete policy per GPU; it does not enable tensor or pipeline parallelism within a policy.
-### RoboTwin Policy Server
+### Generic VLA Session Server
+
+Start the preferred additive WebSocket service without mounting routes into `telefuser serve`:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 TF_MODEL_ZOO_PATH=/path/to/model_zoo \
+ .venv-vla/bin/python -m examples.lingbot_vla_v2.lingbot_vla_v2_vla_server \
+ --parallelism 1 --num-replicas 1 --host 0.0.0.0 --port 8000
+```
+
+The service exposes `GET /healthz` and `/v1/vla/session`. Each `OPEN` reserves one pipeline replica for that session;
+`PREDICT`, `RESET`, and `CLOSE` are sent to the same worker-local `VLASession`. Closing or disconnecting releases the
+replica. The H100 cuDNN SDPA guard is applied inside each LingBot worker before model loading.
+
+This protocol returns semantic `RobotActionChunk` values. The simulator process should deserialize the chunk, pass it
+to `SimulatorChunkRuntime`, and then execute it through its own `SimulatorAdapter`. The inference server tracks
+pending/ready/expired inference only and does not infer simulator execution. The simulator side can call
+`execute_ready_with_report()` (or `execute_ready_async()` from an async client) to obtain a small transport-neutral
+report containing `executed`, `failed`, or `no_action`, the chunk sequence, and the executed step count. The report can
+be forwarded on an existing control channel; it is intentionally not a new WebSocket operation.
-The standalone policy server implements the persistent MessagePack WebSocket protocol used by the upstream
-`WebsocketClientPolicy`. It is isolated from `telefuser serve`: no TeleFuser API routes, service schemas, or other
-model integrations are changed.
+For a minimal P0 check, use one fake adapter test for report status and one async test for non-blocking execution. A
+long-running soak, cross-machine clock comparison, and full simulator episode are not required to validate this
+runtime API.
+
+### Legacy RoboTwin Protocol Compatibility
+
+This compatibility server implements the persistent MessagePack WebSocket protocol used by the upstream
+`WebsocketClientPolicy`. Use it only when the upstream client cannot yet consume the generic VLA session protocol.
+New transport, scheduling, and simulator integrations belong on the generic VLA path. The compatibility endpoint is
+isolated from `telefuser serve`: no TeleFuser API routes, service schemas, or other model integrations are changed.
Install the protocol dependency in the TeleFuser inference environment:
@@ -328,6 +369,14 @@ profile, and returns absolute-position actions in raw RoboTwin order. `--use-len
start with 50 for upstream-equivalent open-loop execution. The adapter accepts episode reset messages but deliberately
rejects runtime checkpoint switching.
+Internally this compatibility endpoint uses the shared `VLAPolicy`, `EmbodimentAdapter`, `VLASessionManager`, and
+`ChunkExecutor` contracts. The wire protocol and the existing `LingBotVlaV2Pipeline` API remain unchanged. The
+declared control rate is intentionally unresolved until the remote RoboTwin loop supplies its actual frequency.
+
+Keep this endpoint until the RTX client has passed end-to-end action delivery, reset, timeout, reconnect, and
+latest-wins parity checks through the generic protocol. After that migration, the compatibility module can be removed
+without changing the model pipeline or the generic VLA service.
+
For split-machine deployment, run the model endpoint and the repository-owned XPolicyLab proxy on the H100 inference
host. The proxy does not load a second model; it translates XPolicyLab observations to the direct TeleFuser protocol:
@@ -370,6 +419,33 @@ The base checkpoint remains marked `unverified_official_6b_base`. This endpoint
action mapping, transport, and simulator execution continuity; it does not establish RoboTwin task success without
an embodiment-validated checkpoint.
+## Local MuJoCo smoke simulation
+
+When the RTX RoboTwin workstation is unavailable, the generic VLA action path can be exercised locally with MuJoCo.
+The setup script installs MuJoCo into the existing TeleFuser `.venv` and stages EGL/OSMesa packages under the ignored
+`.venv-mujoco-libs` directory; it never runs `apt install` or changes system Python. The scene reads the existing
+RoboTwin ALOHA-Agilex URDF and meshes directly from `/data/RoboTwin` and adds a tabletop, cube, and three cameras.
+
+```bash
+bash examples/lingbot_vla_v2/setup_mujoco_local.sh
+
+# Physics + three-camera local smoke (no VLA model required)
+.venv/bin/python examples/lingbot_vla_v2/lingbot_vla_v2_mujoco.py \
+ --mode local --image-size 256 --execute-horizon 8 \
+ --output-dir work_dirs/lingbot_vla_v2/mujoco_local
+
+# Complete local simulator -> generic VLA WebSocket -> simulator loop
+.venv/bin/python examples/lingbot_vla_v2/lingbot_vla_v2_mujoco.py \
+ --mode websocket --server-url ws://127.0.0.1:18080/v1/vla/session \
+ --chunks 2 --execute-horizon 8 \
+ --output-dir work_dirs/lingbot_vla_v2/mujoco_websocket
+```
+
+The adapter uses a semantic 14-dimensional `absolute_qpos` contract, PD torque control, and the existing
+`SimulatorChunkRuntime`. `--mode local` validates model loading, rendering, action mapping, and execution without
+loading LingBot. `--mode websocket` additionally validates the live generic VLA session and requires a running VLA
+server. This is a chain/physics smoke test, not a RoboTwin task-success or checkpoint-quality evaluation.
+
## Validation
The repository includes strict upstream parity, runtime, quantization, structured-service, fault, and AIPerf
diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py
index 9ccb553a..711d0a98 100644
--- a/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py
+++ b/examples/lingbot_vla_v2/lingbot_vla_v2_inference.py
@@ -1,4 +1,8 @@
-"""Run the LingBot-VLA v2 base checkpoint with a RobotWin observation adapter."""
+"""Run the LingBot-VLA v2 base checkpoint with a RobotWin observation adapter.
+
+This is the reference/offline entrypoint for model output checks. Continuous
+simulator control should use the generic VLA WebSocket service instead.
+"""
from __future__ import annotations
diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_mujoco.py b/examples/lingbot_vla_v2/lingbot_vla_v2_mujoco.py
new file mode 100644
index 00000000..36cd24d7
--- /dev/null
+++ b/examples/lingbot_vla_v2/lingbot_vla_v2_mujoco.py
@@ -0,0 +1,283 @@
+"""Run the generic LingBot VLA loop against a local headless MuJoCo scene."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import click
+import numpy as np
+import torch
+from PIL import Image
+
+from telefuser.client import AsyncVLAClient
+from telefuser.integrations.sim import MuJoCoJointBinding, MuJoCoSimulatorAdapter
+from telefuser.pipelines.lingbot_vla_v2.robot_profile import (
+ ROBOTWIN_ACTION_ORDER,
+ ROBOTWIN_ACTION_SPACE,
+ ROBOTWIN_CAMERA_KEYS,
+ ROBOTWIN_OBSERVATION_SPACE,
+)
+from telefuser.vla import RobotActionChunk
+from telefuser.vla.runtime import ChunkStatus, SimulatorChunkRuntime
+
+REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
+DEFAULT_URDF = Path("/data/RoboTwin/assets/embodiments/aloha-agilex/urdf/arx5_description_isaac.urdf")
+LOCAL_MESA_LIBDIR = REPOSITORY_ROOT / ".venv-mujoco-libs/root/usr/lib/x86_64-linux-gnu"
+CAMERA_NAMES = {
+ ROBOTWIN_CAMERA_KEYS[0]: "telefuser_cam_high",
+ ROBOTWIN_CAMERA_KEYS[1]: "telefuser_cam_left",
+ ROBOTWIN_CAMERA_KEYS[2]: "telefuser_cam_right",
+}
+
+
+def _ensure_headless_runtime() -> None:
+ """Re-exec once so the dynamic loader sees repository-local OSMesa."""
+ library_path = str(LOCAL_MESA_LIBDIR)
+ current_paths = os.environ.get("LD_LIBRARY_PATH", "").split(":")
+ if os.environ.get("MUJOCO_GL") == "osmesa" and library_path in current_paths:
+ return
+ if not (LOCAL_MESA_LIBDIR / "libOSMesa.so.8").is_file():
+ raise RuntimeError(
+ "repository-local OSMesa is missing; run examples/lingbot_vla_v2/setup_mujoco_local.sh first"
+ )
+ environment = os.environ.copy()
+ environment["MUJOCO_GL"] = "osmesa"
+ environment["LD_LIBRARY_PATH"] = ":".join(path for path in (library_path, *current_paths) if path)
+ os.execve(sys.executable, [sys.executable, *sys.argv], environment)
+
+
+def _build_smoke_model(urdf_path: Path) -> Any:
+ """Add a small tabletop scene and three fixed cameras to the RoboTwin URDF."""
+ import mujoco
+
+ spec = mujoco.MjSpec.from_file(str(urdf_path))
+ spec.option.timestep = 0.002
+ spec.worldbody.add_light(name="telefuser_key_light", pos=[1.0, 0.0, 2.5], dir=[-0.3, 0.0, -1.0])
+ spec.worldbody.add_geom(
+ name="telefuser_floor",
+ type=mujoco.mjtGeom.mjGEOM_PLANE,
+ size=[3.0, 3.0, 0.05],
+ rgba=[0.24, 0.26, 0.28, 1.0],
+ )
+ spec.worldbody.add_geom(
+ name="telefuser_table",
+ type=mujoco.mjtGeom.mjGEOM_BOX,
+ pos=[0.7, 0.0, 0.66],
+ size=[0.55, 0.7, 0.05],
+ rgba=[0.45, 0.32, 0.22, 1.0],
+ friction=[1.0, 0.01, 0.001],
+ )
+ cube = spec.worldbody.add_body(name="telefuser_cube", pos=[0.65, 0.0, 0.77])
+ cube.add_freejoint(name="telefuser_cube_free")
+ cube.add_geom(
+ name="telefuser_cube_geom",
+ type=mujoco.mjtGeom.mjGEOM_BOX,
+ size=[0.05, 0.05, 0.05],
+ mass=0.1,
+ rgba=[0.85, 0.12, 0.08, 1.0],
+ friction=[1.0, 0.01, 0.001],
+ )
+ spec.worldbody.add_camera(
+ name=CAMERA_NAMES[ROBOTWIN_CAMERA_KEYS[0]],
+ pos=[2.2, 0.0, 1.55],
+ zaxis=[1.55, 0.0, 0.75],
+ fovy=55.0,
+ )
+ spec.worldbody.add_camera(
+ name=CAMERA_NAMES[ROBOTWIN_CAMERA_KEYS[1]],
+ pos=[1.15, 1.25, 1.15],
+ zaxis=[0.5, 0.95, 0.35],
+ fovy=60.0,
+ )
+ spec.worldbody.add_camera(
+ name=CAMERA_NAMES[ROBOTWIN_CAMERA_KEYS[2]],
+ pos=[1.15, -1.25, 1.15],
+ zaxis=[0.5, -0.95, 0.35],
+ fovy=60.0,
+ )
+ return spec.compile()
+
+
+def _joint_bindings() -> tuple[MuJoCoJointBinding, ...]:
+ return (
+ *(MuJoCoJointBinding(ROBOTWIN_ACTION_ORDER[index], (f"fl_joint{index + 1}",)) for index in range(6)),
+ MuJoCoJointBinding(ROBOTWIN_ACTION_ORDER[6], ("fl_joint7", "fl_joint8"), True),
+ *(MuJoCoJointBinding(ROBOTWIN_ACTION_ORDER[index + 7], (f"fr_joint{index + 1}",)) for index in range(6)),
+ MuJoCoJointBinding(ROBOTWIN_ACTION_ORDER[13], ("fr_joint7", "fr_joint8"), True),
+ )
+
+
+def _create_adapter(urdf_path: Path, image_size: int, steps_per_action: int) -> MuJoCoSimulatorAdapter:
+ model = _build_smoke_model(urdf_path)
+ return MuJoCoSimulatorAdapter(
+ model,
+ ROBOTWIN_ACTION_SPACE,
+ ROBOTWIN_OBSERVATION_SPACE,
+ _joint_bindings(),
+ camera_names=CAMERA_NAMES,
+ image_height=image_size,
+ image_width=image_size,
+ steps_per_action=steps_per_action,
+ )
+
+
+def _save_images(images: dict[str, Any], output_dir: Path | None) -> None:
+ if output_dir is None:
+ return
+ output_dir.mkdir(parents=True, exist_ok=True)
+ for name, image in images.items():
+ Image.fromarray(np.asarray(image)).save(output_dir / f"{name.rsplit('.', maxsplit=1)[-1]}.png")
+
+
+def _run_local_smoke(
+ adapter: MuJoCoSimulatorAdapter,
+ *,
+ execute_horizon: int,
+ output_dir: Path | None,
+) -> dict[str, Any]:
+ initial = adapter.reset()
+ targets = initial.state.values.repeat(execute_horizon, 1)
+ targets[:, 0] += torch.linspace(0.01, 0.06, execute_horizon)
+ targets[:, 7] -= torch.linspace(0.01, 0.04, execute_horizon)
+ targets[:, 6] = 0.5
+ targets[:, 13] = 0.5
+ chunk = RobotActionChunk(
+ actions=targets,
+ action_space=ROBOTWIN_ACTION_SPACE,
+ valid_length=execute_horizon,
+ observation_timestamp_ns=initial.state.timestamp_ns,
+ sequence_id=0,
+ episode_id="mujoco-local-smoke",
+ )
+ runtime = SimulatorChunkRuntime(
+ adapter,
+ ROBOTWIN_ACTION_SPACE,
+ chunk.episode_id,
+ execute_horizon=execute_horizon,
+ )
+ status = runtime.accept(chunk, observation_clock_now_ns=initial.state.timestamp_ns)
+ if status is not ChunkStatus.READY:
+ raise RuntimeError(f"local action chunk was not ready: {status.value}")
+ executed = runtime.execute_ready()
+ final = adapter.observe()
+ _save_images(dict(final.images), output_dir)
+ image_std = {name: float(np.asarray(image).std()) for name, image in final.images.items()}
+ if any(value <= 1.0 for value in image_std.values()):
+ raise RuntimeError(f"MuJoCo returned an empty or nearly uniform camera frame: {image_std}")
+ return {
+ "mode": "local",
+ "executed_actions": executed,
+ "simulation_time_s": final.metadata["simulation_time_s"],
+ "state_delta_l2": float(torch.linalg.vector_norm(final.state.values - initial.state.values)),
+ "image_shape": {name: list(np.asarray(image).shape) for name, image in final.images.items()},
+ "image_std": image_std,
+ "runtime_state": runtime.state.value,
+ }
+
+
+async def _run_websocket_loop(
+ adapter: MuJoCoSimulatorAdapter,
+ *,
+ server_url: str,
+ instruction: str,
+ seed: int,
+ chunks: int,
+ execute_horizon: int,
+ output_dir: Path | None,
+) -> dict[str, Any]:
+ episode_id = "mujoco-websocket-smoke"
+ runtime = SimulatorChunkRuntime(
+ adapter,
+ ROBOTWIN_ACTION_SPACE,
+ episode_id,
+ execute_horizon=execute_horizon,
+ )
+ observation = adapter.reset()
+ executed_total = 0
+ async with AsyncVLAClient(server_url) as client:
+ await client.open_session(
+ "mujoco-local",
+ model_id="lingbot-vla-v2",
+ embodiment_id="robotwin",
+ episode_id=episode_id,
+ execute_horizon=execute_horizon,
+ expected_robot_action_space=ROBOTWIN_ACTION_SPACE,
+ expected_robot_observation_space=ROBOTWIN_OBSERVATION_SPACE,
+ )
+ for sequence_id in range(chunks):
+ chunk = await client.predict(
+ "mujoco-local",
+ observation,
+ instruction,
+ sequence_id,
+ seed=seed,
+ )
+ status = runtime.accept(chunk, observation_clock_now_ns=observation.state.timestamp_ns)
+ if status is not ChunkStatus.READY:
+ raise RuntimeError(f"server action chunk was not ready: {status.value}")
+ executed_total += runtime.execute_ready()
+ observation = adapter.observe()
+ _save_images(dict(observation.images), output_dir)
+ return {
+ "mode": "websocket",
+ "chunks": chunks,
+ "executed_actions": executed_total,
+ "simulation_time_s": observation.metadata["simulation_time_s"],
+ "final_state": observation.state.values.tolist(),
+ "runtime_state": runtime.state.value,
+ }
+
+
+@click.command()
+@click.option("--mode", type=click.Choice(("local", "websocket")), default="local", show_default=True)
+@click.option("--urdf", "urdf_path", type=click.Path(path_type=Path, exists=True, dir_okay=False), default=DEFAULT_URDF)
+@click.option("--server-url", default="ws://127.0.0.1:18080/v1/vla/session", show_default=True)
+@click.option("--instruction", default="pick up the red block", show_default=True)
+@click.option("--seed", type=int, default=7, show_default=True)
+@click.option("--chunks", type=click.IntRange(min=1), default=2, show_default=True)
+@click.option("--execute-horizon", type=click.IntRange(min=1), default=8, show_default=True)
+@click.option("--steps-per-action", type=click.IntRange(min=1), default=5, show_default=True)
+@click.option("--image-size", type=click.IntRange(min=64), default=256, show_default=True)
+@click.option("--output-dir", type=click.Path(path_type=Path, file_okay=False), default=None)
+def main(
+ mode: str,
+ urdf_path: Path,
+ server_url: str,
+ instruction: str,
+ seed: int,
+ chunks: int,
+ execute_horizon: int,
+ steps_per_action: int,
+ image_size: int,
+ output_dir: Path | None,
+) -> None:
+ """Validate local physics alone or the complete generic VLA WebSocket loop."""
+ _ensure_headless_runtime()
+ adapter = _create_adapter(urdf_path.resolve(), image_size, steps_per_action)
+ try:
+ if mode == "local":
+ result = _run_local_smoke(adapter, execute_horizon=execute_horizon, output_dir=output_dir)
+ else:
+ result = asyncio.run(
+ _run_websocket_loop(
+ adapter,
+ server_url=server_url,
+ instruction=instruction,
+ seed=seed,
+ chunks=chunks,
+ execute_horizon=execute_horizon,
+ output_dir=output_dir,
+ )
+ )
+ click.echo(json.dumps(result, indent=2, sort_keys=True))
+ finally:
+ adapter.close()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py
index 9095ad17..53f779ea 100644
--- a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py
+++ b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py
@@ -1,4 +1,9 @@
-"""Native TeleFuser service contract for LingBot-VLA v2 action inference."""
+"""Native TeleFuser HTTP service contract for LingBot-VLA v2.
+
+This module is also the pipeline definition loaded by the generic VLA
+WebSocket server. It remains available for ``telefuser serve`` compatibility;
+the WebSocket server is the primary continuous-control entrypoint.
+"""
from __future__ import annotations
@@ -7,11 +12,16 @@
from typing import Any
from telefuser.pipelines.lingbot_vla_v2.pipeline import LingBotVlaV2Pipeline
-from telefuser.pipelines.lingbot_vla_v2.runtime import get_lingbot_vla_v2_pipeline
+from telefuser.pipelines.lingbot_vla_v2.runtime import (
+ configure_lingbot_vla_v2_h100_sdpa,
+ get_lingbot_vla_v2_pipeline,
+)
from telefuser.pipelines.lingbot_vla_v2.service import (
LingBotVlaV2ActionRequest,
predict_lingbot_vla_v2_action,
)
+from telefuser.pipelines.lingbot_vla_v2.vla_policy import create_lingbot_vla_v2_session_manager
+from telefuser.service.vla_replica import VLAReplicaProvider
from telefuser.utils.logging import logger
TF_MODEL_ZOO_PATH = Path(os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo")).expanduser()
@@ -84,6 +94,7 @@ def get_pipeline(parallelism: int = 1) -> LingBotVlaV2Pipeline:
"""Load one policy replica for the native TeleFuser service."""
if parallelism != 1:
raise ValueError("LingBot-VLA v2 supports parallelism=1 per replica; use --num-replicas for a pipeline pool")
+ configure_lingbot_vla_v2_h100_sdpa(PPL_CONFIG["device"])
logger.info(
f"Loading LingBot-VLA v2 service profile quantization={PPL_CONFIG['quantization'] or 'bf16'} "
f"cuda_graph={PPL_CONFIG['cuda_graph']}"
@@ -98,6 +109,11 @@ def get_pipeline(parallelism: int = 1) -> LingBotVlaV2Pipeline:
)
+def get_vla_provider(pipeline: LingBotVlaV2Pipeline) -> VLAReplicaProvider:
+ """Create worker-local semantic VLA sessions around the loaded pipeline."""
+ return VLAReplicaProvider(create_lingbot_vla_v2_session_manager(pipeline))
+
+
def run_structured(
pipeline: LingBotVlaV2Pipeline,
instruction: str,
diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py
index 62905647..ec28f8e1 100644
--- a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py
+++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py
@@ -1,4 +1,9 @@
-"""Serve LingBot-VLA v2 through the upstream RoboTwin policy protocol."""
+"""Legacy compatibility server for the upstream RoboTwin policy protocol.
+
+New simulator integrations should use the generic VLA WebSocket protocol in
+``lingbot_vla_v2_vla_server.py``. This endpoint remains only for unmodified
+upstream ``WebsocketClientPolicy`` clients.
+"""
from __future__ import annotations
@@ -18,44 +23,32 @@
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from telefuser.pipelines.lingbot_vla_v2 import (
+ ROBOTWIN_ACTION_ORDER,
ROBOTWIN_CAMERA_KEYS,
- LingBotVlaV2Observation,
RobotWinProfile,
+ create_lingbot_vla_v2_session_manager,
)
-from telefuser.pipelines.lingbot_vla_v2.action_scheduler import ActionChunkScheduler
from telefuser.pipelines.lingbot_vla_v2.runtime import (
LINGBOT_VLA_V2_QUANTIZATION_CHOICES,
+ configure_lingbot_vla_v2_h100_sdpa,
get_lingbot_vla_v2_pipeline,
)
from telefuser.utils.logging import logger
+from telefuser.vla import RobotObservation, RobotState
+from telefuser.vla.runtime import ActionChunkScheduler
ROBOTWIN_PROTOCOL_VERSION = "1.0"
ROBOTWIN_ACTION_TYPE = "absolute_qpos"
ROBOTWIN_ACTION_DTYPE = "float32"
ROBOTWIN_MAX_REQUEST_BYTES = 16 * 1024 * 1024
-ROBOTWIN_ACTION_ORDER = (
- "left_arm_joint_0",
- "left_arm_joint_1",
- "left_arm_joint_2",
- "left_arm_joint_3",
- "left_arm_joint_4",
- "left_arm_joint_5",
- "left_gripper",
- "right_arm_joint_0",
- "right_arm_joint_1",
- "right_arm_joint_2",
- "right_arm_joint_3",
- "right_arm_joint_4",
- "right_arm_joint_5",
- "right_gripper",
-)
_TRACE_ID_FIELDS = ("request_id", "episode_id")
+_VLA_SESSION_ID_FIELD = "_telefuser_vla_session_id"
class _Pipeline(Protocol):
config: Any
- def __call__(self, observation: LingBotVlaV2Observation, seed: int | None = None) -> Any: ...
+ def __call__(self, observation: Any, seed: int | None = None) -> Any: ...
def close(self) -> None: ...
@@ -130,6 +123,21 @@ def _trace_fields(request: Mapping[str, Any]) -> dict[str, str | int]:
return fields
+def _optional_nonnegative_int(request: Mapping[str, Any], field: str) -> int | None:
+ value = request.get(field)
+ if value is None:
+ return None
+ if isinstance(value, bool):
+ raise ValueError(f"{field} must be a non-negative integer")
+ try:
+ parsed = operator.index(value)
+ except TypeError as error:
+ raise ValueError(f"{field} must be a non-negative integer") from error
+ if parsed < 0:
+ raise ValueError(f"{field} must be a non-negative integer")
+ return parsed
+
+
class RobotWinPolicyAdapter:
"""Translate upstream RoboTwin observations to the TeleFuser VLA SDK."""
@@ -146,6 +154,8 @@ def __init__(
self.profile = profile or pipeline.config.robot_profile
self.use_length = use_length
self._lock = threading.Lock()
+ self._next_sequence: dict[str, int] = {}
+ self._sessions = create_lingbot_vla_v2_session_manager(pipeline, profile=self.profile)
@property
def metadata(self) -> dict[str, Any]:
@@ -173,28 +183,52 @@ def infer(self, request: Mapping[str, Any]) -> dict[str, Any]:
if missing:
raise ValueError(f"RoboTwin observation is missing fields: {missing}")
- observation = LingBotVlaV2Observation(
- task=request["task"],
- state=request["observation.state"],
+ seed = _optional_seed(request)
+ episode_id = str(trace_fields.get("episode_id", "default"))
+ session_id_value = request.get(_VLA_SESSION_ID_FIELD, f"legacy:{episode_id}")
+ if not isinstance(session_id_value, str) or not session_id_value:
+ raise ValueError("internal VLA session ID must be a non-empty string")
+ sequence_id = _optional_nonnegative_int(request, "sequence_id")
+ observation_timestamp_ns = _optional_nonnegative_int(request, "observation_timestamp_ns")
+ if observation_timestamp_ns is None:
+ observation_timestamp_ns = time.time_ns()
+ robot_observation = RobotObservation(
+ state=RobotState(
+ values=torch.tensor(request["observation.state"], dtype=torch.float32, device="cpu"),
+ dimension_names=ROBOTWIN_ACTION_ORDER,
+ timestamp_ns=observation_timestamp_ns,
+ ),
images={key: request[key] for key in ROBOTWIN_CAMERA_KEYS},
)
- seed = _optional_seed(request)
adapter_started_at = time.monotonic()
with self._lock:
lock_wait_ms = (time.monotonic() - adapter_started_at) * 1000.0
- pipeline_started_at = time.monotonic()
- canonical_chunk = self.pipeline(observation, seed=seed)
- pipeline_ms = (time.monotonic() - pipeline_started_at) * 1000.0
- mapping_started_at = time.monotonic()
- action_chunk = self.profile.structure_actions(
- canonical_chunk.canonical_normalized_actions,
+ if sequence_id is None:
+ sequence_id = self._next_sequence.get(session_id_value, 0)
+ self._next_sequence[session_id_value] = sequence_id + 1
+ if session_id_value not in self._sessions.session_ids():
+ self._sessions.open(
+ session_id_value,
+ model_id="lingbot-vla-v2",
+ embodiment_id=self.profile.embodiment_id,
+ episode_id=episode_id,
+ execute_horizon=self.use_length,
+ )
+ vla_timings: dict[str, float] = {}
+ action_chunk = self._sessions.get(session_id_value).predict(
+ robot_observation,
+ request["task"],
+ sequence_id,
+ seed=seed,
+ timings=vla_timings,
)
- action_mapping_ms = (time.monotonic() - mapping_started_at) * 1000.0
- if action_chunk.horizon < self.use_length:
+ pipeline_ms = vla_timings["policy_ms"]
+ action_mapping_ms = vla_timings["decode_actions_ms"] + vla_timings["prepare_actions_ms"]
+ if action_chunk.valid_length < self.use_length:
raise RuntimeError(
- f"policy returned horizon {action_chunk.horizon}, shorter than use_length={self.use_length}"
+ f"policy returned horizon {action_chunk.valid_length}, shorter than use_length={self.use_length}"
)
- actions = np.ascontiguousarray(action_chunk.raw_actions[: self.use_length].numpy(), dtype=np.float32)
+ actions = np.ascontiguousarray(action_chunk.actions[: self.use_length].numpy(), dtype=np.float32)
expected_shape = (self.use_length, self.profile.raw_state_dim)
if actions.shape != expected_shape:
raise RuntimeError(f"mapped actions must have shape {expected_shape}, got {actions.shape}")
@@ -202,8 +236,8 @@ def infer(self, request: Mapping[str, Any]) -> dict[str, Any]:
raise RuntimeError("mapped actions must contain only finite values")
response: dict[str, Any] = {
"action": actions,
- "policy_verified": canonical_chunk.policy_verified,
- "verification_status": canonical_chunk.verification_status,
+ "policy_verified": action_chunk.metadata["policy_verified"],
+ "verification_status": action_chunk.metadata["verification_status"],
"server_timing": {
"lock_wait_ms": lock_wait_ms,
"pipeline_ms": pipeline_ms,
@@ -222,10 +256,25 @@ def _reset(self, request: Mapping[str, Any]) -> dict[str, Any]:
raise ValueError(f"unsupported robot profile: {robot_name!r}")
if request.get("path_to_pi_model") not in (None, ""):
raise ValueError("runtime checkpoint switching is not supported")
+ episode_id = str(request.get("episode_id", "default"))
+ session_id = request.get(_VLA_SESSION_ID_FIELD, f"legacy:{episode_id}")
+ if isinstance(session_id, str) and session_id in self._sessions.session_ids():
+ self._sessions.reset(session_id, episode_id)
+ self._next_sequence[session_id] = 0
return {"action": None}
+ def release_session(self, session_id: str) -> None:
+ """Release semantic session state after a WebSocket disconnect."""
+ with self._lock:
+ if session_id in self._sessions.session_ids():
+ self._sessions.close(session_id)
+ self._next_sequence.pop(session_id, None)
+
def close(self) -> None:
"""Release resources owned by the resident policy."""
+ for session_id in self._sessions.session_ids():
+ self._sessions.close(session_id)
+ self._next_sequence.clear()
self.pipeline.close()
@@ -305,9 +354,12 @@ async def deliver(
trace_fields = _trace_fields(request)
episode_id = trace_fields.get("episode_id", "default")
session_key = f"{connection_key}:{episode_id!r}"
+ request[_VLA_SESSION_ID_FIELD] = session_key
if request.get("reset", False):
for previous_session_key in session_keys - {session_key}:
scheduler.release_session(previous_session_key)
+ await asyncio.to_thread(adapter.release_session, previous_session_key)
+ scheduler.release_session(session_key)
session_keys.intersection_update({session_key})
session_keys.add(session_key)
response_future = scheduler.submit(request, session_key=session_key)
@@ -334,29 +386,14 @@ async def deliver(
task.cancel()
if delivery_tasks:
await asyncio.gather(*delivery_tasks, return_exceptions=True)
+ if session_keys:
+ await asyncio.gather(
+ *(asyncio.to_thread(adapter.release_session, session_key) for session_key in session_keys)
+ )
return app
-def _configure_h100_sdpa_backends(device: str) -> None:
- """Avoid unsupported cuDNN SDPA plans in the isolated H100 policy process."""
- resolved_device = torch.device(device)
- if resolved_device.type != "cuda" or not torch.cuda.is_available():
- return
- if "H100" not in torch.cuda.get_device_name(resolved_device):
- return
-
- if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
- torch.backends.cuda.enable_cudnn_sdp(False)
- if hasattr(torch.backends.cuda, "enable_flash_sdp"):
- torch.backends.cuda.enable_flash_sdp(True)
- if hasattr(torch.backends.cuda, "enable_math_sdp"):
- torch.backends.cuda.enable_math_sdp(True)
- if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"):
- torch.backends.cuda.enable_mem_efficient_sdp(True)
- logger.info("Disabled cuDNN SDPA for the LingBot-VLA v2 H100 policy process")
-
-
@click.command()
@click.option("--model-root", required=True, type=click.Path(exists=True, file_okay=False))
@click.option("--qwen3vl-root", required=True, type=click.Path(exists=True, file_okay=False))
@@ -388,8 +425,8 @@ def main(
cuda_graph: bool,
quantization: str | None,
) -> None:
- """Start one resident LingBot-VLA v2 policy for a RoboTwin client."""
- _configure_h100_sdpa_backends(device)
+ """Start the legacy-compatible policy endpoint for an upstream RoboTwin client."""
+ configure_lingbot_vla_v2_h100_sdpa(device)
pipeline = get_lingbot_vla_v2_pipeline(
model_root,
qwen3vl_root,
diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py
new file mode 100644
index 00000000..6c0614bd
--- /dev/null
+++ b/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py
@@ -0,0 +1,70 @@
+"""Primary generic VLA WebSocket service backed by TeleFuser replicas."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import uvicorn
+from fastapi import FastAPI
+
+from telefuser.service.core.config import ServerConfig
+from telefuser.service.core.pipeline_pool import PipelinePool
+from telefuser.service.security.security_validator import SecurityLevel
+from telefuser.service.vla_session import create_pipeline_pool_vla_session_app
+
+DEFAULT_PIPELINE_FILE = Path(__file__).with_name("lingbot_vla_v2_native_service.py")
+
+
+def create_app(
+ *,
+ pipeline_file: str | Path = DEFAULT_PIPELINE_FILE,
+ parallelism: int = 1,
+ num_replicas: int = 1,
+ skip_validation: bool = False,
+) -> FastAPI:
+ """Load LingBot replicas and create the standalone generic VLA app."""
+ config = ServerConfig(num_replicas=num_replicas, security_level=SecurityLevel.NONE)
+ replica_device_ids = config.resolve_replica_device_ids(parallelism)
+ pool = PipelinePool(
+ num_replicas=num_replicas,
+ replica_device_ids=replica_device_ids,
+ security_level_name=config.security_level.name,
+ config=config,
+ )
+ started = pool.start_all(
+ ppl_file=str(Path(pipeline_file).expanduser().resolve()),
+ parallelism_per_replica=len(replica_device_ids[0]),
+ task="vla_action",
+ skip_validation=skip_validation,
+ vla_provider_factory="get_vla_provider",
+ )
+ if not started:
+ raise RuntimeError("failed to start LingBot-VLA v2 pipeline replicas")
+ return create_pipeline_pool_vla_session_app(pool)
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--host", default="0.0.0.0")
+ parser.add_argument("--port", type=int, default=8000)
+ parser.add_argument("--parallelism", type=int, default=1)
+ parser.add_argument("--num-replicas", type=int, default=1)
+ parser.add_argument("--pipeline-file", type=Path, default=DEFAULT_PIPELINE_FILE)
+ parser.add_argument("--skip-validation", action="store_true")
+ return parser.parse_args()
+
+
+def main() -> None:
+ args = parse_args()
+ app = create_app(
+ pipeline_file=args.pipeline_file,
+ parallelism=args.parallelism,
+ num_replicas=args.num_replicas,
+ skip_validation=args.skip_validation,
+ )
+ uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/lingbot_vla_v2/setup_mujoco_local.sh b/examples/lingbot_vla_v2/setup_mujoco_local.sh
new file mode 100755
index 00000000..c390125b
--- /dev/null
+++ b/examples/lingbot_vla_v2/setup_mujoco_local.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+venv_python="${repo_root}/.venv/bin/python"
+local_lib_root="${repo_root}/.venv-mujoco-libs"
+
+if [[ ! -x "${venv_python}" ]]; then
+ echo "Expected the TeleFuser virtual environment at ${repo_root}/.venv" >&2
+ exit 1
+fi
+if ! command -v uv >/dev/null 2>&1; then
+ echo "uv is required to install into the existing virtual environment" >&2
+ exit 1
+fi
+if ! command -v apt-get >/dev/null 2>&1 || ! command -v dpkg-deb >/dev/null 2>&1; then
+ echo "apt-get and dpkg-deb are required to stage repository-local OSMesa libraries" >&2
+ exit 1
+fi
+
+uv pip install --python "${venv_python}" "mujoco==3.13.0"
+mkdir -p "${local_lib_root}/root"
+pushd "${local_lib_root}" >/dev/null
+apt-get download libegl1 libopengl0 libosmesa6
+for package in ./*.deb; do
+ dpkg-deb -x "${package}" root
+done
+popd >/dev/null
+
+echo "MuJoCo and repository-local OSMesa are ready under ${repo_root}"
diff --git a/telefuser/client/__init__.py b/telefuser/client/__init__.py
index c9e6352c..782ef8e7 100644
--- a/telefuser/client/__init__.py
+++ b/telefuser/client/__init__.py
@@ -17,6 +17,8 @@
from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+
from .tf_client import (
TFClient,
TaskCreationError,
@@ -25,10 +27,24 @@
TeleFuserError,
)
+if TYPE_CHECKING:
+ from .vla import AsyncVLAClient, VLAClientError
+
__all__ = [
"TFClient",
+ "AsyncVLAClient",
"TeleFuserError",
"TaskCreationError",
"TaskFailedError",
"TaskTimeoutError",
+ "VLAClientError",
]
+
+
+def __getattr__(name: str) -> Any:
+ """Load the optional VLA client without changing the existing HTTP client import path."""
+ if name in {"AsyncVLAClient", "VLAClientError"}:
+ from . import vla
+
+ return getattr(vla, name)
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/telefuser/client/vla.py b/telefuser/client/vla.py
new file mode 100644
index 00000000..0fe65cb9
--- /dev/null
+++ b/telefuser/client/vla.py
@@ -0,0 +1,244 @@
+"""Async client for the generic TeleFuser VLA WebSocket protocol."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections.abc import Mapping
+from typing import Any
+
+from telefuser.vla.contracts import ActionSpaceSpec, ObservationSpaceSpec, RobotActionChunk, RobotObservation
+from telefuser.vla.serialization import (
+ VLA_SESSION_PROTOCOL_VERSION,
+ VLA_WIRE_ENCODING,
+ action_space_to_wire,
+ dumps_wire_message,
+ loads_wire_message,
+ observation_space_to_wire,
+ robot_action_chunk_from_wire,
+ robot_observation_to_wire,
+)
+
+DEFAULT_MAX_VLA_MESSAGE_BYTES = 64 * 1024 * 1024
+
+
+class VLAClientError(RuntimeError):
+ """Error returned by the VLA session service."""
+
+ def __init__(self, code: str, message: str, *, request_id: str | int | None = None) -> None:
+ super().__init__(message)
+ self.code = code
+ self.request_id = request_id
+
+
+class AsyncVLAClient:
+ """Maintain one multiplexed connection to a VLA session service."""
+
+ def __init__(
+ self,
+ url: str,
+ *,
+ max_message_bytes: int = DEFAULT_MAX_VLA_MESSAGE_BYTES,
+ open_timeout_s: float = 10.0,
+ ) -> None:
+ if not isinstance(url, str) or not url:
+ raise ValueError("url must be a non-empty string")
+ if not isinstance(max_message_bytes, int) or isinstance(max_message_bytes, bool) or max_message_bytes < 1:
+ raise ValueError("max_message_bytes must be a positive integer")
+ if isinstance(open_timeout_s, bool) or not isinstance(open_timeout_s, (int, float)) or open_timeout_s <= 0:
+ raise ValueError("open_timeout_s must be positive")
+ self.url = url
+ self.max_message_bytes = max_message_bytes
+ self.open_timeout_s = float(open_timeout_s)
+ self.hello: dict[str, Any] | None = None
+ self._websocket: Any = None
+ self._reader: asyncio.Task[None] | None = None
+ self._pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
+ self._sessions: set[str] = set()
+ self._send_lock = asyncio.Lock()
+ self._request_sequence = 0
+
+ async def connect(self) -> dict[str, Any]:
+ """Connect and validate protocol capabilities."""
+ if self._websocket is not None:
+ raise RuntimeError("VLA client is already connected")
+ try:
+ from websockets.asyncio.client import connect
+ except ImportError:
+ from websockets import connect
+
+ websocket = await connect(
+ self.url,
+ max_size=self.max_message_bytes,
+ open_timeout=self.open_timeout_s,
+ )
+ try:
+ hello = loads_wire_message(await websocket.recv(), max_message_bytes=self.max_message_bytes)
+ if hello.get("type") != "HELLO":
+ raise VLAClientError("invalid_message", "VLA service did not send HELLO")
+ if hello.get("protocol_version") != VLA_SESSION_PROTOCOL_VERSION:
+ raise VLAClientError("unsupported_version", "VLA service protocol version is not supported")
+ if hello.get("encoding") != VLA_WIRE_ENCODING:
+ raise VLAClientError("unsupported_encoding", "VLA service wire encoding is not supported")
+ except Exception:
+ await websocket.close()
+ raise
+ self._websocket = websocket
+ self.hello = hello
+ self._reader = asyncio.create_task(self._read_responses())
+ return dict(hello)
+
+ async def open_session(
+ self,
+ session_id: str,
+ *,
+ model_id: str,
+ embodiment_id: str,
+ episode_id: str,
+ execute_horizon: int | None = None,
+ max_observation_age_ns: int | None = None,
+ expected_robot_action_space: ActionSpaceSpec | None = None,
+ expected_robot_observation_space: ObservationSpaceSpec | None = None,
+ ) -> dict[str, Any]:
+ """Open one session and optionally negotiate robot contracts."""
+ payload: dict[str, Any] = {
+ "type": "OPEN",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "session_id": session_id,
+ "model_id": model_id,
+ "embodiment_id": embodiment_id,
+ "episode_id": episode_id,
+ }
+ if execute_horizon is not None:
+ payload["execute_horizon"] = execute_horizon
+ if max_observation_age_ns is not None:
+ payload["max_observation_age_ns"] = max_observation_age_ns
+ if expected_robot_action_space is not None:
+ payload["expected_robot_action_space"] = action_space_to_wire(expected_robot_action_space)
+ if expected_robot_observation_space is not None:
+ payload["expected_robot_observation_space"] = observation_space_to_wire(expected_robot_observation_space)
+ response = await self._request(payload)
+ self._sessions.add(session_id)
+ return response
+
+ async def predict(
+ self,
+ session_id: str,
+ observation: RobotObservation,
+ instruction: str,
+ sequence_id: int,
+ *,
+ seed: int | None = None,
+ request_ttl_ms: float | None = None,
+ observation_clock_now_ns: int | None = None,
+ ) -> RobotActionChunk:
+ """Submit one observation and return its semantic robot action chunk."""
+ payload: dict[str, Any] = {
+ "type": "PREDICT",
+ "session_id": session_id,
+ "sequence_id": sequence_id,
+ "observation_timestamp_ns": observation.state.timestamp_ns,
+ "instruction": instruction,
+ "observation": robot_observation_to_wire(observation),
+ }
+ if seed is not None:
+ payload["seed"] = seed
+ if request_ttl_ms is not None:
+ payload["request_ttl_ms"] = request_ttl_ms
+ if observation_clock_now_ns is not None:
+ payload["observation_clock_now_ns"] = observation_clock_now_ns
+ response = await self._request(payload)
+ chunk = response.get("chunk")
+ if not isinstance(chunk, Mapping):
+ raise VLAClientError("invalid_message", "ACTION_CHUNK response is missing chunk")
+ return robot_action_chunk_from_wire(chunk)
+
+ async def reset_session(self, session_id: str, *, episode_id: str | None = None) -> None:
+ """Reset model history and queued action chunks for one session."""
+ payload: dict[str, Any] = {"type": "RESET", "session_id": session_id}
+ if episode_id is not None:
+ payload["episode_id"] = episode_id
+ await self._request(payload)
+
+ async def close_session(self, session_id: str) -> None:
+ """Close one server-side session while retaining the connection."""
+ await self._request({"type": "CLOSE", "session_id": session_id})
+ self._sessions.discard(session_id)
+
+ async def aclose(self) -> None:
+ """Close owned sessions and the WebSocket connection."""
+ if self._websocket is None:
+ return
+ for session_id in tuple(self._sessions):
+ with contextlib.suppress(Exception):
+ await self.close_session(session_id)
+ websocket = self._websocket
+ self._websocket = None
+ await websocket.close()
+ if self._reader is not None:
+ self._reader.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await self._reader
+ self._reader = None
+ self._fail_pending(VLAClientError("connection_closed", "VLA client connection closed"))
+
+ async def __aenter__(self) -> "AsyncVLAClient":
+ await self.connect()
+ return self
+
+ async def __aexit__(self, *_exc_info: object) -> None:
+ await self.aclose()
+
+ async def _request(self, payload: dict[str, Any]) -> dict[str, Any]:
+ websocket = self._websocket
+ if websocket is None:
+ raise RuntimeError("VLA client is not connected")
+ self._request_sequence += 1
+ request_id = f"vla-{self._request_sequence}"
+ payload["request_id"] = request_id
+ response = asyncio.get_running_loop().create_future()
+ self._pending[request_id] = response
+ try:
+ async with self._send_lock:
+ await websocket.send(dumps_wire_message(payload))
+ return await response
+ except BaseException:
+ self._pending.pop(request_id, None)
+ raise
+
+ async def _read_responses(self) -> None:
+ try:
+ while True:
+ payload = loads_wire_message(await self._websocket.recv(), max_message_bytes=self.max_message_bytes)
+ request_id = payload.get("request_id")
+ if not isinstance(request_id, str):
+ raise VLAClientError("invalid_message", "VLA response requires a string request_id")
+ response = self._pending.pop(request_id, None)
+ if response is None:
+ continue
+ if payload.get("type") == "ERROR":
+ error = payload.get("error")
+ if not isinstance(error, Mapping):
+ response.set_exception(VLAClientError("invalid_message", "invalid VLA error response"))
+ else:
+ response.set_exception(
+ VLAClientError(
+ str(error.get("code", "internal_error")),
+ str(error.get("message", "VLA request failed")),
+ request_id=request_id,
+ )
+ )
+ else:
+ response.set_result(payload)
+ except asyncio.CancelledError:
+ raise
+ except Exception as error:
+ self._fail_pending(
+ error if isinstance(error, VLAClientError) else VLAClientError("connection_closed", str(error))
+ )
+
+ def _fail_pending(self, error: Exception) -> None:
+ for response in self._pending.values():
+ if not response.done():
+ response.set_exception(error)
+ self._pending.clear()
diff --git a/telefuser/integrations/__init__.py b/telefuser/integrations/__init__.py
new file mode 100644
index 00000000..5fd93994
--- /dev/null
+++ b/telefuser/integrations/__init__.py
@@ -0,0 +1 @@
+"""Optional integrations between TeleFuser and external runtimes."""
diff --git a/telefuser/integrations/sim/__init__.py b/telefuser/integrations/sim/__init__.py
new file mode 100644
index 00000000..6999cfa1
--- /dev/null
+++ b/telefuser/integrations/sim/__init__.py
@@ -0,0 +1,12 @@
+"""Simulator-neutral VLA integration interfaces."""
+
+from .base import SimulatorAdapter
+from .mujoco import MuJoCoJointBinding, MuJoCoSimulatorAdapter
+from .robotwin import RoboTwinSimulatorAdapter
+
+__all__ = [
+ "MuJoCoJointBinding",
+ "MuJoCoSimulatorAdapter",
+ "RoboTwinSimulatorAdapter",
+ "SimulatorAdapter",
+]
diff --git a/telefuser/integrations/sim/base.py b/telefuser/integrations/sim/base.py
new file mode 100644
index 00000000..599bae83
--- /dev/null
+++ b/telefuser/integrations/sim/base.py
@@ -0,0 +1,25 @@
+"""Simulator adapter protocol consumed by the generic VLA runtime."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from telefuser.vla.contracts import RobotObservation
+from telefuser.vla.runtime.executor import RobotAction
+
+
+@runtime_checkable
+class SimulatorAdapter(Protocol):
+ """Minimal interface implemented by one simulator connection."""
+
+ def observe(self) -> RobotObservation:
+ """Capture the latest semantically described robot observation."""
+ ...
+
+ def execute(self, action: RobotAction) -> None:
+ """Submit one semantically described robot action."""
+ ...
+
+ def reset(self) -> RobotObservation:
+ """Reset the simulator and return its initial observation."""
+ ...
diff --git a/telefuser/integrations/sim/mujoco.py b/telefuser/integrations/sim/mujoco.py
new file mode 100644
index 00000000..bf32823d
--- /dev/null
+++ b/telefuser/integrations/sim/mujoco.py
@@ -0,0 +1,254 @@
+"""Optional MuJoCo implementation of the generic simulator boundary."""
+
+from __future__ import annotations
+
+import threading
+import time
+from dataclasses import dataclass
+from types import ModuleType
+from typing import Any, Mapping, Sequence
+
+import numpy as np
+import torch
+
+from telefuser.vla.contracts import ActionSpaceSpec, ObservationSpaceSpec, RobotObservation, RobotState
+from telefuser.vla.runtime.executor import RobotAction
+
+
+@dataclass(frozen=True)
+class MuJoCoJointBinding:
+ """Map one semantic action dimension to one or more MuJoCo joints."""
+
+ dimension_name: str
+ joint_names: tuple[str, ...]
+ normalized_to_joint_range: bool = False
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.dimension_name, str) or not self.dimension_name:
+ raise ValueError("dimension_name must be a non-empty string")
+ if not isinstance(self.joint_names, tuple) or not self.joint_names:
+ raise ValueError("joint_names must be a non-empty tuple")
+ if any(not isinstance(name, str) or not name for name in self.joint_names):
+ raise ValueError("joint_names must contain non-empty strings")
+ if len(set(self.joint_names)) != len(self.joint_names):
+ raise ValueError("joint_names must be unique within one binding")
+ if not isinstance(self.normalized_to_joint_range, bool):
+ raise TypeError("normalized_to_joint_range must be a boolean")
+
+
+@dataclass(frozen=True)
+class _ResolvedJoint:
+ qpos_address: int
+ dof_address: int
+ limited: bool
+ lower: float
+ upper: float
+
+
+class MuJoCoSimulatorAdapter:
+ """Apply semantic robot actions to a MuJoCo model with PD control.
+
+ MuJoCo is imported only when this adapter is instantiated, so the package
+ remains optional and cannot affect existing TeleFuser pipelines.
+ """
+
+ def __init__(
+ self,
+ model: Any,
+ action_space: ActionSpaceSpec,
+ observation_space: ObservationSpaceSpec,
+ joint_bindings: Sequence[MuJoCoJointBinding],
+ *,
+ camera_names: Mapping[str, str] | None = None,
+ image_height: int = 256,
+ image_width: int = 256,
+ steps_per_action: int = 5,
+ position_gain: float = 80.0,
+ damping_gain: float = 8.0,
+ max_force: float = 200.0,
+ ) -> None:
+ self._mujoco = self._import_mujoco()
+ if not isinstance(action_space, ActionSpaceSpec):
+ raise TypeError("action_space must be an ActionSpaceSpec")
+ if not isinstance(observation_space, ObservationSpaceSpec):
+ raise TypeError("observation_space must be an ObservationSpaceSpec")
+ if isinstance(steps_per_action, bool) or not isinstance(steps_per_action, int) or steps_per_action < 1:
+ raise ValueError("steps_per_action must be a positive integer")
+ if isinstance(image_height, bool) or not isinstance(image_height, int) or image_height < 1:
+ raise ValueError("image_height must be a positive integer")
+ if isinstance(image_width, bool) or not isinstance(image_width, int) or image_width < 1:
+ raise ValueError("image_width must be a positive integer")
+ for name, value in (
+ ("position_gain", position_gain),
+ ("damping_gain", damping_gain),
+ ("max_force", max_force),
+ ):
+ if isinstance(value, bool) or not isinstance(value, (int, float)) or not np.isfinite(value) or value <= 0:
+ raise ValueError(f"{name} must be positive and finite")
+
+ bindings = tuple(joint_bindings)
+ if tuple(binding.dimension_name for binding in bindings) != action_space.dimension_names:
+ raise ValueError("joint binding order must match action_space.dimension_names")
+ physical_names = [name for binding in bindings for name in binding.joint_names]
+ if len(set(physical_names)) != len(physical_names):
+ raise ValueError("a MuJoCo joint cannot be controlled by multiple action dimensions")
+
+ self.model = model
+ self.data = self._mujoco.MjData(model)
+ self.action_space = action_space
+ self.observation_space = observation_space
+ self.joint_bindings = bindings
+ self.steps_per_action = steps_per_action
+ self.position_gain = float(position_gain)
+ self.damping_gain = float(damping_gain)
+ self.max_force = float(max_force)
+ self._resolved_bindings = tuple(self._resolve_binding(binding) for binding in bindings)
+ self._camera_names = dict(camera_names or {})
+ expected_images = {image.name for image in observation_space.images}
+ if set(self._camera_names) != expected_images:
+ raise ValueError("camera_names keys must exactly match observation_space image names")
+ for camera_name in self._camera_names.values():
+ if self._mujoco.mj_name2id(model, self._mujoco.mjtObj.mjOBJ_CAMERA, camera_name) < 0:
+ raise ValueError(f"MuJoCo model does not contain camera {camera_name!r}")
+ self._renderer = (
+ self._mujoco.Renderer(model, height=image_height, width=image_width) if self._camera_names else None
+ )
+ self._lock = threading.RLock()
+ self._closed = False
+ self._mujoco.mj_forward(self.model, self.data)
+
+ @staticmethod
+ def _import_mujoco() -> ModuleType:
+ try:
+ import mujoco
+ except ImportError as error:
+ raise RuntimeError(
+ "MuJoCo is optional; install it in the simulator virtual environment before using this adapter"
+ ) from error
+ return mujoco
+
+ def observe(self) -> RobotObservation:
+ """Return current joint state and named RGB camera frames."""
+ with self._lock:
+ self._ensure_open()
+ self._mujoco.mj_forward(self.model, self.data)
+ timestamp_ns = time.monotonic_ns()
+ state = torch.tensor(self._read_action_state(), dtype=torch.float32)
+ images: dict[str, np.ndarray] = {}
+ if self._renderer is not None:
+ for observation_name, camera_name in self._camera_names.items():
+ self._renderer.update_scene(self.data, camera=camera_name)
+ images[observation_name] = np.ascontiguousarray(self._renderer.render().copy())
+ observation = RobotObservation(
+ state=RobotState(state, self.action_space.dimension_names, timestamp_ns),
+ images=images,
+ metadata={"simulator": "mujoco", "simulation_time_s": float(self.data.time)},
+ )
+ self.observation_space.validate(observation)
+ return observation
+
+ def execute(self, action: RobotAction) -> None:
+ """Advance physics while driving joints toward one absolute action."""
+ if not isinstance(action, RobotAction):
+ raise TypeError("MuJoCo execute expects RobotAction")
+ self.action_space.require_compatible(action.action_space, context="MuJoCo action space")
+ values = action.values.detach().to(device="cpu", dtype=torch.float64).numpy()
+ if values.shape != (self.action_space.dimension,):
+ raise ValueError(f"MuJoCo action must have shape ({self.action_space.dimension},), got {values.shape}")
+ if not np.isfinite(values).all():
+ raise ValueError("MuJoCo action must contain only finite values")
+
+ with self._lock:
+ self._ensure_open()
+ targets = self._resolve_targets(values)
+ for _ in range(self.steps_per_action):
+ self._mujoco.mj_forward(self.model, self.data)
+ self.data.qfrc_applied.fill(0.0)
+ for target, joints in zip(targets, self._resolved_bindings, strict=True):
+ for joint in joints:
+ position_error = target - float(self.data.qpos[joint.qpos_address])
+ velocity = float(self.data.qvel[joint.dof_address])
+ force = (
+ float(self.data.qfrc_bias[joint.dof_address])
+ + self.position_gain * position_error
+ - self.damping_gain * velocity
+ )
+ self.data.qfrc_applied[joint.dof_address] = np.clip(
+ force,
+ -self.max_force,
+ self.max_force,
+ )
+ self._mujoco.mj_step(self.model, self.data)
+ self.data.qfrc_applied.fill(0.0)
+
+ def reset(self) -> RobotObservation:
+ """Reset model state and return the first rendered observation."""
+ with self._lock:
+ self._ensure_open()
+ self._mujoco.mj_resetData(self.model, self.data)
+ self._mujoco.mj_forward(self.model, self.data)
+ return self.observe()
+
+ def close(self) -> None:
+ """Release the optional offscreen renderer."""
+ with self._lock:
+ if self._closed:
+ return
+ if self._renderer is not None:
+ self._renderer.close()
+ self._closed = True
+
+ def _resolve_binding(self, binding: MuJoCoJointBinding) -> tuple[_ResolvedJoint, ...]:
+ resolved: list[_ResolvedJoint] = []
+ for name in binding.joint_names:
+ joint_id = self._mujoco.mj_name2id(self.model, self._mujoco.mjtObj.mjOBJ_JOINT, name)
+ if joint_id < 0:
+ raise ValueError(f"MuJoCo model does not contain joint {name!r}")
+ joint_type = int(self.model.jnt_type[joint_id])
+ supported_types = {
+ int(self._mujoco.mjtJoint.mjJNT_HINGE),
+ int(self._mujoco.mjtJoint.mjJNT_SLIDE),
+ }
+ if joint_type not in supported_types:
+ raise ValueError(f"MuJoCo joint {name!r} must be a hinge or slide joint")
+ limited = bool(self.model.jnt_limited[joint_id])
+ lower, upper = (float(value) for value in self.model.jnt_range[joint_id])
+ if binding.normalized_to_joint_range and (not limited or upper <= lower):
+ raise ValueError(f"normalized MuJoCo joint {name!r} requires a finite increasing range")
+ resolved.append(
+ _ResolvedJoint(
+ qpos_address=int(self.model.jnt_qposadr[joint_id]),
+ dof_address=int(self.model.jnt_dofadr[joint_id]),
+ limited=limited,
+ lower=lower,
+ upper=upper,
+ )
+ )
+ return tuple(resolved)
+
+ def _resolve_targets(self, values: np.ndarray) -> tuple[float, ...]:
+ targets: list[float] = []
+ for value, binding, joints in zip(values, self.joint_bindings, self._resolved_bindings, strict=True):
+ target = float(value)
+ if binding.normalized_to_joint_range:
+ target = float(np.clip(target, 0.0, 1.0))
+ target = joints[0].lower + target * (joints[0].upper - joints[0].lower)
+ lower = max((joint.lower for joint in joints if joint.limited), default=-np.inf)
+ upper = min((joint.upper for joint in joints if joint.limited), default=np.inf)
+ targets.append(float(np.clip(target, lower, upper)))
+ return tuple(targets)
+
+ def _read_action_state(self) -> np.ndarray:
+ values = np.empty(len(self.joint_bindings), dtype=np.float32)
+ for index, (binding, joints) in enumerate(zip(self.joint_bindings, self._resolved_bindings, strict=True)):
+ joint_values = np.asarray([self.data.qpos[joint.qpos_address] for joint in joints], dtype=np.float64)
+ value = float(joint_values.mean())
+ if binding.normalized_to_joint_range:
+ value = (value - joints[0].lower) / (joints[0].upper - joints[0].lower)
+ value = float(np.clip(value, 0.0, 1.0))
+ values[index] = value
+ return values
+
+ def _ensure_open(self) -> None:
+ if self._closed:
+ raise RuntimeError("MuJoCo simulator adapter is closed")
diff --git a/telefuser/integrations/sim/robotwin.py b/telefuser/integrations/sim/robotwin.py
new file mode 100644
index 00000000..041b3f30
--- /dev/null
+++ b/telefuser/integrations/sim/robotwin.py
@@ -0,0 +1,68 @@
+"""Dependency-free boundary between the VLA runtime and RoboTwin processes."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import numpy as np
+import torch
+
+from telefuser.vla.contracts import ActionSpaceSpec, RobotActionChunk, RobotObservation
+from telefuser.vla.runtime.executor import ChunkExecutor, RobotAction
+
+
+class RoboTwinSimulatorAdapter:
+ """Validate semantic actions before invoking RoboTwin-owned callbacks.
+
+ The callbacks keep RoboTwin, SAPIEN, and Vulkan imports in the remote
+ simulator process rather than making them TeleFuser dependencies.
+ """
+
+ def __init__(
+ self,
+ action_space: ActionSpaceSpec,
+ *,
+ observe_fn: Callable[[], RobotObservation],
+ execute_fn: Callable[[np.ndarray], None],
+ reset_fn: Callable[[], RobotObservation],
+ ) -> None:
+ self.action_space = action_space
+ self._observe_fn = observe_fn
+ self._execute_fn = execute_fn
+ self._reset_fn = reset_fn
+
+ def observe(self) -> RobotObservation:
+ """Read one observation from the RoboTwin process."""
+ observation = self._observe_fn()
+ if not isinstance(observation, RobotObservation):
+ raise TypeError("RoboTwin observe callback must return RobotObservation")
+ return observation
+
+ def execute(self, action: RobotAction) -> None:
+ """Submit one validated float32 action vector to RoboTwin."""
+ if not isinstance(action, RobotAction):
+ raise TypeError("RoboTwin execute expects RobotAction")
+ self.action_space.require_compatible(action.action_space, context="RoboTwin action space")
+ values = action.values.detach().to(device="cpu", dtype=torch.float32).numpy()
+ values = np.ascontiguousarray(values, dtype=np.float32)
+ if values.shape != (self.action_space.dimension,):
+ raise ValueError(f"RoboTwin action must have shape ({self.action_space.dimension},), got {values.shape}")
+ if not np.isfinite(values).all():
+ raise ValueError("RoboTwin action must contain only finite values")
+ self._execute_fn(values)
+
+ def execute_chunk(self, chunk: RobotActionChunk) -> int:
+ """Execute the valid portion of an already prepared chunk."""
+ self.action_space.require_compatible(chunk.action_space, context="RoboTwin action space")
+ count = 0
+ for action in ChunkExecutor.iter_actions(chunk):
+ self.execute(action)
+ count += 1
+ return count
+
+ def reset(self) -> RobotObservation:
+ """Reset RoboTwin and validate its initial observation."""
+ observation = self._reset_fn()
+ if not isinstance(observation, RobotObservation):
+ raise TypeError("RoboTwin reset callback must return RobotObservation")
+ return observation
diff --git a/telefuser/pipelines/lingbot_vla_v2/__init__.py b/telefuser/pipelines/lingbot_vla_v2/__init__.py
index de38aed1..d51d7fc3 100644
--- a/telefuser/pipelines/lingbot_vla_v2/__init__.py
+++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py
@@ -3,7 +3,20 @@
from .data import LingBotVlaV2InputProcessor, LingBotVlaV2Inputs, LingBotVlaV2Observation
from .pipeline import LingBotVlaV2CanonicalActionChunk, LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig
from .policy import LingBotVlaV2PolicyStage
-from .robot_profile import ROBOTWIN_CAMERA_KEYS, LingBotVlaV2ActionChunk, RobotWinProfile
+from .robot_profile import (
+ LINGBOT_VLA_V2_ACTION_SPACE,
+ ROBOTWIN_ACTION_ORDER,
+ ROBOTWIN_ACTION_SPACE,
+ ROBOTWIN_CAMERA_KEYS,
+ ROBOTWIN_OBSERVATION_SPACE,
+ LingBotVlaV2ActionChunk,
+ RobotWinProfile,
+)
+from .vla_policy import (
+ LINGBOT_VLA_V2_MODEL_ID,
+ LingBotVlaV2VLAPolicy,
+ create_lingbot_vla_v2_session_manager,
+)
__all__ = [
"LingBotVlaV2ActionChunk",
@@ -14,6 +27,13 @@
"LingBotVlaV2Pipeline",
"LingBotVlaV2PipelineConfig",
"LingBotVlaV2PolicyStage",
+ "LingBotVlaV2VLAPolicy",
+ "create_lingbot_vla_v2_session_manager",
+ "LINGBOT_VLA_V2_ACTION_SPACE",
+ "LINGBOT_VLA_V2_MODEL_ID",
+ "ROBOTWIN_ACTION_ORDER",
+ "ROBOTWIN_ACTION_SPACE",
"ROBOTWIN_CAMERA_KEYS",
+ "ROBOTWIN_OBSERVATION_SPACE",
"RobotWinProfile",
]
diff --git a/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py b/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py
index 73439d4f..5427a74f 100644
--- a/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py
+++ b/telefuser/pipelines/lingbot_vla_v2/action_scheduler.py
@@ -1,271 +1,5 @@
-"""Bounded latest-wins scheduling for LingBot-VLA v2 action chunks."""
+"""Compatibility import for the generic VLA action scheduler."""
-from __future__ import annotations
+from telefuser.vla.runtime.scheduler import ActionChunkScheduler
-import asyncio
-import math
-import operator
-import time
-from collections import OrderedDict
-from dataclasses import dataclass
-from typing import Any, Callable, Mapping
-
-
-@dataclass(slots=True)
-class _ScheduledAction:
- request: Mapping[str, Any]
- session_key: str
- generation: int
- received_at: float
- deadline_at: float | None
- future: asyncio.Future[dict[str, Any]]
-
-
-class ActionChunkScheduler:
- """Serialize GPU inference while retaining only the newest pending chunk per session."""
-
- def __init__(
- self,
- infer: Callable[[Mapping[str, Any]], dict[str, Any]],
- *,
- max_pending_sessions: int = 32,
- ) -> None:
- if max_pending_sessions < 1:
- raise ValueError("max_pending_sessions must be positive")
- self._infer = infer
- self._max_pending_sessions = max_pending_sessions
- self._pending: OrderedDict[str, _ScheduledAction] = OrderedDict()
- self._latest_generation: dict[str, int] = {}
- self._latest_sequence: dict[str, int] = {}
- self._wake = asyncio.Event()
- self._worker: asyncio.Task[None] | None = None
- self._closed = False
-
- @property
- def metadata(self) -> dict[str, Any]:
- """Describe the additive scheduling controls accepted by the endpoint."""
- return {
- "scheduling": {
- "mode": "latest_wins",
- "max_pending_per_session": 1,
- "max_pending_sessions": self._max_pending_sessions,
- "sequence_field": "sequence_id",
- "ttl_field": "request_ttl_ms",
- "inflight_cancellation": False,
- }
- }
-
- async def start(self) -> None:
- """Start the single inference worker on the current event loop."""
- if self._worker is not None:
- return
- if self._closed:
- raise RuntimeError("action scheduler is closed")
- self._worker = asyncio.create_task(self._run(), name="lingbot-vla-v2-action-scheduler")
-
- def submit(self, request: Mapping[str, Any], *, session_key: str) -> asyncio.Future[dict[str, Any]]:
- """Admit one request and return a future without waiting for inference."""
- if self._worker is None or self._closed:
- raise RuntimeError("action scheduler is not running")
- if not session_key:
- raise ValueError("session_key must be non-empty")
-
- loop = asyncio.get_running_loop()
- future: asyncio.Future[dict[str, Any]] = loop.create_future()
- received_at = time.monotonic()
- ttl_ms = self._optional_ttl_ms(request)
- sequence_id = self._optional_sequence_id(request)
- previous = self._pending.get(session_key)
- if previous is None and len(self._pending) >= self._max_pending_sessions:
- future.set_result(
- self._discarded_response(
- request,
- status="overloaded",
- message="the action scheduler has no free pending-session slot",
- )
- )
- return future
- if sequence_id is not None:
- latest_sequence = self._latest_sequence.get(session_key)
- if latest_sequence is not None and sequence_id <= latest_sequence:
- future.set_result(
- self._discarded_response(
- request,
- status="stale_sequence",
- message=f"sequence_id={sequence_id} is not newer than {latest_sequence}",
- )
- )
- return future
- self._latest_sequence[session_key] = sequence_id
-
- generation = self._latest_generation.get(session_key, 0) + 1
- self._latest_generation[session_key] = generation
-
- deadline_at = None if ttl_ms is None else received_at + ttl_ms / 1000.0
- job = _ScheduledAction(
- request=dict(request),
- session_key=session_key,
- generation=generation,
- received_at=received_at,
- deadline_at=deadline_at,
- future=future,
- )
- if previous is not None:
- self._resolve(
- previous,
- self._discarded_response(
- previous.request,
- status="superseded",
- message="a newer observation replaced this pending action request",
- ),
- )
- self._pending[session_key] = job
- self._wake.set()
- return future
-
- def release_session(self, session_key: str) -> None:
- """Discard pending work and sequence state for a disconnected session."""
- pending = self._pending.pop(session_key, None)
- if pending is not None:
- self._resolve(
- pending,
- self._discarded_response(
- pending.request,
- status="session_closed",
- message="the client session closed before inference",
- ),
- )
- self._latest_generation.pop(session_key, None)
- self._latest_sequence.pop(session_key, None)
-
- async def close(self) -> None:
- """Drain an in-flight call and reject work that has not started."""
- if self._closed:
- return
- self._closed = True
- for job in self._pending.values():
- self._resolve(
- job,
- self._discarded_response(
- job.request,
- status="server_stopping",
- message="the action scheduler is stopping",
- ),
- )
- self._pending.clear()
- self._wake.set()
- if self._worker is not None:
- await self._worker
- self._worker = None
-
- async def _run(self) -> None:
- while True:
- await self._wake.wait()
- if self._closed and not self._pending:
- return
- if not self._pending:
- self._wake.clear()
- continue
-
- _, job = self._pending.popitem(last=False)
- if not self._pending:
- self._wake.clear()
- if job.future.done():
- continue
- discard = self._discard_reason(job)
- if discard is not None:
- self._resolve(job, discard)
- continue
-
- inference_started_at = time.monotonic()
- try:
- response = await asyncio.to_thread(self._infer, job.request)
- except Exception as error:
- if not job.future.done():
- job.future.set_exception(error)
- continue
- inference_ms = (time.monotonic() - inference_started_at) * 1000.0
-
- discard = self._discard_reason(job)
- if discard is not None:
- self._resolve(job, discard)
- continue
- completed_at = time.monotonic()
- response = dict(response)
- server_timing = dict(response.get("server_timing", {}))
- server_timing.update(
- queue_wait_ms=(inference_started_at - job.received_at) * 1000.0,
- infer_ms=inference_ms,
- scheduler_total_ms=(completed_at - job.received_at) * 1000.0,
- )
- response.update(scheduler_status="completed", server_timing=server_timing)
- for field in ("request_id", "episode_id", "sequence_id"):
- if field in job.request:
- response.setdefault(field, job.request[field])
- if job.deadline_at is not None:
- response["request_ttl_ms"] = (job.deadline_at - job.received_at) * 1000.0
- self._resolve(job, response)
-
- def _discard_reason(self, job: _ScheduledAction) -> dict[str, Any] | None:
- if job.deadline_at is not None and time.monotonic() >= job.deadline_at:
- return self._discarded_response(
- job.request,
- status="expired",
- message="the action request exceeded request_ttl_ms",
- )
- if self._latest_generation.get(job.session_key) != job.generation:
- return self._discarded_response(
- job.request,
- status="superseded",
- message="a newer observation superseded this action result",
- )
- return None
-
- @staticmethod
- def _resolve(job: _ScheduledAction, response: dict[str, Any]) -> None:
- if not job.future.done():
- job.future.set_result(response)
-
- @staticmethod
- def _optional_sequence_id(request: Mapping[str, Any]) -> int | None:
- value = request.get("sequence_id")
- if value is None:
- return None
- if isinstance(value, bool):
- raise ValueError("sequence_id must be a non-negative integer")
- try:
- sequence_id = operator.index(value)
- except TypeError as error:
- raise ValueError("sequence_id must be a non-negative integer") from error
- if sequence_id < 0:
- raise ValueError("sequence_id must be a non-negative integer")
- return sequence_id
-
- @staticmethod
- def _optional_ttl_ms(request: Mapping[str, Any]) -> float | None:
- value = request.get("request_ttl_ms")
- if value is None:
- return None
- if isinstance(value, bool) or not isinstance(value, int | float):
- raise ValueError("request_ttl_ms must be a positive finite number")
- ttl_ms = float(value)
- if not math.isfinite(ttl_ms) or ttl_ms <= 0:
- raise ValueError("request_ttl_ms must be a positive finite number")
- return ttl_ms
-
- @staticmethod
- def _discarded_response(
- request: Mapping[str, Any],
- *,
- status: str,
- message: str,
- ) -> dict[str, Any]:
- response: dict[str, Any] = {
- "action": None,
- "scheduler_status": status,
- "error": {"code": status, "message": message},
- }
- for field in ("request_id", "episode_id", "sequence_id"):
- if field in request:
- response[field] = request[field]
- return response
+__all__ = ["ActionChunkScheduler"]
diff --git a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py
index e460b88d..eff67389 100644
--- a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py
+++ b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py
@@ -10,6 +10,17 @@
import torch
+from telefuser.vla.contracts import (
+ ActionSpaceSpec,
+ ImageObservationSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+)
+
ROBOTWIN_CAMERA_KEYS = (
"observation.images.cam_high",
"observation.images.cam_left_wrist",
@@ -19,6 +30,44 @@
CANONICAL_DIM = 55
ARM_SLICE = slice(0, 12)
EFFECTOR_SLICE = slice(28, 30)
+ROBOTWIN_ACTION_ORDER = (
+ "left_arm_joint_0",
+ "left_arm_joint_1",
+ "left_arm_joint_2",
+ "left_arm_joint_3",
+ "left_arm_joint_4",
+ "left_arm_joint_5",
+ "left_gripper",
+ "right_arm_joint_0",
+ "right_arm_joint_1",
+ "right_arm_joint_2",
+ "right_arm_joint_3",
+ "right_arm_joint_4",
+ "right_arm_joint_5",
+ "right_gripper",
+)
+LINGBOT_VLA_V2_ACTION_SPACE = ActionSpaceSpec(
+ representation="canonical_normalized",
+ dimension_names=tuple(f"canonical_action_{index}" for index in range(CANONICAL_DIM)),
+ units=("normalized",) * CANONICAL_DIM,
+ frame=None,
+ control_hz=None,
+ normalized=True,
+ normalization_profile="lingbot_vla_v2_canonical",
+)
+ROBOTWIN_ACTION_SPACE = ActionSpaceSpec(
+ representation="absolute_qpos",
+ dimension_names=ROBOTWIN_ACTION_ORDER,
+ units=("radian",) * 6 + ("normalized_position",) + ("radian",) * 6 + ("normalized_position",),
+ frame="robot_joint",
+ control_hz=None,
+ normalized=False,
+)
+ROBOTWIN_OBSERVATION_SPACE = ObservationSpaceSpec(
+ state_dimension_names=ROBOTWIN_ACTION_ORDER,
+ images=tuple(ImageObservationSpec(name=key) for key in ROBOTWIN_CAMERA_KEYS),
+ allow_extra_images=True,
+)
@dataclass(frozen=True)
@@ -39,6 +88,7 @@ class RobotWinProfile:
"""Map RobotWin observations and actions to LingBot's canonical space."""
name = "robotwin"
+ embodiment_id = "robotwin"
camera_keys = ROBOTWIN_CAMERA_KEYS
canonical_dim = CANONICAL_DIM
raw_state_dim = ROBOTWIN_STATE_DIM
@@ -81,6 +131,50 @@ def action_mask(self) -> torch.Tensor:
mask[EFFECTOR_SLICE] = True
return mask
+ @property
+ def model_action_space(self) -> ActionSpaceSpec:
+ """Return the semantic action space accepted from LingBot-VLA v2."""
+ return LINGBOT_VLA_V2_ACTION_SPACE
+
+ @property
+ def robot_action_space(self) -> ActionSpaceSpec:
+ """Return the semantic action space emitted for RoboTwin."""
+ return ROBOTWIN_ACTION_SPACE
+
+ @property
+ def observation_space(self) -> ObservationSpaceSpec:
+ """Return the RoboTwin state and camera contract."""
+ return ROBOTWIN_OBSERVATION_SPACE
+
+ def encode_observation(self, observation: RobotObservation) -> ModelObservation:
+ """Validate a RoboTwin observation while retaining raw state for the model processor."""
+ self.observation_space.validate(observation)
+ return ModelObservation(
+ state=observation.state.values,
+ images={key: observation.images[key] for key in self.camera_keys},
+ metadata=observation.metadata,
+ )
+
+ def decode_actions(
+ self,
+ actions: ModelActionChunk,
+ robot_state: RobotState,
+ ) -> RobotActionChunk:
+ """Map a semantic LingBot chunk to an absolute RoboTwin joint chunk."""
+ self.model_action_space.require_compatible(actions.action_space, context="LingBot model action space")
+ if robot_state.dimension_names != ROBOTWIN_ACTION_ORDER:
+ raise ValueError("RobotWin state dimension order does not match the robot profile")
+ structured = self.structure_actions(actions.actions[: actions.valid_length])
+ return RobotActionChunk(
+ actions=structured.raw_actions,
+ action_space=self.robot_action_space,
+ valid_length=structured.horizon,
+ observation_timestamp_ns=actions.observation_timestamp_ns,
+ sequence_id=actions.sequence_id,
+ episode_id=actions.episode_id,
+ metadata=actions.metadata,
+ )
+
def normalize_state(self, raw_state: torch.Tensor | Sequence[float]) -> torch.Tensor:
"""Convert one raw 14-D RobotWin state to normalized canonical 55-D space."""
state = torch.as_tensor(raw_state, dtype=torch.float32, device="cpu")
diff --git a/telefuser/pipelines/lingbot_vla_v2/runtime.py b/telefuser/pipelines/lingbot_vla_v2/runtime.py
index 83333399..2bb319d8 100644
--- a/telefuser/pipelines/lingbot_vla_v2/runtime.py
+++ b/telefuser/pipelines/lingbot_vla_v2/runtime.py
@@ -8,6 +8,7 @@
from telefuser.core.config import ModelRuntimeConfig, QuantConfig, QuantKernelBackend, QuantType
from telefuser.core.module_manager import ModuleManager
from telefuser.models.lingbot_vla_v2_loader import load_lingbot_vla_v2
+from telefuser.utils.logging import logger
from .pipeline import LingBotVlaV2Pipeline, LingBotVlaV2PipelineConfig
@@ -19,6 +20,24 @@ def _apply_cuda_runtime_flags(device: torch.device) -> None:
torch.set_float32_matmul_precision("high")
+def configure_lingbot_vla_v2_h100_sdpa(device: str) -> None:
+ """Avoid unsupported cuDNN SDPA plans in a LingBot H100 process."""
+ resolved_device = torch.device(device)
+ if resolved_device.type != "cuda" or not torch.cuda.is_available():
+ return
+ if "H100" not in torch.cuda.get_device_name(resolved_device):
+ return
+ if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
+ torch.backends.cuda.enable_cudnn_sdp(False)
+ if hasattr(torch.backends.cuda, "enable_flash_sdp"):
+ torch.backends.cuda.enable_flash_sdp(True)
+ if hasattr(torch.backends.cuda, "enable_math_sdp"):
+ torch.backends.cuda.enable_math_sdp(True)
+ if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"):
+ torch.backends.cuda.enable_mem_efficient_sdp(True)
+ logger.info("Disabled cuDNN SDPA for the LingBot-VLA v2 H100 policy process")
+
+
def lingbot_vla_v2_quant_config(quantization: str | QuantType | None) -> QuantConfig:
"""Resolve a public LingBot-VLA v2 online-quantization name."""
if quantization is None:
diff --git a/telefuser/pipelines/lingbot_vla_v2/vla_policy.py b/telefuser/pipelines/lingbot_vla_v2/vla_policy.py
new file mode 100644
index 00000000..4b3329f5
--- /dev/null
+++ b/telefuser/pipelines/lingbot_vla_v2/vla_policy.py
@@ -0,0 +1,83 @@
+"""Semantic VLA policy wrapper for the existing LingBot-VLA v2 pipeline."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+from telefuser.vla.contracts import ModelActionChunk, VLACapabilities, VLARequest
+from telefuser.vla.registry import VLARegistry
+from telefuser.vla.session import VLASessionManager
+
+from .data import LingBotVlaV2Observation
+from .robot_profile import LINGBOT_VLA_V2_ACTION_SPACE, RobotWinProfile
+
+LINGBOT_VLA_V2_MODEL_ID = "lingbot-vla-v2"
+
+
+class _LingBotPipeline(Protocol):
+ def __call__(self, observation: LingBotVlaV2Observation, seed: int | None = None) -> Any: ...
+
+
+class LingBotVlaV2VLAPolicy:
+ """Adapt the stable LingBot pipeline API to the shared semantic contract."""
+
+ def __init__(self, pipeline: _LingBotPipeline, *, max_horizon: int = 50) -> None:
+ if not isinstance(max_horizon, int) or isinstance(max_horizon, bool) or max_horizon < 1:
+ raise ValueError("max_horizon must be positive")
+ self.pipeline = pipeline
+ self._capabilities = VLACapabilities(
+ model_id=LINGBOT_VLA_V2_MODEL_ID,
+ output_action_space=LINGBOT_VLA_V2_ACTION_SPACE,
+ max_horizon=max_horizon,
+ stateful=False,
+ supports_seed=True,
+ )
+
+ def capabilities(self) -> VLACapabilities:
+ """Describe the canonical LingBot action output."""
+ return self._capabilities
+
+ def predict(self, request: VLARequest) -> ModelActionChunk:
+ """Run the existing pipeline and attach semantic request provenance."""
+ observation = LingBotVlaV2Observation(
+ task=request.instruction,
+ state=request.observation.state,
+ images=request.observation.images,
+ )
+ output = self.pipeline(observation, seed=request.seed)
+ if output.horizon > self._capabilities.max_horizon:
+ raise RuntimeError(
+ f"LingBot pipeline returned horizon {output.horizon}, exceeding declared maximum "
+ f"{self._capabilities.max_horizon}"
+ )
+ return ModelActionChunk(
+ actions=output.canonical_normalized_actions,
+ action_space=self._capabilities.output_action_space,
+ valid_length=output.horizon,
+ observation_timestamp_ns=request.observation_timestamp_ns,
+ sequence_id=request.sequence_id,
+ episode_id=request.episode_id,
+ metadata={
+ "checkpoint_variant": output.checkpoint_variant,
+ "policy_verified": output.policy_verified,
+ "verification_status": output.verification_status,
+ },
+ )
+
+ def reset(self, episode_id: str) -> None:
+ """Validate RESET for the currently stateless LingBot base policy."""
+ if not isinstance(episode_id, str) or not episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+
+
+def create_lingbot_vla_v2_session_manager(
+ pipeline: _LingBotPipeline,
+ *,
+ profile: RobotWinProfile | None = None,
+) -> VLASessionManager:
+ """Register the shared LingBot policy and RoboTwin embodiment once."""
+ resolved_profile = profile or RobotWinProfile.default()
+ registry = VLARegistry()
+ registry.register_policy(LINGBOT_VLA_V2_MODEL_ID, LingBotVlaV2VLAPolicy(pipeline))
+ registry.register_embodiment(resolved_profile.embodiment_id, resolved_profile)
+ return VLASessionManager(registry)
diff --git a/telefuser/service/core/pipeline_pool.py b/telefuser/service/core/pipeline_pool.py
index 74b982fb..97645a2e 100644
--- a/telefuser/service/core/pipeline_pool.py
+++ b/telefuser/service/core/pipeline_pool.py
@@ -7,10 +7,12 @@
from __future__ import annotations
import asyncio
+import contextlib
import multiprocessing as mp_stdlib
import os
import threading
from contextlib import asynccontextmanager
+from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, AsyncIterator
from telefuser.platforms import current_platform
@@ -30,6 +32,13 @@
from .task_manager import TaskManager
+@dataclass
+class _ReplicaSessionLease:
+ replica_id: int
+ lock: asyncio.Lock
+ closing: bool = False
+
+
class PipelinePool:
"""Pool of pipeline replicas in separate subprocesses.
@@ -55,11 +64,17 @@ def __init__(
self._instance_status: list[str] = []
self._live_count = num_replicas
self._status_lock = threading.Lock()
+ self._session_leases: dict[str, _ReplicaSessionLease] = {}
+ self._opening_sessions: set[str] = set()
+ self._session_leases_lock = asyncio.Lock()
+ self._session_shutdown = False
+ self._session_shutdown_event = asyncio.Event()
self.is_running = False
self._cached_server_metadata: dict[str, Any] = {}
self._cached_supported_tasks: tuple[str, ...] = ()
self._cached_task_contracts: dict[str, Any] = {}
+ self._cached_vla_metadata: dict[str, Any] = {}
def start_all(
self,
@@ -67,13 +82,22 @@ def start_all(
parallelism_per_replica: int,
task: str,
skip_validation: bool,
+ *,
+ vla_provider_factory: str | None = None,
) -> bool:
"""Start all replica subprocesses. Returns True on success."""
device_env_var = current_platform.device_control_env_var
original_cvd = os.environ.get(device_env_var)
try:
- return self._start_all_replicas(ppl_file, parallelism_per_replica, task, skip_validation, device_env_var)
+ return self._start_all_replicas(
+ ppl_file,
+ parallelism_per_replica,
+ task,
+ skip_validation,
+ device_env_var,
+ vla_provider_factory,
+ )
finally:
self._restore_cvd(device_env_var, original_cvd)
@@ -84,6 +108,7 @@ def _start_all_replicas(
task: str,
skip_validation: bool,
device_env_var: str,
+ vla_provider_factory: str | None,
) -> bool:
for i, device_ids in enumerate(self._replica_device_ids):
visible_devices = ",".join(device_ids)
@@ -108,6 +133,7 @@ def _start_all_replicas(
self._security_level_name,
skip_validation,
self._server_config_data,
+ vla_provider_factory,
),
daemon=False,
)
@@ -141,6 +167,12 @@ def _start_all_replicas(
self._cached_server_metadata = payload.get("server_metadata", {})
self._cached_supported_tasks = tuple(payload.get("supported_tasks", []))
self._cached_task_contracts = payload.get("task_contracts", {})
+ self._cached_vla_metadata = payload.get("vla", {})
+ elif payload.get("vla", {}) != self._cached_vla_metadata:
+ logger.error(f"Replica {i} VLA metadata differs from replica 0")
+ p.terminate()
+ self._cleanup_started()
+ return False
handle = ReplicaHandle(
replica_id=i,
@@ -154,6 +186,8 @@ def _start_all_replicas(
self._available.put_nowait(i)
self.is_running = True
+ self._session_shutdown = False
+ self._session_shutdown_event.clear()
logger.info(f"Pipeline pool started: num_replicas={self._num_replicas}")
return True
@@ -174,6 +208,8 @@ def _restore_cvd(env_var: str, original: str | None) -> None:
def _evict_replica(self, idx: int, reason: str) -> None:
"""Evict a dead replica: mark status, shutdown, shrink capacity."""
with self._status_lock:
+ if self._instance_status[idx] == "dead":
+ return
self._instance_status[idx] = "dead"
self._live_count -= 1
live = self._live_count
@@ -227,6 +263,108 @@ async def acquire(self) -> AsyncIterator[ReplicaHandle]:
self._instance_status[idx] = "idle"
self._available.put_nowait(idx)
+ async def open_session(self, session_id: str, *, timeout_s: float | None = None) -> int:
+ """Reserve one healthy replica for a long-lived session."""
+ self._validate_session_id(session_id)
+ if timeout_s is not None and (
+ isinstance(timeout_s, bool) or not isinstance(timeout_s, (int, float)) or timeout_s <= 0
+ ):
+ raise ValueError("timeout_s must be None or a positive number")
+ async with self._session_leases_lock:
+ if self._session_shutdown:
+ raise RuntimeError("pipeline pool is shutting down")
+ if session_id in self._session_leases or session_id in self._opening_sessions:
+ raise ValueError(f"pipeline session is already open: {session_id!r}")
+ self._opening_sessions.add(session_id)
+
+ idx: int | None = None
+ registered = False
+ try:
+ idx = await self._take_available_replica(timeout_s=None if timeout_s is None else float(timeout_s))
+ async with self._session_leases_lock:
+ if self._session_shutdown:
+ raise RuntimeError("pipeline pool is shutting down")
+ self._session_leases[session_id] = _ReplicaSessionLease(idx, asyncio.Lock())
+ self._opening_sessions.discard(session_id)
+ registered = True
+ with self._status_lock:
+ self._instance_status[idx] = "session_reserved"
+ return idx
+ finally:
+ if not registered:
+ async with self._session_leases_lock:
+ self._opening_sessions.discard(session_id)
+ if idx is not None:
+ self._return_available_replica(idx)
+
+ @asynccontextmanager
+ async def acquire_session(self, session_id: str) -> AsyncIterator[ReplicaHandle]:
+ """Acquire the replica pinned to a session, serializing its requests."""
+ lease = await self._get_session_lease(session_id)
+ async with lease.lock:
+ async with self._session_leases_lock:
+ if self._session_leases.get(session_id) is not lease or lease.closing:
+ raise KeyError(f"unknown or closing pipeline session: {session_id!r}")
+
+ idx = lease.replica_id
+ handle = self._handles[idx]
+ if self._handle_is_dead(handle):
+ async with self._session_leases_lock:
+ self._session_leases.pop(session_id, None)
+ self._evict_replica(idx, "session replica process exited")
+ raise ReplicaDeadError(f"Replica {idx} assigned to session {session_id!r} is not alive")
+
+ with self._status_lock:
+ self._instance_status[idx] = "session_busy"
+ try:
+ yield handle
+ except ReplicaDeadError:
+ async with self._session_leases_lock:
+ self._session_leases.pop(session_id, None)
+ self._evict_replica(idx, "died during session execution")
+ raise
+ else:
+ with self._status_lock:
+ if self._instance_status[idx] != "dead":
+ self._instance_status[idx] = "session_reserved"
+
+ async def close_session(self, session_id: str) -> int:
+ """Release a session lease and return its live replica to the idle pool."""
+ lease = await self._get_session_lease(session_id)
+ async with self._session_leases_lock:
+ if self._session_leases.get(session_id) is not lease or lease.closing:
+ raise KeyError(f"unknown or closing pipeline session: {session_id!r}")
+ lease.closing = True
+
+ async with lease.lock:
+ async with self._session_leases_lock:
+ if self._session_leases.get(session_id) is not lease:
+ raise KeyError(f"unknown pipeline session: {session_id!r}")
+ self._session_leases.pop(session_id)
+ self._return_available_replica(lease.replica_id)
+ return lease.replica_id
+
+ async def session_replica_id(self, session_id: str) -> int:
+ """Return the stable replica assignment for an open session."""
+ return (await self._get_session_lease(session_id)).replica_id
+
+ async def session_bindings(self) -> dict[str, int]:
+ """Return a monitoring snapshot of active session-to-replica bindings."""
+ async with self._session_leases_lock:
+ return {session_id: lease.replica_id for session_id, lease in sorted(self._session_leases.items())}
+
+ async def run_vla_operation(
+ self,
+ session_id: str,
+ operation: str,
+ payload: dict[str, Any],
+ *,
+ timeout_s: float | None = None,
+ ) -> dict[str, Any]:
+ """Run one VLA operation on the replica permanently assigned to a session."""
+ async with self.acquire_session(session_id) as handle:
+ return await handle.run_vla_operation(operation, payload, timeout_s=timeout_s)
+
async def run_task_with_stop_event(
self,
task_data: dict,
@@ -240,12 +378,97 @@ async def run_task_with_stop_event(
async def aclose(self) -> None:
"""Shutdown all replicas."""
+ async with self._session_leases_lock:
+ self._session_shutdown = True
+ self._session_shutdown_event.set()
+ leases = tuple(self._session_leases.values())
+ for lease in leases:
+ lease.closing = True
+ for lease in leases:
+ async with lease.lock:
+ pass
+ async with self._session_leases_lock:
+ self._session_leases.clear()
+ self._opening_sessions.clear()
for h in self._handles:
h.shutdown()
self._handles.clear()
self._instance_status.clear()
self.is_running = False
+ async def _get_session_lease(self, session_id: str) -> _ReplicaSessionLease:
+ self._validate_session_id(session_id)
+ async with self._session_leases_lock:
+ try:
+ lease = self._session_leases[session_id]
+ except KeyError as error:
+ raise KeyError(f"unknown pipeline session: {session_id!r}") from error
+ if lease.closing:
+ raise KeyError(f"unknown or closing pipeline session: {session_id!r}")
+ return lease
+
+ async def _take_available_replica(self, timeout_s: float | None) -> int:
+ loop = asyncio.get_running_loop()
+ deadline = None if timeout_s is None else loop.time() + timeout_s
+ while True:
+ if self._session_shutdown:
+ raise RuntimeError("pipeline pool is shutting down")
+ with self._status_lock:
+ if self._live_count == 0:
+ raise RuntimeError("All pipeline replicas are dead; no inference capacity")
+ wait_s = 5.0 if deadline is None else min(5.0, max(0.0, deadline - loop.time()))
+ if wait_s == 0:
+ raise TimeoutError("timed out waiting for an available pipeline replica")
+ available = asyncio.create_task(self._available.get())
+ shutdown = asyncio.create_task(self._session_shutdown_event.wait())
+ try:
+ done, _ = await asyncio.wait(
+ (available, shutdown),
+ timeout=wait_s,
+ return_when=asyncio.FIRST_COMPLETED,
+ )
+ if shutdown in done:
+ raise RuntimeError("pipeline pool is shutting down")
+ if available not in done:
+ if deadline is not None and loop.time() >= deadline:
+ raise TimeoutError("timed out waiting for an available pipeline replica")
+ continue
+ idx = available.result()
+ finally:
+ for task in (available, shutdown):
+ if not task.done():
+ task.cancel()
+ for task in (available, shutdown):
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+ handle = self._handles[idx]
+ if self._handle_is_dead(handle):
+ self._evict_replica(idx, "unavailable while opening session")
+ continue
+ return idx
+
+ def _return_available_replica(self, idx: int) -> None:
+ if idx >= len(self._handles):
+ return
+ handle = self._handles[idx]
+ if self._handle_is_dead(handle):
+ if idx < len(self._instance_status) and self._instance_status[idx] != "dead":
+ self._evict_replica(idx, "session closed after replica exit")
+ return
+ with self._status_lock:
+ self._instance_status[idx] = "idle"
+ self._available.put_nowait(idx)
+
+ @staticmethod
+ def _handle_is_dead(handle: ReplicaHandle) -> bool:
+ process = getattr(handle, "process", None)
+ return handle._dead or process is not None and not process.is_alive()
+
+ @staticmethod
+ def _validate_session_id(session_id: str) -> None:
+ if not isinstance(session_id, str) or not session_id:
+ raise ValueError("session_id must be a non-empty string")
+
# --- Metadata proxy (pool overlay) ---
def server_metadata(self) -> dict:
@@ -274,6 +497,10 @@ def get_task_contract(self, task: str) -> Any:
"""Return the task-level contract for a declared task, if available."""
return self._cached_task_contracts.get(task)
+ def vla_metadata(self) -> dict[str, Any]:
+ """Return component metadata reported consistently by all VLA replicas."""
+ return dict(self._cached_vla_metadata)
+
def pool_status(self) -> list[dict]:
"""Return per-replica status for monitoring."""
with self._status_lock:
diff --git a/telefuser/service/core/replica_worker.py b/telefuser/service/core/replica_worker.py
index 61945089..9162328d 100644
--- a/telefuser/service/core/replica_worker.py
+++ b/telefuser/service/core/replica_worker.py
@@ -31,6 +31,14 @@ class ReplicaDeadError(RuntimeError):
pass
+class ReplicaVLAError(RuntimeError):
+ """Structured application error returned by a live VLA replica."""
+
+ def __init__(self, error_type: str, message: str) -> None:
+ super().__init__(message)
+ self.error_type = error_type
+
+
# ---------------------------------------------------------------------------
# Subprocess entry point
# ---------------------------------------------------------------------------
@@ -48,6 +56,7 @@ def _replica_main(
security_level_name: str,
skip_validation: bool,
server_config_data: dict[str, Any] | None = None,
+ vla_provider_factory_name: str | None = None,
) -> None:
"""Entry point for a replica subprocess.
@@ -91,22 +100,57 @@ def _replica_main(
loop.close()
return
- metadata = _collect_metadata(svc)
+ vla_provider = None
+ if vla_provider_factory_name is not None:
+ try:
+ module = svc._module
+ if module is None or not hasattr(module, vla_provider_factory_name):
+ raise RuntimeError(
+ f"Pipeline file must define optional VLA provider factory {vla_provider_factory_name}(pipeline)"
+ )
+ factory = getattr(module, vla_provider_factory_name)
+ vla_provider = factory(svc.pipeline)
+ if not callable(getattr(vla_provider, "dispatch", None)):
+ raise TypeError("VLA provider must define dispatch(operation, payload)")
+ if not callable(getattr(vla_provider, "metadata", None)):
+ raise TypeError("VLA provider must define metadata()")
+ except Exception as e:
+ conn.send(("error", f"Replica {replica_id} VLA provider init failed: {e}\n{traceback.format_exc()}"))
+ loop.run_until_complete(svc.aclose())
+ conn.close()
+ loop.close()
+ return
+
+ try:
+ metadata = _collect_metadata(svc, vla_provider=vla_provider)
+ except Exception as e:
+ conn.send(("error", f"Replica {replica_id} metadata collection failed: {e}\n{traceback.format_exc()}"))
+ if vla_provider is not None and callable(getattr(vla_provider, "close", None)):
+ vla_provider.close()
+ loop.run_until_complete(svc.aclose())
+ conn.close()
+ loop.close()
+ return
conn.send(("ready", metadata))
logger.info(f"Replica {replica_id} ready, entering task loop")
try:
- loop.run_until_complete(_task_loop(replica_id, svc, conn, cancel_event, logger))
+ loop.run_until_complete(_task_loop(replica_id, svc, conn, cancel_event, logger, vla_provider=vla_provider))
except Exception as e:
logger.error(f"Replica {replica_id} loop crashed: {e}")
finally:
+ if vla_provider is not None and callable(getattr(vla_provider, "close", None)):
+ try:
+ vla_provider.close()
+ except Exception as e:
+ logger.warning(f"Replica {replica_id} VLA provider cleanup failed: {e}")
loop.run_until_complete(svc.aclose())
loop.close()
conn.close()
logger.info(f"Replica {replica_id} exited")
-def _collect_metadata(svc: Any) -> dict:
+def _collect_metadata(svc: Any, *, vla_provider: Any | None = None) -> dict:
"""Collect metadata from initialized PipelineService for pool caching."""
task_contracts: dict[str, Any] = {}
for t in svc.supported_tasks():
@@ -115,11 +159,17 @@ def _collect_metadata(svc: Any) -> dict:
task_contracts[t] = tc.to_metadata()
elif isinstance(tc, dict):
task_contracts[t] = tc
- return {
+ metadata = {
"server_metadata": svc.server_metadata(),
"supported_tasks": list(svc.supported_tasks()),
"task_contracts": task_contracts,
}
+ if vla_provider is not None:
+ vla_metadata = vla_provider.metadata()
+ if not isinstance(vla_metadata, dict):
+ raise TypeError("VLA provider metadata must be a dictionary")
+ metadata["vla"] = vla_metadata
+ return metadata
def _cancel_watcher_fn(
@@ -139,6 +189,8 @@ async def _task_loop(
conn: Connection,
cancel_event: mp_stdlib.Event,
logger: Any,
+ *,
+ vla_provider: Any | None = None,
) -> None:
"""Persistent async task loop inside the replica subprocess."""
loop = asyncio.get_running_loop()
@@ -179,6 +231,18 @@ async def _task_loop(
finally:
cancel_event.set()
forwarder_done.wait(1.0)
+ continue
+
+ if msg[0] == "vla":
+ _, operation, payload = msg
+ if vla_provider is None:
+ conn.send(("vla_error", {"type": "RuntimeError", "message": "replica has no VLA provider"}))
+ continue
+ try:
+ result = await asyncio.to_thread(vla_provider.dispatch, operation, payload)
+ conn.send(("ok", result))
+ except Exception as e:
+ conn.send(("vla_error", {"type": type(e).__name__, "message": str(e)}))
def _recv_with_poll(conn: Connection, timeout: float) -> Any:
@@ -279,6 +343,48 @@ async def run_task(
raise RuntimeError(f"Replica {self.replica_id}: {payload}")
return payload
+ async def run_vla_operation(
+ self,
+ operation: str,
+ payload: dict[str, Any],
+ *,
+ timeout_s: float | None = None,
+ ) -> dict[str, Any]:
+ """Send one session-affine VLA operation to this replica."""
+ if not isinstance(operation, str) or not operation:
+ raise ValueError("operation must be a non-empty string")
+ if timeout_s is not None and (
+ isinstance(timeout_s, bool) or not isinstance(timeout_s, (int, float)) or timeout_s <= 0
+ ):
+ raise ValueError("timeout_s must be None or a positive number")
+ ipc_timeout = (float(timeout_s) if timeout_s is not None else 600.0) + _TASK_IPC_MARGIN_S
+ loop = asyncio.get_running_loop()
+ if not self.process.is_alive():
+ self._dead = True
+ raise ReplicaDeadError(f"Replica {self.replica_id} process is not alive")
+ try:
+ self.conn.send(("vla", operation, payload))
+ result = await loop.run_in_executor(None, self._recv_with_health_check, ipc_timeout)
+ except (EOFError, OSError) as error:
+ self._dead = True
+ raise ReplicaDeadError(f"Replica {self.replica_id} IPC failed: {error}") from error
+ if result is None:
+ self._dead = True
+ raise ReplicaDeadError(
+ f"Replica {self.replica_id} did not respond within {ipc_timeout}s "
+ f"(process alive: {self.process.is_alive()})"
+ )
+ tag, response = result
+ if tag == "vla_error":
+ if not isinstance(response, dict):
+ raise RuntimeError(f"Replica {self.replica_id} returned an invalid VLA error")
+ raise ReplicaVLAError(str(response.get("type", "RuntimeError")), str(response.get("message", "")))
+ if tag == "error":
+ raise RuntimeError(f"Replica {self.replica_id}: {response}")
+ if tag != "ok" or not isinstance(response, dict):
+ raise RuntimeError(f"Replica {self.replica_id} returned an invalid VLA response")
+ return response
+
def _recv_with_health_check(self, total_timeout: float) -> tuple[str, Any] | None:
"""Receive with periodic health checks. Returns (tag, payload) or None."""
from telefuser.utils.logging import logger
diff --git a/telefuser/service/vla_replica.py b/telefuser/service/vla_replica.py
new file mode 100644
index 00000000..a8297551
--- /dev/null
+++ b/telefuser/service/vla_replica.py
@@ -0,0 +1,175 @@
+"""Worker-local adapter between replica RPC and semantic VLA sessions."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from telefuser.vla import VLASessionManager
+from telefuser.vla.serialization import (
+ DEFAULT_MAX_TENSOR_BYTES,
+ action_space_from_wire,
+ action_space_to_wire,
+ observation_space_from_wire,
+ observation_space_to_wire,
+ robot_action_chunk_to_wire,
+ robot_observation_from_wire,
+)
+
+
+class VLAReplicaProvider:
+ """Dispatch VLA lifecycle operations inside one pipeline replica."""
+
+ def __init__(
+ self,
+ sessions: VLASessionManager,
+ *,
+ max_tensor_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+ ) -> None:
+ if not isinstance(max_tensor_bytes, int) or isinstance(max_tensor_bytes, bool) or max_tensor_bytes < 1:
+ raise ValueError("max_tensor_bytes must be a positive integer")
+ self.sessions = sessions
+ self.max_tensor_bytes = max_tensor_bytes
+
+ def metadata(self) -> dict[str, Any]:
+ """Return stable component discovery metadata without opening a session."""
+ return {
+ "model_ids": list(self.sessions.registry.model_ids()),
+ "embodiment_ids": list(self.sessions.registry.embodiment_ids()),
+ }
+
+ def dispatch(self, operation: str, payload: Mapping[str, Any]) -> dict[str, Any]:
+ """Execute one validated worker-local VLA operation."""
+ if not isinstance(operation, str):
+ raise ValueError("VLA replica operation must be a string")
+ if not isinstance(payload, Mapping):
+ raise ValueError("VLA replica payload must be an object")
+ operation = operation.upper()
+ if operation == "OPEN":
+ return self._open(payload)
+ session_id = _required_string(payload, "session_id")
+ if operation == "PREDICT":
+ return self._predict(session_id, payload)
+ if operation == "RESET":
+ self.sessions.reset(session_id, _optional_string(payload, "episode_id"))
+ return {"session_id": session_id}
+ if operation == "CLOSE":
+ self.sessions.close(session_id)
+ return {"session_id": session_id}
+ raise ValueError(f"unsupported VLA replica operation: {operation!r}")
+
+ def close(self) -> None:
+ """Release every worker-local session during replica shutdown."""
+ for session_id in self.sessions.session_ids():
+ self.sessions.close(session_id)
+
+ def _open(self, payload: Mapping[str, Any]) -> dict[str, Any]:
+ session_id = _required_string(payload, "session_id")
+ model_id = _required_string(payload, "model_id")
+ embodiment_id = _required_string(payload, "embodiment_id")
+ policy = self.sessions.registry.get_policy(model_id)
+ embodiment = self.sessions.registry.get_embodiment(embodiment_id)
+ expected_payload = payload.get("expected_robot_action_space")
+ if expected_payload is not None:
+ if not isinstance(expected_payload, Mapping):
+ raise ValueError("expected_robot_action_space must be an object")
+ action_space_from_wire(expected_payload).require_compatible(
+ embodiment.robot_action_space,
+ context="client and embodiment robot action space",
+ )
+ expected_observation_payload = payload.get("expected_robot_observation_space")
+ if expected_observation_payload is not None:
+ if not isinstance(expected_observation_payload, Mapping):
+ raise ValueError("expected_robot_observation_space must be an object")
+ observation_space_from_wire(expected_observation_payload).require_compatible(
+ embodiment.observation_space,
+ context="client and embodiment robot observation space",
+ )
+ self.sessions.open(
+ session_id,
+ model_id=model_id,
+ embodiment_id=embodiment_id,
+ episode_id=_required_string(payload, "episode_id"),
+ execute_horizon=_optional_positive_int(payload, "execute_horizon"),
+ max_observation_age_ns=_optional_positive_int(payload, "max_observation_age_ns"),
+ )
+ capabilities = policy.capabilities()
+ return {
+ "session_id": session_id,
+ "model_id": model_id,
+ "embodiment_id": embodiment_id,
+ "model_action_space": action_space_to_wire(capabilities.output_action_space),
+ "robot_action_space": action_space_to_wire(embodiment.robot_action_space),
+ "robot_observation_space": observation_space_to_wire(embodiment.observation_space),
+ "max_horizon": capabilities.max_horizon,
+ "stateful": capabilities.stateful,
+ "supports_seed": capabilities.supports_seed,
+ }
+
+ def _predict(self, session_id: str, payload: Mapping[str, Any]) -> dict[str, Any]:
+ observation_payload = payload.get("observation")
+ if not isinstance(observation_payload, Mapping):
+ raise ValueError("PREDICT requires an observation object")
+ max_tensor_bytes = payload.get("_max_tensor_bytes", self.max_tensor_bytes)
+ if not isinstance(max_tensor_bytes, int) or isinstance(max_tensor_bytes, bool) or max_tensor_bytes < 1:
+ raise ValueError("_max_tensor_bytes must be a positive integer")
+ observation = robot_observation_from_wire(observation_payload, max_tensor_bytes=max_tensor_bytes)
+ timestamp_ns = _required_nonnegative_int(payload, "observation_timestamp_ns")
+ if timestamp_ns != observation.state.timestamp_ns:
+ raise ValueError("PREDICT observation_timestamp_ns must match observation.state.timestamp_ns")
+ chunk = self.sessions.get(session_id).predict(
+ observation,
+ _required_string(payload, "instruction"),
+ _required_nonnegative_int(payload, "sequence_id"),
+ seed=_optional_integer(payload, "seed"),
+ now_ns=_optional_nonnegative_int(payload, "observation_clock_now_ns"),
+ )
+ return {"session_id": session_id, "chunk": robot_action_chunk_to_wire(chunk)}
+
+
+def _required_string(payload: Mapping[str, Any], field: str) -> str:
+ value = payload.get(field)
+ if not isinstance(value, str) or not value:
+ raise ValueError(f"{field} must be a non-empty string")
+ return value
+
+
+def _optional_string(payload: Mapping[str, Any], field: str) -> str | None:
+ value = payload.get(field)
+ if value is None:
+ return None
+ if not isinstance(value, str) or not value:
+ raise ValueError(f"{field} must be a non-empty string")
+ return value
+
+
+def _required_nonnegative_int(payload: Mapping[str, Any], field: str) -> int:
+ value = _optional_nonnegative_int(payload, field)
+ if value is None:
+ raise ValueError(f"{field} is required")
+ return value
+
+
+def _optional_nonnegative_int(payload: Mapping[str, Any], field: str) -> int | None:
+ value = payload.get(field)
+ if value is None:
+ return None
+ if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+ raise ValueError(f"{field} must be a non-negative integer")
+ return value
+
+
+def _optional_positive_int(payload: Mapping[str, Any], field: str) -> int | None:
+ value = _optional_nonnegative_int(payload, field)
+ if value is not None and value < 1:
+ raise ValueError(f"{field} must be a positive integer")
+ return value
+
+
+def _optional_integer(payload: Mapping[str, Any], field: str) -> int | None:
+ value = payload.get(field)
+ if value is None:
+ return None
+ if not isinstance(value, int) or isinstance(value, bool):
+ raise ValueError(f"{field} must be an integer")
+ return value
diff --git a/telefuser/service/vla_session.py b/telefuser/service/vla_session.py
new file mode 100644
index 00000000..c83b4b55
--- /dev/null
+++ b/telefuser/service/vla_session.py
@@ -0,0 +1,777 @@
+"""Generic versioned WebSocket transport for typed VLA sessions."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import math
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Protocol
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+
+from telefuser.service.core.pipeline_pool import PipelinePool
+from telefuser.service.core.replica_worker import ReplicaDeadError, ReplicaVLAError
+from telefuser.vla import VLASessionManager
+from telefuser.vla.contracts import ActionSpaceSpec, ObservationSpaceSpec, RobotActionChunk
+from telefuser.vla.runtime import ActionChunkStateMachine, ChunkStatus
+from telefuser.vla.runtime.chunk_state import ChunkTicket
+from telefuser.vla.serialization import (
+ DEFAULT_MAX_TENSOR_BYTES,
+ VLA_SESSION_PROTOCOL_VERSION,
+ VLA_WIRE_ENCODING,
+ action_space_from_wire,
+ action_space_to_wire,
+ dumps_wire_message,
+ loads_wire_message,
+ observation_space_from_wire,
+ observation_space_to_wire,
+ robot_action_chunk_from_wire,
+ robot_action_chunk_to_wire,
+ robot_observation_from_wire,
+)
+
+DEFAULT_MAX_VLA_MESSAGE_BYTES = 64 * 1024 * 1024
+
+
+class VLAErrorCode(str, Enum):
+ """Stable error categories returned by the VLA session protocol."""
+
+ INVALID_MESSAGE = "invalid_message"
+ UNSUPPORTED_VERSION = "unsupported_version"
+ SESSION_EXISTS = "session_exists"
+ UNKNOWN_SESSION = "unknown_session"
+ UNKNOWN_COMPONENT = "unknown_component"
+ ACTION_SPACE_MISMATCH = "action_space_mismatch"
+ OBSERVATION_SPACE_MISMATCH = "observation_space_mismatch"
+ OUT_OF_ORDER = "out_of_order"
+ SUPERSEDED = "superseded"
+ EXPIRED = "expired"
+ TIMEOUT = "timeout"
+ REPLICA_UNAVAILABLE = "replica_unavailable"
+ SESSION_UNAVAILABLE = "session_unavailable"
+ INTERNAL_ERROR = "internal_error"
+
+
+class VLAProtocolError(ValueError):
+ """Protocol failure carrying a stable machine-readable error code."""
+
+ def __init__(self, code: VLAErrorCode, message: str) -> None:
+ super().__init__(message)
+ self.code = code
+
+
+class _VLABackend(Protocol):
+ def metadata(self) -> Mapping[str, Any]: ...
+
+ async def open(self, payload: Mapping[str, Any]) -> dict[str, Any]: ...
+
+ async def predict(self, session_id: str, payload: Mapping[str, Any]) -> RobotActionChunk: ...
+
+ async def reset(self, session_id: str, episode_id: str | None) -> None: ...
+
+ async def close(self, session_id: str) -> None: ...
+
+
+class _LocalVLABackend:
+ def __init__(self, sessions: VLASessionManager, max_tensor_bytes: int) -> None:
+ self.sessions = sessions
+ self.max_tensor_bytes = max_tensor_bytes
+
+ def metadata(self) -> Mapping[str, Any]:
+ return {
+ "model_ids": list(self.sessions.registry.model_ids()),
+ "embodiment_ids": list(self.sessions.registry.embodiment_ids()),
+ }
+
+ async def open(self, payload: Mapping[str, Any]) -> dict[str, Any]:
+ session_id = _required_string(payload, "session_id")
+ if session_id in self.sessions.session_ids():
+ raise VLAProtocolError(VLAErrorCode.SESSION_EXISTS, f"VLA session is already open: {session_id!r}")
+ model_id = _required_string(payload, "model_id")
+ embodiment_id = _required_string(payload, "embodiment_id")
+ try:
+ policy = self.sessions.registry.get_policy(model_id)
+ embodiment = self.sessions.registry.get_embodiment(embodiment_id)
+ except KeyError as error:
+ raise VLAProtocolError(VLAErrorCode.UNKNOWN_COMPONENT, str(error)) from error
+ _validate_expected_action_space(payload, embodiment.robot_action_space)
+ _validate_expected_observation_space(payload, embodiment.observation_space)
+ await asyncio.to_thread(
+ self.sessions.open,
+ session_id,
+ model_id=model_id,
+ embodiment_id=embodiment_id,
+ episode_id=_required_string(payload, "episode_id"),
+ execute_horizon=_optional_positive_int(payload, "execute_horizon"),
+ max_observation_age_ns=_optional_positive_int(payload, "max_observation_age_ns"),
+ )
+ capabilities = policy.capabilities()
+ return {
+ "session_id": session_id,
+ "model_id": model_id,
+ "embodiment_id": embodiment_id,
+ "model_action_space": action_space_to_wire(capabilities.output_action_space),
+ "robot_action_space": action_space_to_wire(embodiment.robot_action_space),
+ "robot_observation_space": observation_space_to_wire(embodiment.observation_space),
+ "max_horizon": capabilities.max_horizon,
+ "stateful": capabilities.stateful,
+ "supports_seed": capabilities.supports_seed,
+ }
+
+ async def predict(self, session_id: str, payload: Mapping[str, Any]) -> RobotActionChunk:
+ observation_payload = payload.get("observation")
+ if not isinstance(observation_payload, Mapping):
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "PREDICT requires an observation object")
+ observation = robot_observation_from_wire(observation_payload, max_tensor_bytes=self.max_tensor_bytes)
+ timestamp_ns = _required_nonnegative_int(payload, "observation_timestamp_ns")
+ if timestamp_ns != observation.state.timestamp_ns:
+ raise VLAProtocolError(
+ VLAErrorCode.INVALID_MESSAGE,
+ "PREDICT observation_timestamp_ns must match observation.state.timestamp_ns",
+ )
+ return await asyncio.to_thread(
+ self.sessions.get(session_id).predict,
+ observation,
+ _required_string(payload, "instruction"),
+ _required_nonnegative_int(payload, "sequence_id"),
+ seed=_optional_integer(payload, "seed"),
+ now_ns=_optional_nonnegative_int(payload, "observation_clock_now_ns"),
+ )
+
+ async def reset(self, session_id: str, episode_id: str | None) -> None:
+ await asyncio.to_thread(self.sessions.reset, session_id, episode_id)
+
+ async def close(self, session_id: str) -> None:
+ self.sessions.close(session_id)
+
+
+class _PipelinePoolVLABackend:
+ def __init__(self, pool: PipelinePool, max_tensor_bytes: int) -> None:
+ self.pool = pool
+ self.max_tensor_bytes = max_tensor_bytes
+
+ def metadata(self) -> Mapping[str, Any]:
+ return self.pool.vla_metadata()
+
+ async def open(self, payload: Mapping[str, Any]) -> dict[str, Any]:
+ session_id = _required_string(payload, "session_id")
+ try:
+ await self.pool.open_session(session_id)
+ return await self.pool.run_vla_operation(session_id, "OPEN", dict(payload))
+ except Exception:
+ with contextlib.suppress(KeyError):
+ await self.pool.close_session(session_id)
+ raise
+
+ async def predict(self, session_id: str, payload: Mapping[str, Any]) -> RobotActionChunk:
+ replica_payload = dict(payload)
+ replica_payload["_max_tensor_bytes"] = self.max_tensor_bytes
+ response = await self.pool.run_vla_operation(session_id, "PREDICT", replica_payload)
+ chunk_payload = response.get("chunk")
+ if not isinstance(chunk_payload, Mapping):
+ raise RuntimeError("VLA replica returned an invalid action chunk")
+ return robot_action_chunk_from_wire(chunk_payload, max_tensor_bytes=self.max_tensor_bytes)
+
+ async def reset(self, session_id: str, episode_id: str | None) -> None:
+ payload: dict[str, Any] = {"session_id": session_id}
+ if episode_id is not None:
+ payload["episode_id"] = episode_id
+ await self.pool.run_vla_operation(session_id, "RESET", payload)
+
+ async def close(self, session_id: str) -> None:
+ try:
+ await self.pool.run_vla_operation(session_id, "CLOSE", {"session_id": session_id})
+ finally:
+ with contextlib.suppress(KeyError):
+ await self.pool.close_session(session_id)
+
+
+@dataclass
+class _PredictionJob:
+ message: Mapping[str, Any]
+ ticket: ChunkTicket
+ response: asyncio.Future[dict[str, Any]]
+
+
+class _SessionRuntime:
+ """Run one prediction at a time while retaining only the latest waiting request."""
+
+ def __init__(
+ self,
+ backend: _VLABackend,
+ session_id: str,
+ state: ActionChunkStateMachine,
+ *,
+ stateful: bool,
+ ) -> None:
+ self.backend = backend
+ self.session_id = session_id
+ self.state = state
+ self.stateful = stateful
+ self._waiting: _PredictionJob | None = None
+ self._runner: asyncio.Task[None] | None = None
+ self._accepting = True
+ self._recovering = False
+ self._barrier_active = False
+
+ def submit(self, message: Mapping[str, Any]) -> asyncio.Future[dict[str, Any]]:
+ """Admit a request immediately without waiting for the replica."""
+ loop = asyncio.get_running_loop()
+ response: asyncio.Future[dict[str, Any]] = loop.create_future()
+ if not self._accepting or self._recovering:
+ response.set_exception(
+ VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, "session is not accepting predictions")
+ )
+ return response
+
+ timestamp_ns = _required_nonnegative_int(message, "observation_timestamp_ns")
+ timeout_s = _optional_timeout_s(message)
+ ticket = self.state.submit(
+ _required_nonnegative_int(message, "sequence_id"),
+ timestamp_ns,
+ request_ttl_ms=None if timeout_s is None else timeout_s * 1000.0,
+ observation_clock_now_ns=_optional_nonnegative_int(message, "observation_clock_now_ns"),
+ )
+ admission = self.state.status(ticket)
+ if admission is ChunkStatus.REJECTED:
+ response.set_exception(
+ VLAProtocolError(VLAErrorCode.OUT_OF_ORDER, self.state.reason(ticket) or "action chunk was rejected")
+ )
+ return response
+ if admission is ChunkStatus.EXPIRED:
+ response.set_exception(
+ VLAProtocolError(VLAErrorCode.EXPIRED, self.state.reason(ticket) or "observation expired")
+ )
+ return response
+
+ job = _PredictionJob(message, ticket, response)
+ if self._waiting is not None:
+ self._fail(
+ self._waiting,
+ VLAProtocolError(VLAErrorCode.SUPERSEDED, "a newer observation superseded the waiting request"),
+ )
+ self._waiting = job
+ if self._runner is None:
+ self._runner = asyncio.create_task(self._run())
+ return response
+
+ async def reset(self, episode_id: str | None) -> None:
+ """Apply a barrier, discard earlier results, and reset policy history."""
+ self._accepting = False
+ self._barrier_active = True
+ self.state.reset(episode_id)
+ if self._waiting is not None:
+ self._fail(
+ self._waiting,
+ VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, "request was discarded by RESET"),
+ )
+ self._waiting = None
+ await self._await_runner()
+ await self.backend.reset(self.session_id, episode_id)
+ self.state.reset(episode_id)
+ self._barrier_active = False
+ self._accepting = True
+
+ async def close(self, *, disconnect: bool = False) -> None:
+ """Stop admission, isolate late results, and close the backend session."""
+ self._accepting = False
+ self._barrier_active = True
+ self.state.reset()
+ if self._waiting is not None:
+ self._fail(
+ self._waiting,
+ VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, "request was discarded by CLOSE"),
+ )
+ self._waiting = None
+ await self._await_runner()
+ await self.backend.close(self.session_id)
+ if disconnect:
+ self.state.disconnect()
+
+ async def _run(self) -> None:
+ try:
+ while self._waiting is not None:
+ job = self._waiting
+ self._waiting = None
+ if self.state.status(job.ticket) is not ChunkStatus.PENDING:
+ self._fail(
+ job,
+ VLAProtocolError(
+ VLAErrorCode.SUPERSEDED,
+ self.state.reason(job.ticket) or "action chunk was superseded",
+ ),
+ )
+ continue
+ await self._predict(job)
+ except Exception as error:
+ self._accepting = False
+ if self._waiting is not None:
+ self._fail(self._waiting, error)
+ self._waiting = None
+ finally:
+ self._runner = None
+
+ async def _predict(self, job: _PredictionJob) -> None:
+ self.state.mark_inference_started(job.ticket)
+ predict_task = asyncio.create_task(self.backend.predict(self.session_id, job.message))
+ timeout_s = _optional_timeout_s(job.message)
+ try:
+ if timeout_s is None:
+ chunk = await predict_task
+ else:
+ chunk = await asyncio.wait_for(asyncio.shield(predict_task), timeout=timeout_s)
+ except asyncio.TimeoutError:
+ self._recovering = True
+ self._fail(job, VLAProtocolError(VLAErrorCode.TIMEOUT, "PREDICT exceeded request_ttl_ms"))
+ try:
+ try:
+ chunk = await predict_task
+ except Exception as error:
+ self.state.reject(job.ticket, str(error) or "timed-out VLA inference failed")
+ else:
+ self.state.complete(
+ job.ticket,
+ chunk,
+ observation_clock_now_ns=_optional_nonnegative_int(
+ job.message,
+ "observation_clock_now_ns",
+ ),
+ )
+ if not self._barrier_active:
+ await self.backend.reset(self.session_id, None)
+ self.state.reset()
+ finally:
+ self._recovering = False
+ return
+ except Exception as error:
+ self.state.reject(job.ticket, str(error) or "VLA inference failed")
+ self._fail(job, error)
+ return
+
+ completion = self.state.complete(
+ job.ticket,
+ chunk,
+ observation_clock_now_ns=_optional_nonnegative_int(job.message, "observation_clock_now_ns"),
+ )
+ if completion is ChunkStatus.READY:
+ self._succeed(
+ job,
+ {
+ "type": "ACTION_CHUNK",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": _request_id(job.message),
+ "session_id": self.session_id,
+ "chunk": robot_action_chunk_to_wire(chunk),
+ },
+ )
+ return
+ code = VLAErrorCode.EXPIRED if completion is ChunkStatus.EXPIRED else VLAErrorCode.SUPERSEDED
+ self._fail(job, VLAProtocolError(code, self.state.reason(job.ticket) or "action chunk was discarded"))
+ if self.stateful and not self._barrier_active:
+ await self.backend.reset(self.session_id, None)
+
+ async def _await_runner(self) -> None:
+ if self._runner is not None:
+ await self._runner
+
+ @staticmethod
+ def _succeed(job: _PredictionJob, response: dict[str, Any]) -> None:
+ if not job.response.done():
+ job.response.set_result(response)
+
+ @staticmethod
+ def _fail(job: _PredictionJob, error: Exception) -> None:
+ if not job.response.done():
+ job.response.set_exception(error)
+
+
+def create_vla_session_app(
+ sessions: VLASessionManager,
+ *,
+ max_message_bytes: int = DEFAULT_MAX_VLA_MESSAGE_BYTES,
+ max_tensor_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+) -> FastAPI:
+ """Create an additive WebSocket app for an already-loaded VLA registry."""
+ return _create_vla_session_app(
+ _LocalVLABackend(sessions, max_tensor_bytes),
+ max_message_bytes=max_message_bytes,
+ max_tensor_bytes=max_tensor_bytes,
+ )
+
+
+def create_pipeline_pool_vla_session_app(
+ pool: PipelinePool,
+ *,
+ max_message_bytes: int = DEFAULT_MAX_VLA_MESSAGE_BYTES,
+ max_tensor_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+ close_pool_on_shutdown: bool = True,
+) -> FastAPI:
+ """Create a VLA WebSocket app backed by session-affine pipeline replicas."""
+ app = _create_vla_session_app(
+ _PipelinePoolVLABackend(pool, max_tensor_bytes),
+ max_message_bytes=max_message_bytes,
+ max_tensor_bytes=max_tensor_bytes,
+ )
+ if close_pool_on_shutdown:
+
+ @app.on_event("shutdown")
+ async def close_pool() -> None:
+ await pool.aclose()
+
+ return app
+
+
+def _create_vla_session_app(
+ backend: _VLABackend,
+ *,
+ max_message_bytes: int,
+ max_tensor_bytes: int,
+) -> FastAPI:
+ if not isinstance(max_message_bytes, int) or isinstance(max_message_bytes, bool) or max_message_bytes < 1:
+ raise ValueError("max_message_bytes must be a positive integer")
+ if not isinstance(max_tensor_bytes, int) or isinstance(max_tensor_bytes, bool) or max_tensor_bytes < 1:
+ raise ValueError("max_tensor_bytes must be a positive integer")
+
+ app = FastAPI(title="TeleFuser VLA Session")
+
+ @app.get("/healthz")
+ async def healthz() -> dict[str, str]:
+ return {"status": "ok"}
+
+ @app.websocket("/v1/vla/session")
+ async def session_socket(websocket: WebSocket) -> None:
+ await websocket.accept()
+ owned_sessions: set[str] = set()
+ runtimes: dict[str, _SessionRuntime] = {}
+ delivery_tasks: set[asyncio.Task[None]] = set()
+ send_lock = asyncio.Lock()
+ metadata = backend.metadata()
+ await _locked_send(
+ websocket,
+ send_lock,
+ {
+ "type": "HELLO",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "encoding": VLA_WIRE_ENCODING,
+ "operations": ["OPEN", "PREDICT", "RESET", "CLOSE"],
+ "capabilities": {
+ "observation_contract": True,
+ "concurrent_predict": True,
+ "prediction_queue": "latest-wins",
+ "control_barriers": True,
+ },
+ "model_ids": list(metadata.get("model_ids", [])),
+ "embodiment_ids": list(metadata.get("embodiment_ids", [])),
+ "max_message_bytes": max_message_bytes,
+ "max_tensor_bytes": max_tensor_bytes,
+ },
+ )
+ try:
+ while True:
+ request_id: str | int | None = None
+ try:
+ message = loads_wire_message(
+ await websocket.receive_text(),
+ max_message_bytes=max_message_bytes,
+ )
+ request_id = _request_id(message)
+ operation = message.get("type")
+ if isinstance(operation, str) and operation.upper() == "PREDICT":
+ session_id = _required_string(message, "session_id")
+ if session_id not in owned_sessions:
+ raise VLAProtocolError(
+ VLAErrorCode.UNKNOWN_SESSION,
+ f"session is not open on this connection: {session_id!r}",
+ )
+ response_future = runtimes[session_id].submit(message)
+ task = asyncio.create_task(
+ _deliver_prediction(websocket, send_lock, response_future, request_id=request_id)
+ )
+ delivery_tasks.add(task)
+ task.add_done_callback(delivery_tasks.discard)
+ continue
+ response = await _handle_control_message(
+ backend,
+ message,
+ owned_sessions=owned_sessions,
+ runtimes=runtimes,
+ )
+ except WebSocketDisconnect:
+ break
+ except Exception as error:
+ response = _error_response(error, request_id=request_id)
+ await _locked_send(websocket, send_lock, response)
+ finally:
+ for runtime in tuple(runtimes.values()):
+ with contextlib.suppress(Exception):
+ await runtime.close(disconnect=True)
+ for task in tuple(delivery_tasks):
+ task.cancel()
+ for task in tuple(delivery_tasks):
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+
+ return app
+
+
+async def _handle_control_message(
+ backend: _VLABackend,
+ message: Mapping[str, Any],
+ *,
+ owned_sessions: set[str],
+ runtimes: dict[str, _SessionRuntime],
+) -> dict[str, Any]:
+ operation = message.get("type")
+ if not isinstance(operation, str):
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "message type must be a string")
+ operation = operation.upper()
+ request_id = _request_id(message)
+ if operation == "OPEN":
+ return await _open_session(
+ backend,
+ message,
+ owned_sessions=owned_sessions,
+ runtimes=runtimes,
+ request_id=request_id,
+ )
+
+ session_id = _required_string(message, "session_id")
+ if session_id not in owned_sessions:
+ raise VLAProtocolError(VLAErrorCode.UNKNOWN_SESSION, f"session is not open on this connection: {session_id!r}")
+ if operation == "PREDICT":
+ raise RuntimeError("PREDICT must be scheduled by the connection receiver")
+ if operation == "RESET":
+ episode_id = _optional_string(message, "episode_id")
+ await runtimes[session_id].reset(episode_id)
+ return _success_response("RESET", session_id, request_id)
+ if operation == "CLOSE":
+ await runtimes[session_id].close()
+ owned_sessions.remove(session_id)
+ runtimes.pop(session_id).state.disconnect()
+ return _success_response("CLOSE", session_id, request_id)
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"unsupported VLA operation: {operation!r}")
+
+
+async def _open_session(
+ backend: _VLABackend,
+ message: Mapping[str, Any],
+ *,
+ owned_sessions: set[str],
+ runtimes: dict[str, _SessionRuntime],
+ request_id: str | int | None,
+) -> dict[str, Any]:
+ version = message.get("protocol_version")
+ if version != VLA_SESSION_PROTOCOL_VERSION:
+ raise VLAProtocolError(VLAErrorCode.UNSUPPORTED_VERSION, f"unsupported protocol_version: {version!r}")
+ session_id = _required_string(message, "session_id")
+ if session_id in owned_sessions:
+ raise VLAProtocolError(VLAErrorCode.SESSION_EXISTS, f"VLA session is already open: {session_id!r}")
+ opened = await backend.open(message)
+ owned_sessions.add(session_id)
+ max_horizon = opened.get("max_horizon")
+ if not isinstance(max_horizon, int) or isinstance(max_horizon, bool) or max_horizon < 1:
+ await backend.close(session_id)
+ owned_sessions.remove(session_id)
+ raise RuntimeError("VLA backend returned an invalid max_horizon")
+ state = ActionChunkStateMachine(
+ _required_string(message, "episode_id"),
+ execute_horizon=_optional_positive_int(message, "execute_horizon") or max_horizon,
+ max_observation_age_ns=_optional_positive_int(message, "max_observation_age_ns"),
+ )
+ runtimes[session_id] = _SessionRuntime(
+ backend,
+ session_id,
+ state,
+ stateful=bool(opened.get("stateful", False)),
+ )
+ return {
+ "type": "OPENED",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": request_id,
+ **opened,
+ }
+
+
+def _validate_expected_action_space(payload: Mapping[str, Any], robot_action_space: ActionSpaceSpec) -> None:
+ expected_payload = payload.get("expected_robot_action_space")
+ if expected_payload is None:
+ return
+ if not isinstance(expected_payload, Mapping):
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "expected_robot_action_space must be an object")
+ action_space_from_wire(expected_payload).require_compatible(
+ robot_action_space,
+ context="client and embodiment robot action space",
+ )
+
+
+def _validate_expected_observation_space(
+ payload: Mapping[str, Any],
+ robot_observation_space: ObservationSpaceSpec,
+) -> None:
+ expected_payload = payload.get("expected_robot_observation_space")
+ if expected_payload is None:
+ return
+ if not isinstance(expected_payload, Mapping):
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "expected_robot_observation_space must be an object")
+ observation_space_from_wire(expected_payload).require_compatible(
+ robot_observation_space,
+ context="client and embodiment robot observation space",
+ )
+
+
+async def _send(websocket: WebSocket, payload: Mapping[str, Any]) -> None:
+ await websocket.send_text(dumps_wire_message(payload))
+
+
+async def _locked_send(
+ websocket: WebSocket,
+ lock: asyncio.Lock,
+ payload: Mapping[str, Any],
+) -> None:
+ async with lock:
+ await _send(websocket, payload)
+
+
+async def _deliver_prediction(
+ websocket: WebSocket,
+ lock: asyncio.Lock,
+ response: asyncio.Future[dict[str, Any]],
+ *,
+ request_id: str | int | None,
+) -> None:
+ try:
+ payload = await response
+ except Exception as error:
+ payload = _error_response(error, request_id=request_id)
+ with contextlib.suppress(WebSocketDisconnect, RuntimeError):
+ await _locked_send(websocket, lock, payload)
+
+
+def _error_response(error: Exception, *, request_id: str | int | None) -> dict[str, Any]:
+ code = _error_code(error)
+ message = str(error) if code is not VLAErrorCode.INTERNAL_ERROR else "internal VLA session error"
+ return {
+ "type": "ERROR",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": request_id,
+ "error": {"code": code.value, "message": message},
+ }
+
+
+def _error_code(error: Exception) -> VLAErrorCode:
+ if isinstance(error, VLAProtocolError):
+ return error.code
+ if isinstance(error, ReplicaDeadError):
+ return VLAErrorCode.REPLICA_UNAVAILABLE
+ if isinstance(error, ReplicaVLAError):
+ if error.error_type == "KeyError":
+ message = str(error).lower()
+ if "model_id" in message or "embodiment_id" in message:
+ return VLAErrorCode.UNKNOWN_COMPONENT
+ return VLAErrorCode.UNKNOWN_SESSION
+ if error.error_type in {"ValueError", "TypeError"}:
+ return _value_error_code(str(error))
+ return VLAErrorCode.INTERNAL_ERROR
+ if isinstance(error, KeyError):
+ return VLAErrorCode.UNKNOWN_SESSION
+ if isinstance(error, ValueError):
+ return _value_error_code(str(error))
+ if isinstance(error, RuntimeError) and "replica" in str(error).lower():
+ return VLAErrorCode.REPLICA_UNAVAILABLE
+ return VLAErrorCode.INTERNAL_ERROR
+
+
+def _value_error_code(message: str) -> VLAErrorCode:
+ normalized = message.lower()
+ if "observation space" in normalized or "observation-space" in normalized:
+ return VLAErrorCode.OBSERVATION_SPACE_MISMATCH
+ if "action space" in normalized or "action-space" in normalized:
+ return VLAErrorCode.ACTION_SPACE_MISMATCH
+ if "sequence_id" in normalized or "must increase" in normalized:
+ return VLAErrorCode.OUT_OF_ORDER
+ if "stale" in normalized or "expired" in normalized:
+ return VLAErrorCode.EXPIRED
+ if "already open" in normalized:
+ return VLAErrorCode.SESSION_EXISTS
+ return VLAErrorCode.INVALID_MESSAGE
+
+
+def _success_response(operation: str, session_id: str, request_id: str | int | None) -> dict[str, Any]:
+ return {
+ "type": f"{operation}_ACK",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": request_id,
+ "session_id": session_id,
+ }
+
+
+def _request_id(message: Mapping[str, Any]) -> str | int | None:
+ value = message.get("request_id")
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (str, int)) or isinstance(value, str) and not value:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "request_id must be a non-empty string or integer")
+ return value
+
+
+def _required_string(message: Mapping[str, Any], field: str) -> str:
+ value = message.get(field)
+ if not isinstance(value, str) or not value:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"{field} must be a non-empty string")
+ return value
+
+
+def _optional_string(message: Mapping[str, Any], field: str) -> str | None:
+ value = message.get(field)
+ if value is None:
+ return None
+ if not isinstance(value, str) or not value:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"{field} must be a non-empty string")
+ return value
+
+
+def _required_nonnegative_int(message: Mapping[str, Any], field: str) -> int:
+ value = _optional_nonnegative_int(message, field)
+ if value is None:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"{field} is required")
+ return value
+
+
+def _optional_nonnegative_int(message: Mapping[str, Any], field: str) -> int | None:
+ value = message.get(field)
+ if value is None:
+ return None
+ if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"{field} must be a non-negative integer")
+ return value
+
+
+def _optional_positive_int(message: Mapping[str, Any], field: str) -> int | None:
+ value = _optional_nonnegative_int(message, field)
+ if value is not None and value < 1:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"{field} must be a positive integer")
+ return value
+
+
+def _optional_integer(message: Mapping[str, Any], field: str) -> int | None:
+ value = message.get(field)
+ if value is None:
+ return None
+ if not isinstance(value, int) or isinstance(value, bool):
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"{field} must be an integer")
+ return value
+
+
+def _optional_timeout_s(message: Mapping[str, Any]) -> float | None:
+ value = message.get("request_ttl_ms")
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "request_ttl_ms must be a positive finite number")
+ timeout_ms = float(value)
+ if not math.isfinite(timeout_ms) or timeout_ms <= 0:
+ raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, "request_ttl_ms must be a positive finite number")
+ return timeout_ms / 1000.0
diff --git a/telefuser/vla/__init__.py b/telefuser/vla/__init__.py
new file mode 100644
index 00000000..ce212696
--- /dev/null
+++ b/telefuser/vla/__init__.py
@@ -0,0 +1,63 @@
+"""Public semantic contracts for vision-language-action integrations."""
+
+from .contracts import (
+ ActionChunk,
+ ActionSpaceSpec,
+ ImageObservationSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+ VLACapabilities,
+ VLARequest,
+)
+from .embodiment import EmbodimentAdapter
+from .policy import VLAPolicy
+from .registry import VLARegistry
+from .serialization import (
+ VLA_SESSION_PROTOCOL_VERSION,
+ VLA_WIRE_ENCODING,
+ VLA_WIRE_SCHEMA_VERSION,
+ action_space_from_wire,
+ action_space_to_wire,
+ observation_space_from_wire,
+ observation_space_to_wire,
+ robot_action_chunk_from_wire,
+ robot_action_chunk_to_wire,
+ robot_observation_from_wire,
+ robot_observation_to_wire,
+)
+from .session import VLASession, VLASessionContract, VLASessionManager
+
+__all__ = [
+ "ActionChunk",
+ "ActionSpaceSpec",
+ "EmbodimentAdapter",
+ "ImageObservationSpec",
+ "ModelActionChunk",
+ "ModelObservation",
+ "ObservationSpaceSpec",
+ "RobotActionChunk",
+ "RobotObservation",
+ "RobotState",
+ "VLACapabilities",
+ "VLAPolicy",
+ "VLARegistry",
+ "VLARequest",
+ "VLA_WIRE_ENCODING",
+ "VLA_WIRE_SCHEMA_VERSION",
+ "VLA_SESSION_PROTOCOL_VERSION",
+ "VLASession",
+ "VLASessionContract",
+ "VLASessionManager",
+ "action_space_from_wire",
+ "action_space_to_wire",
+ "observation_space_from_wire",
+ "observation_space_to_wire",
+ "robot_action_chunk_from_wire",
+ "robot_action_chunk_to_wire",
+ "robot_observation_from_wire",
+ "robot_observation_to_wire",
+]
diff --git a/telefuser/vla/contracts.py b/telefuser/vla/contracts.py
new file mode 100644
index 00000000..3846473c
--- /dev/null
+++ b/telefuser/vla/contracts.py
@@ -0,0 +1,334 @@
+"""Semantic contracts shared by VLA policies, embodiments, and simulators."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field, replace
+from types import MappingProxyType
+from typing import Any, Mapping
+
+import torch
+
+
+@dataclass(frozen=True)
+class ImageObservationSpec:
+ """Describe one named image in a robot observation."""
+
+ name: str
+ dtype: str = "uint8"
+ layout: str = "HWC"
+ channels: int = 3
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.name, str) or not self.name:
+ raise ValueError("image observation name must be a non-empty string")
+ if not isinstance(self.dtype, str) or not self.dtype:
+ raise ValueError("image observation dtype must be a non-empty string")
+ if self.layout not in {"HWC", "CHW"}:
+ raise ValueError("image observation layout must be HWC or CHW")
+ if not isinstance(self.channels, int) or isinstance(self.channels, bool) or self.channels < 1:
+ raise ValueError("image observation channels must be a positive integer")
+
+
+@dataclass(frozen=True)
+class ObservationSpaceSpec:
+ """Describe state and named image inputs expected by an embodiment."""
+
+ state_dimension_names: tuple[str, ...]
+ images: tuple[ImageObservationSpec, ...] = ()
+ allow_extra_images: bool = True
+ timestamp_unit: str = "nanosecond"
+ timestamp_clock: str = "observation_source_monotonic"
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.state_dimension_names, tuple) or not self.state_dimension_names:
+ raise ValueError("state_dimension_names must be a non-empty tuple")
+ if any(not isinstance(name, str) or not name for name in self.state_dimension_names):
+ raise ValueError("state_dimension_names must contain non-empty strings")
+ if len(set(self.state_dimension_names)) != len(self.state_dimension_names):
+ raise ValueError("state_dimension_names must be unique")
+ if not isinstance(self.images, tuple) or any(
+ not isinstance(spec, ImageObservationSpec) for spec in self.images
+ ):
+ raise ValueError("images must be a tuple of ImageObservationSpec values")
+ if len({spec.name for spec in self.images}) != len(self.images):
+ raise ValueError("image observation names must be unique")
+ if not isinstance(self.allow_extra_images, bool):
+ raise ValueError("allow_extra_images must be a boolean")
+ if self.timestamp_unit != "nanosecond":
+ raise ValueError("timestamp_unit must be nanosecond")
+ if not isinstance(self.timestamp_clock, str) or not self.timestamp_clock:
+ raise ValueError("timestamp_clock must be a non-empty string")
+
+ def require_compatible(self, actual: "ObservationSpaceSpec", *, context: str = "observation space") -> None:
+ """Reject semantic mismatches between declared observation contracts."""
+ mismatches = []
+ if self.state_dimension_names != actual.state_dimension_names:
+ mismatches.append("state_dimension_names")
+ if {spec.name: spec for spec in self.images} != {spec.name: spec for spec in actual.images}:
+ mismatches.append("images")
+ if self.allow_extra_images != actual.allow_extra_images:
+ mismatches.append("allow_extra_images")
+ if self.timestamp_unit != actual.timestamp_unit:
+ mismatches.append("timestamp_unit")
+ if self.timestamp_clock != actual.timestamp_clock:
+ mismatches.append("timestamp_clock")
+ if mismatches:
+ raise ValueError(f"{context} mismatch in fields: {mismatches}")
+
+ def validate(self, observation: "RobotObservation") -> None:
+ """Validate one observation without imposing image height or width."""
+ if not isinstance(observation, RobotObservation):
+ raise TypeError("observation must be a RobotObservation")
+ if observation.state.dimension_names != self.state_dimension_names:
+ raise ValueError("robot observation space mismatch in state_dimension_names")
+ expected_names = {spec.name for spec in self.images}
+ actual_names = set(observation.images)
+ missing = sorted(expected_names - actual_names)
+ if missing:
+ raise ValueError(f"robot observation space is missing images: {missing}")
+ if not self.allow_extra_images:
+ extra = sorted(actual_names - expected_names)
+ if extra:
+ raise ValueError(f"robot observation space has unexpected images: {extra}")
+ for spec in self.images:
+ image = observation.images[spec.name]
+ shape = getattr(image, "shape", None)
+ dtype = getattr(image, "dtype", None)
+ if shape is None or dtype is None:
+ try:
+ image = torch.as_tensor(image)
+ except (TypeError, ValueError) as error:
+ raise TypeError(f"robot observation image {spec.name!r} must be tensor-like") from error
+ shape = image.shape
+ dtype = image.dtype
+ if len(shape) != 3:
+ raise ValueError(f"robot observation image {spec.name!r} must be rank 3")
+ channel_axis = 2 if spec.layout == "HWC" else 0
+ if shape[channel_axis] != spec.channels:
+ raise ValueError(
+ f"robot observation image {spec.name!r} must have {spec.channels} channels in {spec.layout} layout"
+ )
+ dtype_name = str(dtype).removeprefix("torch.")
+ if dtype_name != spec.dtype:
+ raise ValueError(f"robot observation image {spec.name!r} dtype must be {spec.dtype}, got {dtype_name}")
+
+
+@dataclass(frozen=True)
+class ActionSpaceSpec:
+ """Describe action meaning independently of its tensor shape."""
+
+ representation: str
+ dimension_names: tuple[str, ...]
+ units: tuple[str, ...]
+ frame: str | None
+ control_hz: float | None
+ normalized: bool
+ normalization_profile: str | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.representation, str) or not self.representation:
+ raise ValueError("action representation must be a non-empty string")
+ if not isinstance(self.dimension_names, tuple):
+ raise ValueError("dimension_names must be a tuple")
+ if not self.dimension_names or any(not isinstance(name, str) or not name for name in self.dimension_names):
+ raise ValueError("dimension_names must contain non-empty strings")
+ if len(set(self.dimension_names)) != len(self.dimension_names):
+ raise ValueError("dimension_names must be unique")
+ if not isinstance(self.units, tuple):
+ raise ValueError("units must be a tuple")
+ if len(self.units) != len(self.dimension_names) or any(
+ not isinstance(unit, str) or not unit for unit in self.units
+ ):
+ raise ValueError("units must contain one non-empty value per action dimension")
+ if self.frame is not None and (not isinstance(self.frame, str) or not self.frame):
+ raise ValueError("frame must be None or a non-empty string")
+ if self.control_hz is not None:
+ if isinstance(self.control_hz, bool) or not isinstance(self.control_hz, (int, float)):
+ raise ValueError("control_hz must be None or a positive finite number")
+ if not math.isfinite(self.control_hz) or self.control_hz <= 0:
+ raise ValueError("control_hz must be None or a positive finite number")
+ if not isinstance(self.normalized, bool):
+ raise ValueError("normalized must be a boolean")
+ if self.normalized and (not isinstance(self.normalization_profile, str) or not self.normalization_profile):
+ raise ValueError("normalized action spaces require a normalization_profile")
+ if not self.normalized and self.normalization_profile is not None:
+ raise ValueError("unnormalized action spaces cannot declare a normalization_profile")
+
+ @property
+ def dimension(self) -> int:
+ """Return the action vector width."""
+ return len(self.dimension_names)
+
+ def require_compatible(self, actual: "ActionSpaceSpec", *, context: str = "action space") -> None:
+ """Reject semantic mismatches without inferring meaning from tensor width."""
+ fields = (
+ "representation",
+ "dimension_names",
+ "units",
+ "frame",
+ "normalized",
+ "normalization_profile",
+ )
+ mismatches = [name for name in fields if getattr(self, name) != getattr(actual, name)]
+ if self.control_hz is not None and actual.control_hz is not None and self.control_hz != actual.control_hz:
+ mismatches.append("control_hz")
+ if mismatches:
+ raise ValueError(f"{context} mismatch in fields: {mismatches}")
+
+
+@dataclass(frozen=True)
+class ActionChunk:
+ """A time-ordered action tensor with explicit semantics and provenance."""
+
+ actions: torch.Tensor
+ action_space: ActionSpaceSpec
+ valid_length: int
+ observation_timestamp_ns: int
+ sequence_id: int
+ episode_id: str
+ metadata: Mapping[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.actions, torch.Tensor):
+ raise TypeError("actions must be a torch.Tensor")
+ if self.actions.ndim != 2:
+ raise ValueError(f"actions must have shape [T,A], got {tuple(self.actions.shape)}")
+ if self.actions.shape[1] != self.action_space.dimension:
+ raise ValueError(
+ f"actions width must match the action space dimension {self.action_space.dimension}, "
+ f"got {self.actions.shape[1]}"
+ )
+ if not isinstance(self.valid_length, int) or isinstance(self.valid_length, bool):
+ raise ValueError("valid_length must be an integer within the action tensor horizon")
+ if not 1 <= self.valid_length <= self.actions.shape[0]:
+ raise ValueError("valid_length must be within the action tensor horizon")
+ if not isinstance(self.observation_timestamp_ns, int) or isinstance(self.observation_timestamp_ns, bool):
+ raise ValueError("observation_timestamp_ns must be a non-negative integer")
+ if self.observation_timestamp_ns < 0:
+ raise ValueError("observation_timestamp_ns must be a non-negative integer")
+ if not isinstance(self.sequence_id, int) or isinstance(self.sequence_id, bool):
+ raise ValueError("sequence_id must be a non-negative integer")
+ if self.sequence_id < 0:
+ raise ValueError("sequence_id must be a non-negative integer")
+ if not isinstance(self.episode_id, str) or not self.episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+ object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata)))
+
+ @property
+ def horizon(self) -> int:
+ """Return the allocated action horizon."""
+ return int(self.actions.shape[0])
+
+ def trim(self, length: int) -> "ActionChunk":
+ """Return the same semantic chunk limited to at most ``length`` actions."""
+ if not isinstance(length, int) or isinstance(length, bool) or length < 1:
+ raise ValueError("chunk trim length must be positive")
+ valid_length = min(length, self.valid_length)
+ return replace(self, actions=self.actions[:valid_length], valid_length=valid_length)
+
+
+@dataclass(frozen=True)
+class ModelActionChunk(ActionChunk):
+ """Action chunk expressed in a model-owned action space."""
+
+
+@dataclass(frozen=True)
+class RobotActionChunk(ActionChunk):
+ """Action chunk expressed in an embodiment-owned robot action space."""
+
+
+@dataclass(frozen=True)
+class RobotState:
+ """One robot state vector with explicit dimension order."""
+
+ values: torch.Tensor
+ dimension_names: tuple[str, ...]
+ timestamp_ns: int
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.values, torch.Tensor):
+ raise TypeError("robot state values must be a torch.Tensor")
+ if not isinstance(self.dimension_names, tuple) or any(
+ not isinstance(name, str) or not name for name in self.dimension_names
+ ):
+ raise ValueError("robot state dimension_names must be a tuple of non-empty strings")
+ if self.values.ndim != 1 or self.values.shape[0] != len(self.dimension_names):
+ raise ValueError("robot state values must be one-dimensional and match dimension_names")
+ if len(set(self.dimension_names)) != len(self.dimension_names):
+ raise ValueError("robot state dimension_names must be unique")
+ if not isinstance(self.timestamp_ns, int) or isinstance(self.timestamp_ns, bool) or self.timestamp_ns < 0:
+ raise ValueError("robot state timestamp_ns must be a non-negative integer")
+
+
+@dataclass(frozen=True)
+class RobotObservation:
+ """Simulator observation before embodiment-specific model encoding."""
+
+ state: RobotState
+ images: Mapping[str, Any]
+ metadata: Mapping[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "images", MappingProxyType(dict(self.images)))
+ object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata)))
+
+
+@dataclass(frozen=True)
+class ModelObservation:
+ """Observation representation accepted by one VLA policy."""
+
+ state: Any
+ images: Mapping[str, Any]
+ metadata: Mapping[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ object.__setattr__(self, "images", MappingProxyType(dict(self.images)))
+ object.__setattr__(self, "metadata", MappingProxyType(dict(self.metadata)))
+
+
+@dataclass(frozen=True)
+class VLARequest:
+ """One policy prediction request after embodiment encoding."""
+
+ observation: ModelObservation
+ instruction: str
+ episode_id: str
+ sequence_id: int
+ observation_timestamp_ns: int
+ seed: int | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.instruction, str) or not self.instruction.strip():
+ raise ValueError("instruction must be a non-empty string")
+ if not isinstance(self.episode_id, str) or not self.episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+ if not isinstance(self.sequence_id, int) or isinstance(self.sequence_id, bool) or self.sequence_id < 0:
+ raise ValueError("sequence_id must be a non-negative integer")
+ if (
+ not isinstance(self.observation_timestamp_ns, int)
+ or isinstance(self.observation_timestamp_ns, bool)
+ or self.observation_timestamp_ns < 0
+ ):
+ raise ValueError("observation_timestamp_ns must be a non-negative integer")
+ if self.seed is not None and (not isinstance(self.seed, int) or isinstance(self.seed, bool)):
+ raise ValueError("seed must be an integer or None")
+
+
+@dataclass(frozen=True)
+class VLACapabilities:
+ """Static policy behavior used for session compatibility checks."""
+
+ model_id: str
+ output_action_space: ActionSpaceSpec
+ max_horizon: int
+ stateful: bool = False
+ supports_seed: bool = True
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.model_id, str) or not self.model_id:
+ raise ValueError("model_id must be a non-empty string")
+ if not isinstance(self.max_horizon, int) or isinstance(self.max_horizon, bool) or self.max_horizon < 1:
+ raise ValueError("max_horizon must be positive")
+ if not isinstance(self.stateful, bool) or not isinstance(self.supports_seed, bool):
+ raise ValueError("stateful and supports_seed must be booleans")
diff --git a/telefuser/vla/embodiment.py b/telefuser/vla/embodiment.py
new file mode 100644
index 00000000..a92fa096
--- /dev/null
+++ b/telefuser/vla/embodiment.py
@@ -0,0 +1,36 @@
+"""Embodiment protocol separating model coordinates from robot controls."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from .contracts import (
+ ActionSpaceSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+)
+
+
+@runtime_checkable
+class EmbodimentAdapter(Protocol):
+ """Encode robot observations and decode model actions for one embodiment."""
+
+ @property
+ def embodiment_id(self) -> str: ...
+
+ @property
+ def model_action_space(self) -> ActionSpaceSpec: ...
+
+ @property
+ def robot_action_space(self) -> ActionSpaceSpec: ...
+
+ @property
+ def observation_space(self) -> ObservationSpaceSpec: ...
+
+ def encode_observation(self, observation: RobotObservation) -> ModelObservation: ...
+
+ def decode_actions(self, actions: ModelActionChunk, robot_state: RobotState) -> RobotActionChunk: ...
diff --git a/telefuser/vla/policy.py b/telefuser/vla/policy.py
new file mode 100644
index 00000000..b3bcba9a
--- /dev/null
+++ b/telefuser/vla/policy.py
@@ -0,0 +1,18 @@
+"""Policy protocol for model-independent VLA serving."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from .contracts import ModelActionChunk, VLACapabilities, VLARequest
+
+
+@runtime_checkable
+class VLAPolicy(Protocol):
+ """Adapt one model pipeline to the semantic VLA action contract."""
+
+ def capabilities(self) -> VLACapabilities: ...
+
+ def predict(self, request: VLARequest) -> ModelActionChunk: ...
+
+ def reset(self, episode_id: str) -> None: ...
diff --git a/telefuser/vla/registry.py b/telefuser/vla/registry.py
new file mode 100644
index 00000000..c345fd48
--- /dev/null
+++ b/telefuser/vla/registry.py
@@ -0,0 +1,66 @@
+"""Explicit component registry for VLA policies and embodiments."""
+
+from __future__ import annotations
+
+from typing import TypeVar
+
+from .embodiment import EmbodimentAdapter
+from .policy import VLAPolicy
+
+T = TypeVar("T")
+
+
+class VLARegistry:
+ """Resolve already-loaded policy and embodiment adapters by stable IDs."""
+
+ def __init__(self) -> None:
+ self._policies: dict[str, VLAPolicy] = {}
+ self._embodiments: dict[str, EmbodimentAdapter] = {}
+
+ def register_policy(self, model_id: str, policy: VLAPolicy, *, replace: bool = False) -> None:
+ """Register one loaded policy adapter."""
+ if policy.capabilities().model_id != model_id:
+ raise ValueError("registered model_id must match policy capabilities")
+ self._register(self._policies, model_id, policy, replace=replace)
+
+ def register_embodiment(
+ self,
+ embodiment_id: str,
+ embodiment: EmbodimentAdapter,
+ *,
+ replace: bool = False,
+ ) -> None:
+ """Register one embodiment adapter."""
+ if embodiment.embodiment_id != embodiment_id:
+ raise ValueError("registered embodiment_id must match the adapter")
+ self._register(self._embodiments, embodiment_id, embodiment, replace=replace)
+
+ def get_policy(self, model_id: str) -> VLAPolicy:
+ """Return a registered policy or fail with a stable lookup error."""
+ try:
+ return self._policies[model_id]
+ except KeyError as error:
+ raise KeyError(f"unknown VLA model_id: {model_id!r}") from error
+
+ def get_embodiment(self, embodiment_id: str) -> EmbodimentAdapter:
+ """Return a registered embodiment or fail with a stable lookup error."""
+ try:
+ return self._embodiments[embodiment_id]
+ except KeyError as error:
+ raise KeyError(f"unknown VLA embodiment_id: {embodiment_id!r}") from error
+
+ def model_ids(self) -> tuple[str, ...]:
+ """List registered model IDs in deterministic order."""
+ return tuple(sorted(self._policies))
+
+ def embodiment_ids(self) -> tuple[str, ...]:
+ """List registered embodiment IDs in deterministic order."""
+ return tuple(sorted(self._embodiments))
+
+ @staticmethod
+ def _register(registry: dict[str, T], component_id: str, component: T, *, replace: bool) -> None:
+ if not isinstance(component_id, str) or not component_id:
+ raise ValueError("component ID must be a non-empty string")
+ if component_id in registry and not replace:
+ raise ValueError(f"VLA component is already registered: {component_id!r}")
+ registry[component_id] = component
diff --git a/telefuser/vla/runtime/__init__.py b/telefuser/vla/runtime/__init__.py
new file mode 100644
index 00000000..7ce4c573
--- /dev/null
+++ b/telefuser/vla/runtime/__init__.py
@@ -0,0 +1,31 @@
+"""Runtime utilities for semantic VLA action chunks."""
+
+from .chunk_state import (
+ ActionChunkStateMachine,
+ ChunkStatus,
+ ChunkTicket,
+ DisconnectPolicy,
+ RemainderPolicy,
+ RuntimeState,
+)
+from .executor import ChunkExecutor, RobotAction
+from .safety import ActionSafetyPolicy, BoundedActionSafety, FiniteActionSafety
+from .scheduler import ActionChunkScheduler
+from .simulator import ChunkExecutionReport, SimulatorChunkRuntime
+
+__all__ = [
+ "ActionChunkStateMachine",
+ "ActionSafetyPolicy",
+ "ActionChunkScheduler",
+ "BoundedActionSafety",
+ "ChunkExecutor",
+ "ChunkExecutionReport",
+ "ChunkStatus",
+ "ChunkTicket",
+ "DisconnectPolicy",
+ "FiniteActionSafety",
+ "RobotAction",
+ "RemainderPolicy",
+ "RuntimeState",
+ "SimulatorChunkRuntime",
+]
diff --git a/telefuser/vla/runtime/chunk_state.py b/telefuser/vla/runtime/chunk_state.py
new file mode 100644
index 00000000..ebe03fb6
--- /dev/null
+++ b/telefuser/vla/runtime/chunk_state.py
@@ -0,0 +1,447 @@
+"""Deterministic lifecycle for inferred action chunks."""
+
+from __future__ import annotations
+
+import math
+import time
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass, replace
+from enum import Enum
+from threading import RLock
+
+from ..contracts import RobotActionChunk
+
+
+class ChunkStatus(str, Enum):
+ """Lifecycle state of one submitted action chunk request."""
+
+ PENDING = "pending"
+ READY = "ready"
+ EXECUTING = "executing"
+ EXECUTED = "executed"
+ SUPERSEDED = "superseded"
+ EXPIRED = "expired"
+ REJECTED = "rejected"
+
+
+class RuntimeState(str, Enum):
+ """Externally visible state of the chunk runtime."""
+
+ EMPTY = "empty"
+ PENDING = "pending"
+ READY = "ready"
+ EXECUTING = "executing"
+ HOLDING = "holding"
+ STOPPED = "stopped"
+
+
+class RemainderPolicy(str, Enum):
+ """Treatment of actions beyond one execution horizon."""
+
+ DISCARD = "discard"
+ RETAIN = "retain"
+
+
+class DisconnectPolicy(str, Enum):
+ """Fail-closed state selected when the transport disconnects."""
+
+ HOLD = "hold"
+ STOP = "stop"
+
+
+@dataclass(frozen=True)
+class ChunkTicket:
+ """Identity and independent deadlines for one prediction request."""
+
+ generation: int
+ sequence_id: int
+ observation_timestamp_ns: int
+ received_at: float
+ deadline_at: float | None
+
+
+@dataclass
+class _ChunkRecord:
+ ticket: ChunkTicket
+ status: ChunkStatus
+ inference_started: bool = False
+ inference_completed: bool = False
+ recovery_applied: bool = False
+ reason: str | None = None
+
+
+class ActionChunkStateMachine:
+ """Track pending, ready, and executing chunks independently of transport."""
+
+ def __init__(
+ self,
+ episode_id: str,
+ *,
+ execute_horizon: int,
+ max_observation_age_ns: int | None = None,
+ remainder_policy: RemainderPolicy = RemainderPolicy.DISCARD,
+ disconnect_policy: DisconnectPolicy = DisconnectPolicy.HOLD,
+ stateful_policy: bool = False,
+ recover_stateful_policy: Callable[[str], None] | None = None,
+ clock: Callable[[], float] = time.monotonic,
+ terminal_history: int = 128,
+ ) -> None:
+ if not isinstance(episode_id, str) or not episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+ if not isinstance(execute_horizon, int) or isinstance(execute_horizon, bool) or execute_horizon < 1:
+ raise ValueError("execute_horizon must be a positive integer")
+ if max_observation_age_ns is not None and (
+ not isinstance(max_observation_age_ns, int)
+ or isinstance(max_observation_age_ns, bool)
+ or max_observation_age_ns < 1
+ ):
+ raise ValueError("max_observation_age_ns must be a positive integer")
+ if stateful_policy and recover_stateful_policy is None:
+ raise ValueError("stateful policies require a discard recovery callback")
+ if not isinstance(terminal_history, int) or isinstance(terminal_history, bool) or terminal_history < 1:
+ raise ValueError("terminal_history must be a positive integer")
+ self.episode_id = episode_id
+ self.execute_horizon = execute_horizon
+ self.max_observation_age_ns = max_observation_age_ns
+ self.remainder_policy = RemainderPolicy(remainder_policy)
+ self.disconnect_policy = DisconnectPolicy(disconnect_policy)
+ self.stateful_policy = stateful_policy
+ self._recover_stateful_policy = recover_stateful_policy
+ self._clock = clock
+ self._terminal_history = terminal_history
+ self._generation = 0
+ self._latest_sequence_id: int | None = None
+ self._latest_observation_timestamp_ns: int | None = None
+ self._pending: ChunkTicket | None = None
+ self._ready: tuple[ChunkTicket, RobotActionChunk] | None = None
+ self._executing: tuple[ChunkTicket, RobotActionChunk] | None = None
+ self._remainder: RobotActionChunk | None = None
+ self._disconnected_state: RuntimeState | None = None
+ self._records: OrderedDict[int, _ChunkRecord] = OrderedDict()
+ self._lock = RLock()
+
+ @property
+ def state(self) -> RuntimeState:
+ """Return the current externally visible runtime state."""
+ with self._lock:
+ if self._disconnected_state is not None:
+ return self._disconnected_state
+ if self._executing is not None:
+ return RuntimeState.EXECUTING
+ if self._ready is not None:
+ return RuntimeState.READY
+ if self._pending is not None:
+ return RuntimeState.PENDING
+ return RuntimeState.EMPTY
+
+ @property
+ def latest_sequence_id(self) -> int | None:
+ """Return the newest accepted request sequence."""
+ with self._lock:
+ return self._latest_sequence_id
+
+ def submit(
+ self,
+ sequence_id: int,
+ observation_timestamp_ns: int,
+ *,
+ request_ttl_ms: float | None = None,
+ observation_clock_now_ns: int | None = None,
+ ) -> ChunkTicket:
+ """Admit a prediction request and supersede older pending or ready work."""
+ self._validate_nonnegative_int(sequence_id, "sequence_id")
+ self._validate_nonnegative_int(observation_timestamp_ns, "observation_timestamp_ns")
+ if observation_clock_now_ns is not None:
+ self._validate_nonnegative_int(observation_clock_now_ns, "observation_clock_now_ns")
+ ttl_ms = self._validate_ttl(request_ttl_ms)
+ with self._lock:
+ now = self._now()
+ self._generation += 1
+ ticket = ChunkTicket(
+ generation=self._generation,
+ sequence_id=sequence_id,
+ observation_timestamp_ns=observation_timestamp_ns,
+ received_at=now,
+ deadline_at=None if ttl_ms is None else now + ttl_ms / 1000.0,
+ )
+ record = _ChunkRecord(ticket=ticket, status=ChunkStatus.PENDING)
+ self._records[ticket.generation] = record
+ if self._disconnected_state is not None:
+ self._set_terminal(record, ChunkStatus.REJECTED, "runtime is disconnected")
+ return ticket
+ if self._latest_sequence_id is not None and sequence_id <= self._latest_sequence_id:
+ self._set_terminal(record, ChunkStatus.REJECTED, "sequence_id is not newer than the latest request")
+ return ticket
+ if self._observation_is_stale(observation_timestamp_ns, observation_clock_now_ns):
+ self._set_terminal(record, ChunkStatus.EXPIRED, "observation timestamp is stale")
+ return ticket
+ if (
+ self._latest_observation_timestamp_ns is not None
+ and observation_timestamp_ns < self._latest_observation_timestamp_ns
+ ):
+ self._set_terminal(record, ChunkStatus.REJECTED, "observation timestamp moved backwards")
+ return ticket
+ if self._pending is not None:
+ self._set_terminal(
+ self._record(self._pending),
+ ChunkStatus.SUPERSEDED,
+ "a newer observation superseded the pending request",
+ )
+ if self._ready is not None:
+ self._set_terminal(
+ self._record(self._ready[0]),
+ ChunkStatus.SUPERSEDED,
+ "a newer observation superseded the ready chunk",
+ )
+ self._ready = None
+ self._pending = ticket
+ self._latest_sequence_id = sequence_id
+ if (
+ self._latest_observation_timestamp_ns is None
+ or observation_timestamp_ns > self._latest_observation_timestamp_ns
+ ):
+ self._latest_observation_timestamp_ns = observation_timestamp_ns
+ self._trim_history()
+ return ticket
+
+ def mark_inference_started(self, ticket: ChunkTicket) -> ChunkStatus:
+ """Mark that a policy call may mutate state for this request."""
+ with self._lock:
+ record = self._record(ticket)
+ if record.status is ChunkStatus.PENDING:
+ record.inference_started = True
+ return record.status
+
+ def complete(
+ self,
+ ticket: ChunkTicket,
+ chunk: RobotActionChunk,
+ *,
+ observation_clock_now_ns: int | None = None,
+ ) -> ChunkStatus:
+ """Publish a completed chunk or deterministically discard it."""
+ if observation_clock_now_ns is not None:
+ self._validate_nonnegative_int(observation_clock_now_ns, "observation_clock_now_ns")
+ with self._lock:
+ record = self._record(ticket)
+ record.inference_completed = True
+ if record.status is not ChunkStatus.PENDING or self._pending != ticket:
+ self._recover_discarded_state(record)
+ self._trim_history()
+ return record.status
+ try:
+ self._validate_chunk(ticket, chunk)
+ except Exception:
+ self._pending = None
+ self._set_terminal(record, ChunkStatus.REJECTED, "completed chunk failed validation")
+ self._recover_discarded_state(record)
+ raise
+ now = self._now()
+ if ticket.deadline_at is not None and now >= ticket.deadline_at:
+ self._pending = None
+ self._set_terminal(record, ChunkStatus.EXPIRED, "request_ttl_ms elapsed before completion")
+ self._recover_discarded_state(record)
+ return record.status
+ if self._observation_is_stale(ticket.observation_timestamp_ns, observation_clock_now_ns):
+ self._pending = None
+ self._set_terminal(record, ChunkStatus.EXPIRED, "observation became stale before completion")
+ self._recover_discarded_state(record)
+ return record.status
+ self._pending = None
+ self._ready = (ticket, chunk)
+ record.status = ChunkStatus.READY
+ record.reason = None
+ return record.status
+
+ def reject(self, ticket: ChunkTicket, reason: str) -> ChunkStatus:
+ """Reject pending inference after a policy or transport failure."""
+ if not isinstance(reason, str) or not reason:
+ raise ValueError("reason must be a non-empty string")
+ with self._lock:
+ record = self._record(ticket)
+ record.inference_completed = True
+ if record.status is not ChunkStatus.PENDING or self._pending != ticket:
+ self._recover_discarded_state(record)
+ return record.status
+ self._pending = None
+ self._set_terminal(record, ChunkStatus.REJECTED, reason)
+ self._recover_discarded_state(record)
+ return record.status
+
+ def begin_execution(self) -> RobotActionChunk | None:
+ """Move the newest ready chunk into execution and apply execute_horizon."""
+ with self._lock:
+ if self._disconnected_state is not None:
+ return None
+ if self._executing is not None:
+ raise RuntimeError("an action chunk is already executing")
+ if self._ready is None:
+ return None
+ ticket, chunk = self._ready
+ self._ready = None
+ execution_length = min(self.execute_horizon, chunk.valid_length)
+ execution_chunk = chunk.trim(execution_length)
+ if self.remainder_policy is RemainderPolicy.RETAIN and execution_length < chunk.valid_length:
+ remaining = chunk.actions[execution_length : chunk.valid_length]
+ self._remainder = replace(chunk, actions=remaining, valid_length=int(remaining.shape[0]))
+ else:
+ self._remainder = None
+ self._executing = (ticket, execution_chunk)
+ self._record(ticket).status = ChunkStatus.EXECUTING
+ return execution_chunk
+
+ def finish_execution(self, *, success: bool = True) -> RuntimeState:
+ """Finish the active horizon and optionally expose retained remainder."""
+ if not isinstance(success, bool):
+ raise ValueError("success must be a boolean")
+ with self._lock:
+ if self._executing is None:
+ raise RuntimeError("no action chunk is executing")
+ ticket, _ = self._executing
+ record = self._record(ticket)
+ self._executing = None
+ if not success:
+ self._remainder = None
+ self._set_terminal(record, ChunkStatus.REJECTED, "simulator execution failed")
+ elif self._remainder is not None:
+ self._ready = (ticket, self._remainder)
+ self._remainder = None
+ record.status = ChunkStatus.READY
+ record.reason = None
+ else:
+ record.status = ChunkStatus.EXECUTED
+ record.reason = None
+ self._trim_history()
+ return self.state
+
+ def disconnect(self) -> RuntimeState:
+ """Discard buffered work and enter the configured fail-closed state."""
+ with self._lock:
+ for ticket in self._active_tickets():
+ self._set_terminal(self._record(ticket), ChunkStatus.REJECTED, "transport disconnected")
+ self._pending = None
+ self._ready = None
+ self._executing = None
+ self._remainder = None
+ self._disconnected_state = (
+ RuntimeState.HOLDING if self.disconnect_policy is DisconnectPolicy.HOLD else RuntimeState.STOPPED
+ )
+ self._trim_history()
+ return self._disconnected_state
+
+ def reset(self, episode_id: str | None = None) -> None:
+ """Clear every chunk slot and restart sequence ordering."""
+ if episode_id is not None and (not isinstance(episode_id, str) or not episode_id):
+ raise ValueError("episode_id must be a non-empty string")
+ with self._lock:
+ for ticket in self._active_tickets():
+ record = self._record(ticket)
+ self._set_terminal(record, ChunkStatus.REJECTED, "runtime reset")
+ self._recover_discarded_state(record)
+ if episode_id is not None:
+ self.episode_id = episode_id
+ self._pending = None
+ self._ready = None
+ self._executing = None
+ self._remainder = None
+ self._latest_sequence_id = None
+ self._latest_observation_timestamp_ns = None
+ self._disconnected_state = None
+ self._trim_history()
+
+ def status(self, ticket: ChunkTicket) -> ChunkStatus:
+ """Return the current status of a known ticket."""
+ with self._lock:
+ return self._record(ticket).status
+
+ def reason(self, ticket: ChunkTicket) -> str | None:
+ """Return the terminal reason for a known ticket."""
+ with self._lock:
+ return self._record(ticket).reason
+
+ def _observation_is_stale(self, timestamp_ns: int, observation_clock_now_ns: int | None) -> bool:
+ if self.max_observation_age_ns is None or observation_clock_now_ns is None:
+ return False
+ return observation_clock_now_ns - timestamp_ns > self.max_observation_age_ns
+
+ def _validate_chunk(self, ticket: ChunkTicket, chunk: RobotActionChunk) -> None:
+ if not isinstance(chunk, RobotActionChunk):
+ raise TypeError("completed action must be a RobotActionChunk")
+ if chunk.episode_id != self.episode_id:
+ raise ValueError("completed chunk episode_id does not match the runtime")
+ if chunk.sequence_id != ticket.sequence_id:
+ raise ValueError("completed chunk sequence_id does not match its ticket")
+ if chunk.observation_timestamp_ns != ticket.observation_timestamp_ns:
+ raise ValueError("completed chunk observation timestamp does not match its ticket")
+
+ def _recover_discarded_state(self, record: _ChunkRecord) -> None:
+ if not self.stateful_policy or not record.inference_started or record.recovery_applied:
+ return
+ assert self._recover_stateful_policy is not None
+ self._recover_stateful_policy(self.episode_id)
+ record.recovery_applied = True
+
+ def _set_terminal(self, record: _ChunkRecord, status: ChunkStatus, reason: str) -> None:
+ record.status = status
+ record.reason = reason
+ self._trim_history()
+
+ def _record(self, ticket: ChunkTicket) -> _ChunkRecord:
+ try:
+ return self._records[ticket.generation]
+ except KeyError as error:
+ raise KeyError(f"unknown or expired chunk ticket: {ticket.generation}") from error
+
+ def _active_tickets(self) -> tuple[ChunkTicket, ...]:
+ tickets: list[ChunkTicket] = []
+ if self._pending is not None:
+ tickets.append(self._pending)
+ if self._ready is not None:
+ tickets.append(self._ready[0])
+ if self._executing is not None:
+ tickets.append(self._executing[0])
+ return tuple({ticket.generation: ticket for ticket in tickets}.values())
+
+ def _trim_history(self) -> None:
+ terminal = {
+ ChunkStatus.EXECUTED,
+ ChunkStatus.SUPERSEDED,
+ ChunkStatus.EXPIRED,
+ ChunkStatus.REJECTED,
+ }
+ protected = {ticket.generation for ticket in self._active_tickets()}
+ protected.update(
+ generation
+ for generation, record in self._records.items()
+ if record.inference_started and not record.inference_completed
+ )
+ removable = [
+ generation
+ for generation, record in self._records.items()
+ if record.status in terminal and generation not in protected
+ ]
+ for generation in removable[: -self._terminal_history]:
+ self._records.pop(generation)
+
+ def _now(self) -> float:
+ value = self._clock()
+ if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
+ raise ValueError("chunk runtime clock must return a finite number")
+ return float(value)
+
+ @staticmethod
+ def _validate_nonnegative_int(value: int, name: str) -> None:
+ if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+ raise ValueError(f"{name} must be a non-negative integer")
+
+ @staticmethod
+ def _validate_ttl(value: float | None) -> float | None:
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise ValueError("request_ttl_ms must be a positive finite number")
+ ttl_ms = float(value)
+ if not math.isfinite(ttl_ms) or ttl_ms <= 0:
+ raise ValueError("request_ttl_ms must be a positive finite number")
+ return ttl_ms
diff --git a/telefuser/vla/runtime/executor.py b/telefuser/vla/runtime/executor.py
new file mode 100644
index 00000000..836c592d
--- /dev/null
+++ b/telefuser/vla/runtime/executor.py
@@ -0,0 +1,86 @@
+"""Preparation of semantic robot action chunks for execution."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Iterator
+
+import torch
+
+from ..contracts import ActionSpaceSpec, RobotActionChunk, RobotState
+from .safety import ActionSafetyPolicy, FiniteActionSafety
+
+
+@dataclass(frozen=True)
+class RobotAction:
+ """One executable action extracted from a validated chunk."""
+
+ values: torch.Tensor
+ action_space: ActionSpaceSpec
+ sequence_id: int
+ step_index: int
+ episode_id: str
+
+
+class ChunkExecutor:
+ """Trim, validate, and expose actions without simulator-specific logic."""
+
+ def __init__(
+ self,
+ expected_action_space: ActionSpaceSpec,
+ *,
+ execute_horizon: int | None = None,
+ safety_policy: ActionSafetyPolicy | None = None,
+ max_observation_age_ns: int | None = None,
+ ) -> None:
+ if execute_horizon is not None and (
+ not isinstance(execute_horizon, int) or isinstance(execute_horizon, bool) or execute_horizon < 1
+ ):
+ raise ValueError("execute_horizon must be positive")
+ if max_observation_age_ns is not None and (
+ not isinstance(max_observation_age_ns, int)
+ or isinstance(max_observation_age_ns, bool)
+ or max_observation_age_ns < 1
+ ):
+ raise ValueError("max_observation_age_ns must be positive")
+ self.expected_action_space = expected_action_space
+ self.execute_horizon = execute_horizon
+ self.safety_policy = safety_policy or FiniteActionSafety()
+ self.max_observation_age_ns = max_observation_age_ns
+
+ def prepare(
+ self,
+ chunk: RobotActionChunk,
+ robot_state: RobotState,
+ *,
+ now_ns: int | None = None,
+ ) -> RobotActionChunk:
+ """Return a bounded, safe chunk ready for simulator consumption."""
+ self.expected_action_space.require_compatible(chunk.action_space, context="robot action space")
+ if self.max_observation_age_ns is not None:
+ if now_ns is None:
+ raise ValueError("now_ns is required when max_observation_age_ns is configured")
+ if not isinstance(now_ns, int) or isinstance(now_ns, bool) or now_ns < 0:
+ raise ValueError("now_ns must be a non-negative integer")
+ if now_ns < chunk.observation_timestamp_ns:
+ raise ValueError("now_ns cannot precede the observation timestamp")
+ if now_ns - chunk.observation_timestamp_ns > self.max_observation_age_ns:
+ raise ValueError("robot action chunk is stale")
+ valid_length = chunk.valid_length
+ if self.execute_horizon is not None:
+ valid_length = min(valid_length, self.execute_horizon)
+ prepared = replace(chunk, actions=chunk.actions[:valid_length], valid_length=valid_length)
+ self.safety_policy.validate(prepared, robot_state)
+ return prepared
+
+ @staticmethod
+ def iter_actions(chunk: RobotActionChunk) -> Iterator[RobotAction]:
+ """Yield only the validated portion of a prepared action chunk."""
+ for step_index in range(chunk.valid_length):
+ yield RobotAction(
+ values=chunk.actions[step_index],
+ action_space=chunk.action_space,
+ sequence_id=chunk.sequence_id,
+ step_index=step_index,
+ episode_id=chunk.episode_id,
+ )
diff --git a/telefuser/vla/runtime/safety.py b/telefuser/vla/runtime/safety.py
new file mode 100644
index 00000000..5c749eef
--- /dev/null
+++ b/telefuser/vla/runtime/safety.py
@@ -0,0 +1,61 @@
+"""Safety checks applied after embodiment action decoding."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+import torch
+
+from ..contracts import RobotActionChunk, RobotState
+
+
+@runtime_checkable
+class ActionSafetyPolicy(Protocol):
+ """Validate a robot action chunk before it reaches a simulator."""
+
+ def validate(self, chunk: RobotActionChunk, robot_state: RobotState) -> None: ...
+
+
+class FiniteActionSafety:
+ """Reject non-finite state and action values."""
+
+ def validate(self, chunk: RobotActionChunk, robot_state: RobotState) -> None:
+ """Validate finite values in the executable portion of a chunk."""
+ if not torch.isfinite(robot_state.values).all():
+ raise ValueError("robot state must contain only finite values")
+ if not torch.isfinite(chunk.actions[: chunk.valid_length]).all():
+ raise ValueError("robot actions must contain only finite values")
+
+
+@dataclass(frozen=True)
+class BoundedActionSafety(FiniteActionSafety):
+ """Apply per-dimension bounds and an optional first-step delta limit."""
+
+ lower: torch.Tensor
+ upper: torch.Tensor
+ max_initial_delta: torch.Tensor | None = None
+
+ def validate(self, chunk: RobotActionChunk, robot_state: RobotState) -> None:
+ """Validate finite values, limits, and the transition from current state."""
+ super().validate(chunk, robot_state)
+ dimension = chunk.action_space.dimension
+ if self.lower.shape != (dimension,) or self.upper.shape != (dimension,):
+ raise ValueError("safety bounds must match the robot action dimension")
+ if torch.any(self.lower > self.upper):
+ raise ValueError("safety lower bounds must not exceed upper bounds")
+ actions = chunk.actions[: chunk.valid_length]
+ lower = self.lower.to(device=actions.device, dtype=actions.dtype)
+ upper = self.upper.to(device=actions.device, dtype=actions.dtype)
+ if torch.any(actions < lower) or torch.any(actions > upper):
+ raise ValueError("robot actions exceed configured safety bounds")
+ if self.max_initial_delta is None:
+ return
+ if self.max_initial_delta.shape != (dimension,):
+ raise ValueError("max_initial_delta must match the robot action dimension")
+ if robot_state.values.shape != (dimension,):
+ raise ValueError("robot state dimension must match action delta safety checks")
+ maximum = self.max_initial_delta.to(device=actions.device, dtype=actions.dtype)
+ state = robot_state.values.to(device=actions.device, dtype=actions.dtype)
+ if torch.any(torch.abs(actions[0] - state) > maximum):
+ raise ValueError("first robot action exceeds configured state delta limits")
diff --git a/telefuser/vla/runtime/scheduler.py b/telefuser/vla/runtime/scheduler.py
new file mode 100644
index 00000000..24669f6f
--- /dev/null
+++ b/telefuser/vla/runtime/scheduler.py
@@ -0,0 +1,270 @@
+"""Bounded latest-wins scheduling for action chunk inference."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import operator
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any, Callable, Mapping
+
+
+@dataclass(slots=True)
+class _ScheduledAction:
+ request: Mapping[str, Any]
+ session_key: str
+ generation: int
+ received_at: float
+ deadline_at: float | None
+ future: asyncio.Future[dict[str, Any]]
+
+
+class ActionChunkScheduler:
+ """Serialize inference while retaining only the newest pending chunk per session."""
+
+ def __init__(
+ self,
+ infer: Callable[[Mapping[str, Any]], dict[str, Any]],
+ *,
+ max_pending_sessions: int = 32,
+ ) -> None:
+ if max_pending_sessions < 1:
+ raise ValueError("max_pending_sessions must be positive")
+ self._infer = infer
+ self._max_pending_sessions = max_pending_sessions
+ self._pending: OrderedDict[str, _ScheduledAction] = OrderedDict()
+ self._latest_generation: dict[str, int] = {}
+ self._latest_sequence: dict[str, int] = {}
+ self._wake = asyncio.Event()
+ self._worker: asyncio.Task[None] | None = None
+ self._closed = False
+
+ @property
+ def metadata(self) -> dict[str, Any]:
+ """Describe the additive scheduling controls accepted by the endpoint."""
+ return {
+ "scheduling": {
+ "mode": "latest_wins",
+ "max_pending_per_session": 1,
+ "max_pending_sessions": self._max_pending_sessions,
+ "sequence_field": "sequence_id",
+ "ttl_field": "request_ttl_ms",
+ "inflight_cancellation": False,
+ }
+ }
+
+ async def start(self) -> None:
+ """Start the single inference worker on the current event loop."""
+ if self._worker is not None:
+ return
+ if self._closed:
+ raise RuntimeError("action scheduler is closed")
+ self._worker = asyncio.create_task(self._run(), name="vla-action-chunk-scheduler")
+
+ def submit(self, request: Mapping[str, Any], *, session_key: str) -> asyncio.Future[dict[str, Any]]:
+ """Admit one request and return a future without waiting for inference."""
+ if self._worker is None or self._closed:
+ raise RuntimeError("action scheduler is not running")
+ if not session_key:
+ raise ValueError("session_key must be non-empty")
+
+ loop = asyncio.get_running_loop()
+ future: asyncio.Future[dict[str, Any]] = loop.create_future()
+ received_at = time.monotonic()
+ ttl_ms = self._optional_ttl_ms(request)
+ sequence_id = self._optional_sequence_id(request)
+ previous = self._pending.get(session_key)
+ if previous is None and len(self._pending) >= self._max_pending_sessions:
+ future.set_result(
+ self._discarded_response(
+ request,
+ status="overloaded",
+ message="the action scheduler has no free pending-session slot",
+ )
+ )
+ return future
+ if sequence_id is not None:
+ latest_sequence = self._latest_sequence.get(session_key)
+ if latest_sequence is not None and sequence_id <= latest_sequence:
+ future.set_result(
+ self._discarded_response(
+ request,
+ status="stale_sequence",
+ message=f"sequence_id={sequence_id} is not newer than {latest_sequence}",
+ )
+ )
+ return future
+ self._latest_sequence[session_key] = sequence_id
+
+ generation = self._latest_generation.get(session_key, 0) + 1
+ self._latest_generation[session_key] = generation
+ deadline_at = None if ttl_ms is None else received_at + ttl_ms / 1000.0
+ job = _ScheduledAction(
+ request=dict(request),
+ session_key=session_key,
+ generation=generation,
+ received_at=received_at,
+ deadline_at=deadline_at,
+ future=future,
+ )
+ if previous is not None:
+ self._resolve(
+ previous,
+ self._discarded_response(
+ previous.request,
+ status="superseded",
+ message="a newer observation replaced this pending action request",
+ ),
+ )
+ self._pending[session_key] = job
+ self._wake.set()
+ return future
+
+ def release_session(self, session_key: str) -> None:
+ """Discard pending work and sequence state for a disconnected session."""
+ pending = self._pending.pop(session_key, None)
+ if pending is not None:
+ self._resolve(
+ pending,
+ self._discarded_response(
+ pending.request,
+ status="session_closed",
+ message="the client session closed before inference",
+ ),
+ )
+ self._latest_generation.pop(session_key, None)
+ self._latest_sequence.pop(session_key, None)
+
+ async def close(self) -> None:
+ """Drain an in-flight call and reject work that has not started."""
+ if self._closed:
+ return
+ self._closed = True
+ for job in self._pending.values():
+ self._resolve(
+ job,
+ self._discarded_response(
+ job.request,
+ status="server_stopping",
+ message="the action scheduler is stopping",
+ ),
+ )
+ self._pending.clear()
+ self._wake.set()
+ if self._worker is not None:
+ await self._worker
+ self._worker = None
+
+ async def _run(self) -> None:
+ while True:
+ await self._wake.wait()
+ if self._closed and not self._pending:
+ return
+ if not self._pending:
+ self._wake.clear()
+ continue
+
+ _, job = self._pending.popitem(last=False)
+ if not self._pending:
+ self._wake.clear()
+ if job.future.done():
+ continue
+ discard = self._discard_reason(job)
+ if discard is not None:
+ self._resolve(job, discard)
+ continue
+
+ inference_started_at = time.monotonic()
+ try:
+ response = await asyncio.to_thread(self._infer, job.request)
+ except Exception as error:
+ if not job.future.done():
+ job.future.set_exception(error)
+ continue
+ inference_ms = (time.monotonic() - inference_started_at) * 1000.0
+
+ discard = self._discard_reason(job)
+ if discard is not None:
+ self._resolve(job, discard)
+ continue
+ completed_at = time.monotonic()
+ response = dict(response)
+ server_timing = dict(response.get("server_timing", {}))
+ server_timing.update(
+ queue_wait_ms=(inference_started_at - job.received_at) * 1000.0,
+ infer_ms=inference_ms,
+ scheduler_total_ms=(completed_at - job.received_at) * 1000.0,
+ )
+ response.update(scheduler_status="completed", server_timing=server_timing)
+ for field in ("request_id", "episode_id", "sequence_id"):
+ if field in job.request:
+ response.setdefault(field, job.request[field])
+ if job.deadline_at is not None:
+ response["request_ttl_ms"] = (job.deadline_at - job.received_at) * 1000.0
+ self._resolve(job, response)
+
+ def _discard_reason(self, job: _ScheduledAction) -> dict[str, Any] | None:
+ if job.deadline_at is not None and time.monotonic() >= job.deadline_at:
+ return self._discarded_response(
+ job.request,
+ status="expired",
+ message="the action request exceeded request_ttl_ms",
+ )
+ if self._latest_generation.get(job.session_key) != job.generation:
+ return self._discarded_response(
+ job.request,
+ status="superseded",
+ message="a newer observation superseded this action result",
+ )
+ return None
+
+ @staticmethod
+ def _resolve(job: _ScheduledAction, response: dict[str, Any]) -> None:
+ if not job.future.done():
+ job.future.set_result(response)
+
+ @staticmethod
+ def _optional_sequence_id(request: Mapping[str, Any]) -> int | None:
+ value = request.get("sequence_id")
+ if value is None:
+ return None
+ if isinstance(value, bool):
+ raise ValueError("sequence_id must be a non-negative integer")
+ try:
+ sequence_id = operator.index(value)
+ except TypeError as error:
+ raise ValueError("sequence_id must be a non-negative integer") from error
+ if sequence_id < 0:
+ raise ValueError("sequence_id must be a non-negative integer")
+ return sequence_id
+
+ @staticmethod
+ def _optional_ttl_ms(request: Mapping[str, Any]) -> float | None:
+ value = request.get("request_ttl_ms")
+ if value is None:
+ return None
+ if isinstance(value, bool) or not isinstance(value, int | float):
+ raise ValueError("request_ttl_ms must be a positive finite number")
+ ttl_ms = float(value)
+ if not math.isfinite(ttl_ms) or ttl_ms <= 0:
+ raise ValueError("request_ttl_ms must be a positive finite number")
+ return ttl_ms
+
+ @staticmethod
+ def _discarded_response(
+ request: Mapping[str, Any],
+ *,
+ status: str,
+ message: str,
+ ) -> dict[str, Any]:
+ response: dict[str, Any] = {
+ "action": None,
+ "scheduler_status": status,
+ "error": {"code": status, "message": message},
+ }
+ for field in ("request_id", "episode_id", "sequence_id"):
+ if field in request:
+ response[field] = request[field]
+ return response
diff --git a/telefuser/vla/runtime/simulator.py b/telefuser/vla/runtime/simulator.py
new file mode 100644
index 00000000..e84e0f18
--- /dev/null
+++ b/telefuser/vla/runtime/simulator.py
@@ -0,0 +1,161 @@
+"""Client-side action lifecycle for simulator execution."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..contracts import ActionSpaceSpec, RobotActionChunk
+from .chunk_state import (
+ ActionChunkStateMachine,
+ ChunkStatus,
+ DisconnectPolicy,
+ RemainderPolicy,
+ RuntimeState,
+)
+from .executor import ChunkExecutor
+
+if TYPE_CHECKING:
+ from telefuser.integrations.sim.base import SimulatorAdapter
+
+
+@dataclass(frozen=True)
+class ChunkExecutionReport:
+ """Small, transport-neutral report emitted by simulator-side execution.
+
+ The report deliberately does not become a WebSocket operation. A remote
+ simulator can serialize it on its existing control channel while local
+ callers can inspect it directly. ``sequence_id`` is ``None`` when no
+ ready chunk was available.
+ """
+
+ status: str
+ sequence_id: int | None
+ episode_id: str
+ executed_steps: int
+ reason: str | None = None
+
+ def __post_init__(self) -> None:
+ if self.status not in {"executed", "failed", "no_action"}:
+ raise ValueError("status must be executed, failed, or no_action")
+ if self.sequence_id is not None and (
+ not isinstance(self.sequence_id, int) or isinstance(self.sequence_id, bool) or self.sequence_id < 0
+ ):
+ raise ValueError("sequence_id must be a non-negative integer or None")
+ if not isinstance(self.episode_id, str) or not self.episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+ if not isinstance(self.executed_steps, int) or isinstance(self.executed_steps, bool) or self.executed_steps < 0:
+ raise ValueError("executed_steps must be a non-negative integer")
+ if self.status == "no_action" and self.executed_steps != 0:
+ raise ValueError("no_action reports must have zero executed_steps")
+ if self.status == "failed" and not self.reason:
+ raise ValueError("failed reports require a reason")
+
+
+class SimulatorChunkRuntime:
+ """Own READY/EXECUTING/EXECUTED states beside a simulator connection."""
+
+ def __init__(
+ self,
+ simulator: SimulatorAdapter,
+ action_space: ActionSpaceSpec,
+ episode_id: str,
+ *,
+ execute_horizon: int,
+ remainder_policy: RemainderPolicy = RemainderPolicy.DISCARD,
+ disconnect_policy: DisconnectPolicy = DisconnectPolicy.HOLD,
+ max_observation_age_ns: int | None = None,
+ ) -> None:
+ self.simulator = simulator
+ self.action_space = action_space
+ self.state_machine = ActionChunkStateMachine(
+ episode_id,
+ execute_horizon=execute_horizon,
+ remainder_policy=remainder_policy,
+ disconnect_policy=disconnect_policy,
+ max_observation_age_ns=max_observation_age_ns,
+ )
+
+ @property
+ def state(self) -> RuntimeState:
+ """Return the current simulator-side runtime state."""
+ return self.state_machine.state
+
+ def accept(
+ self,
+ chunk: RobotActionChunk,
+ *,
+ observation_clock_now_ns: int | None = None,
+ ) -> ChunkStatus:
+ """Accept a server-ready chunk without claiming that it was executed."""
+ self.action_space.require_compatible(chunk.action_space, context="simulator action space")
+ ticket = self.state_machine.submit(
+ chunk.sequence_id,
+ chunk.observation_timestamp_ns,
+ observation_clock_now_ns=observation_clock_now_ns,
+ )
+ if self.state_machine.status(ticket) is not ChunkStatus.PENDING:
+ return self.state_machine.status(ticket)
+ self.state_machine.mark_inference_started(ticket)
+ return self.state_machine.complete(
+ ticket,
+ chunk,
+ observation_clock_now_ns=observation_clock_now_ns,
+ )
+
+ def execute_ready(self) -> int:
+ """Execute one configured horizon and return the submitted step count."""
+ report = self._execute_ready(raise_errors=True)
+ return report.executed_steps
+
+ def execute_ready_with_report(self) -> ChunkExecutionReport:
+ """Execute one horizon and return a stable result for remote reporting.
+
+ Unlike :meth:`execute_ready`, this method converts simulator failures
+ into a ``failed`` report after entering the configured fail-closed
+ state. The legacy method remains exception-based for compatibility.
+ """
+ return self._execute_ready(raise_errors=False)
+
+ async def execute_ready_async(self) -> ChunkExecutionReport:
+ """Run execution off the event loop for inference/execution overlap.
+
+ Prediction scheduling remains the caller's responsibility. This
+ additive helper only prevents a synchronous simulator adapter from
+ blocking an async VLA client while a ready chunk is being applied.
+ """
+ return await asyncio.to_thread(self.execute_ready_with_report)
+
+ def _execute_ready(self, *, raise_errors: bool) -> ChunkExecutionReport:
+ chunk = self.state_machine.begin_execution()
+ if chunk is None:
+ return ChunkExecutionReport("no_action", None, self.state_machine.episode_id, 0)
+ executed = 0
+ try:
+ for action in ChunkExecutor.iter_actions(chunk):
+ self.simulator.execute(action)
+ executed += 1
+ except Exception as error:
+ self.state_machine.finish_execution(success=False)
+ self.state_machine.disconnect()
+ report = ChunkExecutionReport(
+ "failed",
+ chunk.sequence_id,
+ chunk.episode_id,
+ executed,
+ reason=str(error) or error.__class__.__name__,
+ )
+ if raise_errors:
+ raise
+ return report
+ self.state_machine.finish_execution(success=True)
+ return ChunkExecutionReport("executed", chunk.sequence_id, chunk.episode_id, executed)
+
+ def reset(self, episode_id: str | None = None) -> None:
+ """Clear buffered action state after the remote session is reset."""
+ self.state_machine.reset(episode_id)
+
+ def disconnect(self) -> RuntimeState:
+ """Discard buffered actions and enter the configured hold or stop state."""
+ return self.state_machine.disconnect()
diff --git a/telefuser/vla/serialization.py b/telefuser/vla/serialization.py
new file mode 100644
index 00000000..6eb63419
--- /dev/null
+++ b/telefuser/vla/serialization.py
@@ -0,0 +1,312 @@
+"""Versioned JSON-safe serialization for public VLA contracts."""
+
+from __future__ import annotations
+
+import base64
+import json
+import math
+from collections.abc import Mapping
+from typing import Any
+
+import numpy as np
+import torch
+
+from .contracts import (
+ ActionSpaceSpec,
+ ImageObservationSpec,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+)
+
+VLA_WIRE_SCHEMA_VERSION = 1
+VLA_WIRE_ENCODING = "json-base64-v1"
+VLA_SESSION_PROTOCOL_VERSION = "1.0"
+DEFAULT_MAX_TENSOR_BYTES = 64 * 1024 * 1024
+
+_DTYPES: dict[str, torch.dtype] = {
+ "bool": torch.bool,
+ "uint8": torch.uint8,
+ "int8": torch.int8,
+ "int16": torch.int16,
+ "int32": torch.int32,
+ "int64": torch.int64,
+ "float16": torch.float16,
+ "bfloat16": torch.bfloat16,
+ "float32": torch.float32,
+ "float64": torch.float64,
+}
+_DTYPE_NAMES = {dtype: name for name, dtype in _DTYPES.items()}
+
+
+def action_space_to_wire(spec: ActionSpaceSpec) -> dict[str, Any]:
+ """Serialize an action-space contract without losing field meaning."""
+ return {
+ "schema": "telefuser.vla.action_space",
+ "schema_version": VLA_WIRE_SCHEMA_VERSION,
+ "representation": spec.representation,
+ "dimension_names": list(spec.dimension_names),
+ "units": list(spec.units),
+ "frame": spec.frame,
+ "control_hz": spec.control_hz,
+ "normalized": spec.normalized,
+ "normalization_profile": spec.normalization_profile,
+ }
+
+
+def action_space_from_wire(payload: Mapping[str, Any]) -> ActionSpaceSpec:
+ """Deserialize and validate an action-space contract."""
+ _require_schema(payload, "telefuser.vla.action_space")
+ dimension_names_value = payload.get("dimension_names")
+ units_value = payload.get("units")
+ if not isinstance(dimension_names_value, list) or not isinstance(units_value, list):
+ raise ValueError("VLA action-space dimension_names and units must be arrays")
+ try:
+ dimension_names = tuple(dimension_names_value)
+ units = tuple(units_value)
+ return ActionSpaceSpec(
+ representation=payload["representation"],
+ dimension_names=dimension_names,
+ units=units,
+ frame=payload.get("frame"),
+ control_hz=payload.get("control_hz"),
+ normalized=payload["normalized"],
+ normalization_profile=payload.get("normalization_profile"),
+ )
+ except (KeyError, TypeError) as error:
+ raise ValueError("invalid VLA action-space payload") from error
+
+
+def observation_space_to_wire(spec: ObservationSpaceSpec) -> dict[str, Any]:
+ """Serialize an observation-space contract."""
+ return {
+ "schema": "telefuser.vla.observation_space",
+ "schema_version": VLA_WIRE_SCHEMA_VERSION,
+ "state_dimension_names": list(spec.state_dimension_names),
+ "images": [
+ {"name": image.name, "dtype": image.dtype, "layout": image.layout, "channels": image.channels}
+ for image in spec.images
+ ],
+ "allow_extra_images": spec.allow_extra_images,
+ "timestamp_unit": spec.timestamp_unit,
+ "timestamp_clock": spec.timestamp_clock,
+ }
+
+
+def observation_space_from_wire(payload: Mapping[str, Any]) -> ObservationSpaceSpec:
+ """Deserialize and validate an observation-space contract."""
+ _require_schema(payload, "telefuser.vla.observation_space")
+ state_names = payload.get("state_dimension_names")
+ images = payload.get("images")
+ if not isinstance(state_names, list) or not isinstance(images, list):
+ raise ValueError("VLA observation-space state_dimension_names and images must be arrays")
+ if any(not isinstance(image, Mapping) for image in images):
+ raise ValueError("VLA observation-space images must contain objects")
+ try:
+ return ObservationSpaceSpec(
+ state_dimension_names=tuple(state_names),
+ images=tuple(
+ ImageObservationSpec(
+ name=image["name"],
+ dtype=image.get("dtype", "uint8"),
+ layout=image.get("layout", "HWC"),
+ channels=image.get("channels", 3),
+ )
+ for image in images
+ ),
+ allow_extra_images=payload.get("allow_extra_images", True),
+ timestamp_unit=payload.get("timestamp_unit", "nanosecond"),
+ timestamp_clock=payload.get("timestamp_clock", "observation_source_monotonic"),
+ )
+ except (KeyError, TypeError) as error:
+ raise ValueError("invalid VLA observation-space payload") from error
+
+
+def tensor_to_wire(value: torch.Tensor | np.ndarray) -> dict[str, Any]:
+ """Serialize a dense CPU copy of a tensor using explicit raw bytes."""
+ tensor = torch.from_numpy(np.asarray(value).copy()) if isinstance(value, np.ndarray) else value
+ if not isinstance(tensor, torch.Tensor):
+ raise TypeError("VLA tensor payload must be a torch.Tensor or numpy.ndarray")
+ dtype_name = _DTYPE_NAMES.get(tensor.dtype)
+ if dtype_name is None:
+ raise ValueError(f"unsupported VLA tensor dtype: {tensor.dtype}")
+ contiguous = tensor.detach().to(device="cpu").contiguous()
+ raw = contiguous.view(torch.uint8).numpy().tobytes()
+ return {
+ "schema": "telefuser.vla.tensor",
+ "schema_version": VLA_WIRE_SCHEMA_VERSION,
+ "dtype": dtype_name,
+ "shape": list(contiguous.shape),
+ "data": base64.b64encode(raw).decode("ascii"),
+ }
+
+
+def tensor_from_wire(
+ payload: Mapping[str, Any],
+ *,
+ max_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+) -> torch.Tensor:
+ """Deserialize a bounded tensor payload into independent CPU storage."""
+ _require_schema(payload, "telefuser.vla.tensor")
+ if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes < 1:
+ raise ValueError("max_bytes must be a positive integer")
+ dtype_name = payload.get("dtype")
+ if dtype_name not in _DTYPES:
+ raise ValueError(f"unsupported VLA tensor dtype: {dtype_name!r}")
+ shape_value = payload.get("shape")
+ if not isinstance(shape_value, list) or any(
+ not isinstance(size, int) or isinstance(size, bool) or size < 0 for size in shape_value
+ ):
+ raise ValueError("VLA tensor shape must contain non-negative integers")
+ encoded = payload.get("data")
+ if not isinstance(encoded, str):
+ raise ValueError("VLA tensor data must be a Base64 string")
+ try:
+ raw = base64.b64decode(encoded, validate=True)
+ except ValueError as error:
+ raise ValueError("VLA tensor data is not valid Base64") from error
+ if len(raw) > max_bytes:
+ raise ValueError(f"VLA tensor exceeds the {max_bytes}-byte limit")
+ dtype = _DTYPES[dtype_name]
+ element_size = torch.empty((), dtype=dtype).element_size()
+ element_count = math.prod(shape_value)
+ if len(raw) != element_count * element_size:
+ raise ValueError("VLA tensor byte length does not match dtype and shape")
+ if element_count == 0:
+ return torch.empty(tuple(shape_value), dtype=dtype)
+ return torch.frombuffer(bytearray(raw), dtype=dtype).reshape(tuple(shape_value)).clone()
+
+
+def robot_observation_to_wire(observation: RobotObservation) -> dict[str, Any]:
+ """Serialize one robot observation and its named camera tensors."""
+ if any(not isinstance(name, str) or not name for name in observation.images):
+ raise ValueError("robot observation image names must be non-empty strings")
+ images = {name: tensor_to_wire(_as_tensor(image)) for name, image in observation.images.items()}
+ payload = {
+ "schema": "telefuser.vla.robot_observation",
+ "schema_version": VLA_WIRE_SCHEMA_VERSION,
+ "state": {
+ "values": tensor_to_wire(observation.state.values),
+ "dimension_names": list(observation.state.dimension_names),
+ "timestamp_ns": observation.state.timestamp_ns,
+ },
+ "images": images,
+ "metadata": dict(observation.metadata),
+ }
+ _require_json_value(payload["metadata"], "observation metadata")
+ return payload
+
+
+def robot_observation_from_wire(
+ payload: Mapping[str, Any],
+ *,
+ max_tensor_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+) -> RobotObservation:
+ """Deserialize one bounded robot observation."""
+ _require_schema(payload, "telefuser.vla.robot_observation")
+ state_payload = payload.get("state")
+ images_payload = payload.get("images")
+ if not isinstance(state_payload, Mapping) or not isinstance(images_payload, Mapping):
+ raise ValueError("robot observation requires state and images objects")
+ try:
+ dimension_names = state_payload.get("dimension_names")
+ if not isinstance(dimension_names, list):
+ raise ValueError("robot state dimension_names must be an array")
+ if any(not isinstance(name, str) or not name for name in images_payload):
+ raise ValueError("robot observation image names must be non-empty strings")
+ state = RobotState(
+ values=tensor_from_wire(state_payload["values"], max_bytes=max_tensor_bytes),
+ dimension_names=tuple(dimension_names),
+ timestamp_ns=state_payload["timestamp_ns"],
+ )
+ images = {name: tensor_from_wire(image, max_bytes=max_tensor_bytes) for name, image in images_payload.items()}
+ metadata = payload.get("metadata", {})
+ _require_json_value(metadata, "observation metadata")
+ return RobotObservation(state=state, images=images, metadata=metadata)
+ except (KeyError, TypeError) as error:
+ raise ValueError("invalid robot observation payload") from error
+
+
+def robot_action_chunk_to_wire(chunk: RobotActionChunk) -> dict[str, Any]:
+ """Serialize a semantically described robot action chunk."""
+ metadata = dict(chunk.metadata)
+ _require_json_value(metadata, "action chunk metadata")
+ return {
+ "schema": "telefuser.vla.robot_action_chunk",
+ "schema_version": VLA_WIRE_SCHEMA_VERSION,
+ "actions": tensor_to_wire(chunk.actions),
+ "action_space": action_space_to_wire(chunk.action_space),
+ "valid_length": chunk.valid_length,
+ "observation_timestamp_ns": chunk.observation_timestamp_ns,
+ "sequence_id": chunk.sequence_id,
+ "episode_id": chunk.episode_id,
+ "metadata": metadata,
+ }
+
+
+def robot_action_chunk_from_wire(
+ payload: Mapping[str, Any],
+ *,
+ max_tensor_bytes: int = DEFAULT_MAX_TENSOR_BYTES,
+) -> RobotActionChunk:
+ """Deserialize and validate a robot action chunk."""
+ _require_schema(payload, "telefuser.vla.robot_action_chunk")
+ try:
+ metadata = payload.get("metadata", {})
+ _require_json_value(metadata, "action chunk metadata")
+ return RobotActionChunk(
+ actions=tensor_from_wire(payload["actions"], max_bytes=max_tensor_bytes),
+ action_space=action_space_from_wire(payload["action_space"]),
+ valid_length=payload["valid_length"],
+ observation_timestamp_ns=payload["observation_timestamp_ns"],
+ sequence_id=payload["sequence_id"],
+ episode_id=payload["episode_id"],
+ metadata=metadata,
+ )
+ except (KeyError, TypeError) as error:
+ raise ValueError("invalid robot action chunk payload") from error
+
+
+def dumps_wire_message(payload: Mapping[str, Any]) -> str:
+ """Encode one deterministic compact JSON WebSocket message."""
+ return json.dumps(dict(payload), sort_keys=True, separators=(",", ":"), allow_nan=False)
+
+
+def loads_wire_message(payload: str, *, max_message_bytes: int) -> dict[str, Any]:
+ """Decode one size-bounded JSON WebSocket message."""
+ if not isinstance(payload, str):
+ raise TypeError("VLA WebSocket messages must be JSON text")
+ if not isinstance(max_message_bytes, int) or isinstance(max_message_bytes, bool) or max_message_bytes < 1:
+ raise ValueError("max_message_bytes must be a positive integer")
+ if len(payload.encode("utf-8")) > max_message_bytes:
+ raise ValueError(f"VLA WebSocket message exceeds the {max_message_bytes}-byte limit")
+ try:
+ decoded = json.loads(payload)
+ except json.JSONDecodeError as error:
+ raise ValueError("VLA WebSocket message is not valid JSON") from error
+ if not isinstance(decoded, dict):
+ raise ValueError("VLA WebSocket message must be a JSON object")
+ return decoded
+
+
+def _require_schema(payload: Mapping[str, Any], expected: str) -> None:
+ if not isinstance(payload, Mapping):
+ raise ValueError(f"{expected} payload must be an object")
+ if payload.get("schema") != expected:
+ raise ValueError(f"expected schema {expected!r}")
+ if payload.get("schema_version") != VLA_WIRE_SCHEMA_VERSION:
+ raise ValueError(f"unsupported VLA schema version: {payload.get('schema_version')!r}")
+
+
+def _require_json_value(value: Any, name: str) -> None:
+ try:
+ json.dumps(value, allow_nan=False)
+ except (TypeError, ValueError) as error:
+ raise ValueError(f"{name} must contain finite JSON values") from error
+
+
+def _as_tensor(value: Any) -> torch.Tensor | np.ndarray:
+ if isinstance(value, (torch.Tensor, np.ndarray)):
+ return value
+ raise TypeError("serialized robot images must be torch.Tensor or numpy.ndarray values")
diff --git a/telefuser/vla/session.py b/telefuser/vla/session.py
new file mode 100644
index 00000000..ac174fcc
--- /dev/null
+++ b/telefuser/vla/session.py
@@ -0,0 +1,223 @@
+"""Session lifecycle for semantic VLA inference and action preparation."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from threading import RLock
+
+from .contracts import ModelActionChunk, RobotActionChunk, RobotObservation, VLARequest
+from .embodiment import EmbodimentAdapter
+from .policy import VLAPolicy
+from .registry import VLARegistry
+from .runtime.executor import ChunkExecutor
+from .runtime.safety import ActionSafetyPolicy
+
+
+@dataclass(frozen=True)
+class VLASessionContract:
+ """Immutable component and action-space selection made during OPEN."""
+
+ session_id: str
+ episode_id: str
+ model_id: str
+ embodiment_id: str
+
+
+class VLASession:
+ """Bind one loaded policy and embodiment for a continuous episode."""
+
+ def __init__(
+ self,
+ contract: VLASessionContract,
+ policy: VLAPolicy,
+ embodiment: EmbodimentAdapter,
+ executor: ChunkExecutor,
+ policy_lock: RLock | None = None,
+ ) -> None:
+ self.contract = contract
+ self.policy = policy
+ self.embodiment = embodiment
+ self.executor = executor
+ self._policy_lock = policy_lock or RLock()
+ self._last_sequence_id: int | None = None
+ self._lock = RLock()
+
+ def predict(
+ self,
+ observation: RobotObservation,
+ instruction: str,
+ sequence_id: int,
+ *,
+ seed: int | None = None,
+ now_ns: int | None = None,
+ timings: dict[str, float] | None = None,
+ ) -> RobotActionChunk:
+ """Run one ordered observation through policy, embodiment, and runtime."""
+ with self._lock:
+ if not isinstance(sequence_id, int) or isinstance(sequence_id, bool) or sequence_id < 0:
+ raise ValueError("sequence_id must be a non-negative integer")
+ if self._last_sequence_id is not None and sequence_id <= self._last_sequence_id:
+ raise ValueError(
+ f"sequence_id must increase within a session: got {sequence_id}, "
+ f"last accepted {self._last_sequence_id}"
+ )
+ capabilities = self.policy.capabilities()
+ if seed is not None and not capabilities.supports_seed:
+ raise ValueError(f"VLA policy {capabilities.model_id!r} does not support seeded inference")
+ self.embodiment.observation_space.validate(observation)
+ stage_started_at = time.monotonic()
+ model_observation = self.embodiment.encode_observation(observation)
+ if timings is not None:
+ timings["encode_observation_ms"] = (time.monotonic() - stage_started_at) * 1000.0
+ request = VLARequest(
+ observation=model_observation,
+ instruction=instruction,
+ episode_id=self.contract.episode_id,
+ sequence_id=sequence_id,
+ observation_timestamp_ns=observation.state.timestamp_ns,
+ seed=seed,
+ )
+ stage_started_at = time.monotonic()
+ with self._policy_lock:
+ model_chunk = self.policy.predict(request)
+ if timings is not None:
+ timings["policy_ms"] = (time.monotonic() - stage_started_at) * 1000.0
+ if not isinstance(model_chunk, ModelActionChunk):
+ raise TypeError("VLA policy must return ModelActionChunk")
+ capabilities.output_action_space.require_compatible(
+ model_chunk.action_space,
+ context="policy output action space",
+ )
+ if model_chunk.valid_length > capabilities.max_horizon:
+ raise ValueError("policy output exceeds its declared maximum horizon")
+ if model_chunk.episode_id != request.episode_id:
+ raise ValueError("policy output episode_id does not match the request")
+ if model_chunk.sequence_id != request.sequence_id:
+ raise ValueError("policy output sequence_id does not match the request")
+ if model_chunk.observation_timestamp_ns != request.observation_timestamp_ns:
+ raise ValueError("policy output observation timestamp does not match the request")
+ stage_started_at = time.monotonic()
+ robot_chunk = self.embodiment.decode_actions(model_chunk, observation.state)
+ if timings is not None:
+ timings["decode_actions_ms"] = (time.monotonic() - stage_started_at) * 1000.0
+ if not isinstance(robot_chunk, RobotActionChunk):
+ raise TypeError("embodiment adapter must return RobotActionChunk")
+ if robot_chunk.episode_id != model_chunk.episode_id:
+ raise ValueError("embodiment output episode_id does not match the model chunk")
+ if robot_chunk.sequence_id != model_chunk.sequence_id:
+ raise ValueError("embodiment output sequence_id does not match the model chunk")
+ if robot_chunk.observation_timestamp_ns != model_chunk.observation_timestamp_ns:
+ raise ValueError("embodiment output observation timestamp does not match the model chunk")
+ stage_started_at = time.monotonic()
+ prepared = self.executor.prepare(robot_chunk, observation.state, now_ns=now_ns)
+ if timings is not None:
+ timings["prepare_actions_ms"] = (time.monotonic() - stage_started_at) * 1000.0
+ self._last_sequence_id = sequence_id
+ return prepared
+
+ def reset(self, episode_id: str | None = None) -> None:
+ """Clear temporal policy and ordering state for the bound episode."""
+ with self._lock:
+ old_episode_id = self.contract.episode_id
+ with self._policy_lock:
+ self.policy.reset(old_episode_id)
+ if episode_id is not None:
+ if not isinstance(episode_id, str) or not episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+ self.contract = VLASessionContract(
+ session_id=self.contract.session_id,
+ episode_id=episode_id,
+ model_id=self.contract.model_id,
+ embodiment_id=self.contract.embodiment_id,
+ )
+ self._last_sequence_id = None
+
+ def close(self) -> None:
+ """Release policy-owned state associated with this session."""
+ with self._lock:
+ with self._policy_lock:
+ self.policy.reset(self.contract.episode_id)
+ self._last_sequence_id = None
+
+
+class VLASessionManager:
+ """Implement OPEN, PREDICT, RESET, and CLOSE without transport concerns."""
+
+ def __init__(self, registry: VLARegistry) -> None:
+ self.registry = registry
+ self._sessions: dict[str, VLASession] = {}
+ self._policy_locks: dict[int, RLock] = {}
+ self._lock = RLock()
+
+ def open(
+ self,
+ session_id: str,
+ *,
+ model_id: str,
+ embodiment_id: str,
+ episode_id: str,
+ execute_horizon: int | None = None,
+ safety_policy: ActionSafetyPolicy | None = None,
+ max_observation_age_ns: int | None = None,
+ ) -> VLASession:
+ """Bind one session to loaded component instances and validate semantics."""
+ if not isinstance(session_id, str) or not session_id:
+ raise ValueError("session_id must be a non-empty string")
+ if not isinstance(episode_id, str) or not episode_id:
+ raise ValueError("episode_id must be a non-empty string")
+ with self._lock:
+ if session_id in self._sessions:
+ raise ValueError(f"VLA session is already open: {session_id!r}")
+ policy = self.registry.get_policy(model_id)
+ embodiment = self.registry.get_embodiment(embodiment_id)
+ embodiment.model_action_space.require_compatible(
+ policy.capabilities().output_action_space,
+ context="policy and embodiment action space",
+ )
+ contract = VLASessionContract(
+ session_id=session_id,
+ episode_id=episode_id,
+ model_id=model_id,
+ embodiment_id=embodiment_id,
+ )
+ session = VLASession(
+ contract,
+ policy,
+ embodiment,
+ ChunkExecutor(
+ embodiment.robot_action_space,
+ execute_horizon=execute_horizon,
+ safety_policy=safety_policy,
+ max_observation_age_ns=max_observation_age_ns,
+ ),
+ self._policy_locks.setdefault(id(policy), RLock()),
+ )
+ self._sessions[session_id] = session
+ return session
+
+ def get(self, session_id: str) -> VLASession:
+ """Return one active session."""
+ with self._lock:
+ try:
+ return self._sessions[session_id]
+ except KeyError as error:
+ raise KeyError(f"unknown VLA session_id: {session_id!r}") from error
+
+ def reset(self, session_id: str, episode_id: str | None = None) -> None:
+ """Reset ordering and policy history for one active session."""
+ self.get(session_id).reset(episode_id)
+
+ def close(self, session_id: str) -> None:
+ """Close one session and release its policy state."""
+ with self._lock:
+ try:
+ session = self._sessions.pop(session_id)
+ except KeyError as error:
+ raise KeyError(f"unknown VLA session_id: {session_id!r}") from error
+ session.close()
+
+ def session_ids(self) -> tuple[str, ...]:
+ """List active session IDs in deterministic order."""
+ with self._lock:
+ return tuple(sorted(self._sessions))
diff --git a/tests/integration/test_vla_replica_process.py b/tests/integration/test_vla_replica_process.py
new file mode 100644
index 00000000..686cb5ef
--- /dev/null
+++ b/tests/integration/test_vla_replica_process.py
@@ -0,0 +1,67 @@
+"""Integration test for the VLA session RPC across a real replica process."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+import pytest
+import torch
+
+from telefuser.service.core.pipeline_pool import PipelinePool
+from telefuser.vla import RobotObservation, RobotState
+from telefuser.vla.serialization import robot_action_chunk_from_wire, robot_observation_to_wire
+
+PIPELINE_FILE = Path(__file__).parents[1] / "unit" / "service" / "fixtures" / "fake_vla_pipeline.py"
+
+
+@pytest.mark.slow
+def test_pipeline_pool_dispatches_vla_lifecycle_across_replica_process() -> None:
+ pool = PipelinePool(
+ num_replicas=1,
+ replica_device_ids=[[]],
+ security_level_name="NONE",
+ )
+ assert pool.start_all(
+ str(PIPELINE_FILE),
+ parallelism_per_replica=1,
+ task="vla_action",
+ skip_validation=True,
+ vla_provider_factory="get_vla_provider",
+ )
+
+ async def scenario() -> None:
+ try:
+ await pool.open_session("session")
+ opened = await pool.run_vla_operation(
+ "session",
+ "OPEN",
+ {
+ "session_id": "session",
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ "episode_id": "episode",
+ },
+ )
+ assert opened["model_id"] == "fake"
+ observation = RobotObservation(RobotState(torch.tensor([1.0]), ("joint",), 100), {})
+ response = await pool.run_vla_operation(
+ "session",
+ "PREDICT",
+ {
+ "session_id": "session",
+ "sequence_id": 1,
+ "observation_timestamp_ns": 100,
+ "instruction": "move",
+ "observation": robot_observation_to_wire(observation),
+ },
+ )
+ chunk = robot_action_chunk_from_wire(response["chunk"])
+ assert chunk.actions.flatten().tolist() == pytest.approx([1.25, 1.5])
+ await pool.run_vla_operation("session", "RESET", {"session_id": "session"})
+ await pool.run_vla_operation("session", "CLOSE", {"session_id": "session"})
+ await pool.close_session("session")
+ finally:
+ await pool.aclose()
+
+ asyncio.run(scenario())
diff --git a/tests/unit/integrations/sim/test_mujoco.py b/tests/unit/integrations/sim/test_mujoco.py
new file mode 100644
index 00000000..7109954e
--- /dev/null
+++ b/tests/unit/integrations/sim/test_mujoco.py
@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+mujoco = pytest.importorskip("mujoco")
+
+from telefuser.integrations.sim import MuJoCoJointBinding, MuJoCoSimulatorAdapter
+from telefuser.vla import ActionSpaceSpec, ObservationSpaceSpec
+from telefuser.vla.runtime import RobotAction
+
+ACTION_NAMES = tuple(f"action_{index}" for index in range(14))
+ACTION_SPACE = ActionSpaceSpec(
+ "absolute_qpos",
+ ACTION_NAMES,
+ ("radian",) * 6 + ("normalized_position",) + ("radian",) * 6 + ("normalized_position",),
+ "robot_joint",
+ None,
+ False,
+)
+OBSERVATION_SPACE = ObservationSpaceSpec(ACTION_NAMES)
+
+
+def _model() -> mujoco.MjModel:
+ joint_names = (
+ *(f"left_{index}" for index in range(6)),
+ "left_gripper_a",
+ "left_gripper_b",
+ *(f"right_{index}" for index in range(6)),
+ "right_gripper_a",
+ "right_gripper_b",
+ )
+ bodies = []
+ for index, name in enumerate(joint_names):
+ joint_type = "slide" if "gripper" in name else "hinge"
+ joint_range = "0 0.05" if "gripper" in name else "-1 1"
+ bodies.append(
+ f''
+ f''
+ ''
+ ""
+ )
+ return mujoco.MjModel.from_xml_string(
+ f'{"".join(bodies)}'
+ )
+
+
+def _bindings() -> tuple[MuJoCoJointBinding, ...]:
+ return (
+ *(MuJoCoJointBinding(ACTION_NAMES[index], (f"left_{index}",)) for index in range(6)),
+ MuJoCoJointBinding(ACTION_NAMES[6], ("left_gripper_a", "left_gripper_b"), True),
+ *(MuJoCoJointBinding(ACTION_NAMES[index + 7], (f"right_{index}",)) for index in range(6)),
+ MuJoCoJointBinding(ACTION_NAMES[13], ("right_gripper_a", "right_gripper_b"), True),
+ )
+
+
+def test_mujoco_adapter_executes_pd_action_and_returns_semantic_state() -> None:
+ adapter = MuJoCoSimulatorAdapter(
+ _model(),
+ ACTION_SPACE,
+ OBSERVATION_SPACE,
+ _bindings(),
+ steps_per_action=20,
+ )
+ initial = adapter.reset()
+ target = torch.full((14,), 0.2)
+ target[6] = 0.5
+ target[13] = 0.75
+
+ adapter.execute(RobotAction(target, ACTION_SPACE, sequence_id=0, step_index=0, episode_id="episode"))
+ observed = adapter.observe()
+
+ assert observed.state.dimension_names == ACTION_NAMES
+ assert observed.state.values.shape == (14,)
+ assert observed.metadata["simulator"] == "mujoco"
+ assert observed.metadata["simulation_time_s"] > 0
+ assert torch.linalg.vector_norm(observed.state.values - initial.state.values) > 0
+ assert 0 < observed.state.values[6] <= 1
+ assert 0 < observed.state.values[13] <= 1
+ adapter.close()
+
+
+def test_mujoco_adapter_rejects_binding_order_mismatch() -> None:
+ bindings = list(_bindings())
+ bindings[0], bindings[1] = bindings[1], bindings[0]
+
+ with pytest.raises(ValueError, match="binding order"):
+ MuJoCoSimulatorAdapter(_model(), ACTION_SPACE, OBSERVATION_SPACE, bindings)
diff --git a/tests/unit/integrations/sim/test_robotwin.py b/tests/unit/integrations/sim/test_robotwin.py
new file mode 100644
index 00000000..fd67aeda
--- /dev/null
+++ b/tests/unit/integrations/sim/test_robotwin.py
@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from telefuser.integrations.sim import RoboTwinSimulatorAdapter
+from telefuser.vla import ActionSpaceSpec, RobotActionChunk, RobotObservation, RobotState
+
+SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), "robot_joint", None, False)
+
+
+def _observation() -> RobotObservation:
+ return RobotObservation(RobotState(torch.zeros(1), ("joint",), 0), {})
+
+
+def test_robotwin_adapter_executes_only_valid_chunk_actions_as_float32() -> None:
+ executed: list[np.ndarray] = []
+ adapter = RoboTwinSimulatorAdapter(
+ SPACE,
+ observe_fn=_observation,
+ reset_fn=_observation,
+ execute_fn=executed.append,
+ )
+ chunk = RobotActionChunk(torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16), SPACE, 2, 0, 1, "episode")
+
+ count = adapter.execute_chunk(chunk)
+
+ assert count == 2
+ assert [action.dtype for action in executed] == [np.float32, np.float32]
+ assert [action.tolist() for action in executed] == [[1.0], [2.0]]
+ assert adapter.observe() == _observation()
+
+
+def test_robotwin_adapter_rejects_semantically_different_action() -> None:
+ other = ActionSpaceSpec("joint_delta", ("joint",), ("radian",), "robot_joint", None, False)
+ adapter = RoboTwinSimulatorAdapter(
+ SPACE,
+ observe_fn=_observation,
+ reset_fn=_observation,
+ execute_fn=lambda _action: None,
+ )
+ chunk = RobotActionChunk(torch.zeros(1, 1), other, 1, 0, 0, "episode")
+
+ with pytest.raises(ValueError, match="representation"):
+ adapter.execute_chunk(chunk)
diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py b/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py
index cae1ede1..fcd5c6fa 100644
--- a/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py
+++ b/tests/unit/pipelines/lingbot_vla_v2/test_action_scheduler.py
@@ -8,6 +8,11 @@
import pytest
from telefuser.pipelines.lingbot_vla_v2.action_scheduler import ActionChunkScheduler
+from telefuser.vla.runtime import ActionChunkScheduler as CommonActionChunkScheduler
+
+
+def test_lingbot_scheduler_import_is_a_compatibility_alias() -> None:
+ assert ActionChunkScheduler is CommonActionChunkScheduler
def test_scheduler_discards_inflight_and_pending_work_when_newer_observation_arrives() -> None:
diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_examples.py b/tests/unit/pipelines/lingbot_vla_v2/test_examples.py
index 3a78bd07..f3017c18 100644
--- a/tests/unit/pipelines/lingbot_vla_v2/test_examples.py
+++ b/tests/unit/pipelines/lingbot_vla_v2/test_examples.py
@@ -2,7 +2,11 @@
from click.testing import CliRunner
-from examples.lingbot_vla_v2 import lingbot_vla_v2_inference, lingbot_vla_v2_native_service
+from examples.lingbot_vla_v2 import (
+ lingbot_vla_v2_inference,
+ lingbot_vla_v2_native_service,
+ lingbot_vla_v2_vla_server,
+)
def test_direct_inference_cli_exposes_cuda_graph() -> None:
@@ -58,3 +62,38 @@ def fake_get_pipeline(model_root: str, qwen3vl_root: str, **kwargs: object) -> o
assert captured["warmup"] is True
assert captured["quantization"] == "fused-fp8-graph"
assert captured["cuda_graph"] is True
+
+
+def test_native_service_exposes_worker_local_vla_provider() -> None:
+ provider = lingbot_vla_v2_native_service.get_vla_provider(object())
+ assert provider.metadata() == {
+ "model_ids": ["lingbot-vla-v2"],
+ "embodiment_ids": ["robotwin"],
+ }
+
+
+def test_generic_vla_server_starts_pool_with_optional_provider(monkeypatch) -> None:
+ captured: dict[str, object] = {}
+ sentinel = object()
+
+ class FakePool:
+ def __init__(self, **kwargs: object) -> None:
+ captured["pool"] = kwargs
+
+ def start_all(self, **kwargs: object) -> bool:
+ captured["start"] = kwargs
+ return True
+
+ monkeypatch.setattr(lingbot_vla_v2_vla_server, "PipelinePool", FakePool)
+ monkeypatch.setattr(
+ lingbot_vla_v2_vla_server,
+ "create_pipeline_pool_vla_session_app",
+ lambda pool: sentinel,
+ )
+ app = lingbot_vla_v2_vla_server.create_app(parallelism=1, num_replicas=1)
+
+ assert app is sentinel
+ start = captured["start"]
+ assert isinstance(start, dict)
+ assert start["vla_provider_factory"] == "get_vla_provider"
+ assert start["task"] == "vla_action"
diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py
index f2879e53..0b084ec5 100644
--- a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py
+++ b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py
@@ -3,7 +3,14 @@
import pytest
import torch
-from telefuser.pipelines.lingbot_vla_v2.robot_profile import RobotWinProfile
+from telefuser.pipelines.lingbot_vla_v2.robot_profile import (
+ LINGBOT_VLA_V2_ACTION_SPACE,
+ ROBOTWIN_ACTION_ORDER,
+ ROBOTWIN_ACTION_SPACE,
+ ROBOTWIN_OBSERVATION_SPACE,
+ RobotWinProfile,
+)
+from telefuser.vla import ModelActionChunk, RobotObservation, RobotState
def _stats() -> dict[str, dict[str, list[float]]]:
@@ -82,3 +89,31 @@ def test_profile_rejects_invalid_state_and_action_shapes() -> None:
profile.normalize_state(torch.zeros(13))
with pytest.raises(ValueError, match="canonical actions must have shape"):
profile.structure_actions(torch.zeros(2, 54))
+
+
+def test_profile_implements_semantic_embodiment_contract() -> None:
+ profile = RobotWinProfile(_stats())
+ observation = RobotObservation(
+ RobotState(torch.zeros(14), ROBOTWIN_ACTION_ORDER, timestamp_ns=123),
+ {key: torch.zeros((2, 2, 3), dtype=torch.uint8) for key in profile.camera_keys},
+ )
+ model_observation = profile.encode_observation(observation)
+ model_chunk = ModelActionChunk(
+ torch.zeros(3, 55),
+ LINGBOT_VLA_V2_ACTION_SPACE,
+ 2,
+ 123,
+ 7,
+ "episode",
+ {"source": "lingbot"},
+ )
+
+ robot_chunk = profile.decode_actions(model_chunk, observation.state)
+
+ assert model_observation.state.shape == (14,)
+ assert profile.observation_space == ROBOTWIN_OBSERVATION_SPACE
+ assert robot_chunk.actions.shape == (2, 14)
+ assert robot_chunk.valid_length == 2
+ assert robot_chunk.action_space == ROBOTWIN_ACTION_SPACE
+ assert robot_chunk.sequence_id == 7
+ assert robot_chunk.metadata["source"] == "lingbot"
diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py
index 22c79d07..3ace5201 100644
--- a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py
+++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py
@@ -198,6 +198,21 @@ def test_websocket_is_persistent_and_uses_upstream_response_fields() -> None:
assert reset_response["episode_id"] == "episode-1"
assert reset_response["server_timing"]["prev_total_ms"] >= 0
+ websocket.send_bytes(
+ server.pack_message(
+ {
+ **_observation(),
+ "request_id": "after-reset",
+ "episode_id": "episode-1",
+ "sequence_id": 0,
+ }
+ )
+ )
+ after_reset = server.unpack_message(websocket.receive_bytes())
+ assert after_reset["scheduler_status"] == "completed"
+ assert after_reset["sequence_id"] == 0
+ assert after_reset["action"].shape == (3, 14)
+
def test_websocket_accepts_overlapping_chunks_and_discards_superseded_actions() -> None:
first_started = threading.Event()
@@ -256,32 +271,3 @@ def test_cli_exposes_isolated_robotwin_server_options() -> None:
assert "--use-length" in result.output
assert "--max-pending-sessions" in result.output
assert "--cuda-graph" in result.output
-
-
-def test_h100_policy_process_disables_only_cudnn_sdpa(monkeypatch) -> None:
- calls: list[tuple[str, bool]] = []
- monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
- monkeypatch.setattr(torch.cuda, "get_device_name", lambda _device: "NVIDIA H100 80GB HBM3")
- monkeypatch.setattr(torch.backends.cuda, "enable_cudnn_sdp", lambda enabled: calls.append(("cudnn", enabled)))
- monkeypatch.setattr(torch.backends.cuda, "enable_flash_sdp", lambda enabled: calls.append(("flash", enabled)))
- monkeypatch.setattr(torch.backends.cuda, "enable_math_sdp", lambda enabled: calls.append(("math", enabled)))
- monkeypatch.setattr(
- torch.backends.cuda,
- "enable_mem_efficient_sdp",
- lambda enabled: calls.append(("mem_efficient", enabled)),
- )
-
- server._configure_h100_sdpa_backends("cuda:0")
-
- assert calls == [("cudnn", False), ("flash", True), ("math", True), ("mem_efficient", True)]
-
-
-def test_non_h100_policy_process_preserves_sdpa_backends(monkeypatch) -> None:
- calls: list[bool] = []
- monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
- monkeypatch.setattr(torch.cuda, "get_device_name", lambda _device: "NVIDIA RTX 4090")
- monkeypatch.setattr(torch.backends.cuda, "enable_cudnn_sdp", calls.append)
-
- server._configure_h100_sdpa_backends("cuda:0")
-
- assert calls == []
diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_runtime.py b/tests/unit/pipelines/lingbot_vla_v2/test_runtime.py
index 912470ea..8d3a6567 100644
--- a/tests/unit/pipelines/lingbot_vla_v2/test_runtime.py
+++ b/tests/unit/pipelines/lingbot_vla_v2/test_runtime.py
@@ -1,9 +1,11 @@
from __future__ import annotations
import pytest
+import torch
from telefuser.core.config import QuantKernelBackend, QuantType
from telefuser.pipelines.lingbot_vla_v2.runtime import (
+ configure_lingbot_vla_v2_h100_sdpa,
get_lingbot_vla_v2_pipeline,
lingbot_vla_v2_quant_config,
)
@@ -54,3 +56,32 @@ def test_quantization_rejects_unknown_name() -> None:
def test_cuda_graph_rejects_cpu_before_loading_models() -> None:
with pytest.raises(ValueError, match="CUDA Graph requires a CUDA device"):
get_lingbot_vla_v2_pipeline("unused", "unused", device="cpu", cuda_graph=True)
+
+
+def test_h100_policy_process_disables_only_cudnn_sdpa(monkeypatch) -> None:
+ calls: list[tuple[str, bool]] = []
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "get_device_name", lambda _device: "NVIDIA H100 80GB HBM3")
+ monkeypatch.setattr(torch.backends.cuda, "enable_cudnn_sdp", lambda enabled: calls.append(("cudnn", enabled)))
+ monkeypatch.setattr(torch.backends.cuda, "enable_flash_sdp", lambda enabled: calls.append(("flash", enabled)))
+ monkeypatch.setattr(torch.backends.cuda, "enable_math_sdp", lambda enabled: calls.append(("math", enabled)))
+ monkeypatch.setattr(
+ torch.backends.cuda,
+ "enable_mem_efficient_sdp",
+ lambda enabled: calls.append(("mem_efficient", enabled)),
+ )
+
+ configure_lingbot_vla_v2_h100_sdpa("cuda:0")
+
+ assert calls == [("cudnn", False), ("flash", True), ("math", True), ("mem_efficient", True)]
+
+
+def test_non_h100_policy_process_preserves_sdpa_backends(monkeypatch) -> None:
+ calls: list[bool] = []
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "get_device_name", lambda _device: "NVIDIA RTX 4090")
+ monkeypatch.setattr(torch.backends.cuda, "enable_cudnn_sdp", calls.append)
+
+ configure_lingbot_vla_v2_h100_sdpa("cuda:0")
+
+ assert calls == []
diff --git a/tests/unit/pipelines/lingbot_vla_v2/test_vla_policy.py b/tests/unit/pipelines/lingbot_vla_v2/test_vla_policy.py
new file mode 100644
index 00000000..d09c5150
--- /dev/null
+++ b/tests/unit/pipelines/lingbot_vla_v2/test_vla_policy.py
@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import torch
+
+from telefuser.pipelines.lingbot_vla_v2 import (
+ LINGBOT_VLA_V2_ACTION_SPACE,
+ LingBotVlaV2CanonicalActionChunk,
+ LingBotVlaV2VLAPolicy,
+)
+from telefuser.vla import ModelObservation, VLARequest
+
+
+class _Pipeline:
+ def __init__(self) -> None:
+ self.seed: int | None = None
+
+ def __call__(self, observation, seed: int | None = None) -> LingBotVlaV2CanonicalActionChunk:
+ self.seed = seed
+ assert observation.task == "pick"
+ return LingBotVlaV2CanonicalActionChunk(
+ canonical_normalized_actions=torch.zeros(4, 55),
+ horizon=4,
+ action_dim=55,
+ checkpoint_variant="base",
+ policy_verified=False,
+ verification_status="unverified",
+ )
+
+
+def test_lingbot_policy_wraps_existing_pipeline_without_changing_its_api() -> None:
+ pipeline = _Pipeline()
+ policy = LingBotVlaV2VLAPolicy(pipeline, max_horizon=4)
+ request = VLARequest(
+ ModelObservation(torch.zeros(14), {"camera": object()}),
+ "pick",
+ "episode",
+ sequence_id=3,
+ observation_timestamp_ns=123,
+ seed=9,
+ )
+
+ chunk = policy.predict(request)
+
+ assert pipeline.seed == 9
+ assert chunk.actions.shape == (4, 55)
+ assert chunk.action_space == LINGBOT_VLA_V2_ACTION_SPACE
+ assert chunk.sequence_id == 3
+ assert chunk.observation_timestamp_ns == 123
+ assert chunk.metadata["verification_status"] == "unverified"
diff --git a/tests/unit/service/fixtures/fake_vla_pipeline.py b/tests/unit/service/fixtures/fake_vla_pipeline.py
new file mode 100644
index 00000000..56523ef0
--- /dev/null
+++ b/tests/unit/service/fixtures/fake_vla_pipeline.py
@@ -0,0 +1,92 @@
+"""CPU-only pipeline fixture for replica-process VLA tests."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from telefuser.service.vla_replica import VLAReplicaProvider
+from telefuser.vla import (
+ ActionSpaceSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+ VLACapabilities,
+ VLARegistry,
+ VLASessionManager,
+)
+
+MODEL_SPACE = ActionSpaceSpec("joint_delta", ("joint",), ("radian",), None, 10.0, False)
+ROBOT_SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), None, 10.0, False)
+OBSERVATION_SPACE = ObservationSpaceSpec(("joint",))
+
+PIPELINE_CONTRACT = {
+ "contract_version": "v1",
+ "pipeline_name": "fake_vla_pipeline",
+ "supported_tasks": ["vla_action"],
+ "supported_media_types": ["structured"],
+ "execution_mode": "serial_single_pipeline",
+ "effective_max_concurrent_tasks": 1,
+ "entrypoints": {"get_pipeline": "get_pipeline", "run_with_file": "run_structured"},
+ "task_contracts": {"vla_action": {"media_type": "structured", "required_inputs": [], "optional_inputs": []}},
+}
+
+
+class _Policy:
+ def capabilities(self) -> VLACapabilities:
+ return VLACapabilities("fake", MODEL_SPACE, max_horizon=2)
+
+ def predict(self, request) -> ModelActionChunk:
+ return ModelActionChunk(
+ torch.tensor([[0.25], [0.5]]),
+ MODEL_SPACE,
+ 2,
+ request.observation_timestamp_ns,
+ request.sequence_id,
+ request.episode_id,
+ )
+
+ def reset(self, episode_id: str) -> None:
+ pass
+
+
+@dataclass
+class _Embodiment:
+ embodiment_id: str = "fake-robot"
+ model_action_space: ActionSpaceSpec = MODEL_SPACE
+ robot_action_space: ActionSpaceSpec = ROBOT_SPACE
+ observation_space: ObservationSpaceSpec = OBSERVATION_SPACE
+
+ def encode_observation(self, observation: RobotObservation) -> ModelObservation:
+ return ModelObservation(observation.state.values, observation.images)
+
+ def decode_actions(self, actions: ModelActionChunk, robot_state: RobotState) -> RobotActionChunk:
+ return RobotActionChunk(
+ actions.actions + robot_state.values,
+ self.robot_action_space,
+ actions.valid_length,
+ actions.observation_timestamp_ns,
+ actions.sequence_id,
+ actions.episode_id,
+ )
+
+
+def get_pipeline(parallelism: int = 1) -> object:
+ if parallelism != 1:
+ raise ValueError("fake VLA pipeline only supports parallelism=1")
+ return object()
+
+
+def run_structured(_pipeline: object, **_kwargs: object) -> dict[str, object]:
+ return {}
+
+
+def get_vla_provider(_pipeline: object) -> VLAReplicaProvider:
+ registry = VLARegistry()
+ registry.register_policy("fake", _Policy())
+ registry.register_embodiment("fake-robot", _Embodiment())
+ return VLAReplicaProvider(VLASessionManager(registry))
diff --git a/tests/unit/service/test_pipeline_pool.py b/tests/unit/service/test_pipeline_pool.py
index e9c7f1cb..6a095d40 100644
--- a/tests/unit/service/test_pipeline_pool.py
+++ b/tests/unit/service/test_pipeline_pool.py
@@ -17,7 +17,7 @@
from telefuser.service.api.schema import TaskRequest
from telefuser.service.core.config import ServerConfig
from telefuser.service.core.pipeline_pool import PipelinePool
-from telefuser.service.core.replica_worker import ReplicaDeadError, ReplicaHandle, _forward_cancel_fn
+from telefuser.service.core.replica_worker import ReplicaDeadError, ReplicaHandle, ReplicaVLAError, _forward_cancel_fn
from telefuser.service.core.task_manager import TaskManager, TaskStatus
_DEVICE_ENV_VAR = current_platform.device_control_env_var
@@ -81,6 +81,31 @@ def test_replica_handle_converts_broken_pipe_to_dead_replica() -> None:
assert handle._dead is True
+def test_replica_handle_runs_vla_operation_and_preserves_remote_error_type() -> None:
+ process = MagicMock()
+ process.is_alive.return_value = True
+ connection = MagicMock()
+ handle = ReplicaHandle(
+ replica_id=0,
+ process=process,
+ conn=connection,
+ cancel_event=threading.Event(),
+ metadata={},
+ )
+ handle._recv_with_health_check = MagicMock(return_value=("ok", {"session_id": "session"}))
+
+ result = asyncio.run(handle.run_vla_operation("RESET", {"session_id": "session"}))
+ assert result == {"session_id": "session"}
+ connection.send.assert_called_once_with(("vla", "RESET", {"session_id": "session"}))
+
+ handle._recv_with_health_check = MagicMock(
+ return_value=("vla_error", {"type": "ValueError", "message": "action space mismatch"})
+ )
+ with pytest.raises(ReplicaVLAError, match="action space mismatch") as error:
+ asyncio.run(handle.run_vla_operation("OPEN", {}))
+ assert error.value.error_type == "ValueError"
+
+
def test_pipeline_pool_evicts_exited_replica_and_uses_remaining_capacity() -> None:
task_manager = MagicMock()
pool = PipelinePool(
diff --git a/tests/unit/service/test_pipeline_pool_sessions.py b/tests/unit/service/test_pipeline_pool_sessions.py
new file mode 100644
index 00000000..95b23c95
--- /dev/null
+++ b/tests/unit/service/test_pipeline_pool_sessions.py
@@ -0,0 +1,186 @@
+"""CPU-only tests for session-affine PipelinePool leases."""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from telefuser.service.core.pipeline_pool import PipelinePool
+from telefuser.service.core.replica_worker import ReplicaDeadError
+
+
+def _pool(num_replicas: int = 2) -> tuple[PipelinePool, list[MagicMock]]:
+ pool = PipelinePool(
+ num_replicas=num_replicas,
+ replica_device_ids=[[str(index)] for index in range(num_replicas)],
+ security_level_name="NONE",
+ )
+ handles = []
+ for replica_id in range(num_replicas):
+ handle = MagicMock()
+ handle.replica_id = replica_id
+ handle._dead = False
+ handle.process.is_alive.return_value = True
+ handles.append(handle)
+ pool._available.put_nowait(replica_id)
+ pool._handles = list(handles)
+ pool._instance_status = ["idle"] * num_replicas
+ return pool, handles
+
+
+def test_sessions_are_pinned_to_distinct_replicas_and_close_returns_capacity() -> None:
+ pool, handles = _pool()
+
+ async def scenario() -> None:
+ assert await pool.open_session("one") == 0
+ assert await pool.open_session("two") == 1
+ assert await pool.session_bindings() == {"one": 0, "two": 1}
+
+ async with pool.acquire_session("one") as first:
+ assert first is handles[0]
+ async with pool.acquire_session("one") as again:
+ assert again is handles[0]
+
+ assert await pool.close_session("one") == 0
+ async with pool.acquire() as released:
+ assert released is handles[0]
+
+ asyncio.run(scenario())
+
+
+def test_same_session_requests_are_serialized() -> None:
+ pool, _ = _pool(num_replicas=1)
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ first_entered = asyncio.Event()
+ release_first = asyncio.Event()
+ entered: list[str] = []
+
+ async def use_session(name: str) -> None:
+ async with pool.acquire_session("one"):
+ entered.append(name)
+ if name == "first":
+ first_entered.set()
+ await release_first.wait()
+
+ first = asyncio.create_task(use_session("first"))
+ await first_entered.wait()
+ second = asyncio.create_task(use_session("second"))
+ await asyncio.sleep(0)
+ assert entered == ["first"]
+ release_first.set()
+ await asyncio.gather(first, second)
+ assert entered == ["first", "second"]
+
+ asyncio.run(scenario())
+
+
+def test_vla_operation_uses_the_session_affine_replica() -> None:
+ pool, handles = _pool(num_replicas=1)
+ handles[0].run_vla_operation = AsyncMock(return_value={"session_id": "one"})
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ result = await pool.run_vla_operation("one", "RESET", {"session_id": "one"})
+ assert result == {"session_id": "one"}
+
+ asyncio.run(scenario())
+ handles[0].run_vla_operation.assert_awaited_once_with(
+ "RESET",
+ {"session_id": "one"},
+ timeout_s=None,
+ )
+
+
+def test_dead_session_replica_invalidates_binding_deterministically() -> None:
+ pool, handles = _pool()
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ handles[0].process.is_alive.return_value = False
+
+ with pytest.raises(ReplicaDeadError, match="not alive"):
+ async with pool.acquire_session("one"):
+ pass
+
+ assert await pool.session_bindings() == {}
+ assert pool._instance_status == ["dead", "idle"]
+ assert pool._live_count == 1
+ with pytest.raises(KeyError, match="unknown pipeline session"):
+ await pool.session_replica_id("one")
+
+ asyncio.run(scenario())
+
+
+def test_replica_failure_during_session_execution_evicts_once() -> None:
+ pool, handles = _pool(num_replicas=1)
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ with pytest.raises(ReplicaDeadError, match="failed"):
+ async with pool.acquire_session("one"):
+ raise ReplicaDeadError("failed")
+ assert await pool.session_bindings() == {}
+ assert pool._live_count == 0
+ handles[0].shutdown.assert_called_once()
+
+ asyncio.run(scenario())
+
+
+def test_pool_shutdown_clears_all_session_bindings() -> None:
+ pool, handles = _pool()
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ await pool.open_session("two")
+ await pool.aclose()
+ assert await pool.session_bindings() == {}
+ with pytest.raises(RuntimeError, match="shutting down"):
+ await pool.open_session("three", timeout_s=0.01)
+
+ asyncio.run(scenario())
+ for handle in handles:
+ handle.shutdown.assert_called_once()
+
+
+def test_pool_shutdown_waits_for_active_session_request() -> None:
+ pool, handles = _pool(num_replicas=1)
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ entered = asyncio.Event()
+ release = asyncio.Event()
+
+ async def use_session() -> None:
+ async with pool.acquire_session("one"):
+ entered.set()
+ await release.wait()
+
+ request = asyncio.create_task(use_session())
+ await entered.wait()
+ shutdown = asyncio.create_task(pool.aclose())
+ await asyncio.sleep(0)
+ assert not shutdown.done()
+ release.set()
+ await asyncio.gather(request, shutdown)
+ assert await pool.session_bindings() == {}
+
+ asyncio.run(scenario())
+ handles[0].shutdown.assert_called_once()
+
+
+def test_pool_shutdown_wakes_session_waiting_for_capacity() -> None:
+ pool, _ = _pool(num_replicas=1)
+
+ async def scenario() -> None:
+ await pool.open_session("one")
+ waiting = asyncio.create_task(pool.open_session("two"))
+ await asyncio.sleep(0)
+ await pool.aclose()
+ with pytest.raises(RuntimeError, match="shutting down"):
+ await waiting
+
+ asyncio.run(scenario())
diff --git a/tests/unit/service/test_vla_replica.py b/tests/unit/service/test_vla_replica.py
new file mode 100644
index 00000000..a38ce82c
--- /dev/null
+++ b/tests/unit/service/test_vla_replica.py
@@ -0,0 +1,167 @@
+"""CPU-only tests for worker-local VLA replica dispatch."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from telefuser.service.vla_replica import VLAReplicaProvider
+from telefuser.vla import (
+ ActionSpaceSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+ VLACapabilities,
+ VLARegistry,
+ VLASessionManager,
+ action_space_to_wire,
+ observation_space_to_wire,
+ robot_action_chunk_from_wire,
+ robot_observation_to_wire,
+)
+
+MODEL_SPACE = ActionSpaceSpec("joint_delta", ("joint",), ("radian",), None, 10.0, False)
+ROBOT_SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), None, 10.0, False)
+OBSERVATION_SPACE = ObservationSpaceSpec(("joint",))
+
+
+class _Policy:
+ def __init__(self) -> None:
+ self.resets: list[str] = []
+
+ def capabilities(self) -> VLACapabilities:
+ return VLACapabilities("fake", MODEL_SPACE, max_horizon=2)
+
+ def predict(self, request) -> ModelActionChunk:
+ return ModelActionChunk(
+ torch.tensor([[0.1], [0.2]]),
+ MODEL_SPACE,
+ 2,
+ request.observation_timestamp_ns,
+ request.sequence_id,
+ request.episode_id,
+ )
+
+ def reset(self, episode_id: str) -> None:
+ self.resets.append(episode_id)
+
+
+@dataclass
+class _Embodiment:
+ embodiment_id: str = "fake-robot"
+ model_action_space: ActionSpaceSpec = MODEL_SPACE
+ robot_action_space: ActionSpaceSpec = ROBOT_SPACE
+ observation_space: ObservationSpaceSpec = OBSERVATION_SPACE
+
+ def encode_observation(self, observation: RobotObservation) -> ModelObservation:
+ return ModelObservation(observation.state.values, observation.images)
+
+ def decode_actions(self, actions: ModelActionChunk, robot_state: RobotState) -> RobotActionChunk:
+ return RobotActionChunk(
+ actions.actions + robot_state.values,
+ ROBOT_SPACE,
+ actions.valid_length,
+ actions.observation_timestamp_ns,
+ actions.sequence_id,
+ actions.episode_id,
+ )
+
+
+def _provider() -> tuple[VLAReplicaProvider, _Policy]:
+ policy = _Policy()
+ registry = VLARegistry()
+ registry.register_policy("fake", policy)
+ registry.register_embodiment("fake-robot", _Embodiment())
+ return VLAReplicaProvider(VLASessionManager(registry)), policy
+
+
+def test_provider_dispatches_complete_worker_local_session() -> None:
+ provider, policy = _provider()
+ assert provider.metadata() == {"model_ids": ["fake"], "embodiment_ids": ["fake-robot"]}
+ opened = provider.dispatch(
+ "OPEN",
+ {
+ "session_id": "session",
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ "episode_id": "episode",
+ "expected_robot_action_space": action_space_to_wire(ROBOT_SPACE),
+ "expected_robot_observation_space": observation_space_to_wire(OBSERVATION_SPACE),
+ },
+ )
+ assert opened["max_horizon"] == 2
+ assert opened["robot_action_space"] == action_space_to_wire(ROBOT_SPACE)
+ assert opened["robot_observation_space"] == observation_space_to_wire(OBSERVATION_SPACE)
+
+ observation = RobotObservation(RobotState(torch.tensor([1.0]), ("joint",), 100), {})
+ response = provider.dispatch(
+ "PREDICT",
+ {
+ "session_id": "session",
+ "sequence_id": 1,
+ "observation_timestamp_ns": 100,
+ "instruction": "move",
+ "seed": 7,
+ "observation": robot_observation_to_wire(observation),
+ },
+ )
+ chunk = robot_action_chunk_from_wire(response["chunk"])
+ assert chunk.actions.flatten().tolist() == pytest.approx([1.1, 1.2])
+
+ provider.dispatch("RESET", {"session_id": "session", "episode_id": "episode-2"})
+ provider.dispatch("CLOSE", {"session_id": "session"})
+ assert policy.resets == ["episode", "episode-2"]
+
+
+def test_provider_rejects_action_contract_mismatch_before_open() -> None:
+ provider, _ = _provider()
+ mismatch = ActionSpaceSpec("velocity", ("joint",), ("radian_per_second",), None, 10.0, False)
+ with pytest.raises(ValueError, match="action space"):
+ provider.dispatch(
+ "OPEN",
+ {
+ "session_id": "session",
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ "episode_id": "episode",
+ "expected_robot_action_space": action_space_to_wire(mismatch),
+ },
+ )
+
+
+def test_provider_rejects_observation_contract_mismatch_before_open() -> None:
+ provider, _ = _provider()
+ mismatch = ObservationSpaceSpec(("other_joint",))
+ with pytest.raises(ValueError, match="observation space"):
+ provider.dispatch(
+ "OPEN",
+ {
+ "session_id": "session",
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ "episode_id": "episode",
+ "expected_robot_observation_space": observation_space_to_wire(mismatch),
+ },
+ )
+
+
+def test_provider_close_releases_all_worker_sessions() -> None:
+ provider, policy = _provider()
+ for session_id in ("one", "two"):
+ provider.dispatch(
+ "OPEN",
+ {
+ "session_id": session_id,
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ "episode_id": session_id,
+ },
+ )
+ provider.close()
+ assert provider.sessions.session_ids() == ()
+ assert policy.resets == ["one", "two"]
diff --git a/tests/unit/service/test_vla_session.py b/tests/unit/service/test_vla_session.py
new file mode 100644
index 00000000..8758c538
--- /dev/null
+++ b/tests/unit/service/test_vla_session.py
@@ -0,0 +1,389 @@
+"""Loopback tests for the generic versioned VLA WebSocket protocol."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from threading import Event
+
+import pytest
+import torch
+from fastapi.testclient import TestClient
+
+from telefuser.service.core.replica_worker import ReplicaDeadError
+from telefuser.service.vla_session import (
+ VLA_SESSION_PROTOCOL_VERSION,
+ create_pipeline_pool_vla_session_app,
+ create_vla_session_app,
+)
+from telefuser.vla import (
+ ActionSpaceSpec,
+ ImageObservationSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+ VLACapabilities,
+ VLARegistry,
+ VLASessionManager,
+ action_space_to_wire,
+ observation_space_to_wire,
+ robot_action_chunk_from_wire,
+ robot_action_chunk_to_wire,
+ robot_observation_to_wire,
+)
+
+MODEL_SPACE = ActionSpaceSpec("joint_delta", ("joint",), ("radian",), None, 10.0, False)
+ROBOT_SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), None, 10.0, False)
+OBSERVATION_SPACE = ObservationSpaceSpec(("joint",), (ImageObservationSpec("front"),))
+
+
+class _Policy:
+ def __init__(
+ self,
+ *,
+ delay_s: float = 0,
+ failure: Exception | None = None,
+ stateful: bool = False,
+ block_first: bool = False,
+ ) -> None:
+ self.delay_s = delay_s
+ self.failure = failure
+ self.stateful = stateful
+ self.resets: list[str] = []
+ self.sequences: list[int] = []
+ self.started = Event()
+ self.release = Event()
+ if not block_first:
+ self.release.set()
+
+ def capabilities(self) -> VLACapabilities:
+ return VLACapabilities("fake", MODEL_SPACE, max_horizon=3, stateful=self.stateful)
+
+ def predict(self, request) -> ModelActionChunk:
+ self.sequences.append(request.sequence_id)
+ if len(self.sequences) == 1:
+ self.started.set()
+ assert self.release.wait(timeout=2)
+ if self.delay_s:
+ time.sleep(self.delay_s)
+ if self.failure is not None:
+ raise self.failure
+ return ModelActionChunk(
+ torch.tensor([[0.1], [0.2], [0.3]]),
+ MODEL_SPACE,
+ 3,
+ request.observation_timestamp_ns,
+ request.sequence_id,
+ request.episode_id,
+ )
+
+ def reset(self, episode_id: str) -> None:
+ self.resets.append(episode_id)
+
+
+@dataclass
+class _Embodiment:
+ embodiment_id: str = "fake-robot"
+ model_action_space: ActionSpaceSpec = MODEL_SPACE
+ robot_action_space: ActionSpaceSpec = ROBOT_SPACE
+ observation_space: ObservationSpaceSpec = OBSERVATION_SPACE
+
+ def encode_observation(self, observation: RobotObservation) -> ModelObservation:
+ return ModelObservation(observation.state.values, observation.images)
+
+ def decode_actions(self, actions: ModelActionChunk, robot_state: RobotState) -> RobotActionChunk:
+ return RobotActionChunk(
+ actions.actions + robot_state.values,
+ self.robot_action_space,
+ actions.valid_length,
+ actions.observation_timestamp_ns,
+ actions.sequence_id,
+ actions.episode_id,
+ )
+
+
+def _manager(policy: _Policy | None = None) -> tuple[VLASessionManager, _Policy]:
+ policy = policy or _Policy()
+ registry = VLARegistry()
+ registry.register_policy("fake", policy)
+ registry.register_embodiment("fake-robot", _Embodiment())
+ return VLASessionManager(registry), policy
+
+
+def _open(**overrides) -> dict:
+ return {
+ "type": "OPEN",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": "open",
+ "session_id": "session",
+ "episode_id": "episode",
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ **overrides,
+ }
+
+
+def _predict(sequence_id: int, **overrides) -> dict:
+ timestamp_ns = overrides.pop("observation_timestamp_ns", 100 + sequence_id)
+ observation = RobotObservation(
+ RobotState(torch.tensor([1.0]), ("joint",), timestamp_ns),
+ {"front": torch.zeros((2, 2, 3), dtype=torch.uint8)},
+ )
+ return {
+ "type": "PREDICT",
+ "request_id": f"predict-{sequence_id}",
+ "session_id": "session",
+ "sequence_id": sequence_id,
+ "observation_timestamp_ns": timestamp_ns,
+ "instruction": "move",
+ "seed": 7,
+ "observation": robot_observation_to_wire(observation),
+ **overrides,
+ }
+
+
+def test_protocol_negotiates_capabilities_and_runs_full_session_lifecycle() -> None:
+ manager, policy = _manager()
+ with TestClient(create_vla_session_app(manager)) as client:
+ assert client.get("/healthz").json() == {"status": "ok"}
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ hello = websocket.receive_json()
+ assert hello["type"] == "HELLO"
+ assert hello["protocol_version"] == "1.0"
+ assert hello["operations"] == ["OPEN", "PREDICT", "RESET", "CLOSE"]
+ assert hello["capabilities"]["prediction_queue"] == "latest-wins"
+ assert hello["capabilities"]["control_barriers"] is True
+ assert hello["model_ids"] == ["fake"]
+ assert hello["embodiment_ids"] == ["fake-robot"]
+
+ websocket.send_json(_open(expected_robot_action_space=action_space_to_wire(ROBOT_SPACE)))
+ opened = websocket.receive_json()
+ assert opened["type"] == "OPENED"
+ assert opened["supports_seed"] is True
+ assert opened["robot_action_space"] == action_space_to_wire(ROBOT_SPACE)
+ assert opened["robot_observation_space"] == observation_space_to_wire(OBSERVATION_SPACE)
+
+ websocket.send_json(_predict(1))
+ response = websocket.receive_json()
+ assert response["type"] == "ACTION_CHUNK"
+ chunk = robot_action_chunk_from_wire(response["chunk"])
+ assert chunk.actions.flatten().tolist() == pytest.approx([1.1, 1.2, 1.3])
+ assert chunk.sequence_id == 1
+
+ websocket.send_json({"type": "RESET", "session_id": "session", "episode_id": "episode-2"})
+ assert websocket.receive_json()["type"] == "RESET_ACK"
+ websocket.send_json(_predict(0, observation_timestamp_ns=200))
+ assert websocket.receive_json()["type"] == "ACTION_CHUNK"
+
+ websocket.send_json({"type": "CLOSE", "session_id": "session"})
+ assert websocket.receive_json()["type"] == "CLOSE_ACK"
+
+ assert manager.session_ids() == ()
+ assert policy.resets == ["episode", "episode-2"]
+
+
+def test_protocol_returns_stable_errors_for_version_components_contract_order_and_age() -> None:
+ manager, _ = _manager()
+ with TestClient(create_vla_session_app(manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open(protocol_version="2.0"))
+ assert websocket.receive_json()["error"]["code"] == "unsupported_version"
+
+ websocket.send_json(_open(model_id="missing"))
+ assert websocket.receive_json()["error"]["code"] == "unknown_component"
+
+ mismatch = ActionSpaceSpec("velocity", ("joint",), ("radian_per_second",), None, 10.0, False)
+ websocket.send_json(_open(expected_robot_action_space=action_space_to_wire(mismatch)))
+ assert websocket.receive_json()["error"]["code"] == "action_space_mismatch"
+
+ observation_mismatch = ObservationSpaceSpec(("joint",), (ImageObservationSpec("wrist"),))
+ websocket.send_json(_open(expected_robot_observation_space=observation_space_to_wire(observation_mismatch)))
+ assert websocket.receive_json()["error"]["code"] == "observation_space_mismatch"
+
+ websocket.send_json(_open(max_observation_age_ns=10))
+ assert websocket.receive_json()["type"] == "OPENED"
+ websocket.send_json(_predict(2, observation_clock_now_ns=102))
+ assert websocket.receive_json()["type"] == "ACTION_CHUNK"
+ websocket.send_json(_predict(2))
+ assert websocket.receive_json()["error"]["code"] == "out_of_order"
+ websocket.send_json(_predict(3, observation_clock_now_ns=1_000))
+ assert websocket.receive_json()["error"]["code"] == "expired"
+
+
+def test_protocol_reports_timeout_recovers_session_and_maps_replica_failure() -> None:
+ manager, _ = _manager(_Policy(delay_s=0.02))
+ with TestClient(create_vla_session_app(manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open())
+ websocket.receive_json()
+ websocket.send_json(_predict(1, request_ttl_ms=1))
+ assert websocket.receive_json()["error"]["code"] == "timeout"
+ websocket.send_json({"type": "RESET", "session_id": "session"})
+ assert websocket.receive_json()["type"] == "RESET_ACK"
+
+ failed_manager, _ = _manager(_Policy(failure=ReplicaDeadError("replica exited")))
+ with TestClient(create_vla_session_app(failed_manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open())
+ websocket.receive_json()
+ websocket.send_json(_predict(1))
+ assert websocket.receive_json()["error"]["code"] == "replica_unavailable"
+
+
+def test_disconnect_closes_connection_owned_sessions() -> None:
+ manager, policy = _manager()
+ with TestClient(create_vla_session_app(manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open())
+ assert websocket.receive_json()["type"] == "OPENED"
+ assert manager.session_ids() == ()
+ assert policy.resets == ["episode"]
+
+
+class _FakePipelinePool:
+ def __init__(self) -> None:
+ self.session_ids: set[str] = set()
+ self.operations: list[tuple[str, str]] = []
+
+ def vla_metadata(self) -> dict:
+ return {"model_ids": ["fake"], "embodiment_ids": ["fake-robot"]}
+
+ async def open_session(self, session_id: str) -> int:
+ if session_id in self.session_ids:
+ raise ValueError(f"pipeline session is already open: {session_id!r}")
+ self.session_ids.add(session_id)
+ return 0
+
+ async def run_vla_operation(self, session_id: str, operation: str, payload: dict) -> dict:
+ assert session_id in self.session_ids
+ self.operations.append((session_id, operation))
+ if operation == "OPEN":
+ return {
+ "session_id": session_id,
+ "model_id": "fake",
+ "embodiment_id": "fake-robot",
+ "model_action_space": action_space_to_wire(MODEL_SPACE),
+ "robot_action_space": action_space_to_wire(ROBOT_SPACE),
+ "robot_observation_space": observation_space_to_wire(OBSERVATION_SPACE),
+ "max_horizon": 3,
+ "stateful": False,
+ "supports_seed": True,
+ }
+ if operation == "PREDICT":
+ chunk = RobotActionChunk(
+ torch.tensor([[1.1], [1.2]]),
+ ROBOT_SPACE,
+ 2,
+ payload["observation_timestamp_ns"],
+ payload["sequence_id"],
+ "episode",
+ )
+ return {"session_id": session_id, "chunk": robot_action_chunk_to_wire(chunk)}
+ return {"session_id": session_id}
+
+ async def close_session(self, session_id: str) -> int:
+ self.session_ids.remove(session_id)
+ return 0
+
+
+def test_pipeline_pool_backend_runs_full_protocol_and_releases_lease() -> None:
+ pool = _FakePipelinePool()
+ app = create_pipeline_pool_vla_session_app(pool, close_pool_on_shutdown=False)
+ with TestClient(app) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ hello = websocket.receive_json()
+ assert hello["model_ids"] == ["fake"]
+ websocket.send_json(_open())
+ assert websocket.receive_json()["type"] == "OPENED"
+ websocket.send_json(_predict(1))
+ response = websocket.receive_json()
+ assert response["type"] == "ACTION_CHUNK"
+ assert robot_action_chunk_from_wire(response["chunk"]).valid_length == 2
+ websocket.send_json({"type": "RESET", "session_id": "session"})
+ assert websocket.receive_json()["type"] == "RESET_ACK"
+ websocket.send_json({"type": "CLOSE", "session_id": "session"})
+ assert websocket.receive_json()["type"] == "CLOSE_ACK"
+
+ assert pool.session_ids == set()
+ assert pool.operations == [
+ ("session", "OPEN"),
+ ("session", "PREDICT"),
+ ("session", "RESET"),
+ ("session", "CLOSE"),
+ ]
+
+
+def test_protocol_keeps_one_inference_and_only_the_latest_waiting_prediction() -> None:
+ policy = _Policy(block_first=True, stateful=True)
+ manager, _ = _manager(policy)
+ with TestClient(create_vla_session_app(manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open())
+ websocket.receive_json()
+ websocket.send_json(_predict(1))
+ assert policy.started.wait(timeout=1)
+ websocket.send_json(_predict(2))
+ websocket.send_json(_predict(3))
+ policy.release.set()
+
+ responses = {response["request_id"]: response for response in (websocket.receive_json() for _ in range(3))}
+ assert responses["predict-1"]["error"]["code"] == "superseded"
+ assert responses["predict-2"]["error"]["code"] == "superseded"
+ assert responses["predict-3"]["type"] == "ACTION_CHUNK"
+ assert policy.resets == ["episode"]
+ assert policy.sequences == [1, 3]
+
+
+def test_reset_is_a_barrier_for_inflight_results_and_restarts_ordering() -> None:
+ policy = _Policy(block_first=True, stateful=True)
+ manager, _ = _manager(policy)
+ with TestClient(create_vla_session_app(manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open())
+ websocket.receive_json()
+ websocket.send_json(_predict(1))
+ assert policy.started.wait(timeout=1)
+ websocket.send_json(
+ {"type": "RESET", "request_id": "reset", "session_id": "session", "episode_id": "episode-2"}
+ )
+ policy.release.set()
+
+ barrier_responses = {
+ response["request_id"]: response for response in (websocket.receive_json() for _ in range(2))
+ }
+ assert barrier_responses["predict-1"]["error"]["code"] == "superseded"
+ assert barrier_responses["reset"]["type"] == "RESET_ACK"
+ websocket.send_json(_predict(0, observation_timestamp_ns=200))
+ assert websocket.receive_json()["type"] == "ACTION_CHUNK"
+
+ assert policy.resets == ["episode", "episode-2"]
+
+
+def test_close_is_a_barrier_for_inflight_results() -> None:
+ policy = _Policy(block_first=True)
+ manager, _ = _manager(policy)
+ with TestClient(create_vla_session_app(manager)) as client:
+ with client.websocket_connect("/v1/vla/session") as websocket:
+ websocket.receive_json()
+ websocket.send_json(_open())
+ websocket.receive_json()
+ websocket.send_json(_predict(1))
+ assert policy.started.wait(timeout=1)
+ websocket.send_json({"type": "CLOSE", "request_id": "close", "session_id": "session"})
+ policy.release.set()
+
+ barrier_responses = {
+ response["request_id"]: response for response in (websocket.receive_json() for _ in range(2))
+ }
+ assert barrier_responses["predict-1"]["error"]["code"] == "superseded"
+ assert barrier_responses["close"]["type"] == "CLOSE_ACK"
+ assert manager.session_ids() == ()
diff --git a/tests/unit/vla/test_chunk_state.py b/tests/unit/vla/test_chunk_state.py
new file mode 100644
index 00000000..d7452f81
--- /dev/null
+++ b/tests/unit/vla/test_chunk_state.py
@@ -0,0 +1,181 @@
+"""Deterministic tests for the generic action chunk state machine."""
+
+from __future__ import annotations
+
+import torch
+
+from telefuser.vla import ActionSpaceSpec, RobotActionChunk
+from telefuser.vla.runtime import (
+ ActionChunkStateMachine,
+ ChunkStatus,
+ DisconnectPolicy,
+ RemainderPolicy,
+ RuntimeState,
+)
+
+SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), None, 10.0, False)
+
+
+class _Clock:
+ def __init__(self) -> None:
+ self.now = 1.0
+
+ def __call__(self) -> float:
+ return self.now
+
+
+def _chunk(sequence_id: int, timestamp_ns: int, values: tuple[float, ...] = (1, 2, 3, 4)) -> RobotActionChunk:
+ return RobotActionChunk(
+ torch.tensor(values).reshape(-1, 1),
+ SPACE,
+ len(values),
+ timestamp_ns,
+ sequence_id,
+ "episode",
+ )
+
+
+def test_latest_sequence_supersedes_pending_and_ready_chunks() -> None:
+ runtime = ActionChunkStateMachine("episode", execute_horizon=2)
+ first = runtime.submit(1, 100)
+ second = runtime.submit(2, 200)
+
+ assert runtime.status(first) is ChunkStatus.SUPERSEDED
+ assert runtime.complete(first, _chunk(1, 100)) is ChunkStatus.SUPERSEDED
+ assert runtime.complete(second, _chunk(2, 200)) is ChunkStatus.READY
+
+ third = runtime.submit(3, 300)
+ assert runtime.status(second) is ChunkStatus.SUPERSEDED
+ assert runtime.state is RuntimeState.PENDING
+ assert runtime.status(third) is ChunkStatus.PENDING
+
+ duplicate = runtime.submit(3, 300)
+ assert runtime.status(duplicate) is ChunkStatus.REJECTED
+
+
+def test_observation_age_and_network_ttl_use_independent_clocks() -> None:
+ clock = _Clock()
+ runtime = ActionChunkStateMachine(
+ "episode",
+ execute_horizon=2,
+ max_observation_age_ns=20,
+ clock=clock,
+ )
+
+ stale = runtime.submit(1, 100, request_ttl_ms=10_000, observation_clock_now_ns=121)
+ assert runtime.status(stale) is ChunkStatus.EXPIRED
+ assert runtime.reason(stale) == "observation timestamp is stale"
+
+ network_expired = runtime.submit(2, 200, request_ttl_ms=10, observation_clock_now_ns=200)
+ clock.now += 0.011
+ assert (
+ runtime.complete(
+ network_expired,
+ _chunk(2, 200),
+ observation_clock_now_ns=200,
+ )
+ is ChunkStatus.EXPIRED
+ )
+ assert runtime.reason(network_expired) == "request_ttl_ms elapsed before completion"
+
+ inference_stale = runtime.submit(3, 300, request_ttl_ms=10_000, observation_clock_now_ns=300)
+ assert (
+ runtime.complete(
+ inference_stale,
+ _chunk(3, 300),
+ observation_clock_now_ns=321,
+ )
+ is ChunkStatus.EXPIRED
+ )
+ assert runtime.reason(inference_stale) == "observation became stale before completion"
+
+
+def test_execute_horizon_can_retain_or_discard_remainder() -> None:
+ retained = ActionChunkStateMachine(
+ "episode",
+ execute_horizon=2,
+ remainder_policy=RemainderPolicy.RETAIN,
+ )
+ ticket = retained.submit(1, 100)
+ assert retained.complete(ticket, _chunk(1, 100)) is ChunkStatus.READY
+
+ first = retained.begin_execution()
+ assert first is not None
+ assert first.actions.flatten().tolist() == [1, 2]
+ assert retained.finish_execution() is RuntimeState.READY
+ second = retained.begin_execution()
+ assert second is not None
+ assert second.actions.flatten().tolist() == [3, 4]
+ assert retained.finish_execution() is RuntimeState.EMPTY
+ assert retained.status(ticket) is ChunkStatus.EXECUTED
+
+ discarded = ActionChunkStateMachine("episode", execute_horizon=2)
+ discarded_ticket = discarded.submit(1, 100)
+ discarded.complete(discarded_ticket, _chunk(1, 100))
+ assert discarded.begin_execution() is not None
+ assert discarded.finish_execution() is RuntimeState.EMPTY
+ assert discarded.status(discarded_ticket) is ChunkStatus.EXECUTED
+
+
+def test_reset_disconnect_and_stateful_discard_recovery_are_explicit() -> None:
+ recovered: list[str] = []
+ runtime = ActionChunkStateMachine(
+ "episode",
+ execute_horizon=2,
+ stateful_policy=True,
+ recover_stateful_policy=recovered.append,
+ )
+ discarded = runtime.submit(1, 100)
+ runtime.mark_inference_started(discarded)
+ newest = runtime.submit(2, 200)
+ assert runtime.complete(discarded, _chunk(1, 100)) is ChunkStatus.SUPERSEDED
+ assert recovered == ["episode"]
+
+ runtime.mark_inference_started(newest)
+ runtime.reset("episode-2")
+ assert recovered == ["episode", "episode"]
+ assert runtime.state is RuntimeState.EMPTY
+ assert runtime.latest_sequence_id is None
+ assert runtime.submit(0, 0).sequence_id == 0
+
+ assert runtime.disconnect() is RuntimeState.HOLDING
+ rejected = runtime.submit(1, 1)
+ assert runtime.status(rejected) is ChunkStatus.REJECTED
+
+ stopped = ActionChunkStateMachine(
+ "episode",
+ execute_horizon=1,
+ disconnect_policy=DisconnectPolicy.STOP,
+ )
+ assert stopped.disconnect() is RuntimeState.STOPPED
+
+
+def test_terminal_history_preserves_inflight_ticket_until_discard_recovery() -> None:
+ recovered: list[str] = []
+ runtime = ActionChunkStateMachine(
+ "episode",
+ execute_horizon=1,
+ stateful_policy=True,
+ recover_stateful_policy=recovered.append,
+ terminal_history=1,
+ )
+ inflight = runtime.submit(1, 100)
+ runtime.mark_inference_started(inflight)
+ runtime.submit(2, 200)
+ for _ in range(3):
+ rejected = runtime.submit(2, 200)
+ assert runtime.status(rejected) is ChunkStatus.REJECTED
+
+ assert runtime.status(inflight) is ChunkStatus.SUPERSEDED
+ assert runtime.complete(inflight, _chunk(1, 100)) is ChunkStatus.SUPERSEDED
+ assert recovered == ["episode"]
+
+
+def test_inference_failure_rejects_pending_ticket() -> None:
+ runtime = ActionChunkStateMachine("episode", execute_horizon=1)
+ ticket = runtime.submit(1, 100)
+ runtime.mark_inference_started(ticket)
+
+ assert runtime.reject(ticket, "policy failed") is ChunkStatus.REJECTED
+ assert runtime.reason(ticket) == "policy failed"
+ assert runtime.state is RuntimeState.EMPTY
diff --git a/tests/unit/vla/test_client.py b/tests/unit/vla/test_client.py
new file mode 100644
index 00000000..a35cbfa1
--- /dev/null
+++ b/tests/unit/vla/test_client.py
@@ -0,0 +1,132 @@
+"""Tests for the generic asynchronous VLA WebSocket client."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+import torch
+
+from telefuser.client.vla import AsyncVLAClient, VLAClientError
+from telefuser.vla import ActionSpaceSpec, RobotActionChunk, RobotObservation, RobotState
+from telefuser.vla.serialization import (
+ VLA_SESSION_PROTOCOL_VERSION,
+ VLA_WIRE_ENCODING,
+ dumps_wire_message,
+ loads_wire_message,
+ robot_action_chunk_to_wire,
+)
+
+SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), None, 10.0, False)
+
+
+class _FakeWebSocket:
+ def __init__(self, *, fail_predict: bool = False) -> None:
+ self.incoming: asyncio.Queue[str] = asyncio.Queue()
+ self.incoming.put_nowait(
+ dumps_wire_message(
+ {
+ "type": "HELLO",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "encoding": VLA_WIRE_ENCODING,
+ }
+ )
+ )
+ self.fail_predict = fail_predict
+ self.delayed_prediction: str | None = None
+ self.closed = False
+
+ async def recv(self) -> str:
+ return await self.incoming.get()
+
+ async def send(self, encoded: str) -> None:
+ request = loads_wire_message(encoded, max_message_bytes=1_000_000)
+ operation = request["type"]
+ response = {
+ "type": f"{operation}_ACK",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": request["request_id"],
+ "session_id": request["session_id"],
+ }
+ if operation == "OPEN":
+ response["type"] = "OPENED"
+ elif operation == "PREDICT":
+ if self.fail_predict:
+ response = {
+ "type": "ERROR",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": request["request_id"],
+ "error": {"code": "expired", "message": "observation expired"},
+ }
+ else:
+ response = {
+ "type": "ACTION_CHUNK",
+ "protocol_version": VLA_SESSION_PROTOCOL_VERSION,
+ "request_id": request["request_id"],
+ "session_id": request["session_id"],
+ "chunk": robot_action_chunk_to_wire(
+ RobotActionChunk(
+ torch.tensor([[float(request["sequence_id"])]]),
+ SPACE,
+ 1,
+ request["observation_timestamp_ns"],
+ request["sequence_id"],
+ "episode",
+ )
+ ),
+ }
+ encoded_response = dumps_wire_message(response)
+ if operation == "PREDICT" and request["sequence_id"] == 1 and not self.fail_predict:
+ self.delayed_prediction = encoded_response
+ return
+ self.incoming.put_nowait(encoded_response)
+ if operation == "PREDICT" and self.delayed_prediction is not None:
+ self.incoming.put_nowait(self.delayed_prediction)
+ self.delayed_prediction = None
+
+ async def close(self) -> None:
+ self.closed = True
+
+
+def _observation(timestamp_ns: int) -> RobotObservation:
+ return RobotObservation(RobotState(torch.tensor([0.0]), ("joint",), timestamp_ns), {})
+
+
+@pytest.mark.asyncio
+async def test_client_negotiates_lifecycle_and_correlates_out_of_order_responses(monkeypatch) -> None:
+ websocket = _FakeWebSocket()
+
+ async def connect(*_args, **_kwargs):
+ return websocket
+
+ monkeypatch.setattr("websockets.asyncio.client.connect", connect)
+ client = AsyncVLAClient("ws://example.test/v1/vla/session")
+ hello = await client.connect()
+ assert hello["type"] == "HELLO"
+ await client.open_session("session", model_id="fake", embodiment_id="robot", episode_id="episode")
+
+ first, second = await asyncio.gather(
+ client.predict("session", _observation(101), "move", 1),
+ client.predict("session", _observation(102), "move", 2),
+ )
+ assert first.sequence_id == 1
+ assert second.sequence_id == 2
+ await client.reset_session("session")
+ await client.close_session("session")
+ await client.aclose()
+ assert websocket.closed is True
+
+
+@pytest.mark.asyncio
+async def test_client_exposes_stable_server_error(monkeypatch) -> None:
+ websocket = _FakeWebSocket(fail_predict=True)
+
+ async def connect(*_args, **_kwargs):
+ return websocket
+
+ monkeypatch.setattr("websockets.asyncio.client.connect", connect)
+ async with AsyncVLAClient("ws://example.test/v1/vla/session") as client:
+ await client.open_session("session", model_id="fake", embodiment_id="robot", episode_id="episode")
+ with pytest.raises(VLAClientError, match="observation expired") as error:
+ await client.predict("session", _observation(101), "move", 1)
+ assert error.value.code == "expired"
diff --git a/tests/unit/vla/test_contracts_runtime.py b/tests/unit/vla/test_contracts_runtime.py
new file mode 100644
index 00000000..75fa3869
--- /dev/null
+++ b/tests/unit/vla/test_contracts_runtime.py
@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from telefuser.vla import ActionSpaceSpec, RobotActionChunk, RobotState
+from telefuser.vla.runtime import BoundedActionSafety, ChunkExecutor
+
+
+def _space(*, representation: str = "joint_position") -> ActionSpaceSpec:
+ return ActionSpaceSpec(
+ representation=representation,
+ dimension_names=("joint_0", "joint_1"),
+ units=("radian", "radian"),
+ frame="robot_joint",
+ control_hz=20.0,
+ normalized=False,
+ )
+
+
+def _chunk(space: ActionSpaceSpec | None = None) -> RobotActionChunk:
+ return RobotActionChunk(
+ actions=torch.tensor([[0.1, 0.2], [0.2, 0.3], [0.3, 0.4]]),
+ action_space=space or _space(),
+ valid_length=3,
+ observation_timestamp_ns=100,
+ sequence_id=4,
+ episode_id="episode",
+ metadata={"source": "test"},
+ )
+
+
+def _state() -> RobotState:
+ return RobotState(torch.zeros(2), ("joint_0", "joint_1"), timestamp_ns=100)
+
+
+def test_action_space_rejects_semantic_mismatch_even_when_dimensions_match() -> None:
+ with pytest.raises(ValueError, match="representation"):
+ _space().require_compatible(_space(representation="joint_delta"))
+
+
+def test_chunk_executor_trims_and_checks_age_and_bounds() -> None:
+ executor = ChunkExecutor(
+ _space(),
+ execute_horizon=2,
+ max_observation_age_ns=50,
+ safety_policy=BoundedActionSafety(torch.full((2,), -1.0), torch.full((2,), 1.0)),
+ )
+
+ prepared = executor.prepare(_chunk(), _state(), now_ns=149)
+
+ assert prepared.actions.shape == (2, 2)
+ assert prepared.valid_length == 2
+ assert [action.step_index for action in executor.iter_actions(prepared)] == [0, 1]
+ with pytest.raises(ValueError, match="stale"):
+ executor.prepare(_chunk(), _state(), now_ns=151)
+
+
+def test_contracts_reject_invalid_tensor_shape_and_timestamp_type() -> None:
+ with pytest.raises(ValueError, match="width"):
+ RobotActionChunk(torch.zeros(2, 3), _space(), 2, 0, 0, "episode")
+ with pytest.raises(ValueError, match="timestamp"):
+ RobotState(torch.zeros(2), ("joint_0", "joint_1"), timestamp_ns="100")
diff --git a/tests/unit/vla/test_serialization.py b/tests/unit/vla/test_serialization.py
new file mode 100644
index 00000000..9068b853
--- /dev/null
+++ b/tests/unit/vla/test_serialization.py
@@ -0,0 +1,113 @@
+"""Tests for stable JSON/Base64 VLA wire contracts."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from telefuser.vla import (
+ ActionSpaceSpec,
+ ImageObservationSpec,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+ action_space_from_wire,
+ action_space_to_wire,
+ observation_space_from_wire,
+ observation_space_to_wire,
+ robot_action_chunk_from_wire,
+ robot_action_chunk_to_wire,
+ robot_observation_from_wire,
+ robot_observation_to_wire,
+)
+from telefuser.vla.serialization import dumps_wire_message, loads_wire_message, tensor_from_wire, tensor_to_wire
+
+SPACE = ActionSpaceSpec("joint_position", ("a", "b"), ("radian", "radian"), "base", 20.0, False)
+OBSERVATION_SPACE = ObservationSpaceSpec(("a", "b"), (ImageObservationSpec("front"),))
+
+
+def test_action_space_observation_and_chunk_round_trip_without_semantic_loss() -> None:
+ observation = RobotObservation(
+ RobotState(torch.tensor([0.1, 0.2]), ("a", "b"), 123),
+ {"front": torch.arange(12, dtype=torch.uint8).reshape(2, 2, 3)},
+ {"source": "fake"},
+ )
+ chunk = RobotActionChunk(torch.tensor([[1.0, 2.0], [3.0, 4.0]]), SPACE, 2, 123, 7, "episode")
+
+ assert action_space_from_wire(action_space_to_wire(SPACE)) == SPACE
+ assert observation_space_from_wire(observation_space_to_wire(OBSERVATION_SPACE)) == OBSERVATION_SPACE
+ restored_observation = robot_observation_from_wire(robot_observation_to_wire(observation))
+ assert restored_observation.state.dimension_names == observation.state.dimension_names
+ assert restored_observation.state.timestamp_ns == observation.state.timestamp_ns
+ assert torch.equal(restored_observation.state.values, observation.state.values)
+ assert torch.equal(restored_observation.images["front"], observation.images["front"])
+ assert dict(restored_observation.metadata) == {"source": "fake"}
+ restored_chunk = robot_action_chunk_from_wire(robot_action_chunk_to_wire(chunk))
+ assert restored_chunk.action_space == SPACE
+ assert torch.equal(restored_chunk.actions, chunk.actions)
+ assert restored_chunk.sequence_id == 7
+
+
+def test_wire_json_is_deterministic_and_size_bounded() -> None:
+ encoded = dumps_wire_message({"z": 1, "a": 2})
+ assert encoded == '{"a":2,"z":1}'
+ assert loads_wire_message(encoded, max_message_bytes=len(encoded)) == {"a": 2, "z": 1}
+ with pytest.raises(ValueError, match="exceeds"):
+ loads_wire_message(encoded, max_message_bytes=len(encoded) - 1)
+ with pytest.raises(ValueError, match="positive integer"):
+ loads_wire_message(encoded, max_message_bytes=0)
+
+
+def test_tensor_wire_rejects_schema_shape_base64_and_byte_limit_errors() -> None:
+ payload = tensor_to_wire(torch.ones((2, 2), dtype=torch.float32))
+
+ invalid_shape = dict(payload, shape=[2, -1])
+ with pytest.raises(ValueError, match="shape"):
+ tensor_from_wire(invalid_shape)
+
+ invalid_data = dict(payload, data="not base64!")
+ with pytest.raises(ValueError, match="Base64"):
+ tensor_from_wire(invalid_data)
+
+ with pytest.raises(ValueError, match="exceeds"):
+ tensor_from_wire(payload, max_bytes=15)
+
+ invalid_schema = dict(payload, schema_version=999)
+ with pytest.raises(ValueError, match="unsupported"):
+ tensor_from_wire(invalid_schema)
+
+
+def test_action_space_wire_requires_arrays_not_ambiguous_strings() -> None:
+ payload = action_space_to_wire(SPACE)
+ payload["dimension_names"] = "ab"
+ with pytest.raises(ValueError, match="must be arrays"):
+ action_space_from_wire(payload)
+
+
+def test_observation_space_validates_named_image_layout_dtype_and_channels() -> None:
+ state = RobotState(torch.tensor([0.1, 0.2]), ("a", "b"), 123)
+ OBSERVATION_SPACE.validate(RobotObservation(state, {"front": torch.zeros((2, 2, 3), dtype=torch.uint8)}))
+
+ with pytest.raises(ValueError, match="dtype"):
+ OBSERVATION_SPACE.validate(RobotObservation(state, {"front": torch.zeros((2, 2, 3))}))
+ with pytest.raises(ValueError, match="channels"):
+ OBSERVATION_SPACE.validate(RobotObservation(state, {"front": torch.zeros((2, 2, 1), dtype=torch.uint8)}))
+
+
+def test_observation_space_compatibility_uses_camera_names_not_declaration_order() -> None:
+ expected = ObservationSpaceSpec(
+ ("a", "b"),
+ (ImageObservationSpec("front"), ImageObservationSpec("wrist")),
+ )
+ reordered = ObservationSpaceSpec(
+ ("a", "b"),
+ (ImageObservationSpec("wrist"), ImageObservationSpec("front")),
+ )
+
+ expected.require_compatible(reordered)
+
+ malformed = observation_space_to_wire(expected)
+ malformed["images"] = [1]
+ with pytest.raises(ValueError, match="must contain objects"):
+ observation_space_from_wire(malformed)
diff --git a/tests/unit/vla/test_session.py b/tests/unit/vla/test_session.py
new file mode 100644
index 00000000..94c7c550
--- /dev/null
+++ b/tests/unit/vla/test_session.py
@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from telefuser.vla import (
+ ActionSpaceSpec,
+ ModelActionChunk,
+ ModelObservation,
+ ObservationSpaceSpec,
+ RobotActionChunk,
+ RobotObservation,
+ RobotState,
+ VLACapabilities,
+ VLARegistry,
+ VLASessionManager,
+)
+
+
+def _space(representation: str) -> ActionSpaceSpec:
+ return ActionSpaceSpec(representation, ("joint",), ("radian",), None, None, False)
+
+
+MODEL_SPACE = _space("joint_delta")
+ROBOT_SPACE = _space("joint_position")
+OBSERVATION_SPACE = ObservationSpaceSpec(("joint",))
+
+
+class _Policy:
+ resets: list[str]
+
+ def __init__(self) -> None:
+ self.resets = []
+
+ def capabilities(self) -> VLACapabilities:
+ return VLACapabilities("fake", MODEL_SPACE, max_horizon=3)
+
+ def predict(self, request):
+ return ModelActionChunk(
+ torch.tensor([[0.1], [0.2], [0.3]]),
+ MODEL_SPACE,
+ 3,
+ request.observation_timestamp_ns,
+ request.sequence_id,
+ request.episode_id,
+ )
+
+ def reset(self, episode_id: str) -> None:
+ self.resets.append(episode_id)
+
+
+@dataclass
+class _Embodiment:
+ embodiment_id: str = "fake-robot"
+ model_action_space: ActionSpaceSpec = MODEL_SPACE
+ robot_action_space: ActionSpaceSpec = ROBOT_SPACE
+ observation_space: ObservationSpaceSpec = OBSERVATION_SPACE
+
+ def encode_observation(self, observation: RobotObservation) -> ModelObservation:
+ return ModelObservation(observation.state.values, observation.images)
+
+ def decode_actions(self, actions: ModelActionChunk, robot_state: RobotState) -> RobotActionChunk:
+ return RobotActionChunk(
+ actions.actions + robot_state.values,
+ self.robot_action_space,
+ actions.valid_length,
+ actions.observation_timestamp_ns,
+ actions.sequence_id,
+ actions.episode_id,
+ )
+
+
+def _manager(policy: _Policy | None = None) -> tuple[VLASessionManager, _Policy]:
+ policy = policy or _Policy()
+ registry = VLARegistry()
+ registry.register_policy("fake", policy)
+ registry.register_embodiment("fake-robot", _Embodiment())
+ return VLASessionManager(registry), policy
+
+
+def _observation() -> RobotObservation:
+ return RobotObservation(RobotState(torch.tensor([1.0]), ("joint",), 100), {})
+
+
+def test_session_runs_full_semantic_dataflow_and_rejects_out_of_order_sequence() -> None:
+ manager, _ = _manager()
+ session = manager.open(
+ "session",
+ model_id="fake",
+ embodiment_id="fake-robot",
+ episode_id="episode",
+ execute_horizon=2,
+ )
+
+ timings: dict[str, float] = {}
+ chunk = session.predict(_observation(), "move", sequence_id=5, seed=7, timings=timings)
+
+ assert torch.allclose(chunk.actions, torch.tensor([[1.1], [1.2]]))
+ assert chunk.action_space == ROBOT_SPACE
+ assert set(timings) == {
+ "decode_actions_ms",
+ "encode_observation_ms",
+ "policy_ms",
+ "prepare_actions_ms",
+ }
+ assert all(value >= 0 for value in timings.values())
+ with pytest.raises(ValueError, match="must increase"):
+ session.predict(_observation(), "move", sequence_id=5)
+
+
+def test_session_reset_allows_new_sequence_and_close_releases_state() -> None:
+ manager, policy = _manager()
+ session = manager.open(
+ "session", model_id="fake", embodiment_id="fake-robot", episode_id="episode", execute_horizon=1
+ )
+ session.predict(_observation(), "move", sequence_id=1)
+
+ manager.reset("session", "episode-2")
+ chunk = session.predict(_observation(), "move", sequence_id=0)
+ manager.close("session")
+
+ assert chunk.episode_id == "episode-2"
+ assert policy.resets == ["episode", "episode-2"]
+ assert manager.session_ids() == ()
+
+
+def test_open_rejects_policy_and_embodiment_action_space_mismatch() -> None:
+ policy = _Policy()
+ registry = VLARegistry()
+ registry.register_policy("fake", policy)
+ registry.register_embodiment(
+ "fake-robot",
+ _Embodiment(model_action_space=_space("velocity")),
+ )
+ manager = VLASessionManager(registry)
+
+ with pytest.raises(ValueError, match="representation"):
+ manager.open("session", model_id="fake", embodiment_id="fake-robot", episode_id="episode")
+
+
+def test_session_rejects_observation_contract_before_policy_inference() -> None:
+ manager, _ = _manager()
+ session = manager.open("session", model_id="fake", embodiment_id="fake-robot", episode_id="episode")
+ observation = RobotObservation(RobotState(torch.tensor([1.0]), ("wrong_joint",), 100), {})
+
+ with pytest.raises(ValueError, match="observation space"):
+ session.predict(observation, "move", sequence_id=1)
diff --git a/tests/unit/vla/test_simulator_runtime.py b/tests/unit/vla/test_simulator_runtime.py
new file mode 100644
index 00000000..59a635f5
--- /dev/null
+++ b/tests/unit/vla/test_simulator_runtime.py
@@ -0,0 +1,110 @@
+"""Tests for simulator-side action execution state."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from telefuser.vla import ActionSpaceSpec, RobotActionChunk
+from telefuser.vla.runtime import ChunkStatus, RemainderPolicy, RuntimeState, SimulatorChunkRuntime
+
+SPACE = ActionSpaceSpec("joint_position", ("joint",), ("radian",), None, 10.0, False)
+
+
+class _Simulator:
+ def __init__(self, *, fail_at: int | None = None) -> None:
+ self.fail_at = fail_at
+ self.actions = []
+
+ def execute(self, action) -> None:
+ if self.fail_at == action.step_index:
+ raise RuntimeError("simulator failed")
+ self.actions.append(action)
+
+
+def _chunk(sequence_id: int = 1) -> RobotActionChunk:
+ return RobotActionChunk(
+ torch.tensor([[1.0], [2.0], [3.0]]),
+ SPACE,
+ 3,
+ 100,
+ sequence_id,
+ "episode",
+ )
+
+
+def test_client_runtime_owns_ready_execution_and_remainder_states() -> None:
+ simulator = _Simulator()
+ runtime = SimulatorChunkRuntime(
+ simulator,
+ SPACE,
+ "episode",
+ execute_horizon=2,
+ remainder_policy=RemainderPolicy.RETAIN,
+ )
+
+ assert runtime.accept(_chunk()) is ChunkStatus.READY
+ assert runtime.state is RuntimeState.READY
+ assert runtime.execute_ready() == 2
+ assert runtime.state is RuntimeState.READY
+ assert runtime.execute_ready() == 1
+ assert runtime.state is RuntimeState.EMPTY
+ assert [action.values.item() for action in simulator.actions] == [1.0, 2.0, 3.0]
+
+
+def test_client_runtime_rejects_mismatch_and_fails_closed_on_execution_error() -> None:
+ simulator = _Simulator(fail_at=1)
+ runtime = SimulatorChunkRuntime(simulator, SPACE, "episode", execute_horizon=3)
+ mismatch = ActionSpaceSpec("velocity", ("joint",), ("radian_per_second",), None, 10.0, False)
+ bad_chunk = RobotActionChunk(torch.tensor([[1.0]]), mismatch, 1, 100, 1, "episode")
+
+ with pytest.raises(ValueError, match="action space"):
+ runtime.accept(bad_chunk)
+
+ assert runtime.accept(_chunk()) is ChunkStatus.READY
+ with pytest.raises(RuntimeError, match="simulator failed"):
+ runtime.execute_ready()
+ assert runtime.state is RuntimeState.HOLDING
+
+
+def test_client_runtime_exposes_transport_neutral_execution_reports() -> None:
+ simulator = _Simulator()
+ runtime = SimulatorChunkRuntime(simulator, SPACE, "episode", execute_horizon=2)
+
+ no_action = runtime.execute_ready_with_report()
+ assert no_action.status == "no_action"
+ assert no_action.sequence_id is None
+ assert no_action.executed_steps == 0
+
+ assert runtime.accept(_chunk()) is ChunkStatus.READY
+ report = runtime.execute_ready_with_report()
+ assert report.status == "executed"
+ assert report.sequence_id == 1
+ assert report.executed_steps == 2
+
+
+@pytest.mark.asyncio
+async def test_client_runtime_can_execute_without_blocking_async_caller() -> None:
+ simulator = _Simulator()
+ runtime = SimulatorChunkRuntime(simulator, SPACE, "episode", execute_horizon=3)
+ assert runtime.accept(_chunk()) is ChunkStatus.READY
+
+ report = await runtime.execute_ready_async()
+
+ assert report.status == "executed"
+ assert report.executed_steps == 3
+ assert runtime.state is RuntimeState.EMPTY
+
+
+def test_client_runtime_reports_simulator_failure_without_raising() -> None:
+ simulator = _Simulator(fail_at=1)
+ runtime = SimulatorChunkRuntime(simulator, SPACE, "episode", execute_horizon=3)
+ assert runtime.accept(_chunk()) is ChunkStatus.READY
+
+ report = runtime.execute_ready_with_report()
+
+ assert report.status == "failed"
+ assert report.sequence_id == 1
+ assert report.executed_steps == 1
+ assert report.reason == "simulator failed"
+ assert runtime.state is RuntimeState.HOLDING