| |
| |
| """Policy server wrapper. |
| |
| Encapsulates a `baseframework` instance plus a :class:`PolicyNormProcessor` |
| that reuses the *training-time* :class:`ComposedModalityTransform` for action |
| un-normalization (no hand-rolled math). The websocket server returns |
| already-unnormalized actions. |
| |
| Client-side responsibilities that REMAIN on the client: |
| - environment-specific adapters (image_history, gripper sticky, action |
| ensembling) |
| - chunk-cache scheduling (`step % chunk_size == 0` triggers a new infer) |
| |
| Exposed API: |
| - ``metadata`` (dict, sent at handshake): ``action_chunk_size``, |
| ``available_unnorm_keys``, ``action_keys``, ``state_keys``. |
| - ``predict_action(examples, unnorm_key=None, **kwargs)`` returns |
| ``{"actions": np.ndarray[B, T, action_dim]}``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from typing import Any, Dict, List, Optional |
|
|
| import numpy as np |
| import torch |
|
|
| from starVLA.model.framework.base_framework import baseframework |
| from starVLA.model.framework.share_tools import read_mode_config |
|
|
| from deployment.model_server.policy_norm_processor import PolicyNormProcessor |
|
|
|
|
| class PolicyServerWrapper: |
| """Wraps a `baseframework` for use as a websocket-server policy.""" |
|
|
| def __init__( |
| self, |
| ckpt_path: str, |
| device: str = "cuda", |
| use_bf16: bool = False, |
| unnorm_key: Optional[str] = None, |
| normalized_action_clip: Optional[float] = None, |
| ) -> None: |
| self._ckpt_path = str(ckpt_path) |
| self._normalized_action_clip = normalized_action_clip |
|
|
| logging.info("PolicyServerWrapper: loading framework from %s", self._ckpt_path) |
| framework = baseframework.from_pretrained(self._ckpt_path) |
| if use_bf16: |
| framework = framework.to(torch.bfloat16) |
| framework = framework.to(device).eval() |
| self._framework = framework |
|
|
| |
| model_cfg, _ = read_mode_config(self._ckpt_path) |
| self._model_cfg = model_cfg |
| vla_data_cfg = model_cfg.get("datasets", {}).get("vla_data", {}) |
| include_state = vla_data_cfg.get("include_state", False) |
| self._uses_state = include_state not in (False, None, 0, "False", "false", "0") |
|
|
| |
| action_model_cfg = model_cfg["framework"]["action_model"] |
| |
| if "action_horizon" in action_model_cfg: |
| self._action_chunk_size = int(action_model_cfg["action_horizon"]) |
| elif "future_action_window_size" in action_model_cfg: |
| self._action_chunk_size = int(action_model_cfg["future_action_window_size"]) + 1 |
| else: |
| raise ValueError( |
| f"PolicyServerWrapper: no action_horizon or future_action_window_size found in model config for {self._ckpt_path}" |
| ) |
| |
| |
| |
| self._default_unnorm_key = unnorm_key |
| self._norm_processors: Dict[str, PolicyNormProcessor] = {} |
|
|
| |
| _, _ns = read_mode_config(self._ckpt_path) |
| self._available_unnorm_keys: List[str] = list(_ns.keys()) |
|
|
| |
| if unnorm_key is not None or len(self._available_unnorm_keys) == 1: |
| default_proc = self._get_processor(unnorm_key) |
| self._default_unnorm_key = default_proc.unnorm_key |
| logging.info( |
| "PolicyServerWrapper ready: action_chunk_size=%d, default_unnorm_key=%s, " |
| "available_unnorm_keys=%s, action_keys=%s, state_keys=%s", |
| self._action_chunk_size, |
| default_proc.unnorm_key, |
| default_proc.available_unnorm_keys, |
| default_proc.action_keys, |
| default_proc.state_keys, |
| ) |
| else: |
| logging.info( |
| "PolicyServerWrapper ready (multi-key): action_chunk_size=%d, " |
| "available_unnorm_keys=%s — clients must pass unnorm_key per request.", |
| self._action_chunk_size, |
| self._available_unnorm_keys, |
| ) |
|
|
| def _get_processor(self, unnorm_key: Optional[str]) -> PolicyNormProcessor: |
| cache_key = unnorm_key if unnorm_key is not None else "__default__" |
| if cache_key not in self._norm_processors: |
| self._norm_processors[cache_key] = PolicyNormProcessor( |
| self._ckpt_path, unnorm_key=unnorm_key |
| ) |
| return self._norm_processors[cache_key] |
|
|
| @property |
| def metadata(self) -> Dict[str, Any]: |
| """Model-invariant metadata; sent to client at websocket handshake.""" |
| base = { |
| "env": "starvla_policy_server", |
| "ckpt_path": self._ckpt_path, |
| "action_chunk_size": self._action_chunk_size, |
| "uses_state": self._uses_state, |
| "available_unnorm_keys": self._available_unnorm_keys, |
| "default_unnorm_key": self._default_unnorm_key, |
| "normalized_action_clip": self._normalized_action_clip, |
| "action_head_attention_mask": False, |
| "inference_rng": "continuous_server_stream", |
| } |
| |
| if self._default_unnorm_key is not None: |
| proc = self._get_processor(self._default_unnorm_key) |
| base["action_keys"] = proc.action_keys |
| base["state_keys"] = proc.state_keys |
| return base |
|
|
| def predict_action( |
| self, |
| examples: List[dict], |
| unnorm_key: Optional[str] = None, |
| **kwargs, |
| ) -> Dict[str, np.ndarray]: |
| """Run the framework, then un-normalize via training-time transforms. |
| |
| Args: |
| examples: list of dicts (each with ``image`` / ``lang`` / optional ``state``). |
| unnorm_key: dataset key for un-normalization stats. ``None`` --> |
| use the wrapper's default (auto-picked at startup). |
| **kwargs: forwarded to the framework's ``predict_action`` |
| (``do_sample``, ``use_ddim``, ``num_ddim_steps``, ...). |
| |
| Returns: |
| ``{"actions": np.ndarray[B, T, D]}`` -- un-normalized. |
| """ |
| effective_key = unnorm_key if unnorm_key is not None else self._default_unnorm_key |
| if effective_key is None: |
| if len(self._available_unnorm_keys) == 1: |
| effective_key = self._available_unnorm_keys[0] |
| else: |
| raise ValueError( |
| f"predict_action: unnorm_key not specified and no default set. " |
| f"Pass one of {self._available_unnorm_keys}." |
| ) |
| proc = self._get_processor(effective_key) |
|
|
| out = self._framework.predict_action(examples=examples, **kwargs) |
| normalized = np.asarray(out["normalized_actions"]) |
| if self._normalized_action_clip is not None: |
| limit = float(self._normalized_action_clip) |
| normalized = np.clip(normalized, -limit, limit) |
|
|
| unnorm = np.stack( |
| [proc.unapply_actions(normalized[b]) for b in range(normalized.shape[0])], |
| axis=0, |
| ) |
| return {"actions": unnorm} |
|
|
