File size: 8,120 Bytes
1bfc3f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """Streaming corpus statistics for first- and second-order primitive design."""
from __future__ import annotations
import json
from collections.abc import Iterable, Sequence
from dataclasses import asdict, dataclass
from pathlib import Path
import numpy as np
from .boundary import source_differences
from .configuration import ActionCodec2Config
from .trajectory import TrajectoryLike, prepare_trajectory
__all__ = [
"DEFAULT_QUANTILES",
"AxisStatistics",
"CorpusStatistics",
"analyze_trajectories",
]
DEFAULT_QUANTILES = (0.001, 0.01, 0.5, 0.99, 0.999)
class _StreamingAxis:
def __init__(self, sample_size: int, rng: np.random.Generator) -> None:
self.sample_size = sample_size
self.rng = rng
self.count = 0
self.minimum = np.inf
self.maximum = -np.inf
self.mean = 0.0
self.m2 = 0.0
self.sample = np.empty(0, dtype=np.float32)
self.priority = np.empty(0, dtype=np.float64)
def update(self, raw: np.ndarray) -> None:
values = np.asarray(raw, dtype=np.float32).reshape(-1)
if not len(values):
return
batch_count = len(values)
batch_mean = float(values.mean(dtype=np.float64))
batch_m2 = float(np.square(values.astype(np.float64) - batch_mean).sum())
total = self.count + batch_count
delta = batch_mean - self.mean
self.mean += delta * batch_count / total
self.m2 += batch_m2 + delta * delta * self.count * batch_count / total
self.count = total
self.minimum = min(self.minimum, float(values.min()))
self.maximum = max(self.maximum, float(values.max()))
priority = self.rng.random(batch_count)
candidates = np.concatenate((self.sample, values))
candidate_priority = np.concatenate((self.priority, priority))
if len(candidates) > self.sample_size:
keep = np.argpartition(candidate_priority, self.sample_size - 1)[: self.sample_size]
candidates = candidates[keep]
candidate_priority = candidate_priority[keep]
self.sample = candidates
self.priority = candidate_priority
def finish(self, quantiles: Sequence[float]) -> AxisStatistics:
if self.count == 0:
return AxisStatistics(0, None, None, None, None, {}, {})
signed = np.quantile(self.sample, quantiles)
absolute = np.quantile(np.abs(self.sample), quantiles)
return AxisStatistics(
samples=self.count,
minimum=self.minimum,
maximum=self.maximum,
mean=self.mean,
std=float(np.sqrt(self.m2 / self.count)),
quantiles={f"{q:g}": float(value) for q, value in zip(quantiles, signed, strict=True)},
absolute_quantiles={
f"{q:g}": float(value) for q, value in zip(quantiles, absolute, strict=True)
},
)
@dataclass(frozen=True)
class AxisStatistics:
"""Statistics for one global codebook coordinate at one signal order."""
samples: int
minimum: float | None
maximum: float | None
mean: float | None
std: float | None
quantiles: dict[str, float]
absolute_quantiles: dict[str, float]
@dataclass(frozen=True)
class CorpusStatistics:
"""Per-axis input, first-difference and optional second-difference measurements."""
dimension: int
episodes: int
frames: int
active_values: int
binary_dims: list[int]
trajectory_arities: dict[str, int]
origin_max_abs: list[float | None]
order_0: list[AxisStatistics]
order_1: list[AxisStatistics]
quantiles: list[float]
quantile_sample_size: int
order_2: list[AxisStatistics] | None = None
def to_dict(self) -> dict[str, object]:
payload = asdict(self)
payload["schema"] = "actioncodec2_first_order_statistics/1"
if self.order_2 is None:
payload.pop("order_2")
else:
payload["schema"] = "actioncodec2_differential_statistics/2"
return payload
def save(self, path: str | Path) -> None:
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8")
def _validate_options(
dimension: int,
binary_dims: Sequence[int],
quantiles: Sequence[float],
sample_size: int,
) -> tuple[set[int], tuple[float, ...]]:
if dimension < 1:
raise ValueError("dimension must be positive")
binary = {int(value) for value in binary_dims}
if any(value < 0 or value >= dimension for value in binary):
raise ValueError("binary_dims contains a dimension outside the codebook")
checked_quantiles = tuple(float(value) for value in quantiles)
if not checked_quantiles or any(not 0.0 <= value <= 1.0 for value in checked_quantiles):
raise ValueError("quantiles must contain values in [0, 1]")
if sample_size < 1:
raise ValueError("quantile_sample_size must be positive")
return binary, checked_quantiles
def analyze_trajectories(
trajectories: Iterable[TrajectoryLike],
*,
dimension: int,
binary_dims: Sequence[int] = (),
quantiles: Sequence[float] = DEFAULT_QUANTILES,
quantile_sample_size: int = 1_000_000,
seed: int = 0,
config: ActionCodec2Config | None = None,
) -> CorpusStatistics:
'Measure variable-width canonical absolute trajectories in one pass.'
binary, checked_quantiles = _validate_options(
dimension, binary_dims, quantiles, quantile_sample_size
)
rng = np.random.default_rng(seed)
order_0 = [_StreamingAxis(quantile_sample_size, rng) for _ in range(dimension)]
order_1 = [_StreamingAxis(quantile_sample_size, rng) for _ in range(dimension)]
order_2 = (
None
if config is None
else [_StreamingAxis(quantile_sample_size, rng) for _ in range(dimension)]
)
if config is not None and config.physical.input_dim != dimension:
raise ValueError("statistics dimension must match config.physical.input_dim")
origin = np.full(dimension, -np.inf, dtype=np.float64)
arities: dict[int, int] = {}
episodes = frames = active_values = 0
for index, raw in enumerate(trajectories):
values, active_dims = prepare_trajectory(
raw, dimension=dimension, index=index, dtype=np.float32
)
episodes += 1
frames += len(values)
active_values += values.size
arities[len(active_dims)] = arities.get(len(active_dims), 0) + 1
for local, global_dim in enumerate(active_dims):
axis = values[:, local]
order_0[global_dim].update(axis)
origin[global_dim] = max(origin[global_dim], abs(float(axis[0])))
if global_dim not in binary and config is None:
delta = np.diff(axis, prepend=np.float32(0.0))
order_1[global_dim].update(delta)
if config is not None:
for dims, first, second in source_differences(raw, config).values():
for local, global_dim in enumerate(dims):
if global_dim not in binary:
order_1[global_dim].update(first[:, local])
order_2[global_dim].update(second[:, local])
if episodes == 0:
raise ValueError("statistics require at least one complete trajectory")
return CorpusStatistics(
dimension=dimension,
episodes=episodes,
frames=frames,
active_values=active_values,
binary_dims=sorted(binary),
trajectory_arities={str(key): value for key, value in sorted(arities.items())},
origin_max_abs=[None if value == -np.inf else float(value) for value in origin],
order_0=[axis.finish(checked_quantiles) for axis in order_0],
order_1=[axis.finish(checked_quantiles) for axis in order_1],
quantiles=list(checked_quantiles),
quantile_sample_size=quantile_sample_size,
order_2=None if order_2 is None else [axis.finish(checked_quantiles) for axis in order_2],
)
|