"""The MDN-RNN: stage two ("M") of the V-M-C world model. The job is one line: given the latent now and the action now, predict the latent next. What makes it interesting is the word *predict*. Why a MIXTURE density and not a regression ------------------------------------------ The obvious model is ``z_{t+1} = f(z_t, a_t)`` trained with MSE. That is a unimodal Gaussian with fixed variance, and it is wrong here for a concrete reason: the transition is genuinely multimodal. Two frames before the ball reaches the floor, whether it bounces off the paddle or off the floor depends on sub-pixel detail the encoder may not have resolved. The true predictive distribution has two lumps. An MSE model cannot represent two lumps, so it predicts the average of them -- a ball in between, which is a place the ball can never be. In a one-step metric this looks fine (small MSE!) and in a rollout it is fatal, because the averaged state is off-manifold and the next step is garbage. Averaging futures is the single classic failure of deterministic dynamics models, and the mixture is the fix: it can put mass on "bounces" and mass on "does not" and commit to one when sampled. ``n_gauss=1`` collapses this to exactly the plain-Gaussian baseline (with a learned per-dim variance), which is why it is worth keeping as a flag: it makes the MDN-vs-Gaussian comparison a one-character change. Why RESIDUAL (``predict_delta``) -------------------------------- z_{t+1} is very close to z_t: the ball moves 0.022 world units per frame, which is a small perturbation of a 16-d code. Predicting z_{t+1} directly means the network must spend most of its capacity re-emitting its own input -- learning the identity map to high precision -- before any of its capacity goes to the physics. Predicting the *delta* hands it the identity for free and leaves the network with only the interesting part: the change. In practice this is worth a large chunk of NLL and it makes early training far better behaved. Why an LSTM and not a feedforward net ------------------------------------- Because a single latent does not contain velocity. The VAE sees one frame; a still image of a ball has a position and no direction of travel (the stage-one probes confirm this: velocity R^2 ~ 0 from mu). The transition is therefore NOT Markov in z. The recurrent state is where velocity has to live -- the model must infer it by integrating consecutive latents -- and ``wm.eval_rnn`` part (c) tests exactly that claim by probing h_t for velocity. Heads ----- MDN over z_{t+1} the dynamics hit logit will this transition be a paddle contact? reward scalar dense shaping reward for this transition The extra heads are cheap (a linear map off the same hidden state) and they are what makes the model usable as an environment for a controller in stage three: a controller needs to dream forward AND be told how well it is doing, without ever touching the real simulator. """ from __future__ import annotations import math from dataclasses import asdict, dataclass from typing import Dict, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F # log-std is clamped to this range. Lower bound: a mixture component can # otherwise shrink its std toward zero on a single point and send the NLL to # -inf (the classic degenerate-mixture blow-up). Upper bound: keeps exp() sane. # exp(-7) ~ 9e-4, which is far below the scale of any real latent motion, and # exp(2) ~ 7.4, wider than the latent space itself. LOGSTD_MIN, LOGSTD_MAX = -7.0, 2.0 LOG_SQRT_2PI = 0.5 * math.log(2.0 * math.pi) @dataclass class RNNConfig: z_dim: int = 16 n_actions: int = 3 hidden: int = 256 n_gauss: int = 5 predict_delta: bool = True ablate_actions: bool = False # zero the action input; see eval part (d) # Optional auxiliary head predicting log(mass) from h. Off by default, so # every checkpoint trained before it existed still loads: the field simply # takes its default and no parameter is created, leaving the state dict # identical. See the note on privileged targets in the Heads section above. mass_head: bool = False # v3 control. True replaces the LSTM with a two-hidden-layer MLP on # [z_t, a_t]; ``h`` becomes that MLP's last hidden layer, so every probe and # eval that reads ``h`` keeps working and the two models are compared on the # same footing. Default False, so every pre-v3 checkpoint loads unchanged. # See the note above MDNRNN for why this is the control that matters. feedforward: bool = False # v3 fix (c). Optional auxiliary head predicting the ball's TRUE (x, y) # from h on every frame, hidden ones included. Exactly the same privileged # -target arrangement as ``mass_head``: the label comes from the simulator # at training time and nothing at dream time reads it. Unlike the mass # head, though, this one is not a fair world-model result -- it TELLS the # recurrent state what to hold, so a model trained with it is a CEILING on # how much permanence the architecture can carry, not evidence that the # self-supervised objective would find it. Default False so every # checkpoint trained before it existed loads with an identical state dict. pos_head: bool = False # v3.1 fix. Two extra outputs on h: "frames since the ball was last at # least half visible" and "frames until it next is", both clipped and # scaled into [0, 1] (see ``wm/clock.py``). UNLIKE every other auxiliary # head here the fair version's target is NOT privileged: it is built from a # frozen degree-2 probe on the model's own input latents, so every bit of # it is information the model could read off z at that frame. What it adds # is a REASON to count -- which the one-step NLL, 20 frames from the exit, # cannot supply. ``--clock-privileged`` swaps the probe for the simulator's # ball_visible column and is the labelled ceiling. Default False, so every # earlier checkpoint loads with an identical state dict. clock_head: bool = False # A linear head on h predicting the TRUE ball_vy. Privileged, exactly like # pos_head, and used only in the privileged clock run: "when does it come # out" and "how fast is it falling" are the same fact twice, and the # ceiling run is allowed to be told both. vy_head: bool = False # v4. Which sequence backbone a checkpoint was trained with. This is the # ONLY thing ``load_rnn`` needs in order to hand back the right class, and # giving it a default means every pre-v4 checkpoint -- whose saved config # has no such key -- still loads as an LSTM with a byte-identical state # dict. ``wm.transformer.TransformerConfig`` sets it to "transformer". arch: str = "lstm" class MDNDynamics(nn.Module): """Everything a dynamics model needs EXCEPT the sequence backbone. Split out of ``MDNRNN`` in v4, when a causal transformer arrived as a second backbone. The heads, the mixture likelihood, the delta bookkeeping and the sampler are not properties of "being an LSTM" -- they are properties of "predicting z_{t+1} as a mixture" -- and the LSTM-vs-transformer comparison is only worth anything if *literally the same* head code and *literally the same* loss run on top of both. So they live here, once, and each backbone supplies only ``forward`` and ``init_hidden``. The refactor is deliberately shallow: every parameter is still an attribute of the model itself (``self.mdn``, ``self.hit_head``, ...) rather than of a nested ``self.heads`` module, so the state-dict keys are unchanged and every checkpoint written before v4 loads without a shim. The contract a subclass must honour, and which the evals rely on so that no call site ever has to ask which model it is holding: ``forward(z, a, h) -> (parts, h)`` parts has "h" (B, T, hidden) and "z_in", plus the head outputs ``step(z_t, a_t, h) -> (parts, h)`` one timestep, T kept as 1 ``init_hidden(B, device)`` a tuple whose ``[0][0]`` is the (B, hidden) carried state and whose ``[1][0]`` is the same shape ``cfg.hidden`` the width of ``parts["h"]`` """ cfg: object # RNNConfig or TransformerConfig; both expose .hidden def _build_heads(self) -> None: """Construct every head. Called by the subclass after it has set cfg.""" c = self.cfg # One linear head emitting all mixture parameters at once: # K logits + K*z means + K*z log-stds self.mdn = nn.Linear(c.hidden, c.n_gauss * (1 + 2 * c.z_dim)) self.hit_head = nn.Linear(c.hidden, 1) self.reward_head = nn.Linear(c.hidden, 1) # Privileged at TRAINING time only, exactly like the reward head: the # target comes from the simulator's state vector, never from anything # the model can see at dream time, and nothing downstream reads it. # Its job is to force the colour -> mass factor to stay explicitly # represented in h rather than being smeared through whatever mixture # of latent directions happens to minimise the one-step NLL. self.mass_head = nn.Linear(c.hidden, 1) if c.mass_head else None # (ball_x, ball_y) from h. See RNNConfig.pos_head: privileged, and a # ceiling rather than a result. Its output is exposed under its own # key ``"ball_pos"`` and NOTHING downstream reads that key -- every # eval reads position out of ``h`` with an external probe, which is # still a fair measurement of what h holds. self.pos_head = nn.Linear(c.hidden, 2) if c.pos_head else None # (frames_since, frames_until) / clip. Exposed under ``"clock"``; # nothing downstream reads it, for the same reason nothing reads # ``ball_pos`` -- every evaluation probes ``h`` from outside. self.clock_head = nn.Linear(c.hidden, 2) if c.clock_head else None self.vy_head = nn.Linear(c.hidden, 1) if c.vy_head else None # Start with small MDN outputs so the initial predicted delta is ~0, # i.e. the model starts life as the identity map. With predict_delta # that is already a decent predictor, so training begins from a sane # place instead of from random jumps. nn.init.zeros_(self.mdn.bias) nn.init.normal_(self.mdn.weight, std=1e-3) def _split(self, raw: torch.Tensor) -> Dict[str, torch.Tensor]: """(B, T, K*(1+2z)) -> logits (B,T,K), mean/logstd (B,T,K,z).""" B, T, _ = raw.shape K, z = self.cfg.n_gauss, self.cfg.z_dim logits, mean, logstd = torch.split(raw, [K, K * z, K * z], dim=-1) return { "logits": logits, "mean": mean.view(B, T, K, z), "logstd": logstd.view(B, T, K, z).clamp(LOGSTD_MIN, LOGSTD_MAX), } def _heads(self, out: torch.Tensor, z: torch.Tensor) -> Dict[str, torch.Tensor]: """Backbone output (B, T, hidden) + the input latents -> ``parts``.""" parts = self._split(self.mdn(out)) parts["hit_logit"] = self.hit_head(out) # (B, T, 1) parts["reward"] = self.reward_head(out) # (B, T, 1) if self.mass_head is not None: parts["log_mass"] = self.mass_head(out) # (B, T, 1) if self.pos_head is not None: parts["ball_pos"] = self.pos_head(out) # (B, T, 2) if self.clock_head is not None: parts["clock"] = self.clock_head(out) # (B, T, 2) if self.vy_head is not None: parts["ball_vy"] = self.vy_head(out) # (B, T, 1) parts["h"] = out # (B, T, hidden) # Stash z so the delta bookkeeping lives in one place. parts["z_in"] = z return parts def step(self, z_t: torch.Tensor, a_t: torch.Tensor, h=None): """One timestep. Returns (parts with the T dim kept at 1, h). The generic implementation: unsqueeze, call ``forward``, hand back what it returns. A backbone whose per-step cost differs from its per-window cost (the transformer) overrides this. """ parts, h = self.forward(z_t.unsqueeze(1), a_t.unsqueeze(1), h) return parts, h # -------------------------------------------------------- delta plumbing def _abs_mean(self, parts: Dict[str, torch.Tensor]) -> torch.Tensor: """Component means in absolute z space: (B, T, K, z).""" if not self.cfg.predict_delta: return parts["mean"] # The MDN models z_{t+1} - z_t, so add the input back. Note the std is # unchanged: a shift does not change the spread. return parts["mean"] + parts["z_in"].unsqueeze(2) def _target_in_model_space( self, parts: Dict[str, torch.Tensor], z_next: torch.Tensor ) -> torch.Tensor: """Map the target into whatever space the MDN is parameterising.""" if not self.cfg.predict_delta: return z_next return z_next - parts["z_in"] # ------------------------------------------------------------------ loss def mdn_nll_per_step( self, parts: Dict[str, torch.Tensor], z_next: torch.Tensor ) -> torch.Tensor: """Per-(batch, time) NLL of ``z_next``. (B, T). Split out of ``mdn_nll`` for v3 fix (b): the emergence-weighted loss needs to multiply each transition by its own weight BEFORE the mean, so the reduction cannot be buried inside the likelihood. ``mdn_nll`` is now exactly ``mdn_nll_per_step(...).mean()`` and is bit-identical to what it was, which ``tests/test_fix_v3.py`` pins. """ y = self._target_in_model_space(parts, z_next).unsqueeze(2) # (B,T,1,z) mean, logstd = parts["mean"], parts["logstd"] var = torch.exp(2.0 * logstd) # (B, T, K, z) elementwise Gaussian log-density. log_p = -0.5 * (y - mean) ** 2 / var - logstd - LOG_SQRT_2PI log_p = log_p.sum(-1) # (B,T,K) log_pi = F.log_softmax(parts["logits"], dim=-1) # (B,T,K) return -torch.logsumexp(log_pi + log_p, dim=-1) # (B,T) def mdn_nll( self, parts: Dict[str, torch.Tensor], z_next: torch.Tensor, weights: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Negative log likelihood of ``z_next``. Scalar (mean over B and T). Per (b, t) the log-likelihood of a diagonal mixture is log sum_k pi_k prod_d N(y_d | mu_kd, sigma_kd) = logsumexp_k [ log pi_k + sum_d log N(y_d | mu_kd, sigma_kd) ] Everything stays in the log domain and the sum over k goes through ``logsumexp``. Doing it naively -- exponentiating the per-component densities and summing -- underflows to exactly 0 as soon as one component is a poor fit in 16 dimensions, and then log(0) = -inf ends the run. This is *the* numerical trap in MDNs. Convention: sum over the z dimensions, mean over batch and time. Same reasoning as ``vae_loss``: a mean over dimensions would silently divide the dynamics term by z_dim relative to the auxiliary heads. ``weights`` (B, T), if given, re-weights individual transitions. It is normalised to mean 1 inside, so the loss stays on the same scale as the unweighted one whatever the weights are -- only the RATIO between transitions is meaningful, and keeping the scale fixed means one learning rate works for every ``--emerge-weight``. Validation always calls this with ``weights=None``, so val NLL stays comparable across runs. """ nll = self.mdn_nll_per_step(parts, z_next) # (B,T) if weights is None: return nll.mean() w = weights.to(nll.dtype) return (nll * w).mean() / w.mean().clamp_min(1e-8) # ---------------------------------------------------------- sampling etc. def most_likely_mean(self, parts: Dict[str, torch.Tensor]) -> torch.Tensor: """Mean of the highest-weight component, in absolute z. (B, T, z). NOT the mixture mean sum_k pi_k mu_k -- that is the averaging failure this whole file exists to avoid. For a deterministic rollout you want a point the model considers likely, and the mixture mean can sit in a valley between two modes. """ k = parts["logits"].argmax(-1) # (B,T) mean = self._abs_mean(parts) # (B,T,K,z) idx = k[..., None, None].expand(-1, -1, 1, mean.shape[-1]) return mean.gather(2, idx).squeeze(2) def sample_next( self, parts: Dict[str, torch.Tensor], temperature: float = 1.0, generator: Optional[torch.Generator] = None, ) -> torch.Tensor: """Draw z_{t+1} ~ mixture, with the paper's temperature. (B, T, z). tau scales the component log-stds by tau and the mixture logits by 1/tau. tau = 1 is the model's honest predictive distribution; tau < 1 sharpens both the component choice and the within-component spread; tau -> 0 degenerates to "the mean of the most likely component", which is what ``temperature=0`` returns exactly (no sampling at all, so the rollout is deterministic and reproducible). Why you would ever want tau != 1: a dream rolled out at tau = 1 accumulates the model's own noise and drifts off-manifold; lowering tau gives a cleaner but less diverse dream. Ha & Schmidhuber use tau > 1 for the opposite reason -- a more uncertain dream makes a controller trained inside it less able to exploit the model's flaws. """ if temperature <= 0.0: return self.most_likely_mean(parts) logits = parts["logits"] / temperature # Gumbel-max: argmax(logits + Gumbel noise) is an exact categorical draw # and needs no loop over the batch. u = torch.rand(logits.shape, device=logits.device, generator=generator) g = -torch.log(-torch.log(u.clamp_min(1e-20)).clamp_min(1e-20)) k = (logits + g).argmax(-1) # (B,T) mean = self._abs_mean(parts) # (B,T,K,z) std = torch.exp(parts["logstd"]) * temperature idx = k[..., None, None].expand(-1, -1, 1, mean.shape[-1]) m = mean.gather(2, idx).squeeze(2) s = std.gather(2, idx).squeeze(2) eps = torch.randn(m.shape, device=m.device, generator=generator) return m + eps * s class MDNRNN(MDNDynamics): """Single-layer LSTM + mixture-density head, as in Ha & Schmidhuber (2018).""" def __init__(self, cfg: Optional[RNNConfig] = None): super().__init__() self.cfg = cfg or RNNConfig() c = self.cfg in_dim = c.z_dim + c.n_actions if c.feedforward: # The memory control (v3). Two hidden layers of 256, so it has # MORE per-step nonlinearity than the LSTM and strictly less # information: its output at time t is a function of (z_t, a_t) # alone. It therefore CANNOT carry the ball's position through an # occlusion, which is exactly the floor every memory test needs. # ``lstm`` stays None and the attribute name is not reused, so the # two state dicts are disjoint and neither can load the other by # accident. self.lstm = None self.ff = nn.Sequential( nn.Linear(in_dim, c.hidden), nn.Tanh(), nn.Linear(c.hidden, c.hidden), nn.Tanh(), ) else: # One layer, exactly as in the paper. Depth is not the bottleneck on # this problem -- the thing that is hard is carrying velocity through # time, which is a recurrence property, not a depth property. self.lstm = nn.LSTM(in_dim, c.hidden, num_layers=1, batch_first=True) self.ff = None self._build_heads() # ------------------------------------------------------------------ core def init_hidden( self, batch: int, device: str | torch.device = "cpu" ) -> Tuple[torch.Tensor, torch.Tensor]: """(h, c) of zeros. For ``feedforward=True`` this is a dummy the model passes through untouched -- the signature is kept so callers do not have to know which kind of model they hold.""" h = torch.zeros(1, batch, self.cfg.hidden, device=device) return h, h.clone() def forward( self, z: torch.Tensor, # (B, T, z_dim) a_onehot: torch.Tensor, # (B, T, n_actions) h: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> Tuple[Dict[str, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: if self.cfg.ablate_actions: # The ablation is done HERE rather than by not building the input, # so the parameter count and every shape stay identical to the main # model and the NLL comparison is apples to apples. a_onehot = torch.zeros_like(a_onehot) x = torch.cat([z, a_onehot], dim=-1) if self.ff is not None: # No recurrence: every timestep is processed independently and the # carried state is passed straight back out untouched, so callers # written for the LSTM (warm-up loops, ``step``, the dream # environment) need no branch of their own. out = self.ff(x) # The CARRIED state (the second return value) has to be the thing a # caller can hand to the controller as ``h_pre`` -- ``wm.dream_env`` # and ``wm.eval_controller`` both read ``h[0][0]``, i.e. "the state # produced after consuming the previous input". For the LSTM that is # the last timestep's output, so for the MLP it is the last # timestep's output too. Returning the incoming dummy instead (what # this did before stage three of v3) would hand a controller a # constant zero vector and silently turn the memory FLOOR into a # z-only policy, which is a different experiment. # # This changes nothing about the model: nothing is fed back in (the # MLP ignores its ``h`` argument), so every per-step output, every # probe of ``parts["h"]`` and every number in README_M3 is # unaffected. Only the convenience handle changes. h = ( out[:, -1].unsqueeze(0).contiguous(), torch.zeros_like(out[:, -1]).unsqueeze(0), ) else: out, h = self.lstm(x, h) return self._heads(out, z), h # ------------------------------------------------------------- combined loss def rnn_loss( model: MDNDynamics, parts: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor], pos_weight: float = 1.0, w_hit: float = 1.0, w_reward: float = 1.0, nll_weights: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: """NLL + weighted BCE(hit) + MSE(reward). Returns (loss, parts_dict). ``pos_weight`` matters more than it looks. Paddle contact is ~1% of transitions, so the unweighted BCE optimum is very close to "always say no", which scores a fine loss and a recall of zero. ``pos_weight = (1-p)/p`` re-balances the two classes so the head is actually forced to find the positives; the cost is that the raw probabilities come out over-confident, which is why eval reports precision/recall/PR-AUC rather than accuracy. ``nll_weights`` (B, T) re-weights the dynamics term per transition -- v3 fix (b) uses it to up-weight the frames where the ball re-emerges from behind the occluder. Only the MDN term is weighted; the hit and reward heads are left alone so that those numbers keep meaning what they did. """ nll = model.mdn_nll(parts, batch["z_next"], weights=nll_weights) hit_logit = parts["hit_logit"].squeeze(-1) bce = F.binary_cross_entropy_with_logits( hit_logit, batch["hit"], pos_weight=torch.as_tensor(pos_weight, device=hit_logit.device), ) rew = parts["reward"].squeeze(-1) mse = F.mse_loss(rew, batch["reward"]) loss = nll + w_hit * bce + w_reward * mse return loss, { "loss": loss.detach(), "nll": nll.detach(), "hit_bce": bce.detach(), "reward_mse": mse.detach(), } # ------------------------------------------------------------------ ckpt i/o def save_rnn(path, model: MDNDynamics, args: Optional[dict] = None, extra: Optional[dict] = None): from pathlib import Path path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) payload = { "model": model.state_dict(), "cfg": asdict(model.cfg), "args": args or {}, } if extra: payload.update(extra) torch.save(payload, path) return path def load_rnn(path, device: str = "cpu") -> Tuple[MDNDynamics, object]: """Load whichever dynamics model wrote this checkpoint. v4 added a second backbone, so this dispatches on the ``arch`` key of the saved config. Checkpoints written before v4 have no such key; they get the dataclass default ``"lstm"`` and load exactly as they always did. The filtering on ``__dataclass_fields__`` is the same trick as before and is what lets a config gain a field without invalidating old checkpoints. """ ck = torch.load(path, map_location=device, weights_only=False) saved = dict(ck["cfg"]) arch = str(saved.get("arch", "lstm")) if arch == "transformer": # Imported lazily: transformer.py imports this module for the shared # heads, so a top-level import here would be circular. from .transformer import TransformerConfig, TransformerDynamics cls, cfg_cls = TransformerDynamics, TransformerConfig elif arch == "lstm": cls, cfg_cls = MDNRNN, RNNConfig else: raise ValueError(f"{path}: unknown arch {arch!r}") fields = set(cfg_cls.__dataclass_fields__) cfg = cfg_cls(**{k: v for k, v in saved.items() if k in fields}) model = cls(cfg).to(device) model.load_state_dict(ck["model"]) model.eval() return model, cfg