add svg_term.py + dynamics model modules
7738b02 verified about 2 months ago
50.1 kB
'''SVG(1) actor gradient for train_online.py -- the follow-up the `svg1` hook was left for.
The actor objective today is -Q(s, a) with a BC anchor. This supplies the alternative
    -gamma^H V(f(s, a))
whose action-gradient is gamma^H (df/da)^T grad_s' V(s'), and which reaches the actor by exactly
the same path as dQ/da because it attaches to the same `action` tensor.
WHY IT IS WORTH A RUN. Driving a 12-chunk controller with this composite on held-out Coffee_D0
failures matched dQ/da from the finished 300k-step armF2 critic (median +0.96 vs +1.00 mm,
p = 0.210 paired over 230 episode-runs) using ZERO online interaction, and destroyed a working
insertion in 1 rollout of 308 against the critic's 5 and the geometric cost's 9.
WHAT THE SAME STUDY SAYS WILL GO WRONG, and what this file does about it.
  OFF-DISTRIBUTION.  Every one of those numbers came from demonstration states with a
                     demonstration warm start. df/da away from that distribution is UNMEASURED,
                     and stale-model failure is the standard way model-based RL dies. Hence
                     `finetune_dynamics`: the model is trained online on the agent's own
                     (s, a, s') transitions, so it tracks where the policy actually goes.
  STALE VALUE.       Same argument for V, handled by the trainer's existing --value-finetune.
                     A frozen offline V is a fixed target the actor can walk away from.
  FAT TAIL.          svg|v_td had the largest radial excursion rate of any arm tested (2.3% of
                     rollouts >3x worse than baseline, against 0.6% for dQ/da). Watch the
                     excursion diagnostics, not just the mean return.
  WRONG AXIS.        cos(d pod_z/da, truth) is ~0.11 at contact, where insertion is decided. This
                     term may improve a gradient on the axis that is not binding.
TWO SCALING TRAPS.
  STATE SCALING.  train_online stores states already pushed through StateScaler. value_model.py
                  documents (and asserts) that this is ALREADY V's input space -- the two scalers
                  are bit-identical. But the dynamics model was fitted on RAW observation.state
                  and normalises internally, so it needs the scaling undone first and its RAW
                  output re-scaled before V sees it. `verify()` pins the round trip; a silent
                  mismatch here would look exactly like "SVG does not work".
  OBJECTIVE SCALING.  q_objective is divided by |Q|'s running magnitude so bc_loss_weight means
                      the same thing across tasks. The SVG term is normalised identically, so
                      `svg_weight` is a mixing ratio and not a contest between incompatible units.
'''
from __future__ import annotations
import os

import importlib
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

ROOT = Path(os.environ.get("HDR_ROOT", "/home/chomeed/workspace/hdr_robot"))
DYNAMICS_DIR = ROOT / "0810_dyanmics_model"


class SvgTerm:
    """gamma^H V(f(s, a)), normalised, plus optional online adaptation of f."""

    def __init__(self, value_fn, scaler, action_dim: int, device, discount: float = 0.99,
                 horizon: int = 8, chunks: int = 1, model_tag: str = "model_residual",
                 finetune_dynamics: bool = False, dynamics_lr: float = 1e-4,
                 dynamics_scratch: bool = False, dynamics=None,
                 value_aware_weight: float = 0.0,
                 identity_scaling: bool = False,
                 feature_space: bool = False, pool_mode: str = "mean", proprio_dim: int = 0,
                 value_pool_mode: str | None = None,
                 static_grad_block: str = "none", goal_slice: tuple[int, int] | None = None,
                 fixed_scale: float = 0.0):
        self.value_fn = value_fn                 # the trainer's ValueEnsemble; may be finetuned
        # NORMALISER FOR THE SVG TERM. Default 0.0 keeps the historical behaviour EXACTLY: divide
        # by the running batch mean|V|. That running quantity is the same number logged as
        # `svg_value`, and it moves with the STATE DISTRIBUTION, not with V -- so as the policy
        # walks into low-V states the SVG term is silently amplified against the FIXED 100x BC
        # anchor, which is a positive feedback loop. A positive value pins the denominator so the
        # SVG/BC balance stays constant, which is the whole point of the arm that uses it.
        self.fixed_scale = float(fixed_scale)
        # fixed_scale < 0 selects BATCHNORM-STYLE AUTO CALIBRATION: average the batch statistic
        # over the first |fixed_scale| updates, then FREEZE it -- the analogue of BN switching from
        # batch statistics to frozen running statistics. It must be a SHORT window: frozen-V arms
        # hold ~0.465 indefinitely, but a finetuned V has already drifted by 5-20k (augv_ft
        # 0.406->0.366, dobs_ft 0.462->0.434), so a long warm-up bakes in part of the collapse.
        # Auto-calibration matters because the intrinsic magnitude is CHECKPOINT-SPECIFIC:
        # clean 0.465, aug 0.467, td 0.339, scratch 0.25 -- one hardcoded constant is wrong for
        # most of them (dobstdfix was pinned at 0.47 against a V whose own value is 0.339).
        self._auto_n = int(-fixed_scale) if fixed_scale < 0 else 0
        self._auto_sum = 0.0
        self._auto_count = 0
        self._auto_frozen = None
        # the divisor ACTUALLY used on the last call -- logged as svg_scale. Differs
        # from svg_value (mean V) whenever a pin or auto-calibration is active, which
        # is exactly when you cannot otherwise tell what the SVG term was divided by.
        self.last_scale = float('nan')
        self.device = device
        # `dynamics` supplies a PRE-BUILT f instead of the checkpoint lookup below -- today that is
        # only online_rl_mimicgen/dino_svg.py's DINO-WM adapter, whose weights live in an HF
        # snapshot rather than 0810_dyanmics_model. Everything AFTER the load (the freeze, the
        # optimiser, the generative/anchored probes) runs identically, which is the point: the
        # adapter's requires_grad_/parameters overrides have to survive the real constructor, not
        # a copy of it.
        #
        # `feature_space` says f returns WORLD-MODEL FEATURES, not a state. Then V consumes
        # pooled features (+ optional raw proprio) instead of a scaled state, and nothing may index
        # f's output positionally -- see value_input() and fit_step().
        self.feature_space = bool(feature_space)
        self.pool_mode = pool_mode
        self.proprio_dim = int(proprio_dim)
        # V's pooling, decoupled from the ACTOR's. pool_mode also builds the --dino-obs observation
        # (see train_online._dino_obs_parts), so "flat" there would silently resize the policy
        # input to 150k dims; V alone may want the unpooled grid (patch-transformer V).
        self.value_pool_mode = value_pool_mode or pool_mode
        self._pool = None
        # SVG(k): roll f forward `chunks` action-chunks before applying V. SVG(1) credits the
        # actor with ONE chunk of consequence and defers the rest to V; on the full task that is
        # one chunk out of ~50, and V has been measured to be nearly flat along the axis that
        # matters. Each extra chunk moves credit out of V and into the model.
        self.chunks = max(1, int(chunks))
        self.discount = discount ** (horizon * self.chunks)
        self.finetune_dynamics = finetune_dynamics

        # StateScaler: scaled = 2 (raw - low)/span - 1, clipped. Inverted below.
        self.low = torch.as_tensor(np.asarray(scaler.low, np.float32), device=device)
        self.span = torch.as_tensor(np.asarray(scaler.scale, np.float32), device=device)

        # Dims whose actor-gradient is blocked; see block_static_grad(). Built once, here, so the
        # per-update path is a single torch.where.
        self._static_mask = None
        if static_grad_block != "none":
            # 1e-5 is StateScaler's collapsed-quantile floor and was calibrated on square_d1,
            # whose ONE constant dim (the peg) falls below it. It is too narrow for a task with a
            # whole immobile OBJECT in the state: three_piece_assembly's base occupies dims 9:23,
            # its pose dims 9:16 are all near-constant, and only dim 15 clears 1e-5 -- leaving
            # dim 12 amplified 71,718x by the 2/span unscaling. Measured effect: actor_grad 9.8
            # unmasked, 4.1 masked at 1e-5, against 0.97 on a healthy threading arm.
            # 'degenerate_wide' raises the floor to 1e-3 (amplification > 2000x). Checked against
            # every task in the campaign: it masks dims 9:16 on tpa and NOTHING on threading
            # (max 583x), square_d1 (69x) or square_d0 (71x), so no existing arm changes.
            # It leaves the base's EEF-RELATIVE dims (16:23) untouched, which is where the
            # actionable signal lives -- only the constant absolute pose is blocked.
            floor = 1e-3 if static_grad_block == "degenerate_wide" else 1e-5
            blocked = (self.span <= floor)
            n_degen = int(blocked.sum())
            if static_grad_block == "goal":
                if goal_slice is None:
                    raise ValueError("--svg-static-grad goal needs a goal_slice for this task")
                blocked = blocked.clone()
                blocked[goal_slice[0]:goal_slice[1]] = True
            if blocked.any():
                self._static_mask = blocked.view(1, -1)
                print(f"[svg] static-grad block '{static_grad_block}': "
                      f"{int(blocked.sum())} of {blocked.numel()} dims "
                      f"({n_degen} degenerate-span, spans "
                      f"{[round(float(x), 8) for x in self.span[blocked][:4]]}...); "
                      f"their actor gradient is zeroed, V's input value is unchanged", flush=True)
            else:
                print(f"[svg] static-grad block '{static_grad_block}': no dims matched", flush=True)

        sys.path[:0] = [str(DYNAMICS_DIR), str(ROOT / "0810"), str(ROOT)]
        state_dim = int(self.low.numel())
        if self.feature_space:
            from online_rl_mimicgen.dino_wm import pool_features
            self._pool = pool_features
            if self.chunks > 1:
                raise ValueError(
                    f"--svg-chunks {self.chunks} with a feature-space f: chunk 2 would have to "
                    "feed predicted FEATURES back in as a raw state, and the world model takes "
                    "pixels-derived context plus a 9-dim proprio state. Only k=1 is defined here.")
        if dynamics is not None:
            self.dynamics = dynamics
            if dynamics_scratch:
                raise ValueError("--svg-dynamics-scratch is not defined for a pre-built f")
        else:
            self._load_dynamics(model_tag, state_dim, action_dim, device, dynamics_scratch)
        self.dynamics_scratch = dynamics_scratch
        self.value_aware_weight = float(value_aware_weight)
        # --svg-on-encoder: f and V live in the pixel encoder's 265-dim feature space and
        # were fitted on RAW features. self.low/self.span are the 37-dim TASK quantiles, so
        # applying them here is both a shape error and, if the shapes happened to match, the
        # wrong transform. In that space the scaler is the identity.
        self.identity_scaling = bool(identity_scaling)
        self.dynamics.requires_grad_(finetune_dynamics)
        self.dynamics.train(finetune_dynamics)
        self.optimiser = (torch.optim.Adam(self.dynamics.parameters(), lr=dynamics_lr)
                          if finetune_dynamics else None)
        # A model exposing mu_scale/infer_noise is the paper's f^(s, a, xi) = mu^ + sigma^ * xi;
        # anything else is a deterministic regressor and gets the additive-residual reading.
        self.generative = hasattr(self.dynamics, "mu_scale")
        # A flow model has infer_noise/sample_anchored but NO mu_scale, so `generative` above is
        # False for it and the deterministic branches would silently mishandle it. `anchored` is
        # the flag those paths actually need.
        self.anchored = hasattr(self.dynamics, "sample_anchored")
        # How many times an ANCHORED f was rolled out with no anchor and fell back to the plain
        # forward. Non-zero means part of this run's actor gradient was taken through RESAMPLED
        # noise -- see _warn_unanchored(). A counter rather than only a one-shot print, so a run
        # that did it on every update is distinguishable from one that did it once.
        self.unanchored_calls = 0
        self._warned_unanchored = False
        self._verified = False
        print(f"[svg] generative={self.generative} "
              f"dynamics={model_tag} finetune={finetune_dynamics} scratch={dynamics_scratch} "
              f"anchored={self.anchored} feature_space={self.feature_space} "
              f"pool={self.pool_mode if self.feature_space else '-'} "
              f"proprio_tail={self.proprio_dim if self.feature_space else '-'} "
              f"value={getattr(value_fn, 'variant', '?')} action_dim={action_dim} "
              f"discount^{horizon}={self.discount:.4f}", flush=True)

    def _load_dynamics(self, model_tag, state_dim, action_dim, device, dynamics_scratch):
        # The MODULE is chosen from the tag, not hardcoded. dynamics_model_flow.pt and
        # dynamics_model_gaussian.pt have different parameter layouts from model_residual, so
        # building a model_residual.Dynamics and loading either into it dies on a key mismatch --
        # which is exactly how the first flow and gaussian runs failed. Matched by longest PREFIX
        # because "model" prefixes every other name and a substring test would swallow them all;
        # dynamics_model_residual_succ2.pt correctly resolves to model_residual.
        module = next((m for m in ("model_gaussian", "model_flow", "model_history",
                                   "model_residual", "model_rnn")
                       if model_tag.startswith(m)), "model")
        payload = torch.load(DYNAMICS_DIR / f"dynamics_{model_tag}.pt", map_location=device)
        # History checkpoints carry their own dims; everything else is a bare state_dict at 66 x 56.
        if isinstance(payload, dict) and "state_dict" in payload:
            state_dim, action_dim = payload["state_dim"], payload["chunk_dim"]
            payload = payload["state_dict"]
        # Hidden width is INFERRED from the checkpoint, not assumed. Dynamics defaults to
        # hidden=512, so a wider checkpoint (h2048's first linear is 2048x122) dies on a shape
        # mismatch in load_state_dict -- silently blocking every model from the capacity sweep.
        first = next((v for k, v in payload.items() if k.endswith("net.0.weight")), None)
        kwargs = {"hidden": int(first.shape[0])} if first is not None else {}
        # STATE WIDTH IS ALSO INFERRED, for the same reason hidden is. --svg-on-encoder runs f in
        # the pixel encoder's 265-dim feature space while the TASK state is 37-dim, so sizing from
        # the task built a 37-dim Dynamics and load_state_dict died on five shape mismatches.
        # mu is the output width; the first linear's input is state+chunk, which gives the chunk.
        mu = payload.get("mu")
        # FLOW MODELS BREAK THIS INFERENCE. `ck_action = first.shape[1] - ck_state` assumes the
        # first linear sees exactly [state, chunk]. A flow head also feeds it a Fourier time
        # embedding (marked by the `freqs` buffer), so the subtraction charges those dims to the
        # action: threadfull is 37 + 56 + 69 = 162, inferred as action=125, and Dynamics(37,125)
        # then builds 37 + 125 + 69 = 231 and dies in load_state_dict. The task's own dims are
        # already correct for these checkpoints, so skip the override rather than mis-attribute.
        if "freqs" in payload:
            mu = None
        if mu is not None and mu.dim() == 1 and first is not None:
            ck_state = int(mu.shape[0])
            ck_action = int(first.shape[1]) - ck_state
            if ck_state != state_dim or ck_action != action_dim:
                print(f"[svg] dynamics checkpoint is {ck_state}x{ck_action}, not the task's "
                      f"{state_dim}x{action_dim} -- sizing f from the checkpoint", flush=True)
                state_dim, action_dim = ck_state, ck_action
        self.dynamics = importlib.import_module(module).Dynamics(
            state_dim, action_dim, **kwargs).to(device)
        # `dynamics_scratch` keeps the ARCHITECTURE and the dims read from the checkpoint but
        # discards its weights, so f is learned online from random init. The checkpoint is still
        # loaded above because history models carry their state/action dims in the payload.
        # ALWAYS load, then discard only the LEARNED WEIGHTS. mu/sigma/fitted are registered
        # buffers holding the demonstration-set input statistics -- the same role StateScaler
        # plays for every other component -- and `loss()` asserts on `fitted`. Dropping them
        # would leave an unnormalised model, which is a different (and broken) experiment from
        # "learn the dynamics online". So: keep the normaliser, randomise the network.
        self.dynamics.load_state_dict(payload)
        if dynamics_scratch:
            for module in self.dynamics.modules():
                if hasattr(module, "reset_parameters"):
                    module.reset_parameters()

    # ---- what V consumes -------------------------------------------------------------------
    def value_input(self, predicted: torch.Tensor, scaled_state: torch.Tensor) -> torch.Tensor:
        """f's output, in V's input space.
        STATE SPACE: the scaled successor state, exactly as before.
        FEATURE SPACE: pooled world-model features, optionally followed by the RAW 9-dim proprio
        of the CURRENT state -- which is what 0814/fit_v_dino.py fit V on. Note the proprio is
        s_t's while the features are the model's prediction for t+8; the world model does not
        predict proprio, so this is the only proprio available. It is constant in `action`, so it
        contributes nothing to dV/da and only shifts V's level.
        """
        if not self.feature_space:
            return self.scaled_state(self.block_static_grad(predicted))
        pooled = self._pool(predicted, self.value_pool_mode)
        if not self.proprio_dim:
            return pooled
        return torch.cat([pooled, self.raw_state(scaled_state)[:, :self.proprio_dim]], dim=-1)

    def block_static_grad(self, predicted: torch.Tensor) -> torch.Tensor:
        """Stop the actor gradient on state dims the action cannot move.
        WHY. The chain is  g = gamma^H (df/da)^T d(scaled)/d(raw) grad_scaled V, and
        d(scaled)/d(raw) = 2/span. A dim whose span hit StateScaler's floor therefore multiplies
        its gradient by 2/1e-6 = 2e6 while a typical dim gets ~10. Measured on square_d1
        (0814/phase_breakdown_square.py): dim 25 is the peg's z, constant because the peg is
        bolted to the table, span 1e-6. Its grad_scaled V is unremarkable (0.031, 10th of 26) and
        f's Jacobian row for it is nearly zero (0.00029) -- exactly as it should be, since no
        action moves a fixed peg -- but the product is 2e6 x that, and dim 25 ALONE reconstructs
        the deployed SVG gradient at cos 0.99997, norm ratio 1.00002. The actor was ascending the
        derivative of an environmental constant. Masking it leaves 0.77% of the gradient, which
        runs weakly POSITIVE in every phase (+0.007 / +0.013 / +0.010 / +0.018 vs a critic).
        StateScaler's own `degenerate_dims` does NOT catch this: it lists a dim only when the
        quantiles collapse AND min/max gives a usable fallback, and a globally constant dim fails
        that second test. The span has to be inspected directly.
        "goal" additionally masks the goal-fixture block. threading's tripod spans 3.4-8.2 mm --
        above any floor, so the degenerate test misses it -- yet the goal block carries 80% of
        that task's gradient despite the tripod being bolted down.
        The VALUE is untouched: `torch.where(mask, x.detach(), x)` returns the same numbers, so V
        still sees the true successor state and only dV/da changes. No checkpoint is invalidated.
        """
        if self._static_mask is None:
            return predicted
        return torch.where(self._static_mask, predicted.detach(), predicted)

    # ---- scaling -------------------------------------------------------------------------
    def raw_state(self, scaled: torch.Tensor) -> torch.Tensor:
        if self.identity_scaling:
            return scaled
        return (scaled + 1.0) * 0.5 * self.span + self.low

    def scaled_state(self, raw: torch.Tensor) -> torch.Tensor:
        if self.identity_scaling:
            return raw
        return torch.clamp(2.0 * (raw - self.low) / self.span - 1.0, -5.0, 5.0)

    # ---- the actor term ------------------------------------------------------------------
    def objective(self, scaled_state: torch.Tensor, action: torch.Tensor,
                  follow: torch.Tensor | None = None,
                  anchor: tuple[torch.Tensor, torch.Tensor] | None = None):
        """(normalised objective to MINIMISE, mean value for logging).
        The dynamics model is used in eval mode for the actor gradient even when it is being
        finetuned -- otherwise the actor's gradient would depend on batchnorm/dropout state that
        the model's own update is concurrently moving.
        `anchor` is (a_data, s_next) for the SAME transitions as `scaled_state`. It is required by
        generative models, which need the noise xi that produced the observed transition; it is
        ignored by point regressors, whose additive residual form already cancels without it.
        """
        if not self._verified:            # once, on a real batch, before it can do any damage
            self.verify(scaled_state)
            self._verified = True
        was_training = self.dynamics.training
        self.dynamics.eval()
        predicted_raw = self._rollout(scaled_state, action, follow, anchor)
        value = self.discount * self.value_fn(self.value_input(predicted_raw, scaled_state))
        self.dynamics.train(was_training)
        scale = self._norm(value)
        return -value.mean() / scale, float(value.detach().mean())

    def _norm(self, t: torch.Tensor) -> torch.Tensor:
        """Denominator for the SVG term.
        fixed_scale > 0  -> that constant.
        fixed_scale < 0  -> BN-style: average the batch statistic for |fixed_scale| updates, then
                            freeze it. Uses the live batch statistic during the warm-up itself.
        fixed_scale == 0 -> the historical running batch mean (default, bit-identical to before).
        """
        if self.fixed_scale > 0.0:
            self.last_scale = float(self.fixed_scale)
            return torch.as_tensor(self.fixed_scale, dtype=t.dtype, device=t.device)
        batch = t.detach().abs().mean().clamp_min(1e-3)
        if self._auto_n > 0:
            if self._auto_frozen is not None:
                self.last_scale = float(self._auto_frozen)
                return torch.as_tensor(self._auto_frozen, dtype=t.dtype, device=t.device)
            self._auto_sum += float(batch); self._auto_count += 1
            self.last_scale = float(batch)
            if self._auto_count >= self._auto_n:
                self._auto_frozen = max(self._auto_sum / self._auto_count, 1e-3)
                print(f"[svg-scale] auto-calibrated to {self._auto_frozen:.4f} "
                      f"over {self._auto_count} updates; frozen for the rest of the run", flush=True)
        self.last_scale = float(batch)
        return batch

    def _warn_unanchored(self) -> None:
        """Loud, once: an anchored f was rolled out with NO anchor and fell back to forward().
        WHY THIS IS NOT A CRASH. The guard in _rollout used to be a bare `if self.anchored:`, and
        objective_terms() -- the --svg-phase-gate path -- has never passed an anchor, so
        `action_data, next_state = anchor` unpacked None and every anchored arm died with
        `TypeError: cannot unpack non-iterable NoneType object` on its first actor update. The
        term is still defined without the anchor, so the run continues rather than losing a GPU
        allocation to a plumbing gap.
        WHY IT IS SHOUTED. The fallback is `self.dynamics(raw_state, action)` = forward(), which
        for a flow integrates the ODE from FRESH NOISE on every call. Two consequences, and
        NEITHER of them moves any logged number:
          * the additive residual stops cancelling, so V is evaluated at a sampled successor
            instead of the observed one and f's prediction BIAS re-enters the gradient -- on this
            task the offline model is 8x worse on policy states (12.25 mm) than on demos;
          * d/da is taken through a Monte-Carlo mean over RESAMPLED noise rather than the flow's
            Jacobian at fixed xi. Noise alone displaces the DINO-WM flow 345.88 against the
            action's 23.19 (14.9x), so what survives is not a degraded gradient, it is noise.
        svg_value, svg_scale, svg_feat_mse and the loss curve all stay healthy while this happens,
        which is why it gets a stderr banner and a counter instead of an info line.
        """
        self.unanchored_calls += 1
        if self._warned_unanchored:
            return
        self._warned_unanchored = True
        print("\n" + "!" * 96 + "\n"
              "[svg] UNANCHORED ROLLOUT ON AN ANCHORED f -- THIS ACTOR GRADIENT IS WRONG\n"
              "  self.anchored is True (f exposes sample_anchored) but _rollout was called with\n"
              "  anchor=None, so it fell back to self.dynamics(s, a) = forward(). For a flow that\n"
              "  integrates from FRESH NOISE every call: the residual no longer cancels and d/da\n"
              "  goes through resampled noise, not the Jacobian at fixed xi. Every SVG metric\n"
              "  will still look healthy -- that is the whole hazard.\n"
              "  The caller that does this today is objective_terms() (--svg-phase-gate), which\n"
              "  takes an `anchor` argument and is not handed one by train_online.py:1145.\n"
              "  Counted in svg_term.unanchored_calls; this banner prints ONCE per run.\n"
              + "!" * 96 + "\n", file=sys.stderr, flush=True)

    def _rollout(self, scaled_state: torch.Tensor, action: torch.Tensor,
                 follow: torch.Tensor | None,
                 anchor: tuple[torch.Tensor, torch.Tensor] | None = None):
        """f applied `self.chunks` times. Only the FIRST action carries the actor's gradient.
        Chunks 2..k use `follow` -- the prior's next chunk as actually recorded (batch["next_base"])
        -- held fixed. That is deliberate: holding subsequent actions constant makes the gradient
        the OPEN-LOOP reach of the action being optimised, rather than a closed-loop quantity that
        also depends on how a future policy would react. With no follow supplied the action is
        repeated, which is the cruder fallback and is only right for k = 1.
        """
        raw_state = self.raw_state(scaled_state)
        if self.anchored and anchor is None:
            # An anchored f with nothing to anchor on. Shout, count it, and fall through to the
            # unanchored forward below -- _warn_unanchored() documents why that is a banner and
            # not a one-liner, and why this is a fallback and not a raise.
            self._warn_unanchored()
        if self.anchored and anchor is not None:
            # FLOW MODELS MUST NOT GO THROUGH self.dynamics(...) HERE. model_flow.forward averages
            # FLOW_EVAL_SAMPLES fresh draws through the ODE, so (a) two calls never share noise and
            # the additive residual stops cancelling, and (b) d/da is taken through a Monte-Carlo
            # mean over RESAMPLED noise rather than the flow's Jacobian at fixed xi. sample_anchored
            # infers xi from the observed transition, holds it fixed, and subtracts a detached
            # solver-error offset so a_theta == a_data reproduces s' exactly. Every offline harness
            # in this repo (jacobian_eval, svg_rescue, task_improve, critic_vs_svg,
            # compare_dynamics) already branches this way; this is the online path catching up.
            action_data, next_target = anchor
            # sample_anchored's 4th argument is s' IN f's OWN OUTPUT SPACE -- the space the
            # DETACHED offset  s' - f(s, a_data, xi)  is subtracted in (model_flow.py:183-185).
            # The two spaces are not interchangeable and the conversion only exists on one side:
            #
            #   STATE SPACE   f emits raw observation units; the anchor carries the SCALED 37-dim
            #                 env successor (train_online builds element 1 from batch[svg_next_key]),
            #                 so it is unscaled here exactly as it always was. Unchanged.
            #   FEATURE SPACE f emits (B, V, 196, 384) encoder features, which NEVER went through
            #                 StateScaler. raw_state() on them is a CATEGORY error, not a units
            #                 error: (scaled+1)*0.5*span+low is defined on 37 env dims and this
            #                 tensor has 384 channels, so 37 vs 384 does not even broadcast. The
            #                 anchor's element 1 is ALREADY the encoder's own t+8 features
            #                 (dino_svg.encode_frames of next_image -- bit-identically the
            #                 `target_feat` tensor _fit_step_features regresses onto), so it is
            #                 passed straight through. Anchoring in any other space would anchor f
            #                 against a quantity it was never trained to produce.
            raw = self.dynamics.sample_anchored(
                raw_state, action, action_data,
                next_target if self.feature_space else self.raw_state(next_target))
        else:
            raw = self.dynamics(raw_state, action)
        nxt = action if follow is None else follow
        for _ in range(self.chunks - 1):
            # Only the first step can be anchored -- there is no observed transition for the
            # imagined states that follow -- so deeper chunks fall back to the plain forward.
            raw = self.dynamics(raw, nxt)
        return raw

    def objective_terms(self, scaled_state: torch.Tensor, action: torch.Tensor,
                        follow: torch.Tensor | None = None,
                        anchor: tuple[torch.Tensor, torch.Tensor] | None = None):
        """Per-sample -gamma^H V(f(s,a)) / |.|, i.e. `objective` before the mean is taken.
        Needed for the phase gate: the mixing weight between the Q term and the SVG term is a
        function of the STATE, so both terms have to survive as vectors until after they are
        blended. Reducing first and blending scalars would apply one global ratio everywhere,
        which is exactly what the measurement says is wrong -- SVG carries signal while the pod is
        in the hand and none after it is released.
        `anchor` is the same (a_data, s'-in-f's-output-space) pair `objective` takes and is
        forwarded unchanged. It defaults to None because the phase-gate call site does not build
        one -- on an anchored f that buys the loud fallback in _warn_unanchored() rather than the
        TypeError this signature used to guarantee, but it is still the WRONG gradient. The fix is
        to hand one in, not to rely on the default.
        """
        if not self._verified:
            self.verify(scaled_state)
            self._verified = True
        was_training = self.dynamics.training
        self.dynamics.eval()
        predicted_raw = self._rollout(scaled_state, action, follow, anchor)
        value = self.discount * self.value_fn(self.value_input(predicted_raw, scaled_state))
        self.dynamics.train(was_training)
        scale = self._norm(value)
        return -value / scale, float(value.detach().mean())

    # ---- SVG(1) with replay, as written in Heess et al. 2015 Alg. 2 ------------------------
    def replay_objective(self, actor, batch, residual_scale, weight_clip=10.0):
        """The paper's line 13, minus r_a.
            v_theta = w * (r_a + gamma * V'_s' * f_a) * pi_theta  |_{eta_k, xi_k}
        Three things distinguish this from the model-based DPG form, each line-for-line from the
        algorithm:
          eta_k   inferred from (s_k, a_k). The REPLAYED action is reproduced as a differentiable
                  function of theta by holding the inferred noise fixed and re-running the
                  reparameterisation -- so the gradient is taken AT the action the data contains,
                  not at whatever the current policy would do instead.
          xi_k    inferred from (s_k, a_k, s_{k+1}) as the additive residual s' - f(s,a). The
                  model here is deterministic, so this is the natural reading, and it matters: V
                  is then evaluated at the TRUE next state when a = a_k, cancelling the model's
                  prediction BIAS and leaving only its derivative. On this task that bias is large
                  -- the offline model is 8x worse on policy states (12.25 mm) than on demos.
          w       the ratio p(a_k|s_k;theta_t) / p(a_k|s_k;theta_k), correcting for the action
                  having been drawn by an older policy. Clipped: the unbounded ratio explodes once
                  the policy moves away from the behaviour that filled the buffer.
        r_a is omitted deliberately -- the reward is sparse terminal success, so dr/da is zero
        almost everywhere and there is nothing to differentiate.
        """
        if self.feature_space and "next_feat_obs" not in batch:
            # xi_k = s' - f(s, a); in feature space s' is the OBSERVED t+8 features, which the
            # buffer does not store for a DINO-WM arm (its `next_feat` column exists only for the
            # latent/seghead arms). The caller must encode next_image once per batch and pass it
            # in as `next_feat_obs` -- done there, not here, so the DINOv2 forward is visible in
            # the update loop's cost rather than hidden inside the objective.
            raise NotImplementedError(
                "--svg-mode replay on a feature-space f needs batch['next_feat_obs'] (the encoded "
                "t+8 features). Pass --vtrace/replay through train_online, which supplies it.")
        state, base = batch["state"], batch["base"]
        replayed = ((batch["action"] - base) / residual_scale).clamp(-1.0 + 1e-6, 1.0 - 1e-6)

        mean, log_std = actor(state, base)
        std = log_std.exp()
        eta = ((torch.atanh(replayed) - mean) / std).detach()      # eta_k | (s_k, a_k)
        residual_theta = torch.tanh(mean + std * eta)
        action_theta = (base + residual_scale * residual_theta).clamp(-1.0, 1.0)

        raw_state = self.raw_state(state)
        # Same reason as fit_dynamics: xi is inferred through a one-step model, so it needs the
        # one-step successor rather than the n-step bootstrap state.
        # FEATURE SPACE: `next_feat_obs` IS already the one-step successor -- next_image is the
        # t+8 frame under every --nstep (fill_nstep rewrites next_state but never next_image), so
        # no next_state_1step equivalent is needed or available.
        raw_next = (batch["next_feat_obs"] if self.feature_space
                    else self.raw_state(batch.get("next_state_1step", batch["next_state"])))
        was_training = self.dynamics.training
        self.dynamics.eval()
        if self.generative:
            # Alg. 2 line 10 as written: xi is STANDARDISED noise under f^ = mu^ + sigma^ * xi,
            # so re-running at a_theta carries the (d sigma^/d a) * xi term. The additive form
            # below is the sigma^-constant special case and drops exactly that term.
            with torch.no_grad():
                xi = self.dynamics.infer_noise(raw_state, batch["action"], raw_next)
                # xi is meant to be N(0, I). A draw past 5 sigma means the model's DENSITY is
                # wrong there, not that the transition was a rare event -- and since xi multiplies
                # sigma^(s, a_theta), an unclipped outlier gets amplified by whatever the scale
                # head does at the new action. Clipping bounds that; the fraction clipped is
                # returned so a run that is mostly clipping is visible rather than silent.
                clipped = (xi.abs() > 5.0).float().mean()
                xi = xi.clamp(-5.0, 5.0)
            mean, scale = self.dynamics.mu_scale(raw_state, action_theta)
            predicted = mean + scale * xi
        else:
            with torch.no_grad():                                   # xi_k | (s_k, a_k, s_{k+1})
                xi = raw_next - self.dynamics(raw_state, batch["action"])
                clipped = torch.zeros((), device=xi.device)
            predicted = self.dynamics(raw_state, action_theta) + xi
        self.dynamics.train(was_training)
        value = self.discount * self.value_fn(self.value_input(predicted, state))

        with torch.no_grad():
            logp_now = actor.log_prob_of(state, base, replayed)
            behaviour_logp = batch["behaviour_logp"]
            # NaN marks a row whose action no policy DENSITY generated: expert demonstrations
            # (30% of every batch here), prior-gated exploration steps, and best-of-N picks. Those
            # rows get w = 1 -- the transition is still a valid (s, a, s') for the model-based
            # gradient, it just carries no importance correction. Without this guard exp(x - NaN)
            # poisons the mean and the actor loss is NaN on the first update.
            valid = torch.isfinite(behaviour_logp)
            ratio = torch.where(
                valid,
                torch.exp(logp_now - torch.nan_to_num(behaviour_logp)).clamp(max=weight_clip),
                torch.ones_like(logp_now),
            )
        weighted = ratio * value
        norm = self._norm(weighted)
        return (-weighted.mean() / norm, float(value.detach().mean()),
                float(ratio.mean()), float(valid.float().mean()), float(clipped))

    # ---- online adaptation of f ----------------------------------------------------------
    def fit_step(self, batch: dict) -> dict[str, float]:
        """One supervised step on the agent's OWN transitions, so f tracks the policy's states.
        Trained in raw units with the model's internal normaliser, exactly as it was fitted
        offline -- training it in scaled units here would silently change the objective.
        """
        if not self.finetune_dynamics:
            return {}
        if getattr(self.dynamics, "is_latent56", False):
            return self._fit_step_latent(batch)
        if self.feature_space:
            return self._fit_step_features(batch)
        raw_state = self.raw_state(batch["state"]).detach()
        # ONE-step successor, not the bootstrap one: under --nstep n, fill_nstep rewrites
        # batch["next_state"] to s_{t+n}, which would train f to predict n chunks ahead while SVG
        # differentiates it as a single chunk. Falls back to next_state when the column is absent
        # (nstep=1 runs and older buffers), where the two are the same state.
        raw_next = self.raw_state(
            batch.get("next_state_1step", batch["next_state"])).detach()
        action = batch["action"].detach()
        self.optimiser.zero_grad()
        loss = self.dynamics.loss(raw_state, action, raw_next)
        # VALUE-AWARE MODEL LEARNING (IterVAML). Prediction accuracy is measurably NOT what makes
        # f useful here -- f's loss and success correlate at +0.045, and f keeps improving while
        # success falls. This adds ||V(f(s,a)) - V(s')||^2, which asks f to be right about the only
        # thing SVG reads out of it. It is ADDITIVE, never a replacement: on its own the term is
        # rank-deficient (every state on a level set of V is an equally good prediction), so f
        # could collapse to one canonical state per value. The prediction loss is the anchor that
        # forbids that, and svg_value_aware_weight controls the trade.
        vaml = None
        if self.value_aware_weight > 0.0:
            pred_s = self.dynamics(raw_state, action)
            with torch.no_grad():
                v_target = self.value_fn(self.value_input(raw_next, None))
            v_pred = self.value_fn(self.value_input(pred_s, None))
            vaml = torch.nn.functional.mse_loss(v_pred, v_target)
            loss = loss + self.value_aware_weight * vaml
        loss.backward()
        self.optimiser.step()
        with torch.no_grad():
            predicted = self.dynamics(raw_state, action)
            if self.identity_scaling:
                # ENCODER SPACE has no millimetres. Dims 9:12 are arbitrary channels of the IMPALA
                # feature vector, so svg_pod_mm would be a real-looking number that means nothing
                # and would land in every comparison table -- the same trap _fit_step_features
                # documents for the map path. Emit the feature MSE under its own name instead.
                error = None
                feat_mse = F.mse_loss(predicted, raw_next)
            else:
                # Dims 9:12 are the MANIPULATED object's position on both tasks -- coffee's pod and
                # threading's needle -- because proprio occupies 0:9 and the first object block
                # starts right after. Confirmed in 0814/state_layout_threading.py, so the metric
                # name is coffee-flavoured but the number is right for either.
                error = (predicted - raw_next)[:, 9:12].norm(dim=-1).mean() * 1000
                feat_mse = None
        out = {"svg_dynamics_loss": float(loss.detach())}
        if error is not None:
            out["svg_pod_mm"] = float(error)
        else:
            out["svg_feat_mse"] = float(feat_mse)
        if vaml is not None:
            out["svg_vaml"] = float(vaml.detach())
        return out

    def _fit_step_latent(self, batch: dict) -> dict[str, float]:
        """One supervised step on the 56-d latent world model.
        Cannot reuse the state branch: it regresses onto raw_state(batch["next_state"]), which under
        --pixel-obs is the 265-d IMPALA encoder output, not z_{t+8}. Cannot reuse the feature branch
        either: that one pools a (B,V,P,d) grid and the latent is a flat vector. The target is
        supplied as `target_feat` by the caller, which encodes next_image.
        """
        action = batch["action"].detach()
        target = batch["target_feat"].detach()
        self.optimiser.zero_grad()
        pred = self.dynamics(None, action)
        loss = F.mse_loss(pred, target)
        loss.backward()
        self.optimiser.step()
        with torch.no_grad():
            # COPY BASELINE: "predict no change" is z_t, the cur slot of the stored context.
            copy_mse = F.mse_loss(self.dynamics.cur_z[:len(target)], target)
            feat_mse = F.mse_loss(self.dynamics(None, action), target)
        out = {"svg_dynamics_loss": float(loss.detach()), "svg_feat_mse": float(feat_mse),
               "svg_copy_feat_mse": float(copy_mse)}
        if float(copy_mse) > 0:
            out["svg_skill_vs_copy"] = 1.0 - float(feat_mse) / float(copy_mse)
        return out

    def _fit_step_features(self, batch: dict) -> dict[str, float]:
        """One supervised step of the world model against the ENCODER's own t+8 features.
        `target_feat` and adapter.ctx_feats are prepared by dino_svg.fit_batch, which owns the
        pixels; this stays free of any DINO import.
        THE METRIC IS NOT svg_pod_mm. That number is (predicted - raw_next)[:, 9:12].norm()*1000 --
        the manipulated object's position error in MILLIMETRES, and it is what this whole campaign
        judged dynamics models by. Feature space has no dims 9:12 and no millimetres, so it emits
        `svg_feat_mse` under its own name. Reporting a feature MSE as a millimetre error would put
        an incomparable number into every comparison table in the study, and --exploration-gate-mm
        would be thresholding it against a distance.
        """
        raw_state = self.raw_state(batch["state"]).detach()
        action = batch["action"].detach()
        target = batch["target_feat"].detach()
        self.optimiser.zero_grad()
        loss = self.dynamics.loss(raw_state, action, target)
        # VALUE-AWARE MODEL LEARNING, feature-space twin of the state branch above. This is the
        # side where it should matter most: f's feature MSE and success are uncorrelated here, and
        # V reads f's output through a POOLING that discards most of what the prediction loss is
        # optimising -- so f spends capacity on directions V cannot see. value_input needs the
        # SCALED state (not raw) because feature-space V appends raw proprio decoded from it.
        vaml = None
        if self.value_aware_weight > 0.0:
            # Reuse the prediction loss() just computed -- a second Predictor forward over the map
            # is what made this arm 25-90x slower than its siblings.
            pred_f = getattr(self.dynamics, "_last_pred", None)
            if pred_f is None:
                pred_f = self.dynamics(raw_state, action)
            # V is not trained by this optimiser (its own fit_step zero_grads first), so computing
            # parameter grads for it is pure waste. Gradient still flows to pred_f, which is what
            # the term is for.
            _vp = [q for q in self.value_fn.parameters() if q.requires_grad]
            for q in _vp:
                q.requires_grad_(False)
            try:
                k = min(len(raw_state), 8)          # a subsample is enough for a regulariser
                sl = slice(0, k)
                with torch.no_grad():
                    v_target = self.value_fn(self.value_input(target[sl], batch["state"][sl]))
                v_pred = self.value_fn(self.value_input(pred_f[sl], batch["state"][sl]))
                vaml = F.mse_loss(v_pred, v_target)
            finally:
                for q in _vp:
                    q.requires_grad_(True)
            loss = loss + self.value_aware_weight * vaml
        loss.backward()
        self.optimiser.step()
        with torch.no_grad():
            error = F.mse_loss(self.dynamics(raw_state, action), target)
            # COPY BASELINE. A raw feature MSE is uninterpretable on its own: it depends on how
            # much the scene actually moved between t and t+8, which differs between the demo
            # distribution f was fitted on and the policy states it sees online. "Predict no
            # change" -- the frame-t features, ctx slot -1 -- is the trivial reference, so
            # skill = 1 - model/copy says whether f is contributing at all. Offline on held-out
            # demos this is 1 - 0.52/1.55 = 0.66; the online path had no equivalent, which made the
            # 0.26 plateau impossible to read as either a model floor or a data floor.
            copy = self.dynamics._context(raw_state)[:, -1]
            copy_mse = F.mse_loss(copy, target)
        out = {"svg_dynamics_loss": float(loss.detach()), "svg_feat_mse": float(error),
               "svg_copy_feat_mse": float(copy_mse)}
        if vaml is not None:
            out["svg_vaml"] = float(vaml.detach())
        if float(copy_mse) > 0:
            out["svg_skill_vs_copy"] = 1.0 - float(error) / float(copy_mse)
        return out

    # ---- a guard, not a comment ----------------------------------------------------------
    def verify(self, scaled_state: torch.Tensor) -> None:
        if self.identity_scaling:
            # Round-trip is exact by construction, and the magnitude floor below is derived from
            # the task scaler's q01/q99 -- a statement about 37-dim poses, not encoder features.
            print(f"[svg] identity scaling (encoder space, dim {scaled_state.shape[-1]}); "
                  f"unscale check skipped", flush=True)
            return
        raw = self.raw_state(scaled_state)
        error = (self.scaled_state(raw) - scaled_state).abs().max().item()
        assert error < 1e-3, f"state unscaling round-trip off by {error}"
        # The magnitude floor catches an unscale that silently collapsed to nothing -- a degenerate
        # or unloaded scaler. It used to be a hardcoded 1.5, which was coffee's HINGE ANGLE in
        # disguise: coffee's raw states peak at 2.204 and the argmax is dim 65, the lid hinge, the
        # one dim that is an angle rather than a position or a quaternion component. Threading has
        # no such dim, peaks at 1.042, and tripped a guard that was never about it. Deriving the
        # floor from the scaler's own q01/q99 keeps the intent -- "the unscaled state should reach
        # the magnitudes the demo data does" -- without smuggling in one task's geometry.
        floor = 0.5 * max(self.low.abs().max().item(),
                          (self.low + self.span).abs().max().item())
        assert raw.abs().max().item() > floor, (
            f"unscaled state peaks at {raw.abs().max().item():.3f}, below {floor:.3f} = half the "
            "demo range implied by the scaler -- StateScaler assumptions changed")
        print(f"[svg] unscale verified (round-trip {error:.2e}, raw range "
              f"[{raw.min().item():.3f}, {raw.max().item():.3f}], floor {floor:.3f})", flush=True)