src/interactive_uncertainty_corr.py

# interactive_uncertainty_corr.py
"""Interactive world-model UI that also steps the live env to expose true
one-step prediction error, for correlating u_r / u_f / u_s with ground truth.

Separate from interactive_uncertainty.py (original open-loop vis unchanged).

True error matches collect_data.py:
  true_error = RMS(z_pred - encode(env_next))
and the latent history is teacher-forced with the env encoding so each step
is a clean one-step prediction (not open-loop accumulation).

Run from ``src/``:
  python interactive_uncertainty_corr.py --tokenizer_ckpt ... --dynamics_ckpt ...
"""
from __future__ import annotations

import argparse
import math
from typing import Any, Dict, Optional, Tuple

import numpy as np
import torch
from aiohttp import web

from interactive_uncertainty import (
    InteractiveServer,
    SessionState,
    _as_2d_packed,
    build_action_from_keys,
    build_parser as _base_parser,
    classify_uncertainty,
    decode_single_packed_frame,
    env_obs_to_frame_chw01,
    frame_to_jpeg_bytes,
    frame_to_uint8_hwc,
    reward_from_reward_head_output,
    sample_one_timestep_packed,
)


def rms_latent_error(z_a: torch.Tensor, z_b: torch.Tensor) -> float:
    return float((z_a.float() - z_b.float()).pow(2).mean().sqrt().item())


class CorrInteractiveServer(InteractiveServer):
    """Same interactive server, but each stepped frame also env.steps and
    reports true_error + raw u_* for a correlation histogram frontend."""

    def _step_env_and_encode(self, st: SessionState, a: torch.Tensor) -> Tuple[torch.Tensor, bool]:
        """Apply action to the live env; return (z_gt packed, done)."""
        env = self._get_or_make_env(st.task)
        act_dim = max(0, int(st.act_dim))
        if act_dim <= 0:
            # No controllable dims — still need a zero step for some envs.
            a_np = np.zeros(env.action_space.shape, dtype=np.float32)
        else:
            real_dim = int(env.action_space.shape[0])
            a_np = a.detach().float().cpu().numpy()[:real_dim].astype(np.float32)
            if a_np.shape[0] < real_dim:
                pad = np.zeros(real_dim, dtype=np.float32)
                pad[: a_np.shape[0]] = a_np
                a_np = pad
        obs, _reward, terminated, truncated, _info = env.step(a_np)
        done = bool(terminated or truncated)
        frame = env_obs_to_frame_chw01(obs, H=self.H, W=self.W).to(self.device)
        z_gt = self._encode_frame_to_packed(frame)
        return _as_2d_packed(z_gt), done

    def _render_step_sync(self, st: SessionState) -> Tuple[Optional[bytes], Dict[str, Any]]:
        if st.reset_requested:
            self._reset_session(st)
            # New episode: clear the cumulative error accumulator.
            st.cum_true_error = 0.0
            st.n_true_error = 0

        # Ensure corr-only fields exist (SessionState is a plain dataclass).
        if not hasattr(st, "last_true_error"):
            st.last_true_error = float("nan")
        if not hasattr(st, "env_done"):
            st.env_done = False
        if not hasattr(st, "cum_true_error"):
            st.cum_true_error = 0.0
        if not hasattr(st, "n_true_error"):
            st.n_true_error = 0

        a_raw = build_action_from_keys(
            st.keys_down, act_dim=st.act_dim, A=16
        ).to(self.device)
        a_raw = (a_raw.clamp(-1, 1) * st.act_mask_1d).to(torch.float32)

        beta = float(st.action_beta)
        if beta > 0.0:
            beta = min(max(beta, 0.0), 0.999)
            st.a_smooth = (beta * st.a_smooth + (1.0 - beta) * a_raw).to(torch.float32)
            a = st.a_smooth
        else:
            a = a_raw

        frame_cur: Optional[torch.Tensor] = None
        stepped: bool = False
        env_done = False

        if not st.paused and st.act_dim >= 0:
            stepped = True
            st.a_hist.append(a)

            past, actions_local, actmask_local = self._build_local_window(st)

            need_h = self.rew_head is not None
            N = self.n_samples_u
            z_prev_1 = st.z_hist[-1]
            past_N = past.expand(N, -1, -1, -1).contiguous()
            actions_N = actions_local.expand(N, -1, -1).contiguous()
            actmask_N = actmask_local.expand(N, -1, -1).contiguous()
            z_prev_N = (
                z_prev_1.unsqueeze(0).expand(N, -1, -1).contiguous()
                if self.tau_init > 0.0 else None
            )
            lang_N = None if st.lang_emb is None else st.lang_emb.expand(N, -1).contiguous()
            result = sample_one_timestep_packed(
                self.dyn,
                past_packed=past_N,
                k_max=self.k_max,
                sched=self.sched,
                actions=actions_N,
                act_mask=actmask_N,
                use_amp=self.use_amp,
                return_h=need_h,
                tau_ctx=self.tau_ctx,
                lang_emb=lang_N,
                z_prev=z_prev_N,
                tau_init=self.tau_init,
                use_kv_cache=self.use_kv_cache,
            )
            if need_h:
                z_next_N, h_N, instability = result
            else:
                z_next_N, instability = result
            st.last_u_f = float(instability)

            z_mean_N = z_next_N.float().mean(dim=0)
            u_s_raw = z_next_N.float().var(dim=0).mean().clamp(min=0).sqrt().item()
            motion = (z_mean_N - z_prev_1.float()).pow(2).mean().sqrt().item()
            st.last_u_s = u_s_raw / max(motion, 1e-3)

            z_next = z_next_N[0]
            h = h_N[0:1] if need_h else None

            # Live env one-step target (same action), then teacher-force history.
            z_gt, env_done = self._step_env_and_encode(st, a)
            st.last_true_error = rms_latent_error(z_next, z_gt)
            st.env_done = env_done
            # Cumulative (running sum) of the per-step error since last reset.
            st.cum_true_error = float(st.cum_true_error) + st.last_true_error
            st.n_true_error = int(st.n_true_error) + 1

            # History uses GT latent so the next predictor call is one-step.
            st.z_hist.append(_as_2d_packed(z_gt.detach()))
            st.step += 1

            cap = int(st.ctx_window) + 1
            if len(st.z_hist) > cap:
                st.z_hist = st.z_hist[-cap:]
                st.a_hist = st.a_hist[-cap:]

            if self.args.uncertainty_overlay and (st.step % max(1, int(self.args.u_every)) == 0):
                # Round-trip on the *predicted* latent (what u_r scores), not GT.
                frame_cur = decode_single_packed_frame(
                    self.decoder,
                    z_packed=_as_2d_packed(z_next.detach()),
                    H=self.H, W=self.W, C=self.C, patch=self.patch,
                    packing_factor=self.args.packing_factor,
                    d_bottleneck=self.d_bottleneck,
                )
                z_recon = self._encode_frame_to_packed(frame_cur)
                diff = z_next.to(torch.float32) - z_recon
                st.last_u_r = float(diff.pow(2).mean().sqrt().item())

                if not st.calib_done:
                    st.calib_f_samples.append(st.last_u_f)
                    st.calib_r_samples.append(st.last_u_r)
                    st.calib_s_samples.append(st.last_u_s)
                    if len(st.calib_f_samples) >= int(self.args.calibration_steps):
                        st.calib_done = True

            if need_h:
                logits_btlk, centers = self.rew_head(h[:, -1:])
                st.last_reward_pred = reward_from_reward_head_output(logits_btlk[0, 0], centers)
                st.cum_reward += st.last_reward_pred

            # Display the WM prediction being scored (not the GT frame).
            if frame_cur is None:
                frame_cur = decode_single_packed_frame(
                    self.decoder,
                    z_packed=_as_2d_packed(z_next.detach()),
                    H=self.H, W=self.W, C=self.C, patch=self.patch,
                    packing_factor=self.args.packing_factor,
                    d_bottleneck=self.d_bottleneck,
                )

            if env_done:
                # Episode ended in the real env — reseed on the next tick.
                st.reset_requested = True

        frame_id = st.step
        need_encode = (st.cached_jpeg is None) or (st.cached_frame_id != frame_id)

        jpeg: Optional[bytes] = None
        if need_encode:
            if frame_cur is None:
                frame_cur = decode_single_packed_frame(
                    self.decoder,
                    z_packed=st.z_hist[-1],
                    H=self.H, W=self.W, C=self.C, patch=self.patch,
                    packing_factor=self.args.packing_factor,
                    d_bottleneck=self.d_bottleneck,
                )
            st.cached_jpeg = frame_to_jpeg_bytes(frame_cur, quality=int(self.args.jpeg_quality))
            st.cached_frame_id = frame_id
            jpeg = st.cached_jpeg
            if self.args.record and stepped:
                st.recorded_frames.append(frame_to_uint8_hwc(frame_cur))
        else:
            jpeg = None

        u_r_state = u_f_state = u_s_state = "off"
        u_suffix = ""
        if self.args.uncertainty_overlay:
            if not st.calib_done:
                u_r_state = u_f_state = u_s_state = "calibrating"
                u_suffix = f" [cal {len(st.calib_f_samples)}/{int(self.args.calibration_steps)}]"
            else:
                u_r_state = classify_uncertainty(st.last_u_r, st.calib_r_samples)
                u_f_state = classify_uncertainty(st.last_u_f, st.calib_f_samples)
                u_s_state = classify_uncertainty(st.last_u_s, st.calib_s_samples)

        te = float(getattr(st, "last_true_error", float("nan")))
        te_str = "nan" if not math.isfinite(te) else f"{te:.3f}"
        cum = float(getattr(st, "cum_true_error", 0.0))
        n_te = int(getattr(st, "n_true_error", 0))
        mean_te = (cum / n_te) if n_te else None
        status = {
            "type": "status",
            "task": st.task,
            "paused": bool(st.paused),
            "act_dim": int(st.act_dim),
            "step": int(st.step),
            "u_r": float(st.last_u_r),
            "u_f": float(st.last_u_f),
            "u_s": float(st.last_u_s),
            "u_r_state": u_r_state,
            "u_f_state": u_f_state,
            "u_s_state": u_s_state,
            "true_error": te if math.isfinite(te) else None,
            "cum_error": cum,
            "mean_error": mean_te,
            "text": (
                f"step={st.step} | "
                f"err={te_str} | "
                f"cum={cum:.2f} | "
                f"u_r={st.last_u_r:.3f} u_f={st.last_u_f:.3f} u_s={st.last_u_s:.3f}{u_suffix} | "
                f"r={st.last_reward_pred:+.2f} R={st.cum_reward:+.2f}"
            ),
        }
        return jpeg, status


def build_parser() -> argparse.ArgumentParser:
    p = _base_parser()
    # Retarget help defaults for this entrypoint.
    for action in p._actions:
        if action.dest == "html":
            action.default = "interactive_uncertainty_corr.html"
        if action.dest == "port":
            action.default = 7862
        if action.dest == "uncertainty_overlay":
            # Always on for this tool; keep flag for CLI compatibility.
            pass
    return p


def main():
    args = build_parser().parse_args()
    # Correlation UI always wants the three predictors.
    args.uncertainty_overlay = True

    server = CorrInteractiveServer(args)
    app = web.Application()
    app.router.add_get("/", server.index)
    app.router.add_get("/ws", server.ws_handler)
    app.router.add_get("/status", server.status)
    app.router.add_get("/healthz", server.healthz)

    print(
        f"[web] corr UI on http://{args.host}:{args.port}  "
        f"(task={args.task}; true_error = RMS(z_pred - z_env))"
    )
    web.run_app(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()