| """Offline temporal proposals that preserve StarVLA's dataset/episode marginals.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| from typing import Any, Sequence |
|
|
| import numpy as np |
|
|
| from retime.sampler import ( |
| LogicalDraw, |
| SampleTrace, |
| ValidAnchorManifest, |
| _stateless_seed, |
| ) |
|
|
|
|
| class TemporalProposal: |
| schema_version = 1 |
| valid_score_kinds = { |
| "synthetic", |
| "random", |
| "shuffle", |
| "action_change", |
| "residual", |
| "likelihood", |
| } |
|
|
| def __init__( |
| self, |
| manifest: ValidAnchorManifest, |
| alpha: float, |
| anchor_probabilities: Sequence[Sequence[Sequence[float]]], |
| score_kind: str, |
| scorer_sha256: str, |
| validation_certificate: dict[str, Any] | None = None, |
| ) -> None: |
| if not 0.0 <= alpha < 1.0: |
| raise ValueError("alpha must satisfy 0 <= alpha < 1") |
| if score_kind not in self.valid_score_kinds: |
| raise ValueError(f"unsupported score_kind: {score_kind}") |
| if len(anchor_probabilities) != len(manifest.datasets): |
| raise ValueError("proposal dataset count does not match the manifest") |
|
|
| normalized_probabilities = [] |
| importance_weights = [] |
| normalized_masses = [] |
| for dataset_index, (dataset, dataset_probabilities) in enumerate( |
| zip(manifest.datasets, anchor_probabilities) |
| ): |
| if len(dataset_probabilities) != len(dataset.episodes): |
| raise ValueError( |
| f"proposal episode count differs for dataset {dataset_index}" |
| ) |
| normalized_dataset = [] |
| dataset_weights = [] |
| normalized_dataset_masses = [] |
| for episode_index, (episode, proposal_mass) in enumerate( |
| zip(dataset.episodes, dataset_probabilities) |
| ): |
| mass = np.asarray(proposal_mass, dtype=np.float64) |
| if mass.shape != (len(episode.anchors),): |
| raise ValueError( |
| f"proposal anchor count differs for dataset {dataset_index}, " |
| f"episode {episode_index}" |
| ) |
| if not np.all(np.isfinite(mass)) or np.any(mass < 0) or mass.sum() <= 0: |
| raise ValueError("proposal masses must be finite, non-negative, and non-zero") |
| residual_probability = mass / mass.sum(dtype=np.float64) |
| baseline_probability = np.full(len(mass), 1.0 / len(mass), dtype=np.float64) |
| safe_probability = (1.0 - alpha) * baseline_probability + alpha * residual_probability |
| weights = baseline_probability / safe_probability |
| if not np.isclose(safe_probability.sum(dtype=np.float64), 1.0, atol=1e-12): |
| raise ValueError("safe temporal proposal does not sum to one") |
| if np.max(weights) > 1.0 / (1.0 - alpha) + 1e-12: |
| raise ValueError("importance weight exceeds the analytical safe-mixture bound") |
| normalized_dataset.append(safe_probability) |
| dataset_weights.append(weights) |
| normalized_dataset_masses.append(mass.copy()) |
| normalized_probabilities.append(tuple(normalized_dataset)) |
| importance_weights.append(tuple(dataset_weights)) |
| normalized_masses.append(tuple(normalized_dataset_masses)) |
|
|
| self.manifest_sha256 = manifest.sha256 |
| self.manifest_support = manifest.to_dict(include_hash=False)["datasets"] |
| self.alpha = float(alpha) |
| self.score_kind = score_kind |
| self.scorer_sha256 = scorer_sha256 |
| self.validation_certificate = validation_certificate |
| self.proposal_masses = tuple(normalized_masses) |
| self.anchor_probabilities = tuple(normalized_probabilities) |
| self.importance_weights = tuple(importance_weights) |
| canonical = json.dumps(self.to_dict(include_hash=False), sort_keys=True, separators=(",", ":")) |
| self.sha256 = hashlib.sha256(canonical.encode("utf-8")).hexdigest() |
|
|
| @classmethod |
| def from_masses( |
| cls, |
| manifest: ValidAnchorManifest, |
| alpha: float, |
| masses: Sequence[Sequence[Sequence[float]]], |
| score_kind: str = "synthetic", |
| scorer_sha256: str = "synthetic-fixture", |
| validation_certificate: dict[str, Any] | None = None, |
| ) -> "TemporalProposal": |
| return cls( |
| manifest=manifest, |
| alpha=alpha, |
| anchor_probabilities=masses, |
| score_kind=score_kind, |
| scorer_sha256=scorer_sha256, |
| validation_certificate=validation_certificate, |
| ) |
|
|
| @classmethod |
| def from_dict( |
| cls, manifest: ValidAnchorManifest, payload: dict[str, Any] |
| ) -> "TemporalProposal": |
| if payload.get("schema_version") != cls.schema_version: |
| raise ValueError( |
| f"unsupported temporal proposal schema: {payload.get('schema_version')}" |
| ) |
| expected_support = payload.get("manifest_support") |
| current_support = json.loads( |
| json.dumps(manifest.to_dict(include_hash=False)["datasets"]) |
| ) |
| if expected_support is not None and expected_support != current_support: |
| raise ValueError("temporal proposal changes dataset/episode/anchor support") |
| proposal = cls( |
| manifest=manifest, |
| alpha=float(payload["alpha"]), |
| anchor_probabilities=payload["proposal_masses"], |
| score_kind=str(payload["score_kind"]), |
| scorer_sha256=str(payload["scorer_sha256"]), |
| validation_certificate=payload.get("validation_certificate"), |
| ) |
| return proposal |
|
|
| @classmethod |
| def load( |
| cls, path: str | Path, manifest: ValidAnchorManifest |
| ) -> "TemporalProposal": |
| with open(path, "r", encoding="utf-8") as proposal_file: |
| return cls.from_dict(manifest, json.load(proposal_file)) |
|
|
| def save(self, path: str | Path) -> None: |
| destination = Path(path) |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| temporary = destination.with_suffix(destination.suffix + ".tmp") |
| with open(temporary, "w", encoding="utf-8") as proposal_file: |
| json.dump(self.to_dict(), proposal_file, indent=2) |
| os.replace(temporary, destination) |
|
|
| def sample( |
| self, manifest: ValidAnchorManifest, seed: int, draw: LogicalDraw |
| ) -> SampleTrace: |
| rng = np.random.default_rng(_stateless_seed(seed, draw.epoch, draw.global_draw_id)) |
| dataset_index = int(rng.choice(len(manifest.datasets), p=manifest._dataset_probabilities)) |
| dataset = manifest.datasets[dataset_index] |
| trajectory_index = int( |
| rng.choice( |
| len(dataset.episodes), p=manifest._episode_probabilities[dataset_index] |
| ) |
| ) |
| episode = dataset.episodes[trajectory_index] |
| if self.alpha == 0.0: |
| |
| anchor_index = int(rng.integers(len(episode.anchors))) |
| else: |
| anchor_index = int( |
| rng.choice( |
| len(episode.anchors), |
| p=self.anchor_probabilities[dataset_index][trajectory_index], |
| ) |
| ) |
| baseline_probability = ( |
| dataset.probability |
| * dataset.episode_probabilities[trajectory_index] |
| / len(episode.anchors) |
| ) |
| proposal_probability = ( |
| dataset.probability |
| * dataset.episode_probabilities[trajectory_index] |
| * self.anchor_probabilities[dataset_index][trajectory_index][anchor_index] |
| ) |
| return SampleTrace( |
| manifest_sha256=manifest.sha256, |
| epoch=draw.epoch, |
| global_draw_id=draw.global_draw_id, |
| dataset_index=dataset_index, |
| dataset_name=dataset.name, |
| trajectory_index=episode.trajectory_index, |
| trajectory_id=episode.trajectory_id, |
| anchor_index=anchor_index, |
| base_index=episode.anchors[anchor_index], |
| baseline_probability=float(baseline_probability), |
| proposal_sha256=self.sha256, |
| proposal_probability=float(proposal_probability), |
| importance_weight=float( |
| self.importance_weights[dataset_index][trajectory_index][anchor_index] |
| ), |
| ) |
|
|
| def to_dict(self, include_hash: bool = True) -> dict[str, Any]: |
| payload = { |
| "schema_version": self.schema_version, |
| "manifest_sha256": self.manifest_sha256, |
| "manifest_support": self.manifest_support, |
| "alpha": self.alpha, |
| "score_kind": self.score_kind, |
| "scorer_sha256": self.scorer_sha256, |
| "validation_certificate": self.validation_certificate, |
| "proposal_masses": [ |
| [probabilities.tolist() for probabilities in dataset] |
| for dataset in self.proposal_masses |
| ], |
| } |
| if include_hash: |
| payload["sha256"] = self.sha256 |
| return payload |
|
|
