Spaces:
Running on Zero
Running on Zero
| # SPDX-License-Identifier: GPL-3.0-only | |
| # Copyright (C) 2026 SLING AI Inc. | |
| """ | |
| selector.py -- Stage 2 of the Goal-Conditioned VPA framework (preprint_261008.pdf, Sec. 4.2). | |
| Implements | |
| * ``NeuroSymbolicSelector`` -- the MLP classifier pi_phi^h over the primitive dictionary | |
| U = {u^(1), ..., u^(N_u)} (Eq. 8): | |
| u_t = argmax_{u in U} pi_phi^h(u | z_t, z_g^(m_t), c_text), | |
| trained with cross-entropy against demonstration primitive labels. | |
| * ``MilestoneTracker`` -- the symbolic milestone pointer m_t and its variance-scaled | |
| completion threshold (Eqs. 9a-9c): | |
| m_{t+1} = m_t + 1 if ||z_t - z_g^(m_t)||_2 < tau and m_t < M, else m_t (9a) | |
| gamma_n = (1 - a_g) gamma_{n-1} + a_g * sqrt( (1/d) sum_j Var(Z^(n)_{.,j}) + eps ) (9b) | |
| tau_n = kappa * sqrt(2d) * gamma_n (9c) | |
| Once E_psi is frozen, gamma is fixed to gamma* and deployment uses tau* = kappa sqrt(2d) gamma*. | |
| Indexing convention: the pointer m_t is stored 1-based exactly as in the paper (m_t in {1..M}); | |
| it is converted to a 0-based index only when gathering z_g^(m_t) from a ``[B, M, d]`` tensor. | |
| This file is self-contained (it imports nothing from the other VPA modules) and can be tested | |
| on its own with ``python selector.py``. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from typing import NamedTuple, Optional, Sequence, Tuple | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch import Tensor | |
| __all__ = ["NeuroSymbolicSelector", "MilestoneTracker", "MilestoneUpdate"] | |
| # --------------------------------------------------------------------------------------------- | |
| # pi_phi^h : neuro-symbolic primitive selector (Eq. 8) | |
| # --------------------------------------------------------------------------------------------- | |
| class NeuroSymbolicSelector(nn.Module): | |
| """MLP classifier pi_phi^h mapping Concat(z_t, z_g^(m_t), c_text) to logits over U (Eq. 8). | |
| The per-decision cost is a single MLP evaluation, independent of the number of milestones M, | |
| because only the *active* milestone latent z_g^(m_t) is fed to the network. | |
| Args: | |
| latent_dim: d, dimension of z_t and z_g^(m). | |
| text_dim: d_c, dimension of c_text. | |
| num_primitives: N_u = |U|. | |
| hidden_dims: widths of the hidden layers of the MLP. | |
| Shapes: | |
| ``forward`` / ``probabilities``: (``[B, d]``, ``[B, d]``, ``[B, d_c]``) -> ``[B, N_u]`` | |
| ``select``: (``[B, d]``, ``[B, d]``, ``[B, d_c]``) -> ``[B]`` (int64) | |
| """ | |
| def __init__( | |
| self, | |
| latent_dim: int, | |
| text_dim: int, | |
| num_primitives: int, | |
| hidden_dims: Sequence[int] = (512, 512), | |
| ) -> None: | |
| super().__init__() | |
| if num_primitives < 2: | |
| raise ValueError("the primitive dictionary must contain at least two primitives") | |
| self.latent_dim: int = latent_dim | |
| self.text_dim: int = text_dim | |
| self.num_primitives: int = num_primitives | |
| layers = [] | |
| in_features = 2 * latent_dim + text_dim | |
| for width in hidden_dims: | |
| layers += [nn.Linear(in_features, width), nn.GELU()] | |
| in_features = width | |
| layers.append(nn.Linear(in_features, num_primitives)) | |
| self.mlp = nn.Sequential(*layers) | |
| def forward(self, z_t: Tensor, z_goal: Tensor, c_text: Tensor) -> Tensor: | |
| """Unnormalised scores of pi_phi^h(u | z_t, z_g^(m_t), c_text). | |
| Args: | |
| z_t: ``[B, d]`` current latent. | |
| z_goal: ``[B, d]`` active milestone latent z_g^(m_t). | |
| c_text: ``[B, d_c]`` instruction embedding. | |
| Returns: | |
| logits: ``[B, N_u]``. | |
| """ | |
| if z_t.shape != z_goal.shape or z_t.shape[-1] != self.latent_dim: | |
| raise ValueError(f"expected z_t and z_goal of shape [B, {self.latent_dim}]") | |
| if c_text.shape != (z_t.shape[0], self.text_dim): | |
| raise ValueError(f"expected c_text of shape [B, {self.text_dim}], got {tuple(c_text.shape)}") | |
| return self.mlp(torch.cat([z_t, z_goal, c_text], dim=-1)) | |
| def probabilities(self, z_t: Tensor, z_goal: Tensor, c_text: Tensor) -> Tensor: | |
| """Categorical distribution pi_phi^h(. | z_t, z_g^(m_t), c_text): ``[B, N_u]``.""" | |
| return F.softmax(self(z_t, z_goal, c_text), dim=-1) | |
| def select(self, z_t: Tensor, z_goal: Tensor, c_text: Tensor) -> Tensor: | |
| """Eq. 8: u_t = argmax_u pi_phi^h(u | ...). Returns int64 tokens ``[B]``. | |
| Softmax is strictly monotone, so the argmax of the logits equals the argmax of pi_phi^h. | |
| The network is invoked through ``__call__`` so that forward hooks (used by the pipeline to | |
| audit the sequential depth) observe exactly one evaluation. | |
| """ | |
| return self(z_t, z_goal, c_text).argmax(dim=-1) | |
| def loss(logits: Tensor, primitive_labels: Tensor) -> Tensor: | |
| """Cross-entropy between pi_phi^h and demonstration primitive labels (Sec. 4.2). | |
| Args: | |
| logits: ``[B, N_u]``. | |
| primitive_labels: ``[B]`` int64 in {0, ..., N_u - 1}. | |
| Returns: | |
| scalar loss (mean over the batch). | |
| """ | |
| return F.cross_entropy(logits, primitive_labels) | |
| # --------------------------------------------------------------------------------------------- | |
| # Milestone pointer m_t and variance-scaled threshold tau (Eqs. 9a-9c) | |
| # --------------------------------------------------------------------------------------------- | |
| class MilestoneUpdate(NamedTuple): | |
| """Result of one application of Eq. 9a (all tensors have leading dimension B).""" | |
| m_t: Tensor # [B] int64, pointer used at this decision step (1-based) | |
| m_next: Tensor # [B] int64, m_{t+1} | |
| distance: Tensor # [B] float, ||z_t - z_g^(m_t)||_2 | |
| reached: Tensor # [B] bool, distance < tau | |
| task_complete: Tensor # [B] bool, sticky flag: test satisfied with m_t = M | |
| class MilestoneTracker(nn.Module): | |
| """Symbolic finite-state machine over milestones with a variance-calibrated threshold. | |
| Training-time calibration (Eqs. 9b, 9c): | |
| ``update_variance_margin(Z)`` is called once per training iteration n with the batch | |
| Z^(n) in R^{B x d} of *online*-encoder latents. Var is the unbiased batch estimator | |
| (1 / (B - 1)), i.e. the same estimator as the covariance C(Z) of Eq. 19. | |
| ``threshold`` then returns tau_n = kappa * sqrt(2d) * gamma_n. | |
| The paper does not specify gamma_0. If ``gamma_bar_init`` is None, gamma_0 is set to the | |
| first instantaneous estimate (so gamma_1 equals that estimate); otherwise gamma_0 = | |
| ``gamma_bar_init``. | |
| Deployment (after E_psi is frozen, Sec. 4.2/4.3): | |
| ``freeze()`` fixes gamma* (and therefore tau* = kappa sqrt(2d) gamma*). ``gamma_bar`` is a | |
| registered buffer, so gamma* is saved and restored with ``state_dict``. | |
| Episode state (not part of ``state_dict``): | |
| ``m`` -- ``[B]`` int64 pointer in {1, ..., M}, monotonically non-decreasing; | |
| ``task_complete`` -- ``[B]`` bool, set once Eq. 9a's test holds with m_t = M. | |
| Args: | |
| latent_dim: d. | |
| kappa: dimensionless tolerance kappa in (0, 1). | |
| alpha_gamma: EMA rate alpha_gamma in (0, 1]. | |
| eps: numerical constant eps of Eq. 18 (must be the same value used by the VICReg loss). | |
| gamma_bar_init: optional gamma_0. | |
| """ | |
| def __init__( | |
| self, | |
| latent_dim: int, | |
| kappa: float, | |
| alpha_gamma: float, | |
| eps: float, | |
| gamma_bar_init: Optional[float] = None, | |
| ) -> None: | |
| super().__init__() | |
| if not 0.0 < kappa < 1.0: | |
| raise ValueError("kappa must lie in (0, 1)") | |
| if not 0.0 < alpha_gamma <= 1.0: | |
| raise ValueError("alpha_gamma must lie in (0, 1]") | |
| if eps <= 0.0: | |
| raise ValueError("eps must be positive") | |
| self.latent_dim: int = int(latent_dim) | |
| self.kappa: float = float(kappa) | |
| self.alpha_gamma: float = float(alpha_gamma) | |
| self.eps: float = float(eps) | |
| init = float("nan") if gamma_bar_init is None else float(gamma_bar_init) | |
| if gamma_bar_init is not None and gamma_bar_init <= 0.0: | |
| raise ValueError("gamma_bar_init must be positive") | |
| # float32 on purpose: the Apple MPS backend has no float64, and the buffer must move with .to("mps"). | |
| self.register_buffer("gamma_bar", torch.tensor(init, dtype=torch.float32)) | |
| self.register_buffer("frozen", torch.tensor(False)) | |
| # Episode state (runtime only). | |
| self.m: Optional[Tensor] = None | |
| self.task_complete: Optional[Tensor] = None | |
| self.num_milestones: Optional[int] = None | |
| # ---------------- threshold calibration (Eqs. 9b, 9c) ---------------- | |
| def instantaneous_margin(self, latents: Tensor) -> Tensor: | |
| """sqrt( (1/d) * sum_j Var(Z_{.,j}) + eps ) for a batch ``Z: [B, d]`` -> scalar tensor.""" | |
| if latents.dim() != 2 or latents.shape[1] != self.latent_dim: | |
| raise ValueError(f"expected latents of shape [B, {self.latent_dim}], got {tuple(latents.shape)}") | |
| if latents.shape[0] < 2: | |
| raise ValueError("the unbiased variance estimator requires B >= 2") | |
| var = latents.detach().float().var(dim=0, unbiased=True) # [d] (float32: MPS-safe) | |
| return torch.sqrt(var.mean() + self.eps) | |
| def update_variance_margin(self, latents: Tensor) -> float: | |
| """Apply Eq. 9b with the online-encoder batch ``Z^(n): [B, d]``; returns gamma_n.""" | |
| if bool(self.frozen): | |
| raise RuntimeError("gamma is frozen (E_psi has been frozen); it can no longer be updated") | |
| inst = self.instantaneous_margin(latents).to(self.gamma_bar.device) | |
| if torch.isnan(self.gamma_bar): | |
| self.gamma_bar.copy_(inst) # gamma_0 := first estimate | |
| self.gamma_bar.copy_((1.0 - self.alpha_gamma) * self.gamma_bar + self.alpha_gamma * inst) | |
| return float(self.gamma_bar) | |
| def is_calibrated(self) -> bool: | |
| return not bool(torch.isnan(self.gamma_bar)) | |
| def threshold(self) -> float: | |
| """tau = kappa * sqrt(2d) * gamma (tau_n during training, tau* once frozen; Eq. 9c).""" | |
| if not self.is_calibrated: | |
| raise RuntimeError("gamma has not been estimated yet; call update_variance_margin or freeze") | |
| return self.kappa * math.sqrt(2.0 * self.latent_dim) * float(self.gamma_bar) | |
| def freeze(self, gamma_bar_star: Optional[float] = None) -> float: | |
| """Fix gamma* (optionally overriding it) for deployment; returns tau*.""" | |
| if gamma_bar_star is not None: | |
| if gamma_bar_star <= 0.0: | |
| raise ValueError("gamma_bar_star must be positive") | |
| self.gamma_bar.fill_(float(gamma_bar_star)) | |
| if not self.is_calibrated: | |
| raise RuntimeError("cannot freeze an uncalibrated tracker; provide gamma_bar_star") | |
| self.frozen.fill_(True) | |
| return self.threshold | |
| def gamma_bar_star(self) -> float: | |
| """The frozen latent scale gamma* (Eq. 9b after freezing).""" | |
| if not bool(self.frozen): | |
| raise RuntimeError("gamma* is only defined after freeze()") | |
| return float(self.gamma_bar) | |
| # ---------------- pointer logic (Eq. 9a) ---------------- | |
| def reset(self, num_milestones: int, batch_size: int, device: Optional[torch.device] = None) -> None: | |
| """Start an episode: m_0 = 1 for every batch element.""" | |
| if num_milestones < 1: | |
| raise ValueError("M must be at least 1") | |
| self.num_milestones = int(num_milestones) | |
| self.m = torch.ones(batch_size, dtype=torch.long, device=device) | |
| self.task_complete = torch.zeros(batch_size, dtype=torch.bool, device=device) | |
| def _require_episode(self) -> Tuple[Tensor, Tensor, int]: | |
| if self.m is None or self.task_complete is None or self.num_milestones is None: | |
| raise RuntimeError("call reset() before using the milestone pointer") | |
| return self.m, self.task_complete, self.num_milestones | |
| def active_goal(self, milestone_latents: Tensor) -> Tensor: | |
| """Gather z_g^(m_t) from ``[B, M, d]`` -> ``[B, d]``.""" | |
| m, _, num_m = self._require_episode() | |
| if milestone_latents.dim() != 3 or milestone_latents.shape[1] != num_m: | |
| raise ValueError(f"expected milestone latents of shape [B, {num_m}, d]") | |
| index = (m - 1).view(-1, 1, 1).expand(-1, 1, milestone_latents.shape[-1]) # [B, 1, d] | |
| return milestone_latents.gather(1, index).squeeze(1) # [B, d] | |
| def completion_test(self, z_t: Tensor, milestone_latents: Tensor) -> Tuple[Tensor, Tensor]: | |
| """Threshold predicate ||z_t - z_g^(m_t)||_2 < tau. | |
| Args: | |
| z_t: ``[B, d]``; milestone_latents: ``[B, M, d]``. | |
| Returns: | |
| distance ``[B]`` and reached ``[B]`` (bool). | |
| """ | |
| distance = torch.linalg.vector_norm(z_t - self.active_goal(milestone_latents), ord=2, dim=-1) | |
| return distance, distance < self.threshold | |
| def step( | |
| self, | |
| z_t: Tensor, | |
| milestone_latents: Tensor, | |
| update_mask: Optional[Tensor] = None, | |
| ) -> MilestoneUpdate: | |
| """Apply Eq. 9a and commit m_{t+1} as the new pointer. | |
| Args: | |
| z_t: ``[B, d]`` current latent. | |
| milestone_latents: ``[B, M, d]``. | |
| update_mask: optional ``[B]`` bool; only elements that are at a decision step are | |
| updated (the others keep m_t and their completion flag). | |
| Returns: | |
| ``MilestoneUpdate`` with m_t, m_{t+1}, distance, reached and the task-complete flag. | |
| """ | |
| m, done, num_m = self._require_episode() | |
| distance, reached = self.completion_test(z_t, milestone_latents) | |
| advance = reached & (m < num_m) | |
| finished = reached & (m == num_m) | |
| if update_mask is not None: | |
| update_mask = update_mask.to(device=m.device, dtype=torch.bool) | |
| advance = advance & update_mask | |
| finished = finished & update_mask | |
| m_t = m.clone() | |
| self.m = m + advance.long() | |
| self.task_complete = done | finished | |
| return MilestoneUpdate(m_t, self.m.clone(), distance, reached, self.task_complete.clone()) | |
| # --------------------------------------------------------------------------------------------- | |
| # Self-test (run: python selector.py) | |
| # --------------------------------------------------------------------------------------------- | |
| _RTOL, _ATOL = 1e-5, 1e-6 # float32 module output vs float64 reference (CLAUDE.md) | |
| def _close(actual: Tensor, ref: Tensor) -> bool: | |
| return torch.allclose(actual.double(), ref.double(), rtol=_RTOL, atol=_ATOL) | |
| def _close_scalar(actual: float, ref: float) -> bool: | |
| return abs(actual - ref) <= _ATOL + _RTOL * abs(ref) | |
| def _expect(exc: type, fn) -> None: | |
| try: | |
| fn() | |
| except exc: | |
| return | |
| raise AssertionError(f"expected {exc.__name__}") | |
| def _ref_selector_logits(sel: NeuroSymbolicSelector, z_t: Tensor, z_g: Tensor, c: Tensor) -> Tensor: | |
| """Float64 reference of the Eq. 8 MLP: input slots written by hand, Linear/exact-GELU by hand.""" | |
| b, d, dc = z_t.shape[0], z_t.shape[1], c.shape[1] | |
| x = torch.zeros(b, 2 * d + dc, dtype=torch.float64) | |
| x[:, :d] = z_t.double() | |
| x[:, d:2 * d] = z_g.double() | |
| x[:, 2 * d:] = c.double() | |
| for layer in sel.mlp: | |
| if isinstance(layer, nn.Linear): | |
| x = x @ layer.weight.detach().double().T + layer.bias.detach().double() | |
| else: | |
| assert isinstance(layer, nn.GELU) and layer.approximate == "none" | |
| x = 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0))) | |
| return x # [B, N_u] | |
| def _ref_softmax(s: Tensor) -> Tensor: | |
| e = torch.exp(s - s.max(dim=-1, keepdim=True).values) | |
| return e / e.sum(dim=-1, keepdim=True) | |
| def _ref_margin(z: Tensor, eps: float) -> float: | |
| """sqrt((1/d) sum_j Var_unbiased(Z_.j) + eps) with explicit float64 sums (Eq. 9b).""" | |
| zd = z.double() | |
| bsz, d = zd.shape | |
| total = 0.0 | |
| for j in range(d): | |
| col = [float(zd[i, j]) for i in range(bsz)] | |
| mean = sum(col) / bsz | |
| total += sum((v - mean) ** 2 for v in col) / (bsz - 1) | |
| return math.sqrt(total / d + eps) | |
| def _ref_pointer_step(m, done, z, goals, tau, num_m, mask): | |
| """Per-element float64 reference of Eq. 9a (lists in, lists out).""" | |
| m_next, done_next, dist, reached = list(m), list(done), [], [] | |
| for i in range(len(m)): | |
| g = goals[i, m[i] - 1].double() | |
| dist_i = math.sqrt(float(((z[i].double() - g) ** 2).sum())) | |
| hit = dist_i < tau | |
| dist.append(dist_i) | |
| reached.append(hit) | |
| if mask[i]: | |
| if hit and m[i] < num_m: | |
| m_next[i] = m[i] + 1 | |
| if hit and m[i] == num_m: | |
| done_next[i] = True | |
| return m_next, done_next, dist, reached | |
| def _self_test() -> None: | |
| torch.manual_seed(0) | |
| gen = torch.Generator().manual_seed(2024) | |
| b, d, dc, nu, num_m = 4, 8, 6, 5, 3 | |
| # ---- Eq. 8: u_t = argmax pi_phi^h(u | z_t, z_g^(m_t), c_text). | |
| sel = NeuroSymbolicSelector(d, dc, nu, hidden_dims=(32, 32)) | |
| z_t, z_g = torch.randn(b, d, generator=gen), torch.randn(b, d, generator=gen) | |
| c = torch.randn(b, dc, generator=gen) | |
| captured = [] | |
| pre = sel.mlp.register_forward_pre_hook(lambda _m, args: captured.append(args[0].detach().clone())) | |
| logits = sel(z_t, z_g, c) | |
| pre.remove() | |
| x_in = captured[0] | |
| assert x_in.shape == (b, 2 * d + dc) | |
| assert torch.equal(x_in[:, :d], z_t) and torch.equal(x_in[:, d:2 * d], z_g) | |
| assert torch.equal(x_in[:, 2 * d:], c) | |
| ref_logits = _ref_selector_logits(sel, z_t, z_g, c) | |
| ref_probs = _ref_softmax(ref_logits) | |
| assert logits.shape == (b, nu) and _close(logits, ref_logits) | |
| probs = sel.probabilities(z_t, z_g, c) | |
| assert _close(probs, ref_probs) and _close(probs.sum(-1), torch.ones(b)) | |
| top2 = ref_probs.topk(2, dim=-1).values | |
| assert bool(((top2[:, 0] - top2[:, 1]) > 1e-4).all()) # no near-ties: argmax is well defined | |
| calls = [] | |
| hook = sel.register_forward_hook(lambda *_: calls.append(1)) | |
| u = sel.select(z_t, z_g, c) | |
| hook.remove() | |
| assert len(calls) == 1 # one evaluation through __call__ (depth audit) | |
| assert u.shape == (b,) and u.dtype == torch.long | |
| assert torch.equal(u, probs.argmax(-1)) and torch.equal(u, ref_probs.argmax(-1)) | |
| sel.loss(logits, torch.randint(0, nu, (b,), generator=gen)).backward() | |
| _expect(ValueError, lambda: sel(z_t, z_g[:, :-1], c)) | |
| _expect(ValueError, lambda: sel(z_t[:, :-1], z_g[:, :-1], c)) | |
| _expect(ValueError, lambda: sel(z_t, z_g, c[:, :-1])) | |
| _expect(ValueError, lambda: sel(z_t, z_g, c[:-1])) | |
| _expect(ValueError, lambda: NeuroSymbolicSelector(d, dc, 1)) | |
| # ---- Eq. 9b / 9c: EMA of the variance margin and tau = kappa sqrt(2d) gamma. | |
| kappa, eps = 0.2, 1e-4 | |
| sizes = [2, 3, 5, 16, 64, 2, 7, 32, 4, 128, 9, 2] # includes the minimum B = 2 | |
| batches = [torch.randn(n, d, generator=gen) * (0.5 + 3.0 * float(torch.rand(1, generator=gen))) | |
| for n in sizes] | |
| for alpha, init in ((0.3, None), (0.3, 0.7), (1.0, None)): | |
| tracker = MilestoneTracker(d, kappa=kappa, alpha_gamma=alpha, eps=eps, gamma_bar_init=init) | |
| if init is None: | |
| _expect(RuntimeError, lambda tr=tracker: tr.threshold) # not calibrated yet | |
| else: # gamma_0 = gamma_bar_init already defines tau_0 | |
| assert _close_scalar(tracker.threshold, kappa * math.sqrt(2.0 * d) * init) | |
| _expect(RuntimeError, lambda tr=tracker: tr.gamma_bar_star) # not frozen yet | |
| gamma_ref = init | |
| for n, z in enumerate(batches): | |
| inst = _ref_margin(z, eps) | |
| gamma_ref = inst if gamma_ref is None else gamma_ref # gamma_0 := first estimate | |
| gamma_ref = (1.0 - alpha) * gamma_ref + alpha * inst | |
| if alpha == 1.0: | |
| assert gamma_ref == inst | |
| if n == 0 and init is not None: | |
| assert _close_scalar(gamma_ref, (1.0 - alpha) * init + alpha * inst) | |
| ret = tracker.update_variance_margin(z) | |
| assert _close_scalar(ret, gamma_ref) and _close_scalar(float(tracker.gamma_bar), gamma_ref) | |
| assert _close_scalar(tracker.threshold, kappa * math.sqrt(2.0 * d) * gamma_ref) | |
| # B < 2 raises and leaves gamma unchanged. | |
| before = float(tracker.gamma_bar) | |
| _expect(ValueError, lambda tr=tracker: tr.update_variance_margin(torch.randn(1, d))) | |
| _expect(ValueError, lambda tr=tracker: tr.instantaneous_margin(torch.randn(1, d))) | |
| _expect(ValueError, lambda tr=tracker: tr.update_variance_margin(torch.randn(4, d + 1))) | |
| assert float(tracker.gamma_bar) == before | |
| # freeze(): gamma* and tau* = kappa sqrt(2d) gamma*; no further updates. | |
| tau_star = tracker.freeze() | |
| tau_ref = kappa * math.sqrt(2.0 * d) * gamma_ref | |
| assert _close_scalar(tau_star, tau_ref) and _close_scalar(tracker.threshold, tau_ref) | |
| assert _close_scalar(tracker.gamma_bar_star, gamma_ref) | |
| _expect(RuntimeError, lambda tr=tracker: tr.update_variance_margin(batches[0])) | |
| assert float(tracker.gamma_bar) == before and tracker.threshold == tau_star | |
| over = MilestoneTracker(d, kappa=kappa, alpha_gamma=0.5, eps=eps) | |
| _expect(RuntimeError, over.freeze) # uncalibrated, no gamma_bar_star given | |
| _expect(ValueError, lambda: over.freeze(gamma_bar_star=-1.0)) | |
| assert _close_scalar(over.freeze(gamma_bar_star=2.5), kappa * math.sqrt(2.0 * d) * 2.5) | |
| # ---- Eq. 9a strict boundary: kappa sqrt(2d) = 0.25 * 4 = 1, so tau* = gamma* = 1.5 exactly. | |
| tau = 1.5 | |
| below = float(torch.nextafter(torch.tensor(tau), torch.tensor(0.0))) # next float32 below tau | |
| for mm in (2, 1): # M = 2: advance test at m < M; M = 1: completion test at m = M | |
| tr = MilestoneTracker(d, kappa=0.25, alpha_gamma=0.5, eps=eps) | |
| assert tr.freeze(gamma_bar_star=tau) == tau | |
| goals = torch.randn(2, mm, d, generator=gen) | |
| goals[..., 0] = 0.0 # offsets along axis 0 are then represented exactly | |
| z = goals[:, 0].clone() | |
| z[0, 0], z[1, 0] = tau, below | |
| tr.reset(mm, 2) | |
| upd = tr.step(z, goals) | |
| assert upd.distance[0].item() == tau and upd.distance[1].item() == below # exact distances | |
| assert upd.reached.tolist() == [False, True] | |
| if mm == 2: | |
| assert upd.m_next.tolist() == [1, 2] and upd.task_complete.tolist() == [False, False] | |
| else: # single-stage task: completes without moving the pointer | |
| assert upd.m_next.tolist() == [1, 1] and upd.task_complete.tolist() == [False, True] | |
| # ---- Eq. 9a randomized run against the per-element reference (tau* = 1.5). | |
| tr = MilestoneTracker(d, kappa=0.25, alpha_gamma=0.5, eps=eps) | |
| tr.freeze(gamma_bar_star=tau) | |
| goals = torch.randn(b, num_m, d, generator=gen) * 5.0 | |
| goals[..., 0] = 0.0 | |
| tr.reset(num_m, b) | |
| m_ref, done_ref = [1] * b, [False] * b | |
| seen = {"advance": 0, "hold": 0, "complete": 0, "masked_reach": 0, "boundary": 0} | |
| for _ in range(40): | |
| kinds = torch.randint(0, 3, (b,), generator=gen).tolist() # 0 inside, 1 outside, 2 boundary | |
| mask = (torch.rand(b, generator=gen) < 0.7).tolist() | |
| z = torch.empty(b, d) | |
| for i in range(b): | |
| g = goals[i, m_ref[i] - 1] | |
| if kinds[i] == 2: | |
| z[i] = g | |
| z[i, 0] = tau | |
| else: | |
| direction = torch.randn(d, generator=gen) | |
| radius = 0.7 if kinds[i] == 0 else 4.0 | |
| z[i] = g + radius * direction / direction.norm() | |
| assert torch.equal(tr.active_goal(goals), torch.stack([goals[i, m_ref[i] - 1] for i in range(b)])) | |
| m_next_ref, done_next_ref, dist_ref, reached_ref = _ref_pointer_step( | |
| m_ref, done_ref, z, goals, tau, num_m, mask) | |
| upd = tr.step(z, goals, update_mask=torch.tensor(mask)) | |
| assert upd.m_t.tolist() == m_ref and upd.m_next.tolist() == m_next_ref | |
| assert upd.reached.tolist() == reached_ref and upd.task_complete.tolist() == done_next_ref | |
| assert _close(upd.distance, torch.tensor(dist_ref, dtype=torch.float64)) | |
| assert torch.equal(tr.m, upd.m_next) and torch.equal(tr.task_complete, upd.task_complete) | |
| for i in range(b): | |
| assert m_ref[i] <= m_next_ref[i] <= num_m and m_next_ref[i] >= 1 # monotone, in range | |
| assert not done_ref[i] or done_next_ref[i] # sticky completion | |
| if not mask[i]: | |
| assert m_next_ref[i] == m_ref[i] and done_next_ref[i] == done_ref[i] | |
| seen["masked_reach"] += int(reached_ref[i]) | |
| seen["advance"] += int(m_next_ref[i] > m_ref[i]) | |
| seen["hold"] += int(mask[i] and m_next_ref[i] == m_ref[i]) | |
| seen["complete"] += int(done_next_ref[i] and not done_ref[i]) | |
| seen["boundary"] += int(kinds[i] == 2 and not reached_ref[i]) | |
| m_ref, done_ref = m_next_ref, done_next_ref | |
| assert all(v > 0 for v in seen.values()), seen | |
| assert all(done_ref) # every element completed and stayed complete | |
| # ---- Errors: pointer before reset(), bad constructor / reset arguments. | |
| fresh = MilestoneTracker(d, kappa=0.25, alpha_gamma=0.5, eps=eps) | |
| fresh.freeze(gamma_bar_star=1.0) | |
| _expect(RuntimeError, lambda: fresh.step(torch.randn(b, d), goals)) | |
| _expect(RuntimeError, lambda: fresh.active_goal(goals)) | |
| _expect(RuntimeError, lambda: fresh.completion_test(torch.randn(b, d), goals)) | |
| _expect(ValueError, lambda: fresh.reset(0, b)) | |
| for kw in (dict(kappa=0.0), dict(kappa=1.0), dict(alpha_gamma=0.0), dict(alpha_gamma=1.5), | |
| dict(eps=0.0), dict(gamma_bar_init=0.0), dict(gamma_bar_init=-1.0)): | |
| args = dict(kappa=0.25, alpha_gamma=0.5, eps=eps) | |
| args.update(kw) | |
| _expect(ValueError, lambda a=args: MilestoneTracker(d, **a)) | |
| print("selector.py self-test passed") | |
| if __name__ == "__main__": | |
| _self_test() | |