From 4e44652aba6fcfc56dff94ba0c79ca0aaee490f3 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Thu, 10 Sep 2026 10:55:28 +0000 Subject: [PATCH 1/5] feat(vla): add semantic action session framework Add model-independent VLA action contracts, policy and embodiment adapters, versioned session transport, deterministic chunk runtime, safety checks, and simulator boundaries. Integrate LingBot VLA v2 through the new abstractions while preserving its existing HTTP and RoboTwin WebSocket contracts, and add additive PipelinePool session leases for replica affinity.\n\nVerification:\n- 59 focused VLA, session lease, LingBot compatibility, and simulator adapter tests passed\n- 50 existing PipelinePool, structured-task, and service-route tests passed\n- ruff check passed for changed Python paths\n- git diff --check passed --- README.md | 17 + docs/en/vla.md | 146 ++++++ examples/lingbot_vla_v2/README.md | 5 + .../lingbot_vla_v2_robotwin_server.py | 123 +++-- telefuser/integrations/__init__.py | 1 + telefuser/integrations/sim/__init__.py | 6 + telefuser/integrations/sim/base.py | 25 + telefuser/integrations/sim/robotwin.py | 68 +++ .../pipelines/lingbot_vla_v2/__init__.py | 14 +- .../lingbot_vla_v2/action_scheduler.py | 272 +---------- .../pipelines/lingbot_vla_v2/robot_profile.py | 88 ++++ .../pipelines/lingbot_vla_v2/vla_policy.py | 66 +++ telefuser/service/core/pipeline_pool.py | 193 ++++++++ telefuser/service/vla_session.py | 384 ++++++++++++++++ telefuser/vla/__init__.py | 53 +++ telefuser/vla/contracts.py | 230 ++++++++++ telefuser/vla/embodiment.py | 32 ++ telefuser/vla/policy.py | 18 + telefuser/vla/registry.py | 66 +++ telefuser/vla/runtime/__init__.py | 28 ++ telefuser/vla/runtime/chunk_state.py | 432 ++++++++++++++++++ telefuser/vla/runtime/executor.py | 86 ++++ telefuser/vla/runtime/safety.py | 61 +++ telefuser/vla/runtime/scheduler.py | 270 +++++++++++ telefuser/vla/serialization.py | 259 +++++++++++ telefuser/vla/session.py | 222 +++++++++ tests/unit/integrations/sim/test_robotwin.py | 46 ++ .../lingbot_vla_v2/test_action_scheduler.py | 5 + .../lingbot_vla_v2/test_robot_profile.py | 35 +- .../lingbot_vla_v2/test_robotwin_server.py | 15 + .../lingbot_vla_v2/test_vla_policy.py | 49 ++ .../service/test_pipeline_pool_sessions.py | 169 +++++++ tests/unit/service/test_vla_session.py | 212 +++++++++ tests/unit/vla/test_chunk_state.py | 171 +++++++ tests/unit/vla/test_contracts_runtime.py | 63 +++ tests/unit/vla/test_serialization.py | 79 ++++ tests/unit/vla/test_session.py | 137 ++++++ 37 files changed, 3840 insertions(+), 306 deletions(-) create mode 100644 docs/en/vla.md create mode 100644 telefuser/integrations/__init__.py create mode 100644 telefuser/integrations/sim/__init__.py create mode 100644 telefuser/integrations/sim/base.py create mode 100644 telefuser/integrations/sim/robotwin.py create mode 100644 telefuser/pipelines/lingbot_vla_v2/vla_policy.py create mode 100644 telefuser/service/vla_session.py create mode 100644 telefuser/vla/__init__.py create mode 100644 telefuser/vla/contracts.py create mode 100644 telefuser/vla/embodiment.py create mode 100644 telefuser/vla/policy.py create mode 100644 telefuser/vla/registry.py create mode 100644 telefuser/vla/runtime/__init__.py create mode 100644 telefuser/vla/runtime/chunk_state.py create mode 100644 telefuser/vla/runtime/executor.py create mode 100644 telefuser/vla/runtime/safety.py create mode 100644 telefuser/vla/runtime/scheduler.py create mode 100644 telefuser/vla/serialization.py create mode 100644 telefuser/vla/session.py create mode 100644 tests/unit/integrations/sim/test_robotwin.py create mode 100644 tests/unit/pipelines/lingbot_vla_v2/test_vla_policy.py create mode 100644 tests/unit/service/test_pipeline_pool_sessions.py create mode 100644 tests/unit/service/test_vla_session.py create mode 100644 tests/unit/vla/test_chunk_state.py create mode 100644 tests/unit/vla/test_contracts_runtime.py create mode 100644 tests/unit/vla/test_serialization.py create mode 100644 tests/unit/vla/test_session.py 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..061e8458 --- /dev/null +++ b/docs/en/vla.md @@ -0,0 +1,146 @@ +# 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. + +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 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. + +## 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 existing LingBot RoboTwin WebSocket endpoint now calls this 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` without changing 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 space; semantic incompatibility is rejected before inference. + +The protocol returns machine-readable error codes for malformed messages, unsupported versions, unknown components +or sessions, action-space mismatches, 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. + +`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. +After a timed-out stateful 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. The generic in-process WebSocket factory accepts a `VLASessionManager`; a deployment backed by +subprocess replicas must use these lease methods when adding its replica RPC adapter. + +## 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 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. + +## Current Boundary + +The generic WebSocket application is intentionally an explicit factory rather than an automatically mounted +`telefuser serve` route. Subprocess replica transport still needs VLA-specific OPEN/PREDICT/RESET/CLOSE RPC commands +before the WebSocket factory can directly own `PipelinePool` leases. The lease primitive is implemented and tested, +but this RPC bridge is outside the current P0 scope. + +The chunk state machine is also transport-neutral. The existing LingBot compatibility server retains its proven +latest-wins scheduler until a later compatibility-preserving integration uses the generic state machine end to end. + +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..e45ecfd6 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -29,6 +29,7 @@ The parity reference uses [Robbyant/lingbot-vla-v2](https://github.com/Robbyant/ | 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) | ## Requirements @@ -328,6 +329,10 @@ 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. + 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: 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..66558509 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_robotwin_server.py @@ -18,44 +18,31 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect from telefuser.pipelines.lingbot_vla_v2 import ( + ROBOTWIN_ACTION_ORDER, ROBOTWIN_CAMERA_KEYS, - LingBotVlaV2Observation, + LingBotVlaV2VLAPolicy, RobotWinProfile, ) -from telefuser.pipelines.lingbot_vla_v2.action_scheduler import ActionChunkScheduler from telefuser.pipelines.lingbot_vla_v2.runtime import ( LINGBOT_VLA_V2_QUANTIZATION_CHOICES, get_lingbot_vla_v2_pipeline, ) from telefuser.utils.logging import logger +from telefuser.vla import RobotObservation, RobotState, VLARegistry, VLASessionManager +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 +117,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 +148,11 @@ 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] = {} + registry = VLARegistry() + registry.register_policy("lingbot-vla-v2", LingBotVlaV2VLAPolicy(pipeline)) + registry.register_embodiment(self.profile.embodiment_id, self.profile) + self._sessions = VLASessionManager(registry) @property def metadata(self) -> dict[str, Any]: @@ -173,28 +180,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 +233,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 +253,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 +351,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,6 +383,10 @@ 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 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..f9e75de2 --- /dev/null +++ b/telefuser/integrations/sim/__init__.py @@ -0,0 +1,6 @@ +"""Simulator-neutral VLA integration interfaces.""" + +from .base import SimulatorAdapter +from .robotwin import RoboTwinSimulatorAdapter + +__all__ = ["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/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..4314c98d 100644 --- a/telefuser/pipelines/lingbot_vla_v2/__init__.py +++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py @@ -3,7 +3,15 @@ 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, + LingBotVlaV2ActionChunk, + RobotWinProfile, +) +from .vla_policy import LingBotVlaV2VLAPolicy __all__ = [ "LingBotVlaV2ActionChunk", @@ -14,6 +22,10 @@ "LingBotVlaV2Pipeline", "LingBotVlaV2PipelineConfig", "LingBotVlaV2PolicyStage", + "LingBotVlaV2VLAPolicy", + "LINGBOT_VLA_V2_ACTION_SPACE", + "ROBOTWIN_ACTION_ORDER", + "ROBOTWIN_ACTION_SPACE", "ROBOTWIN_CAMERA_KEYS", "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..939ad4ee 100644 --- a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py +++ b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py @@ -10,6 +10,15 @@ import torch +from telefuser.vla.contracts import ( + ActionSpaceSpec, + ModelActionChunk, + ModelObservation, + RobotActionChunk, + RobotObservation, + RobotState, +) + ROBOTWIN_CAMERA_KEYS = ( "observation.images.cam_high", "observation.images.cam_left_wrist", @@ -19,6 +28,39 @@ 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, +) @dataclass(frozen=True) @@ -39,6 +81,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 +124,51 @@ 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 + + def encode_observation(self, observation: RobotObservation) -> ModelObservation: + """Validate a RoboTwin observation while retaining raw state for the model processor.""" + if not isinstance(observation, RobotObservation): + raise TypeError("observation must be a RobotObservation") + if observation.state.dimension_names != ROBOTWIN_ACTION_ORDER: + raise ValueError("RobotWin state dimension order does not match the robot profile") + missing = [key for key in self.camera_keys if key not in observation.images] + if missing: + raise ValueError(f"RobotWin observation is missing camera keys: {missing}") + 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/vla_policy.py b/telefuser/pipelines/lingbot_vla_v2/vla_policy.py new file mode 100644 index 00000000..44770d30 --- /dev/null +++ b/telefuser/pipelines/lingbot_vla_v2/vla_policy.py @@ -0,0 +1,66 @@ +"""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 .data import LingBotVlaV2Observation +from .robot_profile import LINGBOT_VLA_V2_ACTION_SPACE + + +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", + 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") diff --git a/telefuser/service/core/pipeline_pool.py b/telefuser/service/core/pipeline_pool.py index 74b982fb..7e40cdbc 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,6 +64,11 @@ 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] = {} @@ -154,6 +168,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 +190,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 +245,96 @@ 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_task_with_stop_event( self, task_data: dict, @@ -240,12 +348,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: diff --git a/telefuser/service/vla_session.py b/telefuser/service/vla_session.py new file mode 100644 index 00000000..081dbfa9 --- /dev/null +++ b/telefuser/service/vla_session.py @@ -0,0 +1,384 @@ +"""Generic versioned WebSocket transport for typed VLA sessions.""" + +from __future__ import annotations + +import asyncio +import contextlib +import math +from collections.abc import Mapping +from enum import Enum +from typing import Any + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect + +from telefuser.service.core.replica_worker import ReplicaDeadError +from telefuser.vla import VLASessionManager +from telefuser.vla.serialization import ( + DEFAULT_MAX_TENSOR_BYTES, + VLA_WIRE_ENCODING, + action_space_from_wire, + action_space_to_wire, + dumps_wire_message, + loads_wire_message, + robot_action_chunk_to_wire, + robot_observation_from_wire, +) + +VLA_SESSION_PROTOCOL_VERSION = "1.0" +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" + OUT_OF_ORDER = "out_of_order" + 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 + + +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.""" + 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() + recovery_tasks: dict[str, asyncio.Task[None]] = {} + await _send( + websocket, + { + "type": "HELLO", + "protocol_version": VLA_SESSION_PROTOCOL_VERSION, + "encoding": VLA_WIRE_ENCODING, + "operations": ["OPEN", "PREDICT", "RESET", "CLOSE"], + "model_ids": list(sessions.registry.model_ids()), + "embodiment_ids": list(sessions.registry.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) + response = await _handle_message( + sessions, + message, + owned_sessions=owned_sessions, + recovery_tasks=recovery_tasks, + max_tensor_bytes=max_tensor_bytes, + ) + except WebSocketDisconnect: + break + except Exception as error: + response = _error_response(error, request_id=request_id) + await _send(websocket, response) + finally: + for task in list(recovery_tasks.values()): + with contextlib.suppress(Exception): + await task + for session_id in tuple(owned_sessions): + with contextlib.suppress(KeyError): + await asyncio.to_thread(sessions.close, session_id) + + return app + + +async def _handle_message( + sessions: VLASessionManager, + message: Mapping[str, Any], + *, + owned_sessions: set[str], + recovery_tasks: dict[str, asyncio.Task[None]], + max_tensor_bytes: int, +) -> 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 _open_session(sessions, message, owned_sessions=owned_sessions, 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": + recovery = recovery_tasks.get(session_id) + if recovery is not None and not recovery.done(): + raise VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, "session is recovering from a timed-out request") + observation_payload = message.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=max_tensor_bytes) + timestamp_ns = _required_nonnegative_int(message, "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", + ) + timeout_s = _optional_timeout_s(message) + predict_task = asyncio.create_task( + asyncio.to_thread( + sessions.get(session_id).predict, + observation, + _required_string(message, "instruction"), + _required_nonnegative_int(message, "sequence_id"), + seed=_optional_integer(message, "seed"), + now_ns=_optional_nonnegative_int(message, "observation_clock_now_ns"), + ) + ) + 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 as error: + recovery_tasks[session_id] = asyncio.create_task( + _recover_timed_out_session(sessions, session_id, predict_task, recovery_tasks) + ) + raise VLAProtocolError(VLAErrorCode.TIMEOUT, "PREDICT exceeded request_ttl_ms") from error + return { + "type": "ACTION_CHUNK", + "protocol_version": VLA_SESSION_PROTOCOL_VERSION, + "request_id": request_id, + "session_id": session_id, + "chunk": robot_action_chunk_to_wire(chunk), + } + if operation == "RESET": + await _await_recovery(session_id, recovery_tasks) + await asyncio.to_thread(sessions.reset, session_id, _optional_string(message, "episode_id")) + return _success_response("RESET", session_id, request_id) + if operation == "CLOSE": + await _await_recovery(session_id, recovery_tasks) + await asyncio.to_thread(sessions.close, session_id) + owned_sessions.remove(session_id) + return _success_response("CLOSE", session_id, request_id) + raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"unsupported VLA operation: {operation!r}") + + +def _open_session( + sessions: VLASessionManager, + message: Mapping[str, Any], + *, + owned_sessions: set[str], + 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 sessions.session_ids(): + raise VLAProtocolError(VLAErrorCode.SESSION_EXISTS, f"VLA session is already open: {session_id!r}") + model_id = _required_string(message, "model_id") + embodiment_id = _required_string(message, "embodiment_id") + try: + policy = sessions.registry.get_policy(model_id) + embodiment = sessions.registry.get_embodiment(embodiment_id) + except KeyError as error: + raise VLAProtocolError(VLAErrorCode.UNKNOWN_COMPONENT, str(error)) from error + expected_payload = message.get("expected_robot_action_space") + if expected_payload is not None: + 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( + embodiment.robot_action_space, + context="client and embodiment robot action space", + ) + sessions.open( + session_id, + model_id=model_id, + embodiment_id=embodiment_id, + episode_id=_required_string(message, "episode_id"), + execute_horizon=_optional_positive_int(message, "execute_horizon"), + max_observation_age_ns=_optional_positive_int(message, "max_observation_age_ns"), + ) + owned_sessions.add(session_id) + capabilities = policy.capabilities() + return { + "type": "OPENED", + "protocol_version": VLA_SESSION_PROTOCOL_VERSION, + "request_id": request_id, + "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), + "max_horizon": capabilities.max_horizon, + "stateful": capabilities.stateful, + "supports_seed": capabilities.supports_seed, + } + + +async def _recover_timed_out_session( + sessions: VLASessionManager, + session_id: str, + predict_task: asyncio.Task[Any], + recovery_tasks: dict[str, asyncio.Task[None]], +) -> None: + try: + with contextlib.suppress(Exception): + await predict_task + with contextlib.suppress(KeyError): + await asyncio.to_thread(sessions.reset, session_id) + finally: + recovery_tasks.pop(session_id, None) + + +async def _await_recovery(session_id: str, recovery_tasks: dict[str, asyncio.Task[None]]) -> None: + recovery = recovery_tasks.get(session_id) + if recovery is not None: + await recovery + + +async def _send(websocket: WebSocket, payload: Mapping[str, Any]) -> None: + await websocket.send_text(dumps_wire_message(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, KeyError): + return VLAErrorCode.UNKNOWN_SESSION + if isinstance(error, ValueError): + message = str(error).lower() + if "action space" in message or "action-space" in message: + return VLAErrorCode.ACTION_SPACE_MISMATCH + if "sequence_id" in message or "must increase" in message: + return VLAErrorCode.OUT_OF_ORDER + if "stale" in message or "expired" in message: + return VLAErrorCode.EXPIRED + if "already open" in message: + return VLAErrorCode.SESSION_EXISTS + return VLAErrorCode.INVALID_MESSAGE + if isinstance(error, RuntimeError) and "replica" in str(error).lower(): + return VLAErrorCode.REPLICA_UNAVAILABLE + return VLAErrorCode.INTERNAL_ERROR + + +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..3b93d620 --- /dev/null +++ b/telefuser/vla/__init__.py @@ -0,0 +1,53 @@ +"""Public semantic contracts for vision-language-action integrations.""" + +from .contracts import ( + ActionChunk, + ActionSpaceSpec, + ModelActionChunk, + ModelObservation, + RobotActionChunk, + RobotObservation, + RobotState, + VLACapabilities, + VLARequest, +) +from .embodiment import EmbodimentAdapter +from .policy import VLAPolicy +from .registry import VLARegistry +from .serialization import ( + VLA_WIRE_ENCODING, + VLA_WIRE_SCHEMA_VERSION, + action_space_from_wire, + action_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", + "ModelActionChunk", + "ModelObservation", + "RobotActionChunk", + "RobotObservation", + "RobotState", + "VLACapabilities", + "VLAPolicy", + "VLARegistry", + "VLARequest", + "VLA_WIRE_ENCODING", + "VLA_WIRE_SCHEMA_VERSION", + "VLASession", + "VLASessionContract", + "VLASessionManager", + "action_space_from_wire", + "action_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..3e83a67f --- /dev/null +++ b/telefuser/vla/contracts.py @@ -0,0 +1,230 @@ +"""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 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..b768830f --- /dev/null +++ b/telefuser/vla/embodiment.py @@ -0,0 +1,32 @@ +"""Embodiment protocol separating model coordinates from robot controls.""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from .contracts import ( + ActionSpaceSpec, + ModelActionChunk, + ModelObservation, + 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: ... + + 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..a305639e --- /dev/null +++ b/telefuser/vla/runtime/__init__.py @@ -0,0 +1,28 @@ +"""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 + +__all__ = [ + "ActionChunkStateMachine", + "ActionSafetyPolicy", + "ActionChunkScheduler", + "BoundedActionSafety", + "ChunkExecutor", + "ChunkStatus", + "ChunkTicket", + "DisconnectPolicy", + "FiniteActionSafety", + "RobotAction", + "RemainderPolicy", + "RuntimeState", +] diff --git a/telefuser/vla/runtime/chunk_state.py b/telefuser/vla/runtime/chunk_state.py new file mode 100644 index 00000000..d1810c06 --- /dev/null +++ b/telefuser/vla/runtime/chunk_state.py @@ -0,0 +1,432 @@ +"""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 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/serialization.py b/telefuser/vla/serialization.py new file mode 100644 index 00000000..9128d0d2 --- /dev/null +++ b/telefuser/vla/serialization.py @@ -0,0 +1,259 @@ +"""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, RobotActionChunk, RobotObservation, RobotState + +VLA_WIRE_SCHEMA_VERSION = 1 +VLA_WIRE_ENCODING = "json-base64-v1" +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 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..b2373392 --- /dev/null +++ b/telefuser/vla/session.py @@ -0,0 +1,222 @@ +"""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") + 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/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_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py index f2879e53..6ee7fec5 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,13 @@ 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, + RobotWinProfile, +) +from telefuser.vla import ModelActionChunk, RobotObservation, RobotState def _stats() -> dict[str, dict[str, list[float]]]: @@ -82,3 +88,30 @@ 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: object() 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 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..02313d27 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() 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/test_pipeline_pool_sessions.py b/tests/unit/service/test_pipeline_pool_sessions.py new file mode 100644 index 00000000..d132d060 --- /dev/null +++ b/tests/unit/service/test_pipeline_pool_sessions.py @@ -0,0 +1,169 @@ +"""CPU-only tests for session-affine PipelinePool leases.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import 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_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_session.py b/tests/unit/service/test_vla_session.py new file mode 100644 index 00000000..cc4970eb --- /dev/null +++ b/tests/unit/service/test_vla_session.py @@ -0,0 +1,212 @@ +"""Loopback tests for the generic versioned VLA WebSocket protocol.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +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_vla_session_app +from telefuser.vla import ( + ActionSpaceSpec, + ModelActionChunk, + ModelObservation, + RobotActionChunk, + RobotObservation, + RobotState, + VLACapabilities, + VLARegistry, + VLASessionManager, + action_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) + + +class _Policy: + def __init__(self, *, delay_s: float = 0, failure: Exception | None = None) -> None: + self.delay_s = delay_s + self.failure = failure + self.resets: list[str] = [] + + def capabilities(self) -> VLACapabilities: + return VLACapabilities("fake", MODEL_SPACE, max_horizon=3) + + def predict(self, request) -> ModelActionChunk: + 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 + + 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["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) + + 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" + + 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"] diff --git a/tests/unit/vla/test_chunk_state.py b/tests/unit/vla/test_chunk_state.py new file mode 100644 index 00000000..71597b0c --- /dev/null +++ b/tests/unit/vla/test_chunk_state.py @@ -0,0 +1,171 @@ +"""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"] 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..d643619d --- /dev/null +++ b/tests/unit/vla/test_serialization.py @@ -0,0 +1,79 @@ +"""Tests for stable JSON/Base64 VLA wire contracts.""" + +from __future__ import annotations + +import pytest +import torch + +from telefuser.vla import ( + ActionSpaceSpec, + RobotActionChunk, + RobotObservation, + RobotState, + action_space_from_wire, + action_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) + + +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 + 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) diff --git a/tests/unit/vla/test_session.py b/tests/unit/vla/test_session.py new file mode 100644 index 00000000..0c954391 --- /dev/null +++ b/tests/unit/vla/test_session.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +from telefuser.vla import ( + ActionSpaceSpec, + ModelActionChunk, + ModelObservation, + 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") + + +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 + + 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") From d6383a45a07f645b9b6eb63448c5078f5ce25592 Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 14 Sep 2026 02:34:28 +0000 Subject: [PATCH 2/5] feat(vla): add replica-backed session serving Route generic VLA OPEN, PREDICT, RESET, and CLOSE operations through session-affine PipelinePool replicas and add a standalone LingBot VLA WebSocket entrypoint. Connect deterministic inference chunk state with a simulator-side execution runtime, standardize replica error handling, and preserve existing HTTP and non-VLA service behavior. Centralize LingBot H100 SDPA setup and VLA session composition while retaining the upstream RoboTwin MessagePack endpoint as a documented compatibility path. Verification: 82 VLA/LingBot/simulator unit tests and 231 non-LiveKit service unit tests passed; ruff check, ruff format --check, and git diff --check passed. --- docs/en/vla.md | 38 +- examples/lingbot_vla_v2/README.md | 38 +- .../lingbot_vla_v2_native_service.py | 13 +- .../lingbot_vla_v2_robotwin_server.py | 35 +- .../lingbot_vla_v2_vla_server.py | 70 ++++ .../pipelines/lingbot_vla_v2/__init__.py | 8 +- telefuser/pipelines/lingbot_vla_v2/runtime.py | 19 + .../pipelines/lingbot_vla_v2/vla_policy.py | 21 +- telefuser/service/core/pipeline_pool.py | 36 +- telefuser/service/core/replica_worker.py | 114 +++++- telefuser/service/vla_replica.py | 164 ++++++++ telefuser/service/vla_session.py | 366 ++++++++++++++---- telefuser/vla/runtime/__init__.py | 2 + telefuser/vla/runtime/chunk_state.py | 15 + telefuser/vla/runtime/simulator.py | 95 +++++ .../pipelines/lingbot_vla_v2/test_examples.py | 41 +- .../lingbot_vla_v2/test_robotwin_server.py | 29 -- .../pipelines/lingbot_vla_v2/test_runtime.py | 31 ++ tests/unit/service/test_pipeline_pool.py | 27 +- .../service/test_pipeline_pool_sessions.py | 19 +- tests/unit/service/test_vla_replica.py | 145 +++++++ tests/unit/service/test_vla_session.py | 79 +++- tests/unit/vla/test_chunk_state.py | 10 + tests/unit/vla/test_simulator_runtime.py | 67 ++++ 24 files changed, 1314 insertions(+), 168 deletions(-) create mode 100644 examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py create mode 100644 telefuser/service/vla_replica.py create mode 100644 telefuser/vla/runtime/simulator.py create mode 100644 tests/unit/service/test_vla_replica.py create mode 100644 tests/unit/vla/test_simulator_runtime.py diff --git a/docs/en/vla.md b/docs/en/vla.md index 061e8458..136c1a37 100644 --- a/docs/en/vla.md +++ b/docs/en/vla.md @@ -32,6 +32,7 @@ The public modules are: - `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. ## LingBot-VLA v2 Compatibility @@ -42,8 +43,10 @@ changed. It labels the normalized canonical `[T,55]` result as a `ModelActionChu 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 existing LingBot RoboTwin WebSocket endpoint now calls this policy, embodiment, session, and runtime path -internally. Its URL, MessagePack request fields, metadata frame, response fields, reset behavior, and latest-wins +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. @@ -83,7 +86,9 @@ clocks on separate machines are not comparable. ## Generic WebSocket Protocol -`create_vla_session_app` exposes `/v1/vla/session` without changing an existing TeleFuser service route. The server +`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 space; semantic incompatibility is rejected before inference. @@ -92,7 +97,8 @@ The protocol returns machine-readable error codes for malformed messages, unsupp or sessions, action-space mismatches, 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. +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. @@ -107,8 +113,9 @@ After a timed-out stateful inference finishes, the session is reset before it ac then stops replicas. This API is additive. Existing HTTP work continues to use `PipelinePool.acquire()`, whose allocation behavior and -callers are unchanged. The generic in-process WebSocket factory accepts a `VLASessionManager`; a deployment backed by -subprocess replicas must use these lease methods when adding its replica RPC adapter. +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 @@ -118,6 +125,10 @@ network-TTL checks, configurable discard/retain handling after `execute_horizon` 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. @@ -132,15 +143,16 @@ Adding another simulator requires a new `SimulatorAdapter`; it must not add mode robot requires an `EmbodimentAdapter`. Adding another VLA requires a `VLAPolicy` wrapper around its maintained pipeline. -## Current Boundary +## Service Boundary -The generic WebSocket application is intentionally an explicit factory rather than an automatically mounted -`telefuser serve` route. Subprocess replica transport still needs VLA-specific OPEN/PREDICT/RESET/CLOSE RPC commands -before the WebSocket factory can directly own `PipelinePool` leases. The lease primitive is implemented and tested, -but this RPC bridge is outside the current P0 scope. +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. -The chunk state machine is also transport-neutral. The existing LingBot compatibility server retains its proven -latest-wins scheduler until a later compatibility-preserving integration uses the generic state machine end to end. +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 e45ecfd6..bf5d3b75 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -25,11 +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 @@ -191,6 +192,10 @@ Compare a deterministic quantized capture with the corresponding TeleFuser BF16 ## Serving +The generic VLA session server is the primary online path for new deployments and simulator integrations. The native +structured service remains the HTTP path for single requests, offline evaluation, and batch clients. The RoboTwin +MessagePack server is retained only when an unmodified upstream `WebsocketClientPolicy` must connect directly. + Start the native structured service: ```bash @@ -229,11 +234,30 @@ 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 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. +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 RTX-side 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 report simulator execution as completed. + +### 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: @@ -333,6 +357,10 @@ Internally this compatibility endpoint uses the shared `VLAPolicy`, `EmbodimentA `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: 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..ec9bd5b9 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_native_service.py @@ -7,11 +7,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 +89,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 +104,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 66558509..8a829020 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,4 @@ -"""Serve LingBot-VLA v2 through the upstream RoboTwin policy protocol.""" +"""Compatibility server for the upstream RoboTwin policy protocol.""" from __future__ import annotations @@ -20,15 +20,16 @@ from telefuser.pipelines.lingbot_vla_v2 import ( ROBOTWIN_ACTION_ORDER, ROBOTWIN_CAMERA_KEYS, - LingBotVlaV2VLAPolicy, RobotWinProfile, + create_lingbot_vla_v2_session_manager, ) 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, VLARegistry, VLASessionManager +from telefuser.vla import RobotObservation, RobotState from telefuser.vla.runtime import ActionChunkScheduler ROBOTWIN_PROTOCOL_VERSION = "1.0" @@ -149,10 +150,7 @@ def __init__( self.use_length = use_length self._lock = threading.Lock() self._next_sequence: dict[str, int] = {} - registry = VLARegistry() - registry.register_policy("lingbot-vla-v2", LingBotVlaV2VLAPolicy(pipeline)) - registry.register_embodiment(self.profile.embodiment_id, self.profile) - self._sessions = VLASessionManager(registry) + self._sessions = create_lingbot_vla_v2_session_manager(pipeline, profile=self.profile) @property def metadata(self) -> dict[str, Any]: @@ -391,25 +389,6 @@ async def deliver( 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)) @@ -441,8 +420,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..5fdc5d0f --- /dev/null +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py @@ -0,0 +1,70 @@ +"""Standalone 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/telefuser/pipelines/lingbot_vla_v2/__init__.py b/telefuser/pipelines/lingbot_vla_v2/__init__.py index 4314c98d..a289312d 100644 --- a/telefuser/pipelines/lingbot_vla_v2/__init__.py +++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py @@ -11,7 +11,11 @@ LingBotVlaV2ActionChunk, RobotWinProfile, ) -from .vla_policy import LingBotVlaV2VLAPolicy +from .vla_policy import ( + LINGBOT_VLA_V2_MODEL_ID, + LingBotVlaV2VLAPolicy, + create_lingbot_vla_v2_session_manager, +) __all__ = [ "LingBotVlaV2ActionChunk", @@ -23,7 +27,9 @@ "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", 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 index 44770d30..4b3329f5 100644 --- a/telefuser/pipelines/lingbot_vla_v2/vla_policy.py +++ b/telefuser/pipelines/lingbot_vla_v2/vla_policy.py @@ -5,9 +5,13 @@ 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 +from .robot_profile import LINGBOT_VLA_V2_ACTION_SPACE, RobotWinProfile + +LINGBOT_VLA_V2_MODEL_ID = "lingbot-vla-v2" class _LingBotPipeline(Protocol): @@ -22,7 +26,7 @@ def __init__(self, pipeline: _LingBotPipeline, *, max_horizon: int = 50) -> None raise ValueError("max_horizon must be positive") self.pipeline = pipeline self._capabilities = VLACapabilities( - model_id="lingbot-vla-v2", + model_id=LINGBOT_VLA_V2_MODEL_ID, output_action_space=LINGBOT_VLA_V2_ACTION_SPACE, max_horizon=max_horizon, stateful=False, @@ -64,3 +68,16 @@ 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 7e40cdbc..97645a2e 100644 --- a/telefuser/service/core/pipeline_pool.py +++ b/telefuser/service/core/pipeline_pool.py @@ -74,6 +74,7 @@ def __init__( 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, @@ -81,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) @@ -98,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) @@ -122,6 +133,7 @@ def _start_all_replicas( self._security_level_name, skip_validation, self._server_config_data, + vla_provider_factory, ), daemon=False, ) @@ -155,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, @@ -335,6 +353,18 @@ async def session_bindings(self) -> dict[str, int]: 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, @@ -467,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..4cb086f0 --- /dev/null +++ b/telefuser/service/vla_replica.py @@ -0,0 +1,164 @@ +"""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, + 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", + ) + 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), + "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 index 081dbfa9..a9e84ece 100644 --- a/telefuser/service/vla_session.py +++ b/telefuser/service/vla_session.py @@ -7,12 +7,16 @@ import math from collections.abc import Mapping from enum import Enum -from typing import Any +from typing import Any, Protocol from fastapi import FastAPI, WebSocket, WebSocketDisconnect -from telefuser.service.core.replica_worker import ReplicaDeadError +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, 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_WIRE_ENCODING, @@ -20,6 +24,7 @@ action_space_to_wire, dumps_wire_message, loads_wire_message, + robot_action_chunk_from_wire, robot_action_chunk_to_wire, robot_observation_from_wire, ) @@ -53,6 +58,130 @@ def __init__(self, code: VLAErrorCode, message: str) -> None: 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) + 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), + "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) + + def create_vla_session_app( sessions: VLASessionManager, *, @@ -60,6 +189,41 @@ def create_vla_session_app( 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: @@ -76,6 +240,8 @@ async def session_socket(websocket: WebSocket) -> None: await websocket.accept() owned_sessions: set[str] = set() recovery_tasks: dict[str, asyncio.Task[None]] = {} + chunk_states: dict[str, ActionChunkStateMachine] = {} + metadata = backend.metadata() await _send( websocket, { @@ -83,8 +249,8 @@ async def session_socket(websocket: WebSocket) -> None: "protocol_version": VLA_SESSION_PROTOCOL_VERSION, "encoding": VLA_WIRE_ENCODING, "operations": ["OPEN", "PREDICT", "RESET", "CLOSE"], - "model_ids": list(sessions.registry.model_ids()), - "embodiment_ids": list(sessions.registry.embodiment_ids()), + "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, }, @@ -99,11 +265,11 @@ async def session_socket(websocket: WebSocket) -> None: ) request_id = _request_id(message) response = await _handle_message( - sessions, + backend, message, owned_sessions=owned_sessions, recovery_tasks=recovery_tasks, - max_tensor_bytes=max_tensor_bytes, + chunk_states=chunk_states, ) except WebSocketDisconnect: break @@ -116,18 +282,21 @@ async def session_socket(websocket: WebSocket) -> None: await task for session_id in tuple(owned_sessions): with contextlib.suppress(KeyError): - await asyncio.to_thread(sessions.close, session_id) + await backend.close(session_id) + state = chunk_states.pop(session_id, None) + if state is not None: + state.disconnect() return app async def _handle_message( - sessions: VLASessionManager, + backend: _VLABackend, message: Mapping[str, Any], *, owned_sessions: set[str], recovery_tasks: dict[str, asyncio.Task[None]], - max_tensor_bytes: int, + chunk_states: dict[str, ActionChunkStateMachine], ) -> dict[str, Any]: operation = message.get("type") if not isinstance(operation, str): @@ -135,7 +304,13 @@ async def _handle_message( operation = operation.upper() request_id = _request_id(message) if operation == "OPEN": - return _open_session(sessions, message, owned_sessions=owned_sessions, request_id=request_id) + return await _open_session( + backend, + message, + owned_sessions=owned_sessions, + chunk_states=chunk_states, + request_id=request_id, + ) session_id = _required_string(message, "session_id") if session_id not in owned_sessions: @@ -144,27 +319,22 @@ async def _handle_message( recovery = recovery_tasks.get(session_id) if recovery is not None and not recovery.done(): raise VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, "session is recovering from a timed-out request") - observation_payload = message.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=max_tensor_bytes) timestamp_ns = _required_nonnegative_int(message, "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", - ) timeout_s = _optional_timeout_s(message) - predict_task = asyncio.create_task( - asyncio.to_thread( - sessions.get(session_id).predict, - observation, - _required_string(message, "instruction"), - _required_nonnegative_int(message, "sequence_id"), - seed=_optional_integer(message, "seed"), - now_ns=_optional_nonnegative_int(message, "observation_clock_now_ns"), - ) + state = chunk_states[session_id] + ticket = 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 = state.status(ticket) + if admission is ChunkStatus.REJECTED: + raise VLAProtocolError(VLAErrorCode.OUT_OF_ORDER, state.reason(ticket) or "action chunk was rejected") + if admission is ChunkStatus.EXPIRED: + raise VLAProtocolError(VLAErrorCode.EXPIRED, state.reason(ticket) or "observation expired") + state.mark_inference_started(ticket) + predict_task = asyncio.create_task(backend.predict(session_id, message)) try: if timeout_s is None: chunk = await predict_task @@ -172,9 +342,29 @@ async def _handle_message( chunk = await asyncio.wait_for(asyncio.shield(predict_task), timeout=timeout_s) except asyncio.TimeoutError as error: recovery_tasks[session_id] = asyncio.create_task( - _recover_timed_out_session(sessions, session_id, predict_task, recovery_tasks) + _recover_timed_out_session( + backend, + session_id, + predict_task, + recovery_tasks, + state, + ticket, + _optional_nonnegative_int(message, "observation_clock_now_ns"), + ) ) raise VLAProtocolError(VLAErrorCode.TIMEOUT, "PREDICT exceeded request_ttl_ms") from error + except Exception as error: + state.reject(ticket, str(error) or "VLA inference failed") + raise + completion = state.complete( + ticket, + chunk, + observation_clock_now_ns=_optional_nonnegative_int(message, "observation_clock_now_ns"), + ) + if completion is ChunkStatus.EXPIRED: + raise VLAProtocolError(VLAErrorCode.EXPIRED, state.reason(ticket) or "action chunk expired") + if completion is not ChunkStatus.READY: + raise VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, state.reason(ticket) or "action chunk discarded") return { "type": "ACTION_CHUNK", "protocol_version": VLA_SESSION_PROTOCOL_VERSION, @@ -184,84 +374,87 @@ async def _handle_message( } if operation == "RESET": await _await_recovery(session_id, recovery_tasks) - await asyncio.to_thread(sessions.reset, session_id, _optional_string(message, "episode_id")) + episode_id = _optional_string(message, "episode_id") + await backend.reset(session_id, episode_id) + chunk_states[session_id].reset(episode_id) return _success_response("RESET", session_id, request_id) if operation == "CLOSE": await _await_recovery(session_id, recovery_tasks) - await asyncio.to_thread(sessions.close, session_id) + await backend.close(session_id) owned_sessions.remove(session_id) + chunk_states.pop(session_id).disconnect() return _success_response("CLOSE", session_id, request_id) raise VLAProtocolError(VLAErrorCode.INVALID_MESSAGE, f"unsupported VLA operation: {operation!r}") -def _open_session( - sessions: VLASessionManager, +async def _open_session( + backend: _VLABackend, message: Mapping[str, Any], *, owned_sessions: set[str], + chunk_states: dict[str, ActionChunkStateMachine], 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 sessions.session_ids(): + if session_id in owned_sessions: raise VLAProtocolError(VLAErrorCode.SESSION_EXISTS, f"VLA session is already open: {session_id!r}") - model_id = _required_string(message, "model_id") - embodiment_id = _required_string(message, "embodiment_id") - try: - policy = sessions.registry.get_policy(model_id) - embodiment = sessions.registry.get_embodiment(embodiment_id) - except KeyError as error: - raise VLAProtocolError(VLAErrorCode.UNKNOWN_COMPONENT, str(error)) from error - expected_payload = message.get("expected_robot_action_space") - if expected_payload is not None: - 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( - embodiment.robot_action_space, - context="client and embodiment robot action space", - ) - sessions.open( - session_id, - model_id=model_id, - embodiment_id=embodiment_id, - episode_id=_required_string(message, "episode_id"), - execute_horizon=_optional_positive_int(message, "execute_horizon"), + 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") + chunk_states[session_id] = 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"), ) - owned_sessions.add(session_id) - capabilities = policy.capabilities() return { "type": "OPENED", "protocol_version": VLA_SESSION_PROTOCOL_VERSION, "request_id": request_id, - "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), - "max_horizon": capabilities.max_horizon, - "stateful": capabilities.stateful, - "supports_seed": capabilities.supports_seed, + **opened, } async def _recover_timed_out_session( - sessions: VLASessionManager, + backend: _VLABackend, session_id: str, predict_task: asyncio.Task[Any], recovery_tasks: dict[str, asyncio.Task[None]], + state: ActionChunkStateMachine, + ticket: ChunkTicket, + observation_clock_now_ns: int | None, ) -> None: try: - with contextlib.suppress(Exception): - await predict_task - with contextlib.suppress(KeyError): - await asyncio.to_thread(sessions.reset, session_id) + try: + chunk = await predict_task + except Exception as error: + state.reject(ticket, str(error) or "timed-out VLA inference failed") + else: + state.complete(ticket, chunk, observation_clock_now_ns=observation_clock_now_ns) + await backend.reset(session_id, None) + state.reset() finally: recovery_tasks.pop(session_id, None) +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", + ) + + async def _await_recovery(session_id: str, recovery_tasks: dict[str, asyncio.Task[None]]) -> None: recovery = recovery_tasks.get(session_id) if recovery is not None: @@ -288,24 +481,37 @@ def _error_code(error: Exception) -> VLAErrorCode: 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): - message = str(error).lower() - if "action space" in message or "action-space" in message: - return VLAErrorCode.ACTION_SPACE_MISMATCH - if "sequence_id" in message or "must increase" in message: - return VLAErrorCode.OUT_OF_ORDER - if "stale" in message or "expired" in message: - return VLAErrorCode.EXPIRED - if "already open" in message: - return VLAErrorCode.SESSION_EXISTS - return VLAErrorCode.INVALID_MESSAGE + 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 "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", diff --git a/telefuser/vla/runtime/__init__.py b/telefuser/vla/runtime/__init__.py index a305639e..d035d85c 100644 --- a/telefuser/vla/runtime/__init__.py +++ b/telefuser/vla/runtime/__init__.py @@ -11,6 +11,7 @@ from .executor import ChunkExecutor, RobotAction from .safety import ActionSafetyPolicy, BoundedActionSafety, FiniteActionSafety from .scheduler import ActionChunkScheduler +from .simulator import SimulatorChunkRuntime __all__ = [ "ActionChunkStateMachine", @@ -25,4 +26,5 @@ "RobotAction", "RemainderPolicy", "RuntimeState", + "SimulatorChunkRuntime", ] diff --git a/telefuser/vla/runtime/chunk_state.py b/telefuser/vla/runtime/chunk_state.py index d1810c06..ebe03fb6 100644 --- a/telefuser/vla/runtime/chunk_state.py +++ b/telefuser/vla/runtime/chunk_state.py @@ -254,6 +254,21 @@ def complete( 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: diff --git a/telefuser/vla/runtime/simulator.py b/telefuser/vla/runtime/simulator.py new file mode 100644 index 00000000..cacf4d82 --- /dev/null +++ b/telefuser/vla/runtime/simulator.py @@ -0,0 +1,95 @@ +"""Client-side action lifecycle for simulator execution.""" + +from __future__ import annotations + +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 + + +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.""" + chunk = self.state_machine.begin_execution() + if chunk is None: + return 0 + executed = 0 + try: + for action in ChunkExecutor.iter_actions(chunk): + self.simulator.execute(action) + executed += 1 + except Exception: + self.state_machine.finish_execution(success=False) + self.state_machine.disconnect() + raise + self.state_machine.finish_execution(success=True) + return 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/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_robotwin_server.py b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py index 02313d27..3ace5201 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robotwin_server.py @@ -271,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/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 index d132d060..95b23c95 100644 --- a/tests/unit/service/test_pipeline_pool_sessions.py +++ b/tests/unit/service/test_pipeline_pool_sessions.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -78,6 +78,23 @@ async def use_session(name: str) -> None: 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() diff --git a/tests/unit/service/test_vla_replica.py b/tests/unit/service/test_vla_replica.py new file mode 100644 index 00000000..f65eaecd --- /dev/null +++ b/tests/unit/service/test_vla_replica.py @@ -0,0 +1,145 @@ +"""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, + RobotActionChunk, + RobotObservation, + RobotState, + VLACapabilities, + VLARegistry, + VLASessionManager, + action_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) + + +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 + + 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), + }, + ) + assert opened["max_horizon"] == 2 + assert opened["robot_action_space"] == action_space_to_wire(ROBOT_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_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 index cc4970eb..64c5eb96 100644 --- a/tests/unit/service/test_vla_session.py +++ b/tests/unit/service/test_vla_session.py @@ -10,7 +10,11 @@ from fastapi.testclient import TestClient from telefuser.service.core.replica_worker import ReplicaDeadError -from telefuser.service.vla_session import VLA_SESSION_PROTOCOL_VERSION, create_vla_session_app +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, ModelActionChunk, @@ -23,6 +27,7 @@ VLASessionManager, action_space_to_wire, robot_action_chunk_from_wire, + robot_action_chunk_to_wire, robot_observation_to_wire, ) @@ -210,3 +215,75 @@ def test_disconnect_closes_connection_owned_sessions() -> None: 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), + "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"), + ] diff --git a/tests/unit/vla/test_chunk_state.py b/tests/unit/vla/test_chunk_state.py index 71597b0c..d7452f81 100644 --- a/tests/unit/vla/test_chunk_state.py +++ b/tests/unit/vla/test_chunk_state.py @@ -169,3 +169,13 @@ def test_terminal_history_preserves_inflight_ticket_until_discard_recovery() -> 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_simulator_runtime.py b/tests/unit/vla/test_simulator_runtime.py new file mode 100644 index 00000000..4ecdf2fc --- /dev/null +++ b/tests/unit/vla/test_simulator_runtime.py @@ -0,0 +1,67 @@ +"""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 From 9504db84239c9858b3598967af9f53dcd65e614f Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 14 Sep 2026 08:38:12 +0000 Subject: [PATCH 3/5] feat(vla): complete generic session integration Add the semantic VLA client, contracts, serialization, embodiment and session APIs. Integrate replica-backed VLA lifecycle handling and cover client, service, pipeline, and process behavior with focused tests. Document the generic session and legacy compatibility boundaries. --- docs/en/vla.md | 31 +- telefuser/client/__init__.py | 16 + telefuser/client/vla.py | 244 +++++++++++ .../pipelines/lingbot_vla_v2/__init__.py | 2 + .../pipelines/lingbot_vla_v2/robot_profile.py | 20 +- telefuser/service/vla_replica.py | 11 + telefuser/service/vla_session.py | 407 +++++++++++++----- telefuser/vla/__init__.py | 10 + telefuser/vla/contracts.py | 104 +++++ telefuser/vla/embodiment.py | 4 + telefuser/vla/serialization.py | 55 ++- telefuser/vla/session.py | 1 + tests/integration/test_vla_replica_process.py | 67 +++ .../lingbot_vla_v2/test_robot_profile.py | 4 +- .../service/fixtures/fake_vla_pipeline.py | 92 ++++ tests/unit/service/test_vla_replica.py | 22 + tests/unit/service/test_vla_session.py | 104 ++++- tests/unit/vla/test_client.py | 132 ++++++ tests/unit/vla/test_serialization.py | 34 ++ tests/unit/vla/test_session.py | 12 + 20 files changed, 1243 insertions(+), 129 deletions(-) create mode 100644 telefuser/client/vla.py create mode 100644 tests/integration/test_vla_replica_process.py create mode 100644 tests/unit/service/fixtures/fake_vla_pipeline.py create mode 100644 tests/unit/vla/test_client.py diff --git a/docs/en/vla.md b/docs/en/vla.md index 136c1a37..4801f03f 100644 --- a/docs/en/vla.md +++ b/docs/en/vla.md @@ -21,6 +21,10 @@ Tensor width is never used to infer action meaning. Every chunk carries an `Acti 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. @@ -28,11 +32,13 @@ The public modules are: - `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 spaces, observations, and chunks. +- `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 @@ -91,18 +97,22 @@ clocks on separate machines are not comparable. `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 space; semantic incompatibility is rejected before inference. +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, action-space mismatches, 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. +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. -After a timed-out stateful inference finishes, the session is reset before it accepts more prediction work. +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 @@ -150,6 +160,11 @@ route. The LingBot example supplies a standalone server that starts `PipelinePoo 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. 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/pipelines/lingbot_vla_v2/__init__.py b/telefuser/pipelines/lingbot_vla_v2/__init__.py index a289312d..d51d7fc3 100644 --- a/telefuser/pipelines/lingbot_vla_v2/__init__.py +++ b/telefuser/pipelines/lingbot_vla_v2/__init__.py @@ -8,6 +8,7 @@ ROBOTWIN_ACTION_ORDER, ROBOTWIN_ACTION_SPACE, ROBOTWIN_CAMERA_KEYS, + ROBOTWIN_OBSERVATION_SPACE, LingBotVlaV2ActionChunk, RobotWinProfile, ) @@ -33,5 +34,6 @@ "ROBOTWIN_ACTION_ORDER", "ROBOTWIN_ACTION_SPACE", "ROBOTWIN_CAMERA_KEYS", + "ROBOTWIN_OBSERVATION_SPACE", "RobotWinProfile", ] diff --git a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py index 939ad4ee..eff67389 100644 --- a/telefuser/pipelines/lingbot_vla_v2/robot_profile.py +++ b/telefuser/pipelines/lingbot_vla_v2/robot_profile.py @@ -12,8 +12,10 @@ from telefuser.vla.contracts import ( ActionSpaceSpec, + ImageObservationSpec, ModelActionChunk, ModelObservation, + ObservationSpaceSpec, RobotActionChunk, RobotObservation, RobotState, @@ -61,6 +63,11 @@ 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) @@ -134,15 +141,14 @@ 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.""" - if not isinstance(observation, RobotObservation): - raise TypeError("observation must be a RobotObservation") - if observation.state.dimension_names != ROBOTWIN_ACTION_ORDER: - raise ValueError("RobotWin state dimension order does not match the robot profile") - missing = [key for key in self.camera_keys if key not in observation.images] - if missing: - raise ValueError(f"RobotWin observation is missing camera keys: {missing}") + self.observation_space.validate(observation) return ModelObservation( state=observation.state.values, images={key: observation.images[key] for key in self.camera_keys}, diff --git a/telefuser/service/vla_replica.py b/telefuser/service/vla_replica.py index 4cb086f0..a8297551 100644 --- a/telefuser/service/vla_replica.py +++ b/telefuser/service/vla_replica.py @@ -10,6 +10,8 @@ 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, ) @@ -75,6 +77,14 @@ def _open(self, payload: Mapping[str, Any]) -> dict[str, Any]: 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, @@ -90,6 +100,7 @@ def _open(self, payload: Mapping[str, Any]) -> dict[str, Any]: "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, diff --git a/telefuser/service/vla_session.py b/telefuser/service/vla_session.py index a9e84ece..c83b4b55 100644 --- a/telefuser/service/vla_session.py +++ b/telefuser/service/vla_session.py @@ -6,6 +6,7 @@ import contextlib import math from collections.abc import Mapping +from dataclasses import dataclass from enum import Enum from typing import Any, Protocol @@ -14,22 +15,24 @@ 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, RobotActionChunk +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, ) -VLA_SESSION_PROTOCOL_VERSION = "1.0" DEFAULT_MAX_VLA_MESSAGE_BYTES = 64 * 1024 * 1024 @@ -42,7 +45,9 @@ class VLAErrorCode(str, Enum): 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" @@ -93,6 +98,7 @@ async def open(self, payload: Mapping[str, Any]) -> dict[str, Any]: 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, @@ -109,6 +115,7 @@ async def open(self, payload: Mapping[str, Any]) -> dict[str, Any]: "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, @@ -182,6 +189,205 @@ async def close(self, session_id: str) -> None: 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, *, @@ -239,16 +445,24 @@ async def healthz() -> dict[str, str]: async def session_socket(websocket: WebSocket) -> None: await websocket.accept() owned_sessions: set[str] = set() - recovery_tasks: dict[str, asyncio.Task[None]] = {} - chunk_states: dict[str, ActionChunkStateMachine] = {} + runtimes: dict[str, _SessionRuntime] = {} + delivery_tasks: set[asyncio.Task[None]] = set() + send_lock = asyncio.Lock() metadata = backend.metadata() - await _send( + 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, @@ -264,39 +478,51 @@ async def session_socket(websocket: WebSocket) -> None: max_message_bytes=max_message_bytes, ) request_id = _request_id(message) - response = await _handle_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, - recovery_tasks=recovery_tasks, - chunk_states=chunk_states, + runtimes=runtimes, ) except WebSocketDisconnect: break except Exception as error: response = _error_response(error, request_id=request_id) - await _send(websocket, response) + await _locked_send(websocket, send_lock, response) finally: - for task in list(recovery_tasks.values()): + 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 - for session_id in tuple(owned_sessions): - with contextlib.suppress(KeyError): - await backend.close(session_id) - state = chunk_states.pop(session_id, None) - if state is not None: - state.disconnect() return app -async def _handle_message( +async def _handle_control_message( backend: _VLABackend, message: Mapping[str, Any], *, owned_sessions: set[str], - recovery_tasks: dict[str, asyncio.Task[None]], - chunk_states: dict[str, ActionChunkStateMachine], + runtimes: dict[str, _SessionRuntime], ) -> dict[str, Any]: operation = message.get("type") if not isinstance(operation, str): @@ -308,7 +534,7 @@ async def _handle_message( backend, message, owned_sessions=owned_sessions, - chunk_states=chunk_states, + runtimes=runtimes, request_id=request_id, ) @@ -316,73 +542,15 @@ async def _handle_message( 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": - recovery = recovery_tasks.get(session_id) - if recovery is not None and not recovery.done(): - raise VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, "session is recovering from a timed-out request") - timestamp_ns = _required_nonnegative_int(message, "observation_timestamp_ns") - timeout_s = _optional_timeout_s(message) - state = chunk_states[session_id] - ticket = 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 = state.status(ticket) - if admission is ChunkStatus.REJECTED: - raise VLAProtocolError(VLAErrorCode.OUT_OF_ORDER, state.reason(ticket) or "action chunk was rejected") - if admission is ChunkStatus.EXPIRED: - raise VLAProtocolError(VLAErrorCode.EXPIRED, state.reason(ticket) or "observation expired") - state.mark_inference_started(ticket) - predict_task = asyncio.create_task(backend.predict(session_id, 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 as error: - recovery_tasks[session_id] = asyncio.create_task( - _recover_timed_out_session( - backend, - session_id, - predict_task, - recovery_tasks, - state, - ticket, - _optional_nonnegative_int(message, "observation_clock_now_ns"), - ) - ) - raise VLAProtocolError(VLAErrorCode.TIMEOUT, "PREDICT exceeded request_ttl_ms") from error - except Exception as error: - state.reject(ticket, str(error) or "VLA inference failed") - raise - completion = state.complete( - ticket, - chunk, - observation_clock_now_ns=_optional_nonnegative_int(message, "observation_clock_now_ns"), - ) - if completion is ChunkStatus.EXPIRED: - raise VLAProtocolError(VLAErrorCode.EXPIRED, state.reason(ticket) or "action chunk expired") - if completion is not ChunkStatus.READY: - raise VLAProtocolError(VLAErrorCode.SESSION_UNAVAILABLE, state.reason(ticket) or "action chunk discarded") - return { - "type": "ACTION_CHUNK", - "protocol_version": VLA_SESSION_PROTOCOL_VERSION, - "request_id": request_id, - "session_id": session_id, - "chunk": robot_action_chunk_to_wire(chunk), - } + raise RuntimeError("PREDICT must be scheduled by the connection receiver") if operation == "RESET": - await _await_recovery(session_id, recovery_tasks) episode_id = _optional_string(message, "episode_id") - await backend.reset(session_id, episode_id) - chunk_states[session_id].reset(episode_id) + await runtimes[session_id].reset(episode_id) return _success_response("RESET", session_id, request_id) if operation == "CLOSE": - await _await_recovery(session_id, recovery_tasks) - await backend.close(session_id) + await runtimes[session_id].close() owned_sessions.remove(session_id) - chunk_states.pop(session_id).disconnect() + 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}") @@ -392,7 +560,7 @@ async def _open_session( message: Mapping[str, Any], *, owned_sessions: set[str], - chunk_states: dict[str, ActionChunkStateMachine], + runtimes: dict[str, _SessionRuntime], request_id: str | int | None, ) -> dict[str, Any]: version = message.get("protocol_version") @@ -408,11 +576,17 @@ async def _open_session( await backend.close(session_id) owned_sessions.remove(session_id) raise RuntimeError("VLA backend returned an invalid max_horizon") - chunk_states[session_id] = ActionChunkStateMachine( + 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, @@ -421,28 +595,6 @@ async def _open_session( } -async def _recover_timed_out_session( - backend: _VLABackend, - session_id: str, - predict_task: asyncio.Task[Any], - recovery_tasks: dict[str, asyncio.Task[None]], - state: ActionChunkStateMachine, - ticket: ChunkTicket, - observation_clock_now_ns: int | None, -) -> None: - try: - try: - chunk = await predict_task - except Exception as error: - state.reject(ticket, str(error) or "timed-out VLA inference failed") - else: - state.complete(ticket, chunk, observation_clock_now_ns=observation_clock_now_ns) - await backend.reset(session_id, None) - state.reset() - finally: - recovery_tasks.pop(session_id, None) - - 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: @@ -455,16 +607,49 @@ def _validate_expected_action_space(payload: Mapping[str, Any], robot_action_spa ) -async def _await_recovery(session_id: str, recovery_tasks: dict[str, asyncio.Task[None]]) -> None: - recovery = recovery_tasks.get(session_id) - if recovery is not None: - await recovery +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" @@ -501,6 +686,8 @@ def _error_code(error: Exception) -> VLAErrorCode: 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: diff --git a/telefuser/vla/__init__.py b/telefuser/vla/__init__.py index 3b93d620..ce212696 100644 --- a/telefuser/vla/__init__.py +++ b/telefuser/vla/__init__.py @@ -3,8 +3,10 @@ from .contracts import ( ActionChunk, ActionSpaceSpec, + ImageObservationSpec, ModelActionChunk, ModelObservation, + ObservationSpaceSpec, RobotActionChunk, RobotObservation, RobotState, @@ -15,10 +17,13 @@ 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, @@ -30,8 +35,10 @@ "ActionChunk", "ActionSpaceSpec", "EmbodimentAdapter", + "ImageObservationSpec", "ModelActionChunk", "ModelObservation", + "ObservationSpaceSpec", "RobotActionChunk", "RobotObservation", "RobotState", @@ -41,11 +48,14 @@ "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", diff --git a/telefuser/vla/contracts.py b/telefuser/vla/contracts.py index 3e83a67f..3846473c 100644 --- a/telefuser/vla/contracts.py +++ b/telefuser/vla/contracts.py @@ -10,6 +10,110 @@ 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.""" diff --git a/telefuser/vla/embodiment.py b/telefuser/vla/embodiment.py index b768830f..a92fa096 100644 --- a/telefuser/vla/embodiment.py +++ b/telefuser/vla/embodiment.py @@ -8,6 +8,7 @@ ActionSpaceSpec, ModelActionChunk, ModelObservation, + ObservationSpaceSpec, RobotActionChunk, RobotObservation, RobotState, @@ -27,6 +28,9 @@ 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/serialization.py b/telefuser/vla/serialization.py index 9128d0d2..6eb63419 100644 --- a/telefuser/vla/serialization.py +++ b/telefuser/vla/serialization.py @@ -11,10 +11,18 @@ import numpy as np import torch -from .contracts import ActionSpaceSpec, RobotActionChunk, RobotObservation, RobotState +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] = { @@ -70,6 +78,51 @@ def action_space_from_wire(payload: Mapping[str, Any]) -> ActionSpaceSpec: 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 diff --git a/telefuser/vla/session.py b/telefuser/vla/session.py index b2373392..ac174fcc 100644 --- a/telefuser/vla/session.py +++ b/telefuser/vla/session.py @@ -65,6 +65,7 @@ def predict( 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: 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/pipelines/lingbot_vla_v2/test_robot_profile.py b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py index 6ee7fec5..0b084ec5 100644 --- a/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py +++ b/tests/unit/pipelines/lingbot_vla_v2/test_robot_profile.py @@ -7,6 +7,7 @@ LINGBOT_VLA_V2_ACTION_SPACE, ROBOTWIN_ACTION_ORDER, ROBOTWIN_ACTION_SPACE, + ROBOTWIN_OBSERVATION_SPACE, RobotWinProfile, ) from telefuser.vla import ModelActionChunk, RobotObservation, RobotState @@ -94,7 +95,7 @@ def test_profile_implements_semantic_embodiment_contract() -> None: profile = RobotWinProfile(_stats()) observation = RobotObservation( RobotState(torch.zeros(14), ROBOTWIN_ACTION_ORDER, timestamp_ns=123), - {key: object() for key in profile.camera_keys}, + {key: torch.zeros((2, 2, 3), dtype=torch.uint8) for key in profile.camera_keys}, ) model_observation = profile.encode_observation(observation) model_chunk = ModelActionChunk( @@ -110,6 +111,7 @@ def test_profile_implements_semantic_embodiment_contract() -> None: 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 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_vla_replica.py b/tests/unit/service/test_vla_replica.py index f65eaecd..a38ce82c 100644 --- a/tests/unit/service/test_vla_replica.py +++ b/tests/unit/service/test_vla_replica.py @@ -12,6 +12,7 @@ ActionSpaceSpec, ModelActionChunk, ModelObservation, + ObservationSpaceSpec, RobotActionChunk, RobotObservation, RobotState, @@ -19,12 +20,14 @@ 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: @@ -53,6 +56,7 @@ 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) @@ -87,10 +91,12 @@ def test_provider_dispatches_complete_worker_local_session() -> None: "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( @@ -128,6 +134,22 @@ def test_provider_rejects_action_contract_mismatch_before_open() -> None: ) +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"): diff --git a/tests/unit/service/test_vla_session.py b/tests/unit/service/test_vla_session.py index 64c5eb96..8758c538 100644 --- a/tests/unit/service/test_vla_session.py +++ b/tests/unit/service/test_vla_session.py @@ -4,6 +4,7 @@ import time from dataclasses import dataclass +from threading import Event import pytest import torch @@ -17,8 +18,10 @@ ) from telefuser.vla import ( ActionSpaceSpec, + ImageObservationSpec, ModelActionChunk, ModelObservation, + ObservationSpaceSpec, RobotActionChunk, RobotObservation, RobotState, @@ -26,6 +29,7 @@ VLARegistry, VLASessionManager, action_space_to_wire, + observation_space_to_wire, robot_action_chunk_from_wire, robot_action_chunk_to_wire, robot_observation_to_wire, @@ -33,18 +37,36 @@ 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) -> None: + 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) + 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: @@ -67,6 +89,7 @@ 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) @@ -131,6 +154,8 @@ def test_protocol_negotiates_capabilities_and_runs_full_session_lifecycle() -> N 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"] @@ -139,6 +164,7 @@ def test_protocol_negotiates_capabilities_and_runs_full_session_lifecycle() -> N 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() @@ -174,6 +200,10 @@ def test_protocol_returns_stable_errors_for_version_components_contract_order_an 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)) @@ -241,6 +271,7 @@ async def run_vla_operation(self, session_id: str, operation: str, payload: dict "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, @@ -287,3 +318,72 @@ def test_pipeline_pool_backend_runs_full_protocol_and_releases_lease() -> None: ("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_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_serialization.py b/tests/unit/vla/test_serialization.py index d643619d..9068b853 100644 --- a/tests/unit/vla/test_serialization.py +++ b/tests/unit/vla/test_serialization.py @@ -7,11 +7,15 @@ 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, @@ -20,6 +24,7 @@ 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: @@ -31,6 +36,7 @@ def test_action_space_observation_and_chunk_round_trip_without_semantic_loss() - 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 @@ -77,3 +83,31 @@ def test_action_space_wire_requires_arrays_not_ambiguous_strings() -> None: 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 index 0c954391..94c7c550 100644 --- a/tests/unit/vla/test_session.py +++ b/tests/unit/vla/test_session.py @@ -9,6 +9,7 @@ ActionSpaceSpec, ModelActionChunk, ModelObservation, + ObservationSpaceSpec, RobotActionChunk, RobotObservation, RobotState, @@ -24,6 +25,7 @@ def _space(representation: str) -> ActionSpaceSpec: MODEL_SPACE = _space("joint_delta") ROBOT_SPACE = _space("joint_position") +OBSERVATION_SPACE = ObservationSpaceSpec(("joint",)) class _Policy: @@ -54,6 +56,7 @@ 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) @@ -135,3 +138,12 @@ def test_open_rejects_policy_and_embodiment_action_space_mismatch() -> None: 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) From bce070282594c50c6616431f873c0f1c63fb196a Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 14 Sep 2026 08:39:11 +0000 Subject: [PATCH 4/5] feat(sim): add local MuJoCo VLA smoke path Add an optional MuJoCo simulator adapter with semantic joint bindings, PD execution, headless camera capture, and lifecycle validation. Provide an isolated local setup script and LingBot VLA WebSocket smoke entrypoint, and mark generic WebSocket as the primary online path while retaining compatibility entries. --- examples/lingbot_vla_v2/README.md | 42 ++- .../lingbot_vla_v2_inference.py | 6 +- .../lingbot_vla_v2/lingbot_vla_v2_mujoco.py | 283 ++++++++++++++++++ .../lingbot_vla_v2_native_service.py | 7 +- .../lingbot_vla_v2_robotwin_server.py | 7 +- .../lingbot_vla_v2_vla_server.py | 2 +- examples/lingbot_vla_v2/setup_mujoco_local.sh | 30 ++ telefuser/integrations/sim/__init__.py | 8 +- telefuser/integrations/sim/mujoco.py | 254 ++++++++++++++++ tests/unit/integrations/sim/test_mujoco.py | 88 ++++++ 10 files changed, 719 insertions(+), 8 deletions(-) create mode 100644 examples/lingbot_vla_v2/lingbot_vla_v2_mujoco.py create mode 100755 examples/lingbot_vla_v2/setup_mujoco_local.sh create mode 100644 telefuser/integrations/sim/mujoco.py create mode 100644 tests/unit/integrations/sim/test_mujoco.py diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index bf5d3b75..673d9ded 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -192,9 +192,18 @@ Compare a deterministic quantized capture with the corresponding TeleFuser BF16 ## Serving -The generic VLA session server is the primary online path for new deployments and simulator integrations. The native -structured service remains the HTTP path for single requests, offline evaluation, and batch clients. The RoboTwin -MessagePack server is retained only when an unmodified upstream `WebsocketClientPolicy` must connect directly. +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: @@ -403,6 +412,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 ec9bd5b9..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 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 8a829020..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 @@ -"""Compatibility server for 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 diff --git a/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py b/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py index 5fdc5d0f..6c0614bd 100644 --- a/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py +++ b/examples/lingbot_vla_v2/lingbot_vla_v2_vla_server.py @@ -1,4 +1,4 @@ -"""Standalone generic VLA WebSocket service backed by TeleFuser replicas.""" +"""Primary generic VLA WebSocket service backed by TeleFuser replicas.""" from __future__ import annotations 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/integrations/sim/__init__.py b/telefuser/integrations/sim/__init__.py index f9e75de2..6999cfa1 100644 --- a/telefuser/integrations/sim/__init__.py +++ b/telefuser/integrations/sim/__init__.py @@ -1,6 +1,12 @@ """Simulator-neutral VLA integration interfaces.""" from .base import SimulatorAdapter +from .mujoco import MuJoCoJointBinding, MuJoCoSimulatorAdapter from .robotwin import RoboTwinSimulatorAdapter -__all__ = ["RoboTwinSimulatorAdapter", "SimulatorAdapter"] +__all__ = [ + "MuJoCoJointBinding", + "MuJoCoSimulatorAdapter", + "RoboTwinSimulatorAdapter", + "SimulatorAdapter", +] 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/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'' + ) + + +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) From 3d2861908278f7603b67cb3c909029897fd70d0c Mon Sep 17 00:00:00 2001 From: HappyDog0713 Date: Mon, 14 Sep 2026 10:34:26 +0000 Subject: [PATCH 5/5] feat(vla): report simulator chunk execution Add a transport-neutral simulator execution report and an async helper for overlapping inference with synchronous simulator adapters. Preserve the legacy execute_ready API and cover success, failure, no-action, and async paths with focused tests. Document the trimmed P0 validation and keep the generic VLA WebSocket as the primary service entrypoint. --- examples/lingbot_vla_v2/README.md | 13 ++++- telefuser/vla/runtime/__init__.py | 3 +- telefuser/vla/runtime/simulator.py | 74 ++++++++++++++++++++++-- tests/unit/vla/test_simulator_runtime.py | 43 ++++++++++++++ 4 files changed, 125 insertions(+), 8 deletions(-) diff --git a/examples/lingbot_vla_v2/README.md b/examples/lingbot_vla_v2/README.md index 673d9ded..bea2d696 100644 --- a/examples/lingbot_vla_v2/README.md +++ b/examples/lingbot_vla_v2/README.md @@ -257,9 +257,16 @@ The service exposes `GET /healthz` and `/v1/vla/session`. Each `OPEN` reserves o `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 RTX-side 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 report simulator execution as completed. +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. + +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 diff --git a/telefuser/vla/runtime/__init__.py b/telefuser/vla/runtime/__init__.py index d035d85c..7ce4c573 100644 --- a/telefuser/vla/runtime/__init__.py +++ b/telefuser/vla/runtime/__init__.py @@ -11,7 +11,7 @@ from .executor import ChunkExecutor, RobotAction from .safety import ActionSafetyPolicy, BoundedActionSafety, FiniteActionSafety from .scheduler import ActionChunkScheduler -from .simulator import SimulatorChunkRuntime +from .simulator import ChunkExecutionReport, SimulatorChunkRuntime __all__ = [ "ActionChunkStateMachine", @@ -19,6 +19,7 @@ "ActionChunkScheduler", "BoundedActionSafety", "ChunkExecutor", + "ChunkExecutionReport", "ChunkStatus", "ChunkTicket", "DisconnectPolicy", diff --git a/telefuser/vla/runtime/simulator.py b/telefuser/vla/runtime/simulator.py index cacf4d82..e84e0f18 100644 --- a/telefuser/vla/runtime/simulator.py +++ b/telefuser/vla/runtime/simulator.py @@ -2,6 +2,8 @@ from __future__ import annotations +import asyncio +from dataclasses import dataclass from typing import TYPE_CHECKING from ..contracts import ActionSpaceSpec, RobotActionChunk @@ -18,6 +20,39 @@ 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.""" @@ -71,20 +106,51 @@ def accept( 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 0 + 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: + except Exception as error: self.state_machine.finish_execution(success=False) self.state_machine.disconnect() - raise + 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 executed + 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.""" diff --git a/tests/unit/vla/test_simulator_runtime.py b/tests/unit/vla/test_simulator_runtime.py index 4ecdf2fc..59a635f5 100644 --- a/tests/unit/vla/test_simulator_runtime.py +++ b/tests/unit/vla/test_simulator_runtime.py @@ -65,3 +65,46 @@ def test_client_runtime_rejects_mismatch_and_fails_closed_on_execution_error() - 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