# SPDX-License-Identifier: GPL-3.0-only # Copyright (C) 2026 SLING AI Inc. """ solver.py -- Stage 4 of the Goal-Conditioned VPA framework (preprint_261008.pdf, Secs. 4.4 and 4.5). Implements * ``standardize_latents`` / ``standardize_proprioception`` -- the fixed standardisation of Sec. 4.5: z_tilde = (z - mu_Z) / gamma* (applied to z_t and z_hat_{t+1}) and S_tilde = (S - mu_S) / sigma_S with dataset statistics. * ``ConditionalVectorField`` -- v_theta : R^{H x d_a} x [0, 1] x R^{d_e} -> R^{H x d_a}. A non-causal transformer over the H chunk positions: every evaluation processes the whole chunk in parallel (Sec. 4.4). * ``FlowMatchingSolver`` -- pi_theta^l: e_t = Concat(Embed(u_t), z_t, z_hat_{t+1}, S_t, c_text) in R^{d_e}, d_e = d_u + 2d + d_s + d_c (Eq. 12) A^(rho) = rho A_t + (1 - rho) xi, xi ~ N(0, I) (Eq. 13) L_solver = E ||v_theta(A^(rho), rho | e_t) - (A_t - xi)||_2^2, rho ~ U[0, 1] (Eq. 14) sampling: A^(0) = xi, A^(rho_{k+1}) = A^(rho_k) + (1/K) v_theta(A^(rho_k), rho_k | e_t), rho_k = k / K, k = 0, ..., K - 1 (explicit Euler, K in {1, 2, 3}) * ``UncertaintyHorizonFilter`` -- the fast-attack / slow-release filter and executed horizon: sigma_bar_{t+1} = max{ sigma_{t+1}, (1 - a_s) sigma_bar_t + a_s sigma_{t+1} } H_t = max( H_min, floor( H_max * exp(-beta * sigma_bar_{t+1}) ) ) (Eq. 15) This file is self-contained (it imports nothing from the other VPA modules) and can be tested on its own with ``python solver.py``. """ from __future__ import annotations import math from typing import Any, List, Optional, Tuple, Union import torch import torch.nn as nn from torch import Tensor __all__ = [ "standardize_latents", "standardize_proprioception", "SinusoidalTimeEmbedding", "ConditionalVectorField", "FlowMatchingSolver", "UncertaintyHorizonFilter", "executed_prefix_mask", ] ALLOWED_INTEGRATION_STEPS = (1, 2, 3) # --------------------------------------------------------------------------------------------- # Fixed standardisation (Sec. 4.5) # --------------------------------------------------------------------------------------------- def standardize_latents(z: Tensor, latent_mean: Tensor, gamma_bar_star: Union[float, Tensor]) -> Tensor: """z_tilde = (z - mu_Z) / gamma* (Sec. 4.5), applied to both z_t and z_hat_{t+1}. Args: z: ``[..., d]`` latents. latent_mean: ``[d]`` mean latent mu_Z over the training data. gamma_bar_star: frozen latent scale gamma* (scalar). Returns: ``[..., d]`` standardised latents. """ return (z - latent_mean) / gamma_bar_star def standardize_proprioception(state: Tensor, proprio_mean: Tensor, proprio_std: Tensor) -> Tensor: """S_tilde = (S - mu_S) / sigma_S with per-dimension dataset statistics (Sec. 4.5). Args: state: ``[..., d_s]``; proprio_mean, proprio_std: ``[d_s]`` (std strictly positive). Returns: ``[..., d_s]``. """ return (state - proprio_mean) / proprio_std # --------------------------------------------------------------------------------------------- # v_theta : conditional vector field # --------------------------------------------------------------------------------------------- class SinusoidalTimeEmbedding(nn.Module): """Sinusoidal features of the flow time rho in [0, 1]. Shapes: ``rho: [B]`` -> ``[B, dim]`` """ def __init__(self, dim: int, scale: float = 1000.0) -> None: super().__init__() if dim % 2 != 0: raise ValueError("time embedding dimension must be even") half = dim // 2 freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32) / half) self.register_buffer("freqs", freqs, persistent=False) self.scale: float = float(scale) def forward(self, rho: Tensor) -> Tensor: """``[B]`` -> ``[B, dim]``.""" args = self.scale * rho.to(self.freqs.dtype).unsqueeze(-1) * self.freqs # [B, dim/2] return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) class ConditionalVectorField(nn.Module): """v_theta(A^(rho), rho | e_t): a non-causal transformer over the H action positions. Token h = W_a a_h + p_h + c, where c = MLP(Concat(e_t, time(rho))) is broadcast to every position; a conditioning token carrying c is also prepended so that every layer can attend to it. No causal mask is used: all H positions are produced jointly in one evaluation. Args: horizon: H (= H_max), chunk length. action_dim: d_a. cond_dim: d_e. width, depth, num_heads, mlp_ratio: transformer hyper-parameters. time_embed_dim: width of the rho embedding. Shapes: (``A_rho [B, H, d_a]``, ``rho [B]`` or scalar, ``e_t [B, d_e]``) -> ``[B, H, d_a]`` """ def __init__( self, horizon: int, action_dim: int, cond_dim: int, width: int = 256, depth: int = 4, num_heads: int = 4, mlp_ratio: float = 4.0, time_embed_dim: int = 128, ) -> None: super().__init__() self.horizon: int = horizon self.action_dim: int = action_dim self.cond_dim: int = cond_dim self.action_in = nn.Linear(action_dim, width) self.pos_embed = nn.Parameter(torch.zeros(1, horizon, width)) nn.init.trunc_normal_(self.pos_embed, std=0.02) self.time_embed = SinusoidalTimeEmbedding(time_embed_dim) self.cond_mlp = nn.Sequential( nn.Linear(cond_dim + time_embed_dim, width), nn.GELU(), nn.Linear(width, width) ) layer = nn.TransformerEncoderLayer( d_model=width, nhead=num_heads, dim_feedforward=int(round(width * mlp_ratio)), dropout=0.0, activation="gelu", batch_first=True, norm_first=True, ) self.blocks = nn.TransformerEncoder(layer, num_layers=depth, enable_nested_tensor=False) self.norm = nn.LayerNorm(width) self.action_out = nn.Linear(width, action_dim) def forward(self, actions: Tensor, rho: Union[float, Tensor], cond: Tensor) -> Tensor: """Evaluate the vector field. Args: actions: ``[B, H, d_a]`` current point A^(rho) on the flow. rho: flow time, a Python float, a 0-dim tensor, or ``[B]``. cond: ``[B, d_e]`` unified conditioning vector e_t. Returns: velocity: ``[B, H, d_a]``. """ b = actions.shape[0] if actions.shape[1:] != (self.horizon, self.action_dim): raise ValueError(f"expected actions [B, {self.horizon}, {self.action_dim}], got {tuple(actions.shape)}") if cond.shape != (b, self.cond_dim): raise ValueError(f"expected e_t of shape [B, {self.cond_dim}], got {tuple(cond.shape)}") rho_t = torch.as_tensor(rho, dtype=actions.dtype, device=actions.device) if rho_t.dim() == 0: rho_t = rho_t.expand(b) c = self.cond_mlp(torch.cat([cond, self.time_embed(rho_t).to(cond.dtype)], dim=-1)) # [B, W] tokens = self.action_in(actions) + self.pos_embed + c.unsqueeze(1) # [B, H, W] x = self.blocks(torch.cat([c.unsqueeze(1), tokens], dim=1)) # [B, 1+H, W] return self.action_out(self.norm(x[:, 1:])) # [B, H, d_a] # --------------------------------------------------------------------------------------------- # pi_theta^l : flow-matching action solver # --------------------------------------------------------------------------------------------- class FlowMatchingSolver(nn.Module): """Conditional Flow-Matching action-chunk solver pi_theta^l (Secs. 4.4, 4.5). Holds the learned primitive embedding Embed : U -> R^{d_u}, the fixed standardisation statistics (mu_Z, gamma*, mu_S, sigma_S; set with ``set_standardization_stats`` after the JEPA modules are frozen), and the vector field v_theta. Args: num_primitives: N_u. primitive_embed_dim: d_u. latent_dim: d. proprio_dim: d_s. text_dim: d_c. horizon: H = H_max. action_dim: d_a. num_integration_steps: K in {1, 2, 3}. **field_kwargs: forwarded to ``ConditionalVectorField``. Shapes: ``build_conditioning``: (``u [B]``, ``z_t [B, d]``, ``z_hat [B, d]``, ``S_t [B, d_s]``, ``c_text [B, d_c]``) -> ``e_t [B, d_e]`` ``flow_matching_loss``: (``A_t [B, H, d_a]``, ``e_t [B, d_e]``) -> scalar ``generate_chunk``: ``e_t [B, d_e]`` -> ``A_hat_t [B, H, d_a]`` """ def __init__( self, num_primitives: int, primitive_embed_dim: int, latent_dim: int, proprio_dim: int, text_dim: int, horizon: int, action_dim: int, num_integration_steps: int = 2, **field_kwargs: Any, ) -> None: super().__init__() if num_integration_steps not in ALLOWED_INTEGRATION_STEPS: raise ValueError(f"K must be one of {ALLOWED_INTEGRATION_STEPS}") self.num_primitives: int = num_primitives self.primitive_embed_dim: int = primitive_embed_dim self.latent_dim: int = latent_dim self.proprio_dim: int = proprio_dim self.text_dim: int = text_dim self.horizon: int = horizon self.action_dim: int = action_dim self.num_integration_steps: int = num_integration_steps self.cond_dim: int = primitive_embed_dim + 2 * latent_dim + proprio_dim + text_dim # d_e self.primitive_embedding = nn.Embedding(num_primitives, primitive_embed_dim) # Embed self.vector_field = ConditionalVectorField(horizon, action_dim, self.cond_dim, **field_kwargs) nan = float("nan") self.register_buffer("latent_mean", torch.full((latent_dim,), nan)) self.register_buffer("gamma_bar_star", torch.tensor(nan)) self.register_buffer("proprio_mean", torch.full((proprio_dim,), nan)) self.register_buffer("proprio_std", torch.full((proprio_dim,), nan)) # ---------------- Sec. 4.5: fixed statistics ---------------- @torch.no_grad() def set_standardization_stats( self, latent_mean: Tensor, gamma_bar_star: float, proprio_mean: Tensor, proprio_std: Tensor, ) -> None: """Store mu_Z ``[d]``, gamma* (scalar), mu_S ``[d_s]`` and sigma_S ``[d_s]``.""" latent_mean = torch.as_tensor(latent_mean, dtype=self.latent_mean.dtype) proprio_mean = torch.as_tensor(proprio_mean, dtype=self.proprio_mean.dtype) proprio_std = torch.as_tensor(proprio_std, dtype=self.proprio_std.dtype) if latent_mean.shape != (self.latent_dim,): raise ValueError(f"latent_mean must have shape [{self.latent_dim}]") if proprio_mean.shape != (self.proprio_dim,) or proprio_std.shape != (self.proprio_dim,): raise ValueError(f"proprioception statistics must have shape [{self.proprio_dim}]") if not gamma_bar_star > 0.0: raise ValueError("gamma_bar_star must be positive") if not bool((proprio_std > 0).all()): raise ValueError("proprio_std must be strictly positive in every dimension") if not (torch.isfinite(latent_mean).all() and torch.isfinite(proprio_mean).all()): raise ValueError("statistics must be finite") self.latent_mean.copy_(latent_mean) self.gamma_bar_star.fill_(float(gamma_bar_star)) self.proprio_mean.copy_(proprio_mean) self.proprio_std.copy_(proprio_std) @property def is_calibrated(self) -> bool: return bool( torch.isfinite(self.latent_mean).all() and torch.isfinite(self.gamma_bar_star) and torch.isfinite(self.proprio_mean).all() and torch.isfinite(self.proprio_std).all() ) # ---------------- Eq. 12: unified conditioning ---------------- def build_conditioning( self, primitive: Tensor, z_t: Tensor, z_hat: Tensor, proprioception: Tensor, c_text: Tensor, ) -> Tensor: """e_t = Concat(Embed(u_t), z_tilde_t, z_hat_tilde_{t+1}, S_tilde_t, c_text) (Eq. 12). Args: primitive: ``[B]`` int64 tokens u_t. z_t: ``[B, d]`` current latent (raw; standardised here). z_hat: ``[B, d]`` predicted latent z_hat_{t+1} (raw; standardised here). proprioception: ``[B, d_s]`` raw proprioceptive state S_t (standardised here). c_text: ``[B, d_c]`` instruction embedding. Returns: e_t: ``[B, d_e]`` with d_e = d_u + 2d + d_s + d_c. """ if not self.is_calibrated: raise RuntimeError("standardisation statistics are not set; call set_standardization_stats") b = z_t.shape[0] expected = { "z_t": (z_t, (b, self.latent_dim)), "z_hat": (z_hat, (b, self.latent_dim)), "proprioception": (proprioception, (b, self.proprio_dim)), "c_text": (c_text, (b, self.text_dim)), } for name, (tensor, shape) in expected.items(): if tuple(tensor.shape) != shape: raise ValueError(f"{name} must have shape {list(shape)}, got {list(tensor.shape)}") if primitive.shape != (b,): raise ValueError(f"primitive must have shape [{b}]") e_t = torch.cat( [ self.primitive_embedding(primitive.long()), # [B, d_u] standardize_latents(z_t, self.latent_mean, self.gamma_bar_star), # [B, d] standardize_latents(z_hat, self.latent_mean, self.gamma_bar_star), # [B, d] standardize_proprioception(proprioception, self.proprio_mean, self.proprio_std), # [B, d_s] c_text, # [B, d_c] ], dim=-1, ) return e_t # ---------------- Eqs. 13-14: training objective ---------------- def flow_matching_loss( self, actions: Tensor, cond: Tensor, generator: Optional[torch.Generator] = None ) -> Tensor: """L_solver = E_{t, xi, rho ~ U[0,1]} ||v_theta(A^(rho), rho | e_t) - (A_t - xi)||_2^2 (Eq. 14). ||.||_2^2 is summed over the H x d_a chunk entries; the expectation is the batch mean. Args: actions: ``[B, H, d_a]`` ground-truth chunks A_t. cond: ``[B, d_e]`` conditioning vectors e_t. Returns: scalar loss. """ b = actions.shape[0] xi = torch.randn(actions.shape, generator=generator, device=actions.device, dtype=actions.dtype) rho = torch.rand(b, generator=generator, device=actions.device, dtype=actions.dtype) # U[0, 1) rho_b = rho.view(b, 1, 1) a_rho = rho_b * actions + (1.0 - rho_b) * xi # Eq. 13 velocity = self.vector_field(a_rho, rho, cond) return (velocity - (actions - xi)).pow(2).sum(dim=(1, 2)).mean() # ---------------- Sec. 4.4: K-step Euler sampling ---------------- def generate_chunk( self, cond: Tensor, noise: Optional[Tensor] = None, generator: Optional[torch.Generator] = None, ) -> Tensor: """Sample A_hat_t by integrating v_theta from rho = 0 to 1 with K explicit Euler steps. A^(0) = xi ~ N(0, I); A^(rho_{k+1}) = A^(rho_k) + (1/K) v_theta(A^(rho_k), rho_k | e_t), rho_k = k / K for k = 0, ..., K - 1. Exactly K sequential evaluations of v_theta, each over the full ``[B, H, d_a]`` chunk; the loop length does not depend on H. Args: cond: ``[B, d_e]`` conditioning vectors e_t. noise: optional ``[B, H, d_a]`` initial sample xi (drawn from N(0, I) if omitted). generator: optional RNG for xi. Returns: A_hat_t: ``[B, H, d_a]``. """ b = cond.shape[0] shape = (b, self.horizon, self.action_dim) if noise is None: noise = torch.randn(shape, generator=generator, device=cond.device, dtype=cond.dtype) elif tuple(noise.shape) != shape: raise ValueError(f"noise must have shape {list(shape)}") k_steps = self.num_integration_steps dt = 1.0 / k_steps a = noise for k in range(k_steps): rho_k = k / k_steps a = a + dt * self.vector_field(a, rho_k, cond) return a # --------------------------------------------------------------------------------------------- # Uncertainty-adaptive executed horizon (Eq. 15) # --------------------------------------------------------------------------------------------- class UncertaintyHorizonFilter: """Fast-attack / slow-release filter on sigma and the executed horizon H_t (Eq. 15). sigma_bar_{t+1} = max{ sigma_{t+1}, (1 - alpha_sigma) sigma_bar_t + alpha_sigma sigma_{t+1} } H_t = max( H_min, floor( H_max * exp(-beta * sigma_bar_{t+1}) ) ) The filter state sigma_bar starts at 0 for each episode, so the first decision step uses sigma_bar = sigma exactly (fast attack). Args: h_min: H_min >= 1. h_max: H_max = H (chunk length), H_min <= H_max. beta: sensitivity beta > 0. alpha_sigma: release rate alpha_sigma in (0, 1]. Shapes: ``update``: ``sigma [B]`` (+ optional ``mask [B]`` bool) -> (``sigma_bar [B]``, ``H_t [B]`` int64) """ def __init__(self, h_min: int, h_max: int, beta: float, alpha_sigma: float) -> None: if not 1 <= h_min <= h_max: raise ValueError("require 1 <= H_min <= H_max") if beta <= 0.0: raise ValueError("beta must be positive") if not 0.0 < alpha_sigma <= 1.0: raise ValueError("alpha_sigma must lie in (0, 1]") self.h_min: int = int(h_min) self.h_max: int = int(h_max) self.beta: float = float(beta) self.alpha_sigma: float = float(alpha_sigma) self.sigma_bar: Optional[Tensor] = None def reset(self, batch_size: int, device: Optional[torch.device] = None) -> None: """Start an episode with sigma_bar = 0.""" self.sigma_bar = torch.zeros(batch_size, device=device) def filter(self, sigma: Tensor, mask: Optional[Tensor] = None) -> Tensor: """Asymmetric filter; commits and returns sigma_bar_{t+1} ``[B]``. Args: sigma: ``[B]`` raw scores sigma_{t+1} >= 0. mask: optional ``[B]`` bool; elements outside the mask keep their previous sigma_bar. """ prev = self.sigma_bar if self.sigma_bar is not None else torch.zeros_like(sigma) prev = prev.to(device=sigma.device, dtype=sigma.dtype) new = torch.maximum(sigma, (1.0 - self.alpha_sigma) * prev + self.alpha_sigma * sigma) if mask is not None: new = torch.where(mask.to(device=sigma.device, dtype=torch.bool), new, prev) self.sigma_bar = new return new def horizon(self, sigma_bar: Tensor) -> Tensor: """Eq. 15: H_t = max(H_min, floor(H_max * exp(-beta * sigma_bar))) -> ``[B]`` int64.""" raw = torch.floor(self.h_max * torch.exp(-self.beta * sigma_bar)).long() return torch.clamp(raw, min=self.h_min) def update(self, sigma: Tensor, mask: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]: """Filter sigma and compute H_t: returns (``sigma_bar [B]``, ``H_t [B]`` int64).""" sigma_bar = self.filter(sigma, mask) return sigma_bar, self.horizon(sigma_bar) def executed_prefix_mask(executed_horizon: Tensor, horizon: int) -> Tensor: """Boolean mask of the executed prefix: ``H_t [B]`` -> ``[B, H]`` with True for h < H_t.""" steps = torch.arange(horizon, device=executed_horizon.device) return steps.unsqueeze(0) < executed_horizon.unsqueeze(1) # --------------------------------------------------------------------------------------------- # Self-test (run: python solver.py) # --------------------------------------------------------------------------------------------- class _CountingField(nn.Module): """Wraps v_theta and counts evaluations (used by the self-test only).""" def __init__(self, field: nn.Module) -> None: super().__init__() self.field = field self.calls: List[float] = [] def forward(self, actions: Tensor, rho: Union[float, Tensor], cond: Tensor) -> Tensor: # Record rho exactly as passed: torch.as_tensor(float) would round to float32 (1/3 != 1/3). self.calls.append(rho if isinstance(rho, float) else float(rho.reshape(-1)[0])) return self.field(actions, rho, cond) _RTOL, _ATOL = 1e-5, 1e-6 # float32 module output vs float64 reference (CLAUDE.md) def _close(actual: Any, ref: Any) -> bool: a = torch.as_tensor(actual).detach().double() r = torch.as_tensor(ref).detach().double() return torch.allclose(a, r, rtol=_RTOL, atol=_ATOL) def _expect(exc: type, fn: Any) -> None: try: fn() except exc: return raise AssertionError(f"expected {exc.__name__}") class _ZeroField(nn.Module): """Test stand-in v_theta == 0 (self-test only).""" def forward(self, actions: Tensor, rho: Union[float, Tensor], cond: Tensor) -> Tensor: return torch.zeros_like(actions) class _AffineField(nn.Module): """Test stand-in v(A, rho, e) = a A + b rho + reshape(e W), easy to evaluate in float64. Shapes: (``[B, H, d_a]``, scalar or ``[B]``, ``[B, d_e]``) -> ``[B, H, d_a]`` """ def __init__(self, horizon: int, action_dim: int, cond_dim: int, generator: torch.Generator) -> None: super().__init__() self.a, self.b = 0.7, 0.3 self.register_buffer("w", 0.1 * torch.randn(cond_dim, horizon * action_dim, generator=generator)) def forward(self, actions: Tensor, rho: Union[float, Tensor], cond: Tensor) -> Tensor: bsz = actions.shape[0] rho_t = torch.as_tensor(rho, dtype=actions.dtype) rho_t = rho_t.expand(bsz) if rho_t.dim() == 0 else rho_t return self.a * actions + self.b * rho_t.view(bsz, 1, 1) + (cond @ self.w).view(actions.shape) class _RecordingField(nn.Module): """Wraps v_theta and records its last input and output (self-test only).""" def __init__(self, field: nn.Module) -> None: super().__init__() self.field = field self.record: Optional[Tuple[Tensor, Tensor, Tensor]] = None def forward(self, actions: Tensor, rho: Union[float, Tensor], cond: Tensor) -> Tensor: out = self.field(actions, rho, cond) self.record = (actions.detach().clone(), torch.as_tensor(rho).detach().clone(), out) return out def _ref_conditioning(solver: "FlowMatchingSolver", u: Tensor, z_t: Tensor, z_hat: Tensor, s_t: Tensor, c: Tensor) -> Tensor: """Float64 reference of Eq. 12 with the Sec. 4.5 standardisation, written slot by slot.""" du, d, ds, dc = solver.primitive_embed_dim, solver.latent_dim, solver.proprio_dim, solver.text_dim bsz = z_t.shape[0] emb = solver.primitive_embedding.weight.detach().double() mu_z, g = solver.latent_mean.double(), float(solver.gamma_bar_star) mu_s, sd_s = solver.proprio_mean.double(), solver.proprio_std.double() e = torch.zeros(bsz, du + 2 * d + ds + dc, dtype=torch.float64) for k in range(bsz): for j in range(du): e[k, j] = emb[int(u[k]), j] for j in range(d): e[k, du + j] = (float(z_t[k, j]) - float(mu_z[j])) / g e[k, du + d + j] = (float(z_hat[k, j]) - float(mu_z[j])) / g for j in range(ds): e[k, du + 2 * d + j] = (float(s_t[k, j]) - float(mu_s[j])) / float(sd_s[j]) for j in range(dc): e[k, du + 2 * d + ds + j] = float(c[k, j]) return e def _self_test() -> None: torch.manual_seed(0) gen = torch.Generator().manual_seed(99) b, nu, du, d, ds, dc, h, da = 3, 4, 5, 6, 7, 8, 10, 2 field_kwargs = dict(width=32, depth=2, num_heads=4, time_embed_dim=16) def make(k: int = 2, horizon: int = h) -> FlowMatchingSolver: return FlowMatchingSolver(nu, du, d, ds, dc, horizon, da, num_integration_steps=k, **field_kwargs) lat_mean, gamma_star = torch.randn(d, generator=gen), 2.5 s_mean, s_std = torch.randn(ds, generator=gen), torch.rand(ds, generator=gen) + 0.5 u = torch.randint(0, nu, (b,), generator=gen) z_t, z_hat = torch.randn(b, d, generator=gen), torch.randn(b, d, generator=gen) s_t, c = torch.randn(b, ds, generator=gen), torch.randn(b, dc, generator=gen) # ---- Eq. 12 + Sec. 4.5: unified conditioning and standardisation. solver = make() _expect(RuntimeError, lambda: solver.build_conditioning(u, z_t, z_hat, s_t, c)) # stats not set bad_std_zero, bad_std_neg = s_std.clone(), s_std.clone() bad_std_zero[2], bad_std_neg[4] = 0.0, -0.5 _expect(ValueError, lambda: solver.set_standardization_stats(lat_mean, gamma_star, s_mean, bad_std_zero)) _expect(ValueError, lambda: solver.set_standardization_stats(lat_mean, gamma_star, s_mean, bad_std_neg)) _expect(ValueError, lambda: solver.set_standardization_stats(lat_mean, 0.0, s_mean, s_std)) _expect(ValueError, lambda: solver.set_standardization_stats(lat_mean, -1.0, s_mean, s_std)) _expect(ValueError, lambda: solver.set_standardization_stats(lat_mean[:-1], gamma_star, s_mean, s_std)) _expect(ValueError, lambda: solver.set_standardization_stats(lat_mean, gamma_star, s_mean[:-1], s_std)) nan_mean = lat_mean.clone() nan_mean[0] = float("nan") _expect(ValueError, lambda: solver.set_standardization_stats(nan_mean, gamma_star, s_mean, s_std)) assert not solver.is_calibrated # rejected statistics were not installed solver.set_standardization_stats(lat_mean, gamma_star, s_mean, s_std) e_t = solver.build_conditioning(u, z_t, z_hat, s_t, c) assert solver.cond_dim == du + 2 * d + ds + dc and e_t.shape == (b, du + 2 * d + ds + dc) ref_e = _ref_conditioning(solver, u, z_t, z_hat, s_t, c) assert torch.equal(e_t[:, :du], solver.primitive_embedding.weight[u]) # Embed(u_t) assert _close(e_t[:, du:du + d], ref_e[:, du:du + d]) # (z_t - mu_Z) / gamma* assert _close(e_t[:, du + d:du + 2 * d], ref_e[:, du + d:du + 2 * d]) # (z_hat - mu_Z) / gamma* assert _close(e_t[:, du + 2 * d:du + 2 * d + ds], ref_e[:, du + 2 * d:du + 2 * d + ds]) # S_tilde assert torch.equal(e_t[:, -dc:], c) # c_text, not standardised assert _close(e_t, ref_e) _expect(ValueError, lambda: solver.build_conditioning(u, z_t[:, :-1], z_hat, s_t, c)) _expect(ValueError, lambda: solver.build_conditioning(u, z_t, z_hat[:-1], s_t, c)) _expect(ValueError, lambda: solver.build_conditioning(u, z_t, z_hat, s_t[:, :-1], c)) _expect(ValueError, lambda: solver.build_conditioning(u, z_t, z_hat, s_t, c[:, :-1])) _expect(ValueError, lambda: solver.build_conditioning(u[:-1], z_t, z_hat, s_t, c)) e_t = e_t.detach() # ---- Eqs. 13-14: flow-matching loss with a fixed generator. a_true = torch.randn(b, h, da, generator=gen) def replay(seed: int) -> Tuple[Tensor, Tensor]: g = torch.Generator().manual_seed(seed) xi_r = torch.randn(a_true.shape, generator=g) # same draw order as flow_matching_loss return xi_r, torch.rand(b, generator=g) real_field = solver.vector_field solver.vector_field = _ZeroField() loss0 = solver.flow_matching_loss(a_true, e_t, generator=torch.Generator().manual_seed(7)) xi, rho = replay(7) ref0 = sum(float(a_true[k, i, j] - xi[k, i, j]) ** 2 for k in range(b) for i in range(h) for j in range(da)) / b assert _close(loss0, ref0) rec = _RecordingField(real_field) solver.vector_field = rec loss = solver.flow_matching_loss(a_true, e_t, generator=torch.Generator().manual_seed(7)) a_in, rho_in, vel = rec.record vel_ref = vel.detach() # reference reads values only; the loss keeps its graph assert torch.equal(rho_in, rho) and bool(((rho >= 0) & (rho < 1)).all()) a_rho_ref = torch.zeros(b, h, da, dtype=torch.float64) for k in range(b): r = float(rho[k]) a_rho_ref[k] = r * a_true[k].double() + (1.0 - r) * xi[k].double() # Eq. 13 assert _close(a_in, a_rho_ref) ref_loss = sum((float(vel_ref[k, i, j]) - (float(a_true[k, i, j]) - float(xi[k, i, j]))) ** 2 for k in range(b) for i in range(h) for j in range(da)) / b # Eq. 14 assert _close(loss, ref_loss) loss.backward() assert all(p.grad is not None for p in real_field.parameters()) solver.vector_field = real_field # ---- Euler (Sec. 4.4) for every K, against hand-written loops. for k in ALLOWED_INTEGRATION_STEPS: sk = make(k) sk.set_standardization_stats(lat_mean, gamma_star, s_mean, s_std) ek = sk.build_conditioning(u, z_t, z_hat, s_t, c).detach() xi = torch.randn(b, h, da, generator=gen) real_k = sk.vector_field # (a) affine stand-in field vs a float64 Euler loop. affine = _AffineField(h, da, sk.cond_dim, gen) counting = _CountingField(affine) sk.vector_field = counting with torch.no_grad(): chunk = sk.generate_chunk(ek, noise=xi) assert counting.calls == [j / k for j in range(k)] a_ref = xi.double() bias = (ek.double() @ affine.w.double()).view(b, h, da) for j in range(k): a_ref = a_ref + (1.0 / k) * (affine.a * a_ref + affine.b * (j / k) + bias) assert chunk.shape == (b, h, da) and _close(chunk, a_ref) # (b) real transformer vs a hand-written float32 loop from the same xi. field = _CountingField(real_k) sk.vector_field = field with torch.no_grad(): chunk = sk.generate_chunk(ek, noise=xi) a_loop = xi.clone() for j in range(k): a_loop = a_loop + (1.0 / k) * real_k(a_loop, j / k, ek) assert field.calls == [j / k for j in range(k)] and _close(chunk, a_loop) # (c) noise drawn from a generator == the same noise passed explicitly. with torch.no_grad(): c1 = sk.generate_chunk(ek, generator=torch.Generator().manual_seed(11)) c2 = sk.generate_chunk(ek, noise=torch.randn(b, h, da, generator=torch.Generator().manual_seed(11))) assert torch.equal(c1, c2) _expect(ValueError, lambda s=sk, e=ek: s.generate_chunk(e, noise=torch.randn(b, h + 1, da))) for bad_k in (0, 4): _expect(ValueError, lambda kk=bad_k: make(kk)) # ---- Horizon independence: K evaluations of v_theta for every H. for hh in (4, 16, 64): for k in ALLOWED_INTEGRATION_STEPS: sk = make(k, horizon=hh) sk.set_standardization_stats(lat_mean, gamma_star, s_mean, s_std) ek = sk.build_conditioning(u, z_t, z_hat, s_t, c).detach() counting = _CountingField(sk.vector_field) sk.vector_field = counting with torch.no_grad(): chunk = sk.generate_chunk(ek, generator=torch.Generator().manual_seed(hh + k)) assert len(counting.calls) == k and chunk.shape == (b, hh, da) # ---- Eq. 15 filter: fast attack, slow release, masked elements frozen. alpha = 0.25 seq = [0.0, 0.2, 1.0, 1.0, 1.0, 1.0, 0.3, 0.3, 0.1, 0.0, 0.0, 0.6] # rise, plateau, decay, rise mask1 = [True, True, False, True, False, True, True, False, True, True, False, True] filt = UncertaintyHorizonFilter(h_min=3, h_max=h, beta=1.0, alpha_sigma=alpha) filt.reset(3) assert torch.equal(filt.sigma_bar, torch.zeros(3)) ref_prev = [0.0, 0.0, 0.0] prev_mod = filt.sigma_bar.clone() rises, decays = [], [] for t, s in enumerate(seq): sig = [s, s, 0.5 * s] mask = [True, mask1[t], True] sb, ht = filt.update(torch.tensor(sig), mask=torch.tensor(mask)) ref = [max(sig[i], (1 - alpha) * ref_prev[i] + alpha * sig[i]) if mask[i] else ref_prev[i] for i in range(3)] assert _close(sb, torch.tensor(ref, dtype=torch.float64)) assert torch.equal(ht, filt.horizon(sb)) if t == 0: assert float(sb[0]) == float(torch.tensor(s)) # sigma_bar_0 = 0 -> sigma_bar = sigma if sig[0] > ref_prev[0]: # rise: applied immediately, bit-exact rises.append(t) assert float(sb[0]) == float(torch.tensor(sig[0])) if sig[0] < ref_prev[0]: # decay: gradual, stays above sigma decays.append(t) assert float(sb[0]) > sig[0] and float(sb[0]) < float(prev_mod[0]) if not mask[1]: assert float(sb[1]) == float(prev_mod[1]) # masked element keeps its value ref_prev, prev_mod = ref, sb.clone() assert {1, 2, 11} <= set(rises) and {6, 7, 8, 9, 10} <= set(decays) filt.reset(3) assert torch.equal(filt.sigma_bar, torch.zeros(3)) instant = UncertaintyHorizonFilter(h_min=3, h_max=h, beta=1.0, alpha_sigma=1.0) instant.reset(1) for s in seq: assert torch.equal(instant.filter(torch.tensor([s])), torch.tensor([s])) # ---- Eq. 15 horizon: H_t = max(H_min, floor(H_max exp(-beta sigma_bar))). h_min, beta = 3, 1.0 hf = UncertaintyHorizonFilter(h_min=h_min, h_max=h, beta=beta, alpha_sigma=alpha) def ref_h(sb_val: float) -> int: return max(h_min, math.floor(h * math.exp(-beta * sb_val))) assert hf.horizon(torch.tensor([0.0])).tolist() == [h] # sigma_bar = 0 -> H_max assert hf.horizon(torch.tensor([50.0])).tolist() == [h_min] # large sigma_bar -> H_min assert math.floor(h * math.exp(-1.05)) == h_min and math.floor(h * math.exp(-1.3)) == h_min - 1 assert hf.horizon(torch.tensor([1.05, 1.3])).tolist() == [h_min, h_min] # floor = H_min, clamped grid = torch.linspace(0.0, 3.0, 31)[1:] for v in grid.tolist(): x = h * math.exp(-beta * v) assert abs(x - round(x)) > 1e-4, v # floor is robust to float32 ht = hf.horizon(grid) assert ht.dtype == torch.long and ht.tolist() == [ref_h(v) for v in grid.tolist()] assert bool((ht[1:] <= ht[:-1]).all()) and bool(((ht >= h_min) & (ht <= h)).all()) hf.reset(len(grid)) sb_u, ht_u = hf.update(grid) assert torch.equal(ht_u, hf.horizon(sb_u)) h_exec = torch.tensor([h_min, 6, h]) pm = executed_prefix_mask(h_exec, h) assert pm.shape == (3, h) assert pm.tolist() == [[step < int(h_exec[i]) for step in range(h)] for i in range(3)] for kw in (dict(h_min=0), dict(h_min=h + 1), dict(beta=0.0), dict(beta=-1.0), dict(alpha_sigma=0.0), dict(alpha_sigma=1.5)): args = dict(h_min=3, h_max=h, beta=1.0, alpha_sigma=0.25) args.update(kw) _expect(ValueError, lambda a=args: UncertaintyHorizonFilter(**a)) print("solver.py self-test passed") if __name__ == "__main__": _self_test()