ACSSVAE C2 sync WaveSemanticHybridCodec/V5_ACSSVAE_C2_CausalProtectedReal/V5AwareStructuredSemanticVAE.py

WaveSemanticHybridCodec/V5_ACSSVAE_C2_CausalProtectedReal/V5AwareStructuredSemanticVAE.py

@@ -0,0 +1,622 @@
+"""
+V5-Aware Analytic-Centered Structured Semantic VAE (AC-SSVAE)
+================================================================
+Stage-2 feature extractor for ICWDS.
+
+Design goals
+------------
+1. Consume WaveSystemSetParserV5 soft wave-system outputs instead of hard labels.
+2. Use an analytic physical descriptor as the prior center, and learn only bounded
+   semantic corrections.
+3. Canonicalize each wave system before shape encoding so the free latent does not
+   waste capacity on location, direction, scale, or energy.
+4. Keep a small causal shape latent. The residual decoder is constructed as
+       Delta(z, s) = F(z, s) - F(0, s)
+   so z=0 exactly returns the semantic base.
+5. Expose deterministic decoding APIs for quantization, intervention, and later
+   30-byte packet design.
+
+Internal semantic vector (9D)
+-----------------------------
+[log_energy,
+ peak_frequency_01,
+ sin_peak_direction,
+ cos_peak_direction,
+ frequency_spread_01,
+ direction_spread_over_pi,
+ f_theta_correlation,
+ frequency_skew,
+ direction_skew]
+
+The packet-facing direction is still one circular quantity. sin/cos are only the
+neural internal representation.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, asdict
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class V5AwareACSSVAEConfig:
+    n_freqs: int = 47
+    n_dirs: int = 72
+    n_slots: int = 6
+    semantic_dim: int = 9
+    shape_dim: int = 4
+    context_dim: int = 24
+
+    canonical_h: int = 32
+    canonical_w: int = 48
+    canonical_f_extent: float = 3.0
+    canonical_t_extent: float = 3.0
+
+    encoder_width: int = 48
+    hidden_dim: int = 160
+    decoder_width: int = 64
+
+    # Bounded semantic correction magnitudes around analytic center.
+    # [logE, f, sin/cos correction, spreads, rho, skews]
+    delta_loge: float = 0.35
+    delta_f: float = 0.08
+    delta_angle_rad: float = math.radians(18.0)
+    delta_log_spread: float = 0.45
+    delta_rho_logit: float = 0.75
+    delta_skew: float = 0.90
+
+    min_spread_f: float = 0.010
+    max_spread_f: float = 0.45
+    min_spread_t: float = 0.015
+    max_spread_t: float = 0.95
+    max_abs_rho: float = 0.92
+    max_abs_skew: float = 3.0
+
+    shape_log_residual_scale: float = 2.25
+    energy_log_den: float = math.log1p(47 * 72)
+    eps: float = 1e-6
+
+    def to_dict(self):
+        return asdict(self)
+
+
+class ConvNormAct(nn.Module):
+    def __init__(self, ci: int, co: int, stride: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(ci, co, 3, stride=stride, padding=1, bias=False)
+        self.norm = nn.GroupNorm(min(8, co), co)
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.norm(self.conv(x)))
+
+
+class ResBlock(nn.Module):
+    def __init__(self, ch: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(ch, ch, 3, padding=1, bias=False),
+            nn.GroupNorm(min(8, ch), ch),
+            nn.GELU(),
+            nn.Conv2d(ch, ch, 3, padding=1, bias=False),
+            nn.GroupNorm(min(8, ch), ch),
+        )
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(x + self.net(x))
+
+
+def _wrap_angle(x: torch.Tensor) -> torch.Tensor:
+    return torch.atan2(torch.sin(x), torch.cos(x))
+
+
+def _atanh_safe(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
+    x = x.clamp(-1 + eps, 1 - eps)
+    return 0.5 * (torch.log1p(x) - torch.log1p(-x))
+
+
+def reparameterize(mu: torch.Tensor, logvar: torch.Tensor, stochastic: bool) -> torch.Tensor:
+    if not stochastic:
+        return mu
+    std = torch.exp(0.5 * logvar)
+    return mu + std * torch.randn_like(std)
+
+
+@torch.no_grad()
+def analytic_semantics(
+    energy: torch.Tensor,
+    mask_prob: torch.Tensor,
+    core_prob: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+    energy_log_den: Optional[float] = None,
+) -> torch.Tensor:
+    """Compute stable 9D physical descriptors from soft wave-system masks.
+
+    Parameters
+    ----------
+    energy: [B,1,H,W] or [B,H,W], non-negative display-domain energy.
+    mask_prob: [B,K,H,W].
+    core_prob: optional [B,K,H,W]. Core weighting sharpens peak estimation.
+    """
+    if energy.ndim == 3:
+        energy = energy[:, None]
+    B, K, H, W = mask_prob.shape
+    dtype, device = energy.dtype, energy.device
+    if energy_log_den is None:
+        energy_log_den = math.log1p(H * W)
+
+    E = torch.nan_to_num(energy, nan=0.0, posinf=1.0, neginf=0.0).clamp_min(0)
+    M = torch.nan_to_num(mask_prob, nan=0.0).clamp(0, 1)
+    part = E * M
+    mass = part.sum(dim=(-2, -1)).clamp_min(eps)
+    loge = torch.log1p(mass) / float(energy_log_den)
+
+    fgrid = torch.linspace(0, 1, H, device=device, dtype=dtype).view(1, 1, H, 1)
+    theta = torch.linspace(0, 2 * math.pi, W + 1, device=device, dtype=dtype)[:W]
+    tgrid = theta.view(1, 1, 1, W)
+
+    peak_src = part
+    if core_prob is not None:
+        core = torch.nan_to_num(core_prob, nan=0.0).clamp(0, 1)
+        peak_src = part * (0.25 + 0.75 * core)
+
+    # Smooth soft-peak estimator. It is more stable than hard argmax on noisy spectra.
+    pnorm = peak_src / peak_src.amax(dim=(-2, -1), keepdim=True).clamp_min(eps)
+    peak_w = torch.softmax((12.0 * pnorm).flatten(2), dim=-1).view(B, K, H, W)
+    fp = (peak_w * fgrid).sum(dim=(-2, -1))
+    sx = (peak_w * torch.sin(tgrid)).sum(dim=(-2, -1))
+    cx = (peak_w * torch.cos(tgrid)).sum(dim=(-2, -1))
+    theta_p = torch.atan2(sx, cx)
+    sinp = torch.sin(theta_p)
+    cosp = torch.cos(theta_p)
+
+    w = part / mass[:, :, None, None]
+    df = fgrid - fp[:, :, None, None]
+    dt = _wrap_angle(tgrid - theta_p[:, :, None, None]) / math.pi
+
+    sf = torch.sqrt((w * df.square()).sum(dim=(-2, -1)).clamp_min(eps))
+    st = torch.sqrt((w * dt.square()).sum(dim=(-2, -1)).clamp_min(eps))
+
+    zf = df / sf[:, :, None, None].clamp_min(1e-3)
+    zt = dt / st[:, :, None, None].clamp_min(1e-3)
+    rho = (w * zf * zt).sum(dim=(-2, -1)).clamp(-0.95, 0.95)
+    skew_f = (w * zf.pow(3)).sum(dim=(-2, -1)).clamp(-4.0, 4.0)
+    skew_t = (w * zt.pow(3)).sum(dim=(-2, -1)).clamp(-4.0, 4.0)
+
+    return torch.stack([
+        loge, fp, sinp, cosp, sf, st, rho, skew_f, skew_t
+    ], dim=-1)
+
+
+class Canonicalizer(nn.Module):
+    """Differentiably center/scale each slot in frequency-direction coordinates."""
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        u = torch.linspace(-cfg.canonical_f_extent, cfg.canonical_f_extent, cfg.canonical_h)
+        v = torch.linspace(-cfg.canonical_t_extent, cfg.canonical_t_extent, cfg.canonical_w)
+        self.register_buffer("u", u.view(1, 1, cfg.canonical_h, 1))
+        self.register_buffer("v", v.view(1, 1, 1, cfg.canonical_w))
+
+    def forward(
+        self,
+        energy: torch.Tensor,
+        mask_prob: torch.Tensor,
+        core_prob: torch.Tensor,
+        support_prob: torch.Tensor,
+        semantics: torch.Tensor,
+    ) -> torch.Tensor:
+        if energy.ndim == 3:
+            energy = energy[:, None]
+        B, K, H, W = mask_prob.shape
+        eps = self.cfg.eps
+
+        E = energy[:, None].expand(B, K, 1, H, W)
+        M = mask_prob[:, :, None]
+        C = core_prob[:, :, None]
+        S = support_prob[:, :, None]
+        part = E * M
+        cpart = E * C
+        spart = E * S
+
+        peak = part.amax(dim=(-2, -1), keepdim=True).clamp_min(eps)
+        channels = torch.cat([
+            part / peak,
+            cpart / peak,
+            spart / peak,
+            M,
+            C,
+            S,
+        ], dim=2).reshape(B * K, 6, H, W)
+
+        sem = semantics.reshape(B * K, -1)
+        fp = sem[:, 1].clamp(0, 1)
+        theta = torch.atan2(sem[:, 2], sem[:, 3]) % (2 * math.pi)
+        sf = sem[:, 4].clamp(self.cfg.min_spread_f, self.cfg.max_spread_f)
+        st = sem[:, 5].clamp(self.cfg.min_spread_t, self.cfg.max_spread_t)
+
+        uf = self.u.to(channels.dtype)
+        vt = self.v.to(channels.dtype)
+        f_sample = fp[:, None, None, None] + sf[:, None, None, None] * uf
+        theta_sample = theta[:, None, None, None] + (math.pi * st[:, None, None, None]) * vt
+        theta_frac = torch.remainder(theta_sample, 2 * math.pi) / (2 * math.pi)
+
+        # Tile direction three times and sample from the middle copy.
+        tiled = torch.cat([channels, channels, channels], dim=-1)
+        y = 2.0 * f_sample - 1.0
+        xpix = (1.0 + theta_frac) * W - 0.5
+        x = 2.0 * xpix / max(3 * W - 1, 1) - 1.0
+        grid = torch.stack([
+            x.expand(-1, -1, self.cfg.canonical_h, self.cfg.canonical_w).squeeze(1),
+            y.expand(-1, -1, self.cfg.canonical_h, self.cfg.canonical_w).squeeze(1),
+        ], dim=-1)
+        out = F.grid_sample(
+            tiled,
+            grid,
+            mode="bilinear",
+            padding_mode="zeros",
+            align_corners=True,
+        )
+        return out.reshape(B, K, 6, self.cfg.canonical_h, self.cfg.canonical_w)
+
+
+class SlotEncoder(nn.Module):
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        w = cfg.encoder_width
+        self.image_net = nn.Sequential(
+            ConvNormAct(6, w, 1),
+            ResBlock(w),
+            ConvNormAct(w, w * 2, 2),
+            ResBlock(w * 2),
+            ConvNormAct(w * 2, w * 3, 2),
+            ResBlock(w * 3),
+            ConvNormAct(w * 3, w * 4, 2),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+        )
+        self.context_net = nn.Sequential(
+            nn.Linear(cfg.context_dim + cfg.semantic_dim + 1, cfg.hidden_dim // 2),
+            nn.LayerNorm(cfg.hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(cfg.hidden_dim // 2, cfg.hidden_dim // 2),
+            nn.GELU(),
+        )
+        image_dim = w * 4
+        self.fuse = nn.Sequential(
+            nn.Linear(image_dim + cfg.hidden_dim // 2, cfg.hidden_dim),
+            nn.LayerNorm(cfg.hidden_dim),
+            nn.GELU(),
+            nn.Linear(cfg.hidden_dim, cfg.hidden_dim),
+            nn.GELU(),
+        )
+        self.sem_mu = nn.Linear(cfg.hidden_dim, cfg.semantic_dim)
+        self.sem_logvar = nn.Linear(cfg.hidden_dim, cfg.semantic_dim)
+        self.shape_mu = nn.Linear(cfg.hidden_dim, cfg.shape_dim)
+        self.shape_logvar = nn.Linear(cfg.hidden_dim, cfg.shape_dim)
+
+        # Start from the analytic center and nearly-zero free shape code.
+        nn.init.zeros_(self.sem_mu.weight)
+        nn.init.zeros_(self.sem_mu.bias)
+        nn.init.constant_(self.sem_logvar.bias, -4.0)
+        nn.init.normal_(self.shape_mu.weight, std=0.01)
+        nn.init.zeros_(self.shape_mu.bias)
+        nn.init.constant_(self.shape_logvar.bias, -3.0)
+
+    def forward(
+        self,
+        canonical: torch.Tensor,
+        analytic_sem: torch.Tensor,
+        context: torch.Tensor,
+        exist_prob: torch.Tensor,
+    ) -> Dict[str, torch.Tensor]:
+        B, K = canonical.shape[:2]
+        img = self.image_net(canonical.reshape(B * K, *canonical.shape[2:]))
+        ctx_in = torch.cat([
+            context.reshape(B * K, -1),
+            analytic_sem.reshape(B * K, -1),
+            exist_prob.reshape(B * K, 1),
+        ], dim=-1)
+        ctx = self.context_net(ctx_in)
+        h = self.fuse(torch.cat([img, ctx], dim=-1))
+        sem_mu = self.sem_mu(h).reshape(B, K, -1)
+        sem_logvar = self.sem_logvar(h).clamp(-8.0, 3.0).reshape(B, K, -1)
+        shape_mu = self.shape_mu(h).reshape(B, K, -1)
+        shape_logvar = self.shape_logvar(h).clamp(-8.0, 3.0).reshape(B, K, -1)
+        return {
+            "sem_delta_mu": sem_mu,
+            "sem_delta_logvar": sem_logvar,
+            "shape_mu": shape_mu,
+            "shape_logvar": shape_logvar,
+        }
+
+
+class SemanticRenderer(nn.Module):
+    """Parameter-free generalized elliptical renderer from corrected semantics."""
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        f = torch.linspace(0, 1, cfg.n_freqs).view(1, 1, cfg.n_freqs, 1)
+        t = torch.linspace(0, 2 * math.pi, cfg.n_dirs + 1)[:cfg.n_dirs].view(1, 1, 1, cfg.n_dirs)
+        self.register_buffer("fgrid", f)
+        self.register_buffer("tgrid", t)
+
+    def forward(self, sem: torch.Tensor) -> torch.Tensor:
+        cfg = self.cfg
+        eps = cfg.eps
+        loge, fp = sem[..., 0], sem[..., 1]
+        theta = torch.atan2(sem[..., 2], sem[..., 3])
+        sf = sem[..., 4].clamp(cfg.min_spread_f, cfg.max_spread_f)
+        st = sem[..., 5].clamp(cfg.min_spread_t, cfg.max_spread_t)
+        rho = sem[..., 6].clamp(-cfg.max_abs_rho, cfg.max_abs_rho)
+        skew_f = sem[..., 7].clamp(-cfg.max_abs_skew, cfg.max_abs_skew)
+        skew_t = sem[..., 8].clamp(-cfg.max_abs_skew, cfg.max_abs_skew)
+
+        df = (self.fgrid - fp[:, :, None, None]) / sf[:, :, None, None].clamp_min(1e-3)
+        dt = _wrap_angle(self.tgrid - theta[:, :, None, None])
+        dt = dt / (math.pi * st[:, :, None, None].clamp_min(1e-3))
+        den = (1.0 - rho.square()).clamp_min(0.08)
+        q = (df.square() + dt.square() - 2.0 * rho[:, :, None, None] * df * dt) / den[:, :, None, None]
+        shape = torch.exp(-0.5 * q.clamp_max(60.0))
+        asym = torch.exp(
+            0.28 * skew_f[:, :, None, None] * torch.tanh(df)
+            + 0.28 * skew_t[:, :, None, None] * torch.tanh(dt)
+        ).clamp(0.15, 6.0)
+        shape = shape * asym
+        shape = shape / shape.sum(dim=(-2, -1), keepdim=True).clamp_min(eps)
+        mass = torch.expm1(loge * cfg.energy_log_den).clamp_min(0.0)
+        return shape * mass[:, :, None, None]
+
+
+class ShapeResidualNet(nn.Module):
+    def __init__(self, cfg: V5AwareACSSVAEConfig):
+        super().__init__()
+        self.cfg = cfg
+        w = cfg.decoder_width
+        self.fc = nn.Sequential(
+            nn.Linear(cfg.semantic_dim + cfg.shape_dim, cfg.hidden_dim),
+            nn.GELU(),
+            nn.Linear(cfg.hidden_dim, w * 6 * 9),
+            nn.GELU(),
+        )
+        self.net = nn.Sequential(
+            ConvNormAct(w, w, 1),
+            ResBlock(w),
+            nn.Upsample(size=(12, 18), mode="bilinear", align_corners=False),
+            ConvNormAct(w, w, 1),
+            ResBlock(w),
+            nn.Upsample(size=(24, 36), mode="bilinear", align_corners=False),
+            ConvNormAct(w, w // 2, 1),
+            ResBlock(w // 2),
+            nn.Upsample(size=(48, 72), mode="bilinear", align_corners=False),
+            ConvNormAct(w // 2, w // 2, 1),
+            nn.Conv2d(w // 2, 1, 3, padding=1),
+        )
+
+    def field(self, sem: torch.Tensor, zshape: torch.Tensor) -> torch.Tensor:
+        B, K = sem.shape[:2]
+        x = torch.cat([sem, zshape], dim=-1).reshape(B * K, -1)
+        h = self.fc(x).view(B * K, self.cfg.decoder_width, 6, 9)
+        out = self.net(h)[:, :, : self.cfg.n_freqs, : self.cfg.n_dirs]
+        return out.reshape(B, K, self.cfg.n_freqs, self.cfg.n_dirs)
+
+    def forward(self, sem: torch.Tensor, zshape: torch.Tensor) -> torch.Tensor:
+        raw = self.field(sem, zshape)
+        zero = self.field(sem, torch.zeros_like(zshape))
+        return self.cfg.shape_log_residual_scale * torch.tanh(raw - zero)
+
+
+class V5AwareStructuredSemanticVAE(nn.Module):
+    def __init__(self, cfg: Optional[V5AwareACSSVAEConfig] = None):
+        super().__init__()
+        self.cfg = cfg or V5AwareACSSVAEConfig()
+        self.canonicalizer = Canonicalizer(self.cfg)
+        self.encoder = SlotEncoder(self.cfg)
+        self.semantic_renderer = SemanticRenderer(self.cfg)
+        self.shape_decoder = ShapeResidualNet(self.cfg)
+
+    def apply_semantic_delta(self, base: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        c = self.cfg
+        out = base.clone()
+        out[..., 0] = (base[..., 0] + c.delta_loge * torch.tanh(delta[..., 0])).clamp_min(0.0)
+        out[..., 1] = (base[..., 1] + c.delta_f * torch.tanh(delta[..., 1])).clamp(0.0, 1.0)
+
+        theta0 = torch.atan2(base[..., 2], base[..., 3])
+        # Use both sin/cos correction channels to form one stable tangent-angle correction.
+        dtheta = c.delta_angle_rad * torch.tanh(0.7071 * (delta[..., 2] - delta[..., 3]))
+        theta = theta0 + dtheta
+        out[..., 2] = torch.sin(theta)
+        out[..., 3] = torch.cos(theta)
+
+        out[..., 4] = (
+            base[..., 4].clamp_min(c.min_spread_f)
+            * torch.exp(c.delta_log_spread * torch.tanh(delta[..., 4]))
+        ).clamp(c.min_spread_f, c.max_spread_f)
+        out[..., 5] = (
+            base[..., 5].clamp_min(c.min_spread_t)
+            * torch.exp(c.delta_log_spread * torch.tanh(delta[..., 5]))
+        ).clamp(c.min_spread_t, c.max_spread_t)
+        out[..., 6] = torch.tanh(
+            _atanh_safe(base[..., 6], c.eps)
+            + c.delta_rho_logit * torch.tanh(delta[..., 6])
+        ).clamp(-c.max_abs_rho, c.max_abs_rho)
+        out[..., 7] = (
+            base[..., 7] + c.delta_skew * torch.tanh(delta[..., 7])
+        ).clamp(-c.max_abs_skew, c.max_abs_skew)
+        out[..., 8] = (
+            base[..., 8] + c.delta_skew * torch.tanh(delta[..., 8])
+        ).clamp(-c.max_abs_skew, c.max_abs_skew)
+        return out
+
+    def decode_from_codes(
+        self,
+        semantics: torch.Tensor,
+        shape_code: torch.Tensor,
+        exist_prob: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        base = self.semantic_renderer(semantics)
+        residual_log = self.shape_decoder(semantics, shape_code)
+        # Multiplicative residual in log domain, followed by exact energy renormalization.
+        part_hat = (base + self.cfg.eps) * torch.exp(residual_log)
+        target_mass = torch.expm1(semantics[..., 0] * self.cfg.energy_log_den).clamp_min(0.0)
+        part_hat = part_hat * (
+            target_mass[:, :, None, None]
+            / part_hat.sum(dim=(-2, -1), keepdim=True).clamp_min(self.cfg.eps)
+        )
+        if exist_prob is not None:
+            part_hat = part_hat * exist_prob[:, :, None, None]
+            base = base * exist_prob[:, :, None, None]
+        return {
+            "semantic_base": base,
+            "shape_residual_log": residual_log,
+            "part_hat": part_hat,
+        }
+
+    def forward(
+        self,
+        energy: torch.Tensor,
+        mask_prob: torch.Tensor,
+        core_prob: torch.Tensor,
+        support_prob: torch.Tensor,
+        exist_prob: torch.Tensor,
+        slot_context: torch.Tensor,
+        stochastic: bool = True,
+        shape_code_override: Optional[torch.Tensor] = None,
+        semantic_delta_override: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        analytic = analytic_semantics(
+            energy,
+            mask_prob,
+            core_prob=core_prob,
+            eps=self.cfg.eps,
+            energy_log_den=self.cfg.energy_log_den,
+        )
+        canonical = self.canonicalizer(
+            energy,
+            mask_prob,
+            core_prob,
+            support_prob,
+            analytic,
+        )
+        enc = self.encoder(canonical, analytic, slot_context, exist_prob)
+        sem_delta = (
+            semantic_delta_override
+            if semantic_delta_override is not None
+            else reparameterize(enc["sem_delta_mu"], enc["sem_delta_logvar"], stochastic)
+        )
+        zshape = (
+            shape_code_override
+            if shape_code_override is not None
+            else reparameterize(enc["shape_mu"], enc["shape_logvar"], stochastic)
+        )
+        corrected = self.apply_semantic_delta(analytic, sem_delta)
+        dec = self.decode_from_codes(corrected, zshape, exist_prob=exist_prob)
+        return {
+            "analytic_semantics": analytic,
+            "corrected_semantics": corrected,
+            "semantic_delta": sem_delta,
+            "shape_code": zshape,
+            "canonical": canonical,
+            **enc,
+            **dec,
+        }
+
+    @property
+    def token_dim_per_wave(self) -> int:
+        return self.cfg.semantic_dim + self.cfg.shape_dim
+
+    def num_params(self) -> int:
+        return sum(p.numel() for p in self.parameters())
+
+
+def kl_standard_normal(mu: torch.Tensor, logvar: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+    kl = -0.5 * (1.0 + logvar - mu.square() - logvar.exp()).sum(dim=-1)
+    if weight is None:
+        return kl.mean()
+    w = weight.float()
+    return (kl * w).sum() / w.sum().clamp_min(1.0)
+
+
+def covariance_penalty(x: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Off-diagonal covariance penalty for [B,K,D] latent tensors."""
+    D = x.shape[-1]
+    flat = x.reshape(-1, D)
+    if weight is not None:
+        w = weight.reshape(-1).float()
+        keep = w > 0.25
+        flat = flat[keep]
+    if flat.shape[0] < max(4, D):
+        return x.new_zeros(())
+    flat = flat - flat.mean(dim=0, keepdim=True)
+    cov = flat.T @ flat / max(flat.shape[0] - 1, 1)
+    off = cov - torch.diag(torch.diag(cov))
+    return off.square().mean()
+
+
+def semantic_distance(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Per-slot interpretable semantic distance with circular direction handling."""
+    d = []
+    d.append((pred[..., 0] - target[..., 0]).abs())
+    d.append(2.0 * (pred[..., 1] - target[..., 1]).abs())
+    tp = torch.atan2(pred[..., 2], pred[..., 3])
+    tt = torch.atan2(target[..., 2], target[..., 3])
+    d.append(_wrap_angle(tp - tt).abs() / math.pi)
+    d.append(2.0 * (pred[..., 4] - target[..., 4]).abs())
+    d.append(1.5 * (pred[..., 5] - target[..., 5]).abs())
+    d.append(0.5 * (pred[..., 6] - target[..., 6]).abs())
+    d.append(0.15 * (pred[..., 7] - target[..., 7]).abs())
+    d.append(0.15 * (pred[..., 8] - target[..., 8]).abs())
+    return torch.stack(d, dim=-1).mean(dim=-1)
+
+
+def marginal_l1(pred: torch.Tensor, target: torch.Tensor, eps: float = 1e-6) -> Tuple[torch.Tensor, torch.Tensor]:
+    p = pred.clamp_min(0)
+    t = target.clamp_min(0)
+    pf = p.sum(dim=-1); tf = t.sum(dim=-1)
+    pt = p.sum(dim=-2); tt = t.sum(dim=-2)
+    pf = pf / pf.sum(dim=-1, keepdim=True).clamp_min(eps)
+    tf = tf / tf.sum(dim=-1, keepdim=True).clamp_min(eps)
+    pt = pt / pt.sum(dim=-1, keepdim=True).clamp_min(eps)
+    tt = tt / tt.sum(dim=-1, keepdim=True).clamp_min(eps)
+    return (pf - tf).abs().mean(), (pt - tt).abs().mean()
+
+
+def reconstruction_terms(pred: torch.Tensor, target: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> Dict[str, torch.Tensor]:
+    w = weight[:, :, None, None].float()
+    denom = w.sum().clamp_min(1.0)
+    mse_slot = (pred - target).square().mean(dim=(-2, -1))
+    log_slot = (
+        torch.log1p(20.0 * pred.clamp_min(0))
+        - torch.log1p(20.0 * target.clamp_min(0))
+    ).abs().mean(dim=(-2, -1))
+    mass_p = pred.sum(dim=(-2, -1))
+    mass_t = target.sum(dim=(-2, -1))
+    energy_slot = (mass_p - mass_t).abs() / mass_t.clamp_min(eps)
+    mse = (mse_slot * weight).sum() / denom
+    log_l1 = (log_slot * weight).sum() / denom
+    energy = (energy_slot * weight).sum() / denom
+    freq, direction = marginal_l1(pred * w, target * w, eps=eps)
+    return {
+        "mse": mse,
+        "log_l1": log_l1,
+        "energy": energy,
+        "freq": freq,
+        "direction": direction,
+    }
+
+
+__all__ = [
+    "V5AwareACSSVAEConfig",
+    "V5AwareStructuredSemanticVAE",
+    "analytic_semantics",
+    "kl_standard_normal",
+    "covariance_penalty",
+    "semantic_distance",
+    "reconstruction_terms",
+    "marginal_l1",
+]