"""MDLM ELBO loss with SUBS parameterisation. The craftax twin is src/diffusion/loss.py:compute_loss. Computes continuous-time loss on masked positions only, with analytic SUBS weighting clipped for numerical stability. """ from __future__ import annotations from collections.abc import Callable import torch import torch.nn.functional as F from torch import Tensor from src.diffusion.schedules import get_schedule_deriv_for _MAX_WEIGHT: float = 1000.0 # matches loss_weight_clip default; craftax twin identical _WEIGHT_DENOM_EPS: float = 1e-5 # floor for 1 - alpha_t; craftax _EPS identical def mdlm_loss( logits: Tensor, x0: Tensor, zt: Tensor, t: Tensor, mask_token: int, pad_token: int, schedule_fn: Callable[[Tensor], Tensor], schedule_deriv_fn: Callable[[Tensor], Tensor] | None = None, weight_clip: float = _MAX_WEIGHT, label_smoothing: float = 0.0, reduction: str = "mean", ) -> Tensor: """Monte-Carlo estimate of the continuous-time MDLM NELBO. Per sample: ``w(t) * sum_masked(CE) / L`` with the analytic weight ``w(t) = -alpha'(t) / (1 - alpha_t)`` clipped at *weight_clip*, then the batch mean. This is the estimator stated by MDLM eq (10) and Shi et al. eq (4) under a constant per-token normalisation. ``reduction="none"`` returns the per-sample ``[B]`` vector instead of the batch mean (the ablation suite weighs samples by advantage, exactly as the craftax twin's ``compute_loss`` does internally). Replaces a flat average over all masked tokens in the batch — the MaskGIT loss of Shi et al. App. eq (28), which is not a likelihood bound — and the opt-in ``use_importance_weighting`` path, which divided by the realised masked count (a ``1/(1-alpha_t)`` distortion of the weight). Args: logits: Model output. Shape ``[B, L, vocab]``. x0: Clean action sequences. Shape ``[B, L]``, int64. zt: Noisy sequences. Shape ``[B, L]``, int64. t: Per-sample diffusion time in [0, 1]. Shape ``[B]``. mask_token: MASK token ID. pad_token: PAD token ID. schedule_fn: Noise schedule returning alpha(t). schedule_deriv_fn: Analytic d(alpha)/dt; resolved from *schedule_fn* via the registry when ``None``. weight_clip: Upper clamp for w(t) (default 1000). label_smoothing: Smoothing epsilon for cross-entropy. reduction: ``"mean"`` (default, scalar) or ``"none"`` (``[B]``). Returns: Scalar loss, or the per-sample ``[B]`` vector under ``reduction="none"``. Zero(s) when no masked positions exist, and differentiable in ``logits`` even then — see the mask comment below. """ if logits.ndim != 3 or x0.shape != zt.shape or x0.shape != logits.shape[:2]: raise ValueError( "mdlm_loss expects logits [B, L, V] with x0/zt [B, L]; got " f"{tuple(logits.shape)}, {tuple(x0.shape)}, {tuple(zt.shape)}" ) if reduction not in ("mean", "none"): raise ValueError(f"Unknown reduction: {reduction!r}") B, L, V = logits.shape # Mask: compute loss only on masked, non-PAD positions. # # An all-False mask is a legitimate draw, not an error: at a t where # alpha(t) is near 1 nothing gets masked. It is handled by the arithmetic # below rather than by an early return, because the value is not the only # thing that matters — the result has to stay differentiable in `logits`. # `ce * is_masked.float()` gives exactly zero while keeping `logits` in # the graph; a freshly allocated zero tensor gives the same number with # no graph, and any caller that back-propagates it raises "element 0 of # tensors does not require grad and does not have a grad_fn". That is # reachable whenever the caller's other loss terms cannot carry the graph # either, which is the case for every ablation that freezes the goal # head's input path. The craftax twin has always computed this zero # arithmetically. is_masked = (zt == mask_token) & (x0 != pad_token) # [B, L] # Per-position cross-entropy # Clamp targets to valid vocab range — out-of-range positions (PAD, # MASK) will be zeroed out by is_masked anyway. safe_targets = x0.clamp(0, V - 1) # [B, L] ce = F.cross_entropy( logits.reshape(-1, V), safe_targets.reshape(-1), reduction="none", label_smoothing=label_smoothing, ) # [B*L] ce = ce.reshape(B, L) # [B, L] # Zero out non-masked positions ce = ce * is_masked.float() # [B, L] # NELBO weight w(t) = -alpha'(t) / (1 - alpha_t), analytic derivative if schedule_deriv_fn is None: schedule_deriv_fn = get_schedule_deriv_for(schedule_fn) alpha_t = schedule_fn(t) # [B] w_t = (-schedule_deriv_fn(t)) / torch.clamp( 1.0 - alpha_t, min=_WEIGHT_DENOM_EPS, ) # [B] w_t = torch.clamp(w_t, max=weight_clip) # [B] # Constant per-token normalisation (1/L), NOT the realised masked count per_sample = w_t * ce.sum(dim=1) / L # [B] if reduction == "none": return per_sample # `sum / max(B, 1)`, not `mean`: identical for every non-empty batch, and # zero rather than NaN for an empty one, which the removed early return # also happened to cover. return per_sample.sum() / max(B, 1) def auxiliary_goal_loss( goal_pred: Tensor, global_obs: Tensor, pad_value: float = -1.0, ) -> Tensor: """MSE loss for auxiliary staircase-coordinate prediction. Args: goal_pred: Predicted normalised staircase coords. Shape ``[B, 2]``. global_obs: Full map glyphs. Shape ``[B, 21, 79]``, int. pad_value: Coordinate value used when staircase is not visible. Returns: Scalar MSE loss over samples where the staircase is visible. Exactly zero when no staircase is visible in the batch, whatever *goal_pred* holds, and differentiable in *goal_pred* whenever *goal_pred* itself is. """ targets = find_staircase_from_glyphs(global_obs) # [B, 2] targets = targets.to(goal_pred.device, dtype=goal_pred.dtype) # Only supervise where staircase is visible valid = targets[:, 0] != pad_value # [B] if not valid.any(): # An empty selection: exactly 0.0, in the graph, zero gradient. # # Three constraints meet here. The caller adds this term to the ELBO # term and back-propagates the sum, so a detached constant silently # drops this term from the graph -- that was `goal_pred.new_tensor(0.0)`, # removed in `0cfc632` because it left every arm with a frozen goal # head unable to back-propagate at all. Its replacement multiplied by # the empty `valid` mask, which keeps the graph but returns NaN for a # non-finite `goal_pred`, since `nan * False` is `nan` -- while the # supervised branch below excludes exactly those rows. Indexing the # same way that branch does satisfies all three: `goal_pred[valid]` is # empty, so the sum is exactly zero whatever `goal_pred` holds, and it # is still a function of `goal_pred`, so the graph survives. return goal_pred[valid].sum() diff = (goal_pred[valid] - targets[valid]) ** 2 # [N, 2] return diff.mean() def find_staircase_from_glyphs(global_obs: Tensor) -> Tensor: """Locate the staircase '>' in the global glyph map. Searches for NLE staircase-down glyph (character code 62 = '>'). Returns normalised (row/H, col/W) coordinates per batch element, or (-1, -1) when the staircase is not visible. Args: global_obs: Glyph map. Shape ``[B, H, W]`` or ``[H, W]``, int. Returns: Normalised coordinates. Shape ``[B, 2]`` (float32). """ if global_obs.ndim == 2: global_obs = global_obs.unsqueeze(0) B, H, W = global_obs.shape # NLE staircase-down glyphs: ord('>') = 62, plus NLE tile variants # 2310 (S_dnstair), 2368 (S_dnstairs), 2383 (S_vodoor). is_stair = ( (global_obs == 62) | (global_obs == 2310) | (global_obs == 2368) | (global_obs == 2383) ) # Vectorised over the batch. The previous form looped over B # calling `is_stair[b].nonzero()`, and `nonzero` needs its output size # on the host, so every sample forced a device sync: 2048 syncs per # gradient step at `dagger_batch_size: 2048`, which dominated the step. # # `nonzero` returns indices in row-major order, so `positions[0]` is the # lowest flat index that is set. Taking the minimum flat index over the # masked positions reproduces that exactly, with no host round-trip. flat = is_stair.reshape(B, H * W) idx = torch.arange(H * W, device=global_obs.device, dtype=torch.int32) masked_idx = torch.where(flat, idx, torch.full_like(idx, H * W)) first = masked_idx.min(dim=1).values # [B]; == H*W when no staircase found = first < H * W row = (first // W).float() / max(1, H - 1) col = (first % W).float() / max(1, W - 1) coords = torch.stack( ( torch.where(found, row, -1.0), torch.where(found, col, -1.0), ), dim=1, ) return coords.to(torch.float32)