| |
| |
| """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) |
| |
| Embodiments whose DataConfig declares a ``robot_command_layout`` (currently |
| xtrainer2) additionally get the raw-robot conversions here, so the client never |
| has to reproduce training-time normalization or the quaternion order: |
| |
| - ``state`` arrives in robot coordinates (16-D, xyzw) and is reordered to the |
| training layout, then normalized with the training transforms (20-D) |
| - returned actions are un-normalized *and* reordered back to xyzw |
| |
| Exposed API: |
| - ``metadata`` (dict, sent at handshake): ``action_chunk_size``, |
| ``available_unnorm_keys``, ``action_keys``, ``state_keys``, plus the |
| embodiment's raw-robot contract. |
| - ``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 |
| from deployment.model_server import xtrainer2_command |
|
|
|
|
| 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] |
|
|
| @staticmethod |
| def _uses_xtrainer2_layout(proc: PolicyNormProcessor) -> bool: |
| """Whether this embodiment speaks the raw xtrainer2 robot layout.""" |
| layout = getattr(proc._data_config, "robot_command_layout", None) |
| return layout == xtrainer2_command.COMMAND_LAYOUT |
|
|
| @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", |
| |
| |
| "image_format": "HWC_RGB_uint8", |
| "image_key": "image", |
| } |
| |
| 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 |
| base["camera_order"] = list(getattr(proc._data_config, "video_keys", [])) |
| base["model_action_dim"] = proc.action_dim |
| base["model_state_dim"] = proc.state_dim |
| base["server_normalizes_state"] = self._server_normalizes_state(proc) |
| if base["server_normalizes_state"]: |
| base["state_dim"] = proc.raw_state_dim |
| if self._uses_xtrainer2_layout(proc): |
| base["robot_command_layout"] = xtrainer2_command.COMMAND_LAYOUT |
| base["action_dim"] = proc.raw_action_dim |
| base["model_quaternion_format"] = xtrainer2_command.MODEL_QUATERNION_FORMAT |
| base["robot_quaternion_format"] = xtrainer2_command.SOURCE_QUATERNION_FORMAT |
| base["quaternion_starts"] = list(xtrainer2_command.QUATERNION_STARTS) |
| action_mode = str( |
| getattr(self, "_model_cfg", {}) |
| .get("datasets", {}) |
| .get("vla_data", {}) |
| .get("action_mode", "abs") |
| ) |
| base["action_mode"] = action_mode |
| base["server_composes_relative_eef"] = ( |
| xtrainer2_command.is_relative_action_mode(action_mode) |
| ) |
| return base |
|
|
| @staticmethod |
| def _server_normalizes_state(proc: PolicyNormProcessor) -> bool: |
| """Whether ``state`` arrives raw and must be normalized server-side.""" |
| return bool(getattr(proc._data_config, "server_normalizes_state", False)) |
|
|
| @classmethod |
| def _prepare_examples(cls, examples: List[dict], proc: PolicyNormProcessor) -> List[dict]: |
| """Turn raw-robot observations into the tensors the model was trained on. |
| |
| Leaves examples untouched for embodiments that do not opt in, so clients |
| that already send normalized state keep working. |
| """ |
| if not cls._server_normalizes_state(proc): |
| return examples |
|
|
| to_model_layout = ( |
| xtrainer2_command.robot_state_to_model |
| if cls._uses_xtrainer2_layout(proc) |
| else None |
| ) |
|
|
| prepared: List[dict] = [] |
| for example in examples: |
| if "state" not in example or example["state"] is None: |
| prepared.append(example) |
| continue |
| raw_state = np.asarray(example["state"], dtype=np.float32) |
| if raw_state.shape[-1] != proc.raw_state_dim: |
| raise ValueError( |
| f"state must be the robot's raw {proc.raw_state_dim}-D vector " |
| f"(state_keys={proc.state_keys}); got shape {raw_state.shape}. " |
| "The server applies the training-time normalization itself." |
| ) |
| if to_model_layout is not None: |
| raw_state = to_model_layout(raw_state) |
| example = dict(example) |
| |
| |
| example["_raw_state_wxyz"] = raw_state |
| example["state"] = proc.apply_state(raw_state) |
| prepared.append(example) |
| return prepared |
|
|
| 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``). For embodiments that declare |
| ``server_normalizes_state``, ``state`` is the robot's raw vector |
| and is normalized here. |
| 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) |
| examples = self._prepare_examples(examples, proc) |
|
|
| 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, |
| ) |
| if self._uses_xtrainer2_layout(proc): |
| if xtrainer2_command.is_relative_action_mode( |
| self._model_cfg.get("datasets", {}).get("vla_data", {}).get("action_mode") |
| ): |
| composed = [] |
| for batch_index in range(unnorm.shape[0]): |
| state_wxyz = examples[batch_index].get("_raw_state_wxyz") |
| if state_wxyz is None: |
| raise ValueError( |
| "action_mode=rel requires a raw robot state so relative " |
| "EE commands can be composed back to absolute poses" |
| ) |
| composed.append( |
| xtrainer2_command.compose_relative_eef_wxyz(unnorm[batch_index], state_wxyz) |
| ) |
| unnorm = np.stack(composed, axis=0) |
| unnorm = xtrainer2_command.model_actions_to_robot(unnorm) |
| return {"actions": unnorm} |
|
|
