src/uncertainty.py

# uncertainty.py
"""
Multi-sample variance uncertainty estimation for the Dreamer 4 dynamics model.

The diffusion-based dynamics model starts each prediction from random noise z ~ N(0,1)
and integrates to a prediction via Euler ODE steps. Different noise seeds produce
different predictions. In well-covered regions, predictions converge regardless of seed.
In poorly-covered regions, predictions diverge. This variance is a free, reward-free
uncertainty estimator that requires no architectural changes.
"""
from typing import Dict, Any, Optional

import torch
from torch.amp import autocast

from model import (
    Dynamics, Encoder, Decoder,
    temporal_patchify, pack_bottleneck_to_spatial,
)
from train_dynamics import (
    sample_one_timestep_packed,
    decode_packed_to_frames,
)


# ---------------------------------------------------------------------------
# Unified sampler + scorers (used by curiosity MPC to support multiple signals)
# ---------------------------------------------------------------------------

@torch.no_grad()
def sample_predictions_for_actions(
    dyn: Dynamics,
    *,
    past_packed: torch.Tensor,                     # (1, t, Sz, Dz)
    candidate_actions: torch.Tensor,               # (K, t+1, A)
    k_max: int,
    sched: Dict[str, Any],
    act_mask: Optional[torch.Tensor] = None,
    tau_ctx: float = 0.1,
    lang_emb: Optional[torch.Tensor] = None,
    n_samples: int = 2,
    use_kv_cache: bool = False,
) -> torch.Tensor:
    """
    Run N independent diffusion samples for each of K candidate action sequences.
    Returns (K, N, Sz, Dz) float32.

    With `use_kv_cache=True` the t context tokens are run through the
    transformer once and reused across all K denoising steps — the same
    optimization as in `train_dynamics.sample_one_timestep_packed`. Default
    off; enable for the per-env-step uncertainty logging in collect_data.py,
    where this is the dominant per-step cost.
    """
    K = candidate_actions.shape[0]
    _, t, Sz, Dz = past_packed.shape

    dtype = next(dyn.parameters()).dtype
    past_packed = past_packed.to(dtype)
    candidate_actions = candidate_actions.to(dtype)

    past_KN = past_packed.expand(K, -1, -1, -1).unsqueeze(1).expand(-1, n_samples, -1, -1, -1)
    past_KN = past_KN.reshape(K * n_samples, t, Sz, Dz)

    T_act, A = candidate_actions.shape[1], candidate_actions.shape[2]
    actions_KN = candidate_actions.unsqueeze(1).expand(-1, n_samples, -1, -1)
    actions_KN = actions_KN.reshape(K * n_samples, T_act, A)

    lang_KN = None if lang_emb is None else lang_emb.expand(K * n_samples, -1)

    predictions = sample_one_timestep_packed(
        dyn,
        past_packed=past_KN,
        k_max=k_max,
        sched=sched,
        actions=actions_KN,
        act_mask=act_mask,
        tau_ctx=tau_ctx,
        lang_emb=lang_KN,
        use_kv_cache=use_kv_cache,
    )  # (K*N, Sz, Dz)
    return predictions.float().reshape(K, n_samples, Sz, Dz)


class CrossSeedScorer:
    """
    Score K candidates by per-element variance across N diffusion seeds — the
    inter-seed denoising-variance predictor (u_s). Raw cross-seed variance is
    motion-invariant only in aggregate, so it confounds with scene motion.
    """

    def score_components(self, predictions_KN: torch.Tensor, z_prev_K: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Return the per-candidate score components. Only "u_r_norm" (here the
        inter-seed variance) is returned; the `u_r` and `motion` keys do not
        apply to this scorer."""
        return {
            "u_r_norm": predictions_KN.float().var(dim=1).mean(dim=(1, 2)),  # (K,)
        }

    def __call__(self, predictions_KN: torch.Tensor, z_prev_K: torch.Tensor) -> torch.Tensor:
        # Hot path: CEM scoring. predictions_KN: (K, N, Sz, Dz); z_prev_K unused.
        return self.score_components(predictions_KN, z_prev_K)["u_r_norm"]


class URNormScorer:
    """
    Score K candidates by the tokenizer round-trip residual of the mean predicted
    latent, normalized by the predicted latent-space step motion:

        u_r      = RMS( z_pred - encode(decode(z_pred)) )
        motion   = RMS( z_pred - z_prev )
        u_r_norm = u_r / max(motion, eps)

    This is the motion-normalized tokenizer round-trip residual signal (u_norm)
    and is label-free.
    """

    def __init__(
        self,
        encoder: Encoder,
        decoder: Decoder,
        *,
        patch: int,
        packing_factor: int,
        n_spatial: int,
        H: int,
        W: int,
        C: int = 3,
        motion_eps: float = 1e-3,
    ):
        self.encoder = encoder
        self.decoder = decoder
        self.patch = patch
        self.packing_factor = packing_factor
        self.n_spatial = n_spatial
        self.H = H
        self.W = W
        self.C = C
        self.motion_eps = motion_eps

    @torch.no_grad()
    def score_components(self, predictions_KN: torch.Tensor, z_prev_K: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Same compute as __call__, but returns the raw `u_r`, `motion`, and
        the ratio `u_r_norm` separately.

        Returns dict with three (K,)-shaped float tensors:
          - "u_r":      RMS( z_pred - encode(decode(z_pred)) ) — round-trip residual
          - "motion":   RMS( z_pred - z_prev )                 — predicted step
          - "u_r_norm": u_r / max(motion, motion_eps)          — what __call__ returns
        """
        # Collapse over N via mean to get one predicted latent per candidate.
        z_pred_K = predictions_KN.float().mean(dim=1)                        # (K, Sz, Dz)

        motion_K = (z_pred_K - z_prev_K.float()).pow(2).mean(dim=(1, 2)).sqrt()  # (K,)

        enc_dtype = next(self.encoder.parameters()).dtype
        z_pred_in = z_pred_K.unsqueeze(1).to(enc_dtype)                       # (K, 1, Sz, Dz)
        with autocast(device_type=z_pred_K.device.type, dtype=torch.bfloat16):
            frames = decode_packed_to_frames(
                self.decoder,
                z_packed=z_pred_in,
                H=self.H, W=self.W, C=self.C,
                patch=self.patch,
                packing_factor=self.packing_factor,
            )                                                                 # (K, 1, C, H, W)
            patches = temporal_patchify(frames, self.patch)                   # (K, 1, Np, Dp)
            z_recon_btLd, _ = self.encoder(patches)                           # (K, 1, L, D_b)
        z_recon_K = pack_bottleneck_to_spatial(
            z_recon_btLd, n_spatial=self.n_spatial, k=self.packing_factor,
        )[:, 0].float()                                                        # (K, Sz, Dz)

        u_r_K = (z_pred_K - z_recon_K).pow(2).mean(dim=(1, 2)).sqrt()          # (K,)
        u_r_norm_K = u_r_K / motion_K.clamp(min=self.motion_eps)
        return {"u_r": u_r_K, "motion": motion_K, "u_r_norm": u_r_norm_K}

    @torch.no_grad()
    def __call__(self, predictions_KN: torch.Tensor, z_prev_K: torch.Tensor) -> torch.Tensor:
        # Hot path: CEM scoring. Returns just the u_r_norm ratio scalar.
        return self.score_components(predictions_KN, z_prev_K)["u_r_norm"]