| from __future__ import annotations |
|
|
| import hashlib |
| from pathlib import Path |
| from typing import Any |
|
|
| try: |
| import numpy as np |
| except ImportError: |
| np = None |
|
|
|
|
| CONTEXT_HASH_WIDTH = 8 |
| OBSERVATION_EMBED_DIM = 32 |
| OBJECT_LAYOUT_EMBED_DIM = 64 |
| CHART_FEATURE_MODES = ( |
| "base", |
| "base_context", |
| "base_context_obs", |
| "base_context_obj", |
| "base_context_obs_obj", |
| ) |
| _EMBEDDING_CACHE: dict[str, Any] = {} |
|
|
|
|
| def build_chart_feature( |
| base_action: Any, |
| metadata: dict[str, Any] | None = None, |
| *, |
| mode: str = "base", |
| ) -> Any: |
| """Build a deployment-visible chart feature vector. |
| |
| `base` preserves the original behavior: the chart token is the flattened |
| base action chunk. `base_context` appends stable hashes of task/language |
| metadata that are visible at proposal time. `base_context_obs` additionally |
| appends a precomputed observation embedding referenced by metadata. |
| `base_context_obj` appends a precomputed RGB object-layout embedding, and |
| `base_context_obs_obj` appends both. These modes intentionally do not read |
| outcomes, labels, hidden chart branches, or evaluator-only fields. |
| """ |
|
|
| if np is None: |
| raise ImportError("build_chart_feature requires numpy") |
| if mode not in CHART_FEATURE_MODES: |
| raise ValueError(f"unknown chart feature mode: {mode}") |
| base = np.asarray(base_action, dtype=np.float32).reshape(-1) |
| if mode == "base": |
| return base |
| metadata = metadata or {} |
| context = np.asarray( |
| [ |
| *_stable_hash_features(str(metadata.get("task_id", "")), CONTEXT_HASH_WIDTH), |
| *_stable_hash_features(str(metadata.get("instruction", "")).lower(), CONTEXT_HASH_WIDTH), |
| min(len(str(metadata.get("instruction", ""))) / 128.0, 4.0), |
| min(len(str(metadata.get("instruction", "")).split()) / 32.0, 4.0), |
| ], |
| dtype=np.float32, |
| ) |
| if mode in {"base_context_obs", "base_context_obs_obj"}: |
| obs = _load_embedding( |
| metadata.get("observation_embedding_path"), |
| dim=OBSERVATION_EMBED_DIM, |
| chart_root=metadata.get("_chart_root"), |
| ) |
| context = np.concatenate([context, obs.astype(np.float32, copy=False)]) |
| if mode in {"base_context_obj", "base_context_obs_obj"}: |
| obj = _load_embedding( |
| metadata.get("object_embedding_path"), |
| dim=OBJECT_LAYOUT_EMBED_DIM, |
| chart_root=metadata.get("_chart_root"), |
| ) |
| context = np.concatenate([context, obj.astype(np.float32, copy=False)]) |
| return np.concatenate([base, context]).astype(np.float32, copy=False) |
|
|
|
|
| def chart_feature_dim(base_action: Any, *, mode: str = "base") -> int: |
| return int(build_chart_feature(base_action, {}, mode=mode).reshape(-1).shape[0]) |
|
|
|
|
| def _stable_hash_features(text: str, width: int) -> list[float]: |
| digest = hashlib.sha256(text.encode("utf-8")).digest() |
| return [float(digest[index] / 127.5 - 1.0) for index in range(width)] |
|
|
|
|
| def _load_embedding(value: Any, *, dim: int, chart_root: Any = None) -> Any: |
| if np is None: |
| raise ImportError("build_chart_feature requires numpy") |
| if not value: |
| return np.zeros(dim, dtype=np.float32) |
| path_text, dataset, row_index = _parse_embedding_ref(str(value)) |
| path = Path(path_text) |
| if not path.is_absolute() and chart_root: |
| path = Path(str(chart_root)) / path |
| cache_key = str(path.resolve()) |
| if cache_key not in _EMBEDDING_CACHE: |
| with np.load(path, allow_pickle=False) as data: |
| _EMBEDDING_CACHE[cache_key] = np.asarray(data[dataset], dtype=np.float32) |
| matrix = _EMBEDDING_CACHE[cache_key] |
| vector = np.asarray(matrix[int(row_index)], dtype=np.float32).reshape(-1) |
| if vector.shape[0] == dim: |
| return vector |
| output = np.zeros(dim, dtype=np.float32) |
| width = min(dim, vector.shape[0]) |
| output[:width] = vector[:width] |
| return output |
|
|
|
|
| def _parse_embedding_ref(value: str) -> tuple[str, str, int]: |
| if "#" not in value: |
| return value, "embeddings", 0 |
| path_text, ref = value.split("#", 1) |
| parts = [part for part in ref.split("/") if part] |
| if len(parts) != 2: |
| raise ValueError(f"invalid observation embedding ref: {value}") |
| return path_text, parts[0], int(parts[1]) |
|
|
