| """Deterministic hierarchical sampling primitives for InfoFlow experiments.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| from typing import Any, Iterator, Sequence |
|
|
| import numpy as np |
| from torch.utils.data import Sampler |
|
|
|
|
| def _identity_files(dataset_root: Path) -> list[Path]: |
| files = list((dataset_root / "data").rglob("*.parquet")) |
| files.extend((dataset_root / "videos").rglob("*.mp4")) |
| meta_root = dataset_root / "meta" |
| if meta_root.exists(): |
| for pattern in ("*.json", "*.jsonl", "*.parquet"): |
| files.extend(meta_root.rglob(pattern)) |
| return sorted(set(files), key=lambda path: path.relative_to(dataset_root).as_posix()) |
|
|
|
|
| def compute_dataset_content_fingerprint(dataset_root: str | Path) -> dict[str, Any]: |
| """Record the immutable LeRobot data and metadata file inventory.""" |
| root = Path(dataset_root) |
| entries = [] |
| total_bytes = 0 |
| for file_path in _identity_files(root): |
| relative_path = file_path.relative_to(root).as_posix() |
| file_size = file_path.stat().st_size |
| entries.append({"path": relative_path, "bytes": file_size}) |
| total_bytes += file_size |
| return { |
| "file_count": len(entries), |
| "total_bytes": total_bytes, |
| "files": entries, |
| } |
|
|
|
|
| def validate_dataset_content_inventory( |
| dataset_root: str | Path, fingerprint: dict[str, Any] |
| ) -> None: |
| """Check that the current file inventory matches the recorded dataset inventory.""" |
| root = Path(dataset_root) |
| current_entries = [ |
| { |
| "path": path.relative_to(root).as_posix(), |
| "bytes": path.stat().st_size, |
| } |
| for path in _identity_files(root) |
| ] |
| if current_entries != fingerprint.get("files"): |
| raise RuntimeError( |
| f"Dataset file inventory differs from its content fingerprint: {root}" |
| ) |
| if len(current_entries) != fingerprint.get("file_count"): |
| raise RuntimeError("Dataset content fingerprint has an inconsistent file count") |
| if sum(entry["bytes"] for entry in current_entries) != fingerprint.get("total_bytes"): |
| raise RuntimeError("Dataset content fingerprint has an inconsistent byte count") |
|
|
|
|
| def _normalize_probabilities(values: Sequence[float], name: str) -> tuple[float, ...]: |
| probabilities = np.asarray(values, dtype=np.float64) |
| if probabilities.ndim != 1 or probabilities.size == 0: |
| raise ValueError(f"{name} must be a non-empty one-dimensional sequence") |
| if not np.all(np.isfinite(probabilities)) or np.any(probabilities <= 0): |
| raise ValueError(f"{name} must contain finite, strictly positive values") |
| probabilities /= probabilities.sum(dtype=np.float64) |
| return tuple(float(value) for value in probabilities) |
|
|
|
|
| def _stateless_seed(seed: int, epoch: int, global_draw_id: int) -> int: |
| payload = f"{int(seed)}:{int(epoch)}:{int(global_draw_id)}".encode("utf-8") |
| return int.from_bytes(hashlib.sha256(payload).digest()[:8], "little", signed=False) |
|
|
|
|
| @dataclass(frozen=True) |
| class EpisodeSupport: |
| trajectory_index: int |
| trajectory_id: int | str |
| anchors: tuple[int, ...] |
|
|
| def __post_init__(self) -> None: |
| if not self.anchors: |
| raise ValueError("episode support must contain at least one anchor") |
| if len(set(self.anchors)) != len(self.anchors): |
| raise ValueError("episode anchors must be unique") |
|
|
|
|
| @dataclass(frozen=True) |
| class DatasetSupport: |
| name: str |
| probability: float |
| episodes: tuple[EpisodeSupport, ...] |
| episode_probabilities: tuple[float, ...] |
|
|
| def __post_init__(self) -> None: |
| if not self.episodes: |
| raise ValueError(f"dataset {self.name!r} has no valid episodes") |
| if len(self.episodes) != len(self.episode_probabilities): |
| raise ValueError("episode support and probability lengths differ") |
| normalized = _normalize_probabilities(self.episode_probabilities, "episode_probabilities") |
| if not np.allclose(normalized, self.episode_probabilities, rtol=0, atol=1e-12): |
| raise ValueError("episode probabilities must already sum to one") |
|
|
|
|
| @dataclass(frozen=True) |
| class LogicalDraw: |
| epoch: int |
| global_draw_id: int |
|
|
| def __post_init__(self) -> None: |
| if self.epoch < 0 or self.global_draw_id < 0: |
| raise ValueError("epoch and global_draw_id must be non-negative") |
|
|
|
|
| @dataclass(frozen=True) |
| class SampleTrace: |
| manifest_sha256: str |
| epoch: int |
| global_draw_id: int |
| dataset_index: int |
| dataset_name: str |
| trajectory_index: int |
| trajectory_id: int | str |
| anchor_index: int |
| base_index: int |
| baseline_probability: float |
| proposal_sha256: str | None = None |
| proposal_probability: float | None = None |
| importance_weight: float = 1.0 |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return asdict(self) |
|
|
|
|
| class ValidAnchorManifest: |
| schema_version = 1 |
|
|
| def __init__(self, datasets: Sequence[DatasetSupport], identity: dict[str, Any] | None = None): |
| if not datasets: |
| raise ValueError("manifest must contain at least one dataset") |
| dataset_probabilities = _normalize_probabilities( |
| [dataset.probability for dataset in datasets], "dataset_probabilities" |
| ) |
| if not np.allclose( |
| dataset_probabilities, |
| [dataset.probability for dataset in datasets], |
| rtol=0, |
| atol=1e-12, |
| ): |
| raise ValueError("dataset probabilities must already sum to one") |
| self.datasets = tuple(datasets) |
| self.identity = identity or {} |
| self._dataset_probabilities = np.asarray( |
| [dataset.probability for dataset in self.datasets], dtype=np.float64 |
| ) |
| self._episode_probabilities = tuple( |
| np.asarray(dataset.episode_probabilities, dtype=np.float64) |
| for dataset in self.datasets |
| ) |
| 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_hierarchy( |
| cls, |
| dataset_names: Sequence[str], |
| dataset_weights: Sequence[float], |
| episodes: Sequence[Sequence[tuple[int | str, Sequence[int], float]]], |
| identity: dict[str, Any] | None = None, |
| ) -> "ValidAnchorManifest": |
| if len(dataset_names) != len(dataset_weights) or len(dataset_names) != len(episodes): |
| raise ValueError("dataset names, weights, and episode supports must have equal lengths") |
| dataset_probabilities = _normalize_probabilities(dataset_weights, "dataset_weights") |
| supports = [] |
| for dataset_index, (name, dataset_episodes) in enumerate(zip(dataset_names, episodes)): |
| if not dataset_episodes: |
| raise ValueError(f"dataset {name!r} has no valid episodes") |
| episode_probabilities = _normalize_probabilities( |
| [entry[2] for entry in dataset_episodes], f"episode_weights[{dataset_index}]" |
| ) |
| episode_supports = tuple( |
| EpisodeSupport( |
| trajectory_index=trajectory_index, |
| trajectory_id=trajectory_id, |
| anchors=tuple(int(anchor) for anchor in anchors), |
| ) |
| for trajectory_index, (trajectory_id, anchors, _) in enumerate(dataset_episodes) |
| ) |
| supports.append( |
| DatasetSupport( |
| name=str(name), |
| probability=dataset_probabilities[dataset_index], |
| episodes=episode_supports, |
| episode_probabilities=episode_probabilities, |
| ) |
| ) |
| return cls(supports, identity=identity) |
|
|
| @classmethod |
| def from_dict(cls, payload: dict[str, Any]) -> "ValidAnchorManifest": |
| datasets = [] |
| for dataset_payload in payload["datasets"]: |
| episodes = tuple( |
| EpisodeSupport( |
| trajectory_index=int(episode["trajectory_index"]), |
| trajectory_id=episode["trajectory_id"], |
| anchors=tuple(int(anchor) for anchor in episode["anchors"]), |
| ) |
| for episode in dataset_payload["episodes"] |
| ) |
| datasets.append( |
| DatasetSupport( |
| name=str(dataset_payload["name"]), |
| probability=float(dataset_payload["probability"]), |
| episodes=episodes, |
| episode_probabilities=tuple( |
| float(value) for value in dataset_payload["episode_probabilities"] |
| ), |
| ) |
| ) |
| return cls(datasets, identity=payload.get("identity", {})) |
|
|
| @classmethod |
| def load(cls, path: str | Path) -> "ValidAnchorManifest": |
| with open(path, "r", encoding="utf-8") as manifest_file: |
| return cls.from_dict(json.load(manifest_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 manifest_file: |
| json.dump(self.to_dict(), manifest_file, indent=2) |
| os.replace(temporary, destination) |
|
|
| def sample(self, 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(self.datasets), p=self._dataset_probabilities)) |
| dataset = self.datasets[dataset_index] |
| trajectory_index = int( |
| rng.choice(len(dataset.episodes), p=self._episode_probabilities[dataset_index]) |
| ) |
| episode = dataset.episodes[trajectory_index] |
| anchor_index = int(rng.integers(len(episode.anchors))) |
| baseline_probability = ( |
| dataset.probability |
| * dataset.episode_probabilities[trajectory_index] |
| / len(episode.anchors) |
| ) |
| return SampleTrace( |
| manifest_sha256=self.sha256, |
| epoch=int(draw.epoch), |
| global_draw_id=int(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), |
| ) |
|
|
| def to_dict(self, include_hash: bool = True) -> dict[str, Any]: |
| payload = { |
| "schema_version": self.schema_version, |
| "identity": self.identity, |
| "datasets": [asdict(dataset) for dataset in self.datasets], |
| } |
| if include_hash: |
| payload["sha256"] = self.sha256 |
| return payload |
|
|
|
|
| class StatelessDrawSampler(Sampler[LogicalDraw]): |
| """Yield rank-disjoint logical draw ids; sampling itself stays inside the dataset.""" |
|
|
| def __init__( |
| self, |
| num_local_draws: int, |
| rank: int = 0, |
| world_size: int = 1, |
| start_global_draw_id: int = 0, |
| ) -> None: |
| if num_local_draws < 0: |
| raise ValueError("num_local_draws must be non-negative") |
| if world_size <= 0 or rank < 0 or rank >= world_size: |
| raise ValueError("rank must satisfy 0 <= rank < world_size") |
| self.num_local_draws = int(num_local_draws) |
| self.rank = int(rank) |
| self.world_size = int(world_size) |
| self.start_global_draw_id = int(start_global_draw_id) |
| if self.start_global_draw_id < 0: |
| raise ValueError("start_global_draw_id must be non-negative") |
| self.epoch = 0 |
| self._resume_epoch = 0 |
|
|
| def set_resume_position(self, epoch: int, start_global_draw_id: int) -> None: |
| """Resume at an inclusive global draw cursor within one epoch.""" |
| if epoch < 0 or start_global_draw_id < 0: |
| raise ValueError("resume epoch and draw id must be non-negative") |
| self.epoch = int(epoch) |
| self._resume_epoch = int(epoch) |
| self.start_global_draw_id = int(start_global_draw_id) |
|
|
| def set_epoch(self, epoch: int) -> None: |
| if epoch < 0: |
| raise ValueError("epoch must be non-negative") |
| self.epoch = int(epoch) |
| if self.epoch != self._resume_epoch: |
| self.start_global_draw_id = 0 |
|
|
| def _first_local_draw_id(self) -> int: |
| return max( |
| 0, |
| (self.start_global_draw_id - self.rank + self.world_size - 1) |
| // self.world_size, |
| ) |
|
|
| def __iter__(self) -> Iterator[LogicalDraw]: |
| for local_draw_id in range(self._first_local_draw_id(), self.num_local_draws): |
| yield LogicalDraw( |
| epoch=self.epoch, |
| global_draw_id=local_draw_id * self.world_size + self.rank, |
| ) |
|
|
| def __len__(self) -> int: |
| return max(0, self.num_local_draws - self._first_local_draw_id()) |
|
|
