perception.py

# SPDX-License-Identifier: GPL-3.0-only
# Copyright (C) 2026 SLING AI Inc.
"""
perception.py -- Stage 1 of the Goal-Conditioned VPA framework (preprint_261008.pdf, Sec. 4.1).

Implements
    * ``VisionEncoder``      -- the shared (Siamese) vision encoder E_psi that maps observation
                                frames I_t and milestone frames I_g^(m) into the latent space
                                Z subset R^d (Eq. 7):
                                    z_t = E_psi(I_t),   z_g^(m) = E_psi(I_g^(m)),  m = 1..M.
    * ``MomentumEncoder``    -- the momentum (exponential-moving-average) copy E_psi_bar that
                                produces the stop-gradient target z_{t+1} = sg(E_psi_bar(I_{t+1}))
                                of the JEPA objective (Eq. 11, Fig. 3).
    * ``TextEncoderWrapper`` -- the lightweight *frozen* text encoder (e.g. CLIP) that embeds the
                                natural-language instruction x into c_text in R^{d_c}.

Layout convention
    The paper writes I_t in R^{H_img x W_img x C}. PyTorch convolutions are channels-first, so
    every image tensor in this module is ``[Batch, C, H_img, W_img]``; a milestone sequence is
    ``[Batch, M, C, H_img, W_img]``. With V >= 2 camera views (``VisionEncoder(num_views=V)``), a
    view axis follows the batch (and milestone) axes: ``[Batch, V, C, H_img, W_img]`` and
    ``[Batch, M, V, C, H_img, W_img]``. With one camera there is no view axis.

This file is self-contained (it imports nothing from the other VPA modules) and can be tested
on its own with ``python perception.py``.
"""

from __future__ import annotations

import copy
from types import SimpleNamespace
from typing import Any, Callable, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
from torch import Tensor

__all__ = [
    "PatchEmbedding",
    "TransformerBlock",
    "VisionEncoder",
    "MomentumEncoder",
    "TextEncoderWrapper",
]

ImageSize = Union[int, Tuple[int, int]]


def _pair(x: ImageSize) -> Tuple[int, int]:
    return (x, x) if isinstance(x, int) else (int(x[0]), int(x[1]))


# ---------------------------------------------------------------------------------------------
# Vision Transformer building blocks
# ---------------------------------------------------------------------------------------------
class PatchEmbedding(nn.Module):
    """Non-overlapping patch tokenizer of a Vision Transformer.

    Args:
        image_size: (H_img, W_img) or a single int for square frames.
        patch_size: side length P of the square patches; must divide H_img and W_img.
        in_channels: number of image channels C.
        width: token width D_vit.

    Shapes:
        input  ``images``: ``[B, C, H_img, W_img]``
        output ``tokens``: ``[B, N_patch, D_vit]`` with N_patch = (H_img / P) * (W_img / P)
    """

    def __init__(self, image_size: ImageSize, patch_size: int, in_channels: int, width: int) -> None:
        super().__init__()
        h, w = _pair(image_size)
        if h % patch_size != 0 or w % patch_size != 0:
            raise ValueError(f"image size {(h, w)} is not divisible by patch size {patch_size}")
        self.image_size: Tuple[int, int] = (h, w)
        self.patch_size: int = patch_size
        self.num_patches: int = (h // patch_size) * (w // patch_size)
        self.proj = nn.Conv2d(in_channels, width, kernel_size=patch_size, stride=patch_size)

    def forward(self, images: Tensor) -> Tensor:
        """``[B, C, H_img, W_img]`` -> ``[B, N_patch, D_vit]``."""
        x = self.proj(images)  # [B, D_vit, H_img/P, W_img/P]
        return x.flatten(2).transpose(1, 2)  # [B, N_patch, D_vit]


class TransformerBlock(nn.Module):
    """Pre-norm ViT block: x <- x + MHSA(LN(x));  x <- x + MLP(LN(x)).

    Only LayerNorm is used (no BatchNorm), so every sample in a batch is processed independently.
    This is what makes Siamese batching of observations and milestones exact (see
    ``VisionEncoder.encode_siamese``).

    Shapes:
        input/output ``x``: ``[B, N_tok, D_vit]``
    """

    def __init__(self, width: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
        super().__init__()
        if width % num_heads != 0:
            raise ValueError(f"width {width} must be divisible by num_heads {num_heads}")
        hidden = int(round(width * mlp_ratio))
        self.norm1 = nn.LayerNorm(width)
        self.attn = nn.MultiheadAttention(width, num_heads, dropout=0.0, batch_first=True)
        self.norm2 = nn.LayerNorm(width)
        self.mlp = nn.Sequential(nn.Linear(width, hidden), nn.GELU(), nn.Linear(hidden, width))

    def forward(self, x: Tensor) -> Tensor:
        """``[B, N_tok, D_vit]`` -> ``[B, N_tok, D_vit]``."""
        h = self.norm1(x)
        x = x + self.attn(h, h, h, need_weights=False)[0]
        x = x + self.mlp(self.norm2(x))
        return x


# ---------------------------------------------------------------------------------------------
# E_psi : shared Siamese vision encoder (Eq. 7)
# ---------------------------------------------------------------------------------------------
class VisionEncoder(nn.Module):
    """Shared vision encoder E_psi (Vision Transformer backbone, Sec. 4.1, Eq. 7).

    A single set of weights psi embeds both the current observation and every milestone goal
    frame, i.e. the encoder is used as a Siamese network:

        z_t     = E_psi(I_t)          in R^d
        z_g^(m) = E_psi(I_g^(m))      in R^d,   m = 1, ..., M

    Single camera (``num_views = 1``, the paper's setting): the latent is the linear projection of
    the final [CLS] token after the last LayerNorm.

    Several cameras (``num_views = V >= 2``; a decision where the paper is silent). The paper writes
    I_t as one frame. With V cameras, I_t is the ordered tuple of the V views (I_t^1, ..., I_t^V),
    and every milestone I_g^(m) is the tuple of the same views in the same order. One ViT body
    (shared by all views, as it is shared by observations and milestones) yields the final
    LayerNorm'd [CLS] token f(I^v) of each view, and the head maps their concatenation in view order:

        E_psi(I) = W_head Concat(f(I^1), ..., f(I^V)) + b_head,    W_head in R^{d x V D_vit}

    So E_psi is still a single map from an observation to R^d, and everything downstream of z
    (Eqs. 8-15, 18, 19) is unchanged. With V = 1 this is exactly the single-camera encoder: no extra
    module, the same parameter names and shapes, and the same frame shape ``[C, H_img, W_img]``.

    Args:
        image_size: (H_img, W_img) or int.
        patch_size: ViT patch size P.
        in_channels: image channels C.
        width: ViT token width D_vit.
        depth: number of transformer blocks.
        num_heads: attention heads per block.
        latent_dim: dimension d of the latent space Z.
        mlp_ratio: hidden width multiplier of the block MLPs.
        num_views: number of camera views V per frame (1 = single camera).

    Shapes (``frame_shape`` = ``[C, H_img, W_img]`` if V = 1, ``[V, C, H_img, W_img]`` if V >= 2):
        ``forward``:          ``[B, *frame_shape]`` -> ``[B, d]``
        ``encode_milestones``: ``[B, M, *frame_shape]`` -> ``[B, M, d]``
        ``encode_siamese``:   (``[B, *frame_shape]``, ``[B, M, *frame_shape]``)
                              -> (``[B, d]``, ``[B, M, d]``)
    """

    def __init__(
        self,
        image_size: ImageSize = 224,
        patch_size: int = 16,
        in_channels: int = 3,
        width: int = 384,
        depth: int = 6,
        num_heads: int = 6,
        latent_dim: int = 256,
        mlp_ratio: float = 4.0,
        num_views: int = 1,
    ) -> None:
        super().__init__()
        if int(num_views) != num_views or int(num_views) < 1:
            raise ValueError(f"num_views must be an integer >= 1, got {num_views}")
        self.in_channels: int = in_channels
        self.latent_dim: int = latent_dim
        self.num_views: int = int(num_views)
        self.patch_embed = PatchEmbedding(image_size, patch_size, in_channels, width)
        self.image_size: Tuple[int, int] = self.patch_embed.image_size
        self.cls_token = nn.Parameter(torch.zeros(1, 1, width))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, width))
        self.blocks = nn.ModuleList(
            [TransformerBlock(width, num_heads, mlp_ratio) for _ in range(depth)]
        )
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(self.num_views * width, latent_dim)  # V = 1: Linear(D_vit, d) as before
        self._init_weights()

    def _init_weights(self) -> None:
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.trunc_normal_(module.weight, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

    @property
    def frame_shape(self) -> Tuple[int, ...]:
        """Shape of one observation: ``(C, H_img, W_img)`` if V = 1, ``(V, C, H_img, W_img)`` if V >= 2."""
        c, (h, w) = self.in_channels, self.image_size
        return (c, h, w) if self.num_views == 1 else (self.num_views, c, h, w)

    def _check_frames(self, images: Tensor) -> None:
        expected = self.frame_shape
        if images.dim() != 1 + len(expected) or tuple(images.shape[1:]) != expected:
            layout = "[B, C, H, W]" if self.num_views == 1 else "[B, V, C, H, W]"
            raise ValueError(
                f"expected frames {layout} = [B, {', '.join(str(s) for s in expected)}], "
                f"got {tuple(images.shape)}"
            )

    def _cls_features(self, images: Tensor) -> Tensor:
        """ViT body on single-view frames: ``[N, C, H_img, W_img]`` -> ``[N, D_vit]``.

        Returns the final LayerNorm'd [CLS] token. A plain method, not a submodule, so a
        multi-view forward is still one evaluation of E_psi for the depth audit (Prop. 5.1).
        """
        x = self.patch_embed(images)  # [N, N_patch, D_vit]
        cls = self.cls_token.expand(x.shape[0], -1, -1)  # [N, 1, D_vit]
        x = torch.cat([cls, x], dim=1) + self.pos_embed  # [N, N_patch + 1, D_vit]
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        return x[:, 0]  # [N, D_vit]

    def forward(self, images: Tensor) -> Tensor:
        """Encode a batch of observations (Eq. 7).

        Args:
            images: ``[B, C, H_img, W_img]`` (V = 1) or ``[B, V, C, H_img, W_img]`` (V >= 2) float
                frames (observations and/or milestone goals).

        Returns:
            z: ``[B, d]`` latent embeddings E_psi(I).
        """
        self._check_frames(images)
        if self.num_views == 1:
            return self.head(self._cls_features(images))  # [B, d]
        b = images.shape[0]
        # all B * V views through the shared body in one batch (row b * V + v is view v of sample b)
        cls = self._cls_features(images.flatten(0, 1))  # [B * V, D_vit]
        return self.head(cls.reshape(b, self.num_views * cls.shape[-1]))  # Concat in view order -> [B, d]

    def encode_milestones(self, milestones: Tensor) -> Tensor:
        """Encode an ordered milestone sequence (computed once per episode, Sec. 4.1).

        Args:
            milestones: ``[B, M, *frame_shape]`` -- (I_g^(1), ..., I_g^(M)) per batch element.

        Returns:
            z_g: ``[B, M, d]`` -- (z_g^(1), ..., z_g^(M)).
        """
        if milestones.dim() != 2 + len(self.frame_shape):
            layout = "[B, M, C, H, W]" if self.num_views == 1 else "[B, M, V, C, H, W]"
            raise ValueError(f"expected milestones of shape {layout}, got {tuple(milestones.shape)}")
        b, m = milestones.shape[:2]
        z = self(milestones.reshape(b * m, *milestones.shape[2:]))  # [B*M, d]
        return z.reshape(b, m, self.latent_dim)

    def encode_siamese(self, observation: Tensor, milestones: Tensor) -> Tuple[Tensor, Tensor]:
        """Siamese batching: encode observations and milestones in ONE forward pass of E_psi.

        Both streams are concatenated along the batch axis, passed through the shared weights,
        and split again. Because the network contains no cross-sample operation (LayerNorm only),
        the result is identical to encoding the two streams separately.

        Args:
            observation: ``[B, *frame_shape]`` -- current frames I_t.
            milestones:  ``[B, M, *frame_shape]`` -- milestone frames I_g^(1..M).

        Returns:
            z_t: ``[B, d]``;  z_g: ``[B, M, d]``.
        """
        n = len(self.frame_shape)
        if observation.dim() != 1 + n or milestones.dim() != 2 + n:
            if self.num_views == 1:
                raise ValueError("expected observation [B, C, H, W] and milestones [B, M, C, H, W]")
            raise ValueError("expected observation [B, V, C, H, W] and milestones [B, M, V, C, H, W]")
        b, m = milestones.shape[:2]
        if observation.shape[0] != b:
            raise ValueError(f"batch mismatch: observation {observation.shape[0]} vs milestones {b}")
        frames = torch.cat([observation, milestones.reshape(b * m, *milestones.shape[2:])], dim=0)
        z = self(frames)  # [B + B*M, d]
        z_t = z[:b]  # [B, d]
        z_g = z[b:].reshape(b, m, self.latent_dim)  # [B, M, d]
        return z_t, z_g


# ---------------------------------------------------------------------------------------------
# E_psi_bar : momentum target encoder (Eq. 11, Fig. 3)
# ---------------------------------------------------------------------------------------------
class MomentumEncoder(nn.Module):
    """Momentum (EMA) copy E_psi_bar of the online encoder, used only to produce JEPA targets.

    Update rule after each optimizer step (psi_bar <- momentum * psi_bar + (1 - momentum) * psi).
    The paper does not fix the momentum coefficient; it is a constructor argument here.
    The target branch receives no gradients: ``forward`` runs under ``torch.no_grad`` and returns a
    detached tensor, which realises the stop-gradient sg(.) of Eq. 11.

    Shapes:
        ``forward``: ``[B, *frame_shape]`` -> ``[B, d]`` (detached); ``frame_shape`` is the online
        encoder's (``[C, H_img, W_img]``, or ``[V, C, H_img, W_img]`` with V camera views).
    """

    def __init__(self, online_encoder: VisionEncoder, momentum: float = 0.996) -> None:
        super().__init__()
        if not 0.0 <= momentum < 1.0:
            raise ValueError("momentum must lie in [0, 1)")
        self.momentum: float = float(momentum)
        self.encoder: VisionEncoder = copy.deepcopy(online_encoder)
        self.encoder.requires_grad_(False)

    @torch.no_grad()
    def update(self, online_encoder: VisionEncoder) -> None:
        """psi_bar <- momentum * psi_bar + (1 - momentum) * psi (parameters and buffers)."""
        for p_bar, p in zip(self.encoder.parameters(), online_encoder.parameters(), strict=True):
            p_bar.mul_(self.momentum).add_(p.detach(), alpha=1.0 - self.momentum)
        for b_bar, b in zip(self.encoder.buffers(), online_encoder.buffers(), strict=True):
            b_bar.copy_(b)

    @torch.no_grad()
    def forward(self, images: Tensor) -> Tensor:
        """``[B, *frame_shape]`` -> ``[B, d]`` target latents sg(E_psi_bar(I))."""
        return self.encoder(images).detach()


# ---------------------------------------------------------------------------------------------
# Frozen text encoder (c_text)
# ---------------------------------------------------------------------------------------------
class TextEncoderWrapper(nn.Module):
    """Frozen, lightweight text encoder producing the instruction embedding c_text in R^{d_c}.

    By default it loads a pre-trained CLIP text tower with its projection head
    (``transformers.CLIPTextModelWithProjection`` + ``CLIPTokenizer``) and uses the projected
    pooled output (``text_embeds``) as c_text, so d_c = ``config.projection_dim``
    (512 for ``openai/clip-vit-base-patch32``). All parameters are frozen and the module is kept
    in eval mode permanently; c_text is computed once per episode (Sec. 4.1).

    For offline use or tests, a ``text_model`` and ``tokenizer`` can be injected. The injected
    model must accept ``(input_ids, attention_mask)`` and return an object with a
    ``text_embeds`` attribute of shape ``[B, d_c]``; it must expose ``config.projection_dim``.

    Shapes:
        ``forward``:       list of B strings -> ``[B, d_c]``
        ``encode_tokens``: ``input_ids [B, L]`` (+ ``attention_mask [B, L]``) -> ``[B, d_c]``
    """

    def __init__(
        self,
        model_name_or_path: str = "openai/clip-vit-base-patch32",
        text_model: Optional[nn.Module] = None,
        tokenizer: Optional[Callable[..., Any]] = None,
        max_length: int = 77,
    ) -> None:
        super().__init__()
        if text_model is None or tokenizer is None:
            # Imported lazily so the rest of this module works without `transformers` installed.
            from transformers import CLIPTextModelWithProjection, CLIPTokenizer

            if text_model is None:
                text_model = CLIPTextModelWithProjection.from_pretrained(model_name_or_path)
            if tokenizer is None:
                tokenizer = CLIPTokenizer.from_pretrained(model_name_or_path)
        self.text_model: nn.Module = text_model
        self.tokenizer: Callable[..., Any] = tokenizer
        self.max_length: int = int(max_length)
        self.embed_dim: int = int(text_model.config.projection_dim)  # d_c
        self.text_model.requires_grad_(False)
        self.text_model.eval()

    def train(self, mode: bool = True) -> "TextEncoderWrapper":
        """The text encoder is frozen: it stays in eval mode regardless of ``mode``."""
        super().train(mode)
        self.text_model.eval()
        return self

    def _device(self) -> torch.device:
        try:
            return next(self.text_model.parameters()).device
        except StopIteration:
            return torch.device("cpu")

    @torch.no_grad()
    def encode_tokens(self, input_ids: Tensor, attention_mask: Optional[Tensor] = None) -> Tensor:
        """Embed pre-tokenised instructions.

        Args:
            input_ids:      ``[B, L]`` int64 token ids.
            attention_mask: ``[B, L]`` (optional) 1 for real tokens, 0 for padding.

        Returns:
            c_text: ``[B, d_c]``.
        """
        out = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
        c_text = out.text_embeds if hasattr(out, "text_embeds") else out
        if c_text.dim() != 2 or c_text.shape[-1] != self.embed_dim:
            raise RuntimeError(f"text model returned {tuple(c_text.shape)}, expected [B, {self.embed_dim}]")
        return c_text

    @torch.no_grad()
    def forward(self, instructions: Union[str, Sequence[str]]) -> Tensor:
        """Embed natural-language instructions x.

        Args:
            instructions: one string or a sequence of B strings.

        Returns:
            c_text: ``[B, d_c]`` (B = 1 for a single string).
        """
        texts = [instructions] if isinstance(instructions, str) else list(instructions)
        if len(texts) == 0:
            raise ValueError("at least one instruction is required")
        tokens = self.tokenizer(
            texts, padding=True, truncation=True, max_length=self.max_length, return_tensors="pt"
        )
        device = self._device()
        attention_mask = tokens.get("attention_mask")
        return self.encode_tokens(
            tokens["input_ids"].to(device),
            attention_mask.to(device) if attention_mask is not None else None,
        )


# ---------------------------------------------------------------------------------------------
# Self-test (run: python perception.py)
# ---------------------------------------------------------------------------------------------
class _ToyTextModel(nn.Module):
    """Tiny stand-in with the CLIPTextModelWithProjection interface (offline testing only)."""

    def __init__(self, vocab: int = 100, width: int = 32, projection_dim: int = 24) -> None:
        super().__init__()
        self.config = SimpleNamespace(projection_dim=projection_dim)
        self.embed = nn.Embedding(vocab, width)
        self.proj = nn.Linear(width, projection_dim, bias=False)

    def forward(self, input_ids: Tensor, attention_mask: Optional[Tensor] = None) -> SimpleNamespace:
        h = self.embed(input_ids)
        if attention_mask is not None:
            h = h * attention_mask.unsqueeze(-1)
        return SimpleNamespace(text_embeds=self.proj(h.mean(1)))


def _toy_tokenizer(texts: Sequence[str], max_length: int = 16, **_: Any) -> dict:
    ids = torch.zeros(len(texts), max_length, dtype=torch.long)
    mask = torch.zeros(len(texts), max_length, dtype=torch.long)
    for i, text in enumerate(texts):
        codes = [ord(ch) % 100 for ch in text][:max_length]
        ids[i, : len(codes)] = torch.tensor(codes, dtype=torch.long)
        mask[i, : len(codes)] = 1
    return {"input_ids": ids, "attention_mask": mask}


_RTOL, _ATOL = 1e-5, 1e-6  # float32 module output vs float64 reference (CLAUDE.md)


def _close(actual: Tensor, ref: Tensor) -> bool:
    return torch.allclose(actual.double(), ref.double(), rtol=_RTOL, atol=_ATOL)


def _expect_value_error(fn: Callable[[], Any]) -> None:
    try:
        fn()
    except ValueError:
        return
    raise AssertionError("expected ValueError")


def _ref_layer_norm(x: Tensor, weight: Tensor, bias: Tensor, eps: float) -> Tensor:
    """LayerNorm over the last axis, written out: (x - mean) / sqrt(biased var + eps) * w + b."""
    mean = x.sum(-1, keepdim=True) / x.shape[-1]
    var = ((x - mean) ** 2).sum(-1, keepdim=True) / x.shape[-1]
    return (x - mean) / torch.sqrt(var + eps) * weight + bias


def _ref_vit_cls(enc: VisionEncoder, images: Tensor) -> Tensor:
    """Independent float64 reference of the ViT body on single-view frames ``[B, C, H, W]``.

    Re-derives the ViT forward from its definition with float64 copies of the module weights:
    explicit patch loops, per-head softmax attention, exact erf-GELU, hand-written LayerNorm.
    Returns the final LayerNorm'd [CLS] token, ``[B, D_vit]`` float64.
    """
    f64 = {k: v.detach().double() for k, v in enc.state_dict().items()}
    x = images.double()
    bsz, _, h, w = x.shape
    p = enc.patch_embed.patch_size
    w_patch = f64["patch_embed.proj.weight"].reshape(f64["patch_embed.proj.weight"].shape[0], -1)
    tokens = []
    for i in range(h // p):  # row-major patch order
        for j in range(w // p):
            patch = x[:, :, i * p:(i + 1) * p, j * p:(j + 1) * p].reshape(bsz, -1)  # [B, C*P*P]
            tokens.append(patch @ w_patch.T + f64["patch_embed.proj.bias"])
    t = torch.stack(tokens, dim=1)  # [B, N_patch, D]
    t = torch.cat([f64["cls_token"].expand(bsz, -1, -1), t], dim=1) + f64["pos_embed"]
    width = t.shape[-1]
    for n, block in enumerate(enc.blocks):
        pre = f"blocks.{n}."
        heads = block.attn.num_heads
        dh = width // heads
        hn = _ref_layer_norm(t, f64[pre + "norm1.weight"], f64[pre + "norm1.bias"], block.norm1.eps)
        w_in, b_in = f64[pre + "attn.in_proj_weight"], f64[pre + "attn.in_proj_bias"]
        q = hn @ w_in[:width].T + b_in[:width]
        k = hn @ w_in[width:2 * width].T + b_in[width:2 * width]
        v = hn @ w_in[2 * width:].T + b_in[2 * width:]
        head_out = []
        for a in range(heads):
            sl = slice(a * dh, (a + 1) * dh)
            scores = q[..., sl] @ k[..., sl].transpose(1, 2) / dh ** 0.5  # [B, N, N]
            e = torch.exp(scores - scores.max(dim=-1, keepdim=True).values)
            head_out.append((e / e.sum(dim=-1, keepdim=True)) @ v[..., sl])
        attn = torch.cat(head_out, dim=-1) @ f64[pre + "attn.out_proj.weight"].T + f64[pre + "attn.out_proj.bias"]
        t = t + attn
        hn = _ref_layer_norm(t, f64[pre + "norm2.weight"], f64[pre + "norm2.bias"], block.norm2.eps)
        hid = hn @ f64[pre + "mlp.0.weight"].T + f64[pre + "mlp.0.bias"]
        hid = 0.5 * hid * (1.0 + torch.erf(hid / 2.0 ** 0.5))  # exact GELU
        t = t + hid @ f64[pre + "mlp.2.weight"].T + f64[pre + "mlp.2.bias"]
    t = _ref_layer_norm(t, f64["norm.weight"], f64["norm.bias"], enc.norm.eps)
    return t[:, 0]  # [B, D_vit]


def _ref_vit_forward(enc: VisionEncoder, images: Tensor) -> Tensor:
    """Independent float64 reference of z = E_psi(I) (Eq. 7).

    ``[B, C, H, W]`` (one camera): the head applied to the [CLS] token.
    ``[B, V, C, H, W]`` (V cameras): the [CLS] tokens of the V views, each from its own body pass,
    concatenated in view order, then the head. Returns ``[B, d]`` float64.
    """
    w_head = enc.head.weight.detach().double()
    b_head = enc.head.bias.detach().double()
    if images.dim() == 4:
        feats = _ref_vit_cls(enc, images)  # [B, D_vit]
    else:
        feats = torch.cat([_ref_vit_cls(enc, images[:, v]) for v in range(images.shape[1])], dim=-1)  # [B, V*D_vit]
    return feats @ w_head.T + b_head  # [B, d]


def _self_test() -> None:
    torch.manual_seed(0)
    gen = torch.Generator().manual_seed(1234)
    b, m, c, h, w, d = 2, 3, 3, 32, 32, 16  # B != M so a B/M transposition cannot pass
    enc = VisionEncoder(image_size=(h, w), patch_size=8, in_channels=c, width=32, depth=2,
                        num_heads=4, latent_dim=d).eval()
    obs = torch.randn(b, c, h, w, generator=gen)
    goals = torch.randn(b, m, c, h, w, generator=gen)

    # ---- Eq. 7: z_t = E_psi(I_t), z_g^(m) = E_psi(I_g^(m)) against the float64 reference.
    with torch.no_grad():
        z_t = enc(obs)
        z_g = enc.encode_milestones(goals)
        zs_t, zs_g = enc.encode_siamese(obs, goals)
    assert z_t.shape == (b, d) and z_g.shape == (b, m, d)
    assert zs_t.shape == (b, d) and zs_g.shape == (b, m, d)
    ref_t = _ref_vit_forward(enc, obs)
    assert _close(z_t, ref_t)
    for bi in range(b):
        for mi in range(m):
            ref_g = _ref_vit_forward(enc, goals[bi, mi].unsqueeze(0))[0]
            assert _close(z_g[bi, mi], ref_g)
            assert _close(zs_g[bi, mi], ref_g)
    assert _close(zs_t, ref_t)
    # Siamese batching (one pass, shared psi) == encoding the two streams separately.
    assert torch.allclose(zs_t, z_t, rtol=_RTOL, atol=_ATOL)
    assert torch.allclose(zs_g, z_g, rtol=_RTOL, atol=_ATOL)
    # B = 1: no cross-sample statistics.
    with torch.no_grad():
        z1_t, z1_g = enc.encode_siamese(obs[:1], goals[:1])
    assert z1_t.shape == (1, d) and z1_g.shape == (1, m, d)
    assert _close(z1_t, ref_t[:1]) and torch.allclose(z1_g, z_g[:1], rtol=_RTOL, atol=_ATOL)

    # ---- Malformed frames raise ValueError.
    _expect_value_error(lambda: enc(torch.randn(b, c + 1, h, w)))          # wrong C
    _expect_value_error(lambda: enc(torch.randn(b, c, h + 8, w)))          # wrong H
    _expect_value_error(lambda: enc(torch.randn(b, c, h, w + 8)))          # wrong W
    _expect_value_error(lambda: enc(torch.randn(c, h, w)))                 # 3-D
    _expect_value_error(lambda: enc(torch.randn(b, m, c, h, w)))           # 5-D into forward
    _expect_value_error(lambda: enc.encode_milestones(torch.randn(b, c, h, w)))          # 4-D
    _expect_value_error(lambda: enc.encode_milestones(torch.randn(b, m, c, h, w + 8)))   # frame size
    _expect_value_error(lambda: enc.encode_siamese(torch.randn(b + 1, c, h, w), goals))  # batch mismatch
    _expect_value_error(lambda: enc.encode_siamese(torch.randn(c, h, w), goals))         # 3-D obs
    _expect_value_error(lambda: enc.encode_siamese(obs, torch.randn(b, c, h, w)))        # 4-D milestones

    # ---- Momentum target E_psi_bar (Eq. 11): EMA psi_bar <- mu psi_bar + (1 - mu) psi.
    mu = 0.9
    target = MomentumEncoder(enc, momentum=mu)
    assert not any(p.requires_grad for p in target.parameters())
    for p_bar, p in zip(target.encoder.parameters(), enc.parameters(), strict=True):
        assert torch.equal(p_bar, p) and p_bar.data_ptr() != p.data_ptr()  # copy, not alias
    psi_bar_ref = [p.detach().double().clone() for p in enc.parameters()]
    for _ in range(5):
        with torch.no_grad():  # simulated optimizer step on the online encoder
            for p in enc.parameters():
                p.add_(0.1 * torch.randn(p.shape, generator=gen))
        psi_before = [p.detach().clone() for p in enc.parameters()]
        target.update(enc)
        for i, p in enumerate(enc.parameters()):
            psi_bar_ref[i] = mu * psi_bar_ref[i] + (1.0 - mu) * p.detach().double()
        for p_bar, ref, p, p_old in zip(
            target.encoder.parameters(), psi_bar_ref, enc.parameters(), psi_before, strict=True
        ):
            assert _close(p_bar, ref)
            assert torch.equal(p, p_old)  # update() never writes to the online encoder
    # E_psi_bar is the Eq. 7 map evaluated with the EMA weights.
    assert _close(target(obs), _ref_vit_forward(target.encoder, obs))
    # Stop-gradient sg(.): no graph from the target branch, gradients only reach psi.
    obs_g = obs.clone().requires_grad_(True)
    z_tgt = target(obs_g)
    assert not z_tgt.requires_grad and z_tgt.grad_fn is None
    enc.zero_grad(set_to_none=True)
    loss = ((enc(obs_g) - z_tgt) ** 2).sum()
    loss.backward()
    assert all(p.grad is not None for p in enc.parameters())
    assert all(p.grad is None for p in target.parameters())
    _expect_value_error(lambda: MomentumEncoder(enc, momentum=1.0))
    _expect_value_error(lambda: MomentumEncoder(enc, momentum=-0.1))

    # ---- Frozen text encoder c_text in R^{d_c} (offline stand-in with the CLIP interface).
    text = TextEncoderWrapper(text_model=_ToyTextModel(), tokenizer=_toy_tokenizer)
    for mode in (None, True):
        text.train() if mode is None else text.train(mode)
        assert not text.text_model.training
        assert not any(mod.training for mod in text.text_model.modules())
    assert not any(p.requires_grad for p in text.text_model.parameters())
    texts = ["move the cup", "throw the cup away", "push the block"]
    with torch.enable_grad():
        c_text = text(texts)
        c_one = text("move the cup")
    assert c_text.shape == (len(texts), text.embed_dim) and c_one.shape == (1, text.embed_dim)
    assert not c_text.requires_grad and c_text.grad_fn is None
    tokens = _toy_tokenizer(texts, padding=True, truncation=True, max_length=text.max_length,
                            return_tensors="pt")
    with torch.no_grad():
        direct = text.text_model(input_ids=tokens["input_ids"],
                                 attention_mask=tokens["attention_mask"]).text_embeds
    assert torch.equal(c_text, direct)  # wrapper passes ids and mask through unchanged
    assert torch.allclose(c_one[0], c_text[0], rtol=_RTOL, atol=_ATOL)  # B=1 vs B=3 BLAS rounding
    for i in range(len(texts)):
        for j in range(i + 1, len(texts)):
            assert not torch.allclose(c_text[i], c_text[j])
    _expect_value_error(lambda: text([]))

    # ---- Several cameras (num_views = V = 2): one shared ViT body for every view, the head maps the
    # concatenated [CLS] tokens in view order. Own generator, so the checks above are unaffected.
    gen2 = torch.Generator().manual_seed(4321)
    v = 2
    enc2 = VisionEncoder(image_size=(h, w), patch_size=8, in_channels=c, width=32, depth=2,
                         num_heads=4, latent_dim=d, num_views=v).eval()
    assert enc.num_views == 1 and enc.frame_shape == (c, h, w)
    assert enc2.num_views == v and enc2.frame_shape == (v, c, h, w)
    assert enc.head.in_features == 32 and enc2.head.in_features == v * 32
    # the same parameters as the single-camera encoder, except the head's input width
    shapes1 = {k: tuple(p.shape) for k, p in enc.named_parameters()}
    shapes2 = {k: tuple(p.shape) for k, p in enc2.named_parameters()}
    assert shapes1.keys() == shapes2.keys()
    assert all(shapes1[k] == shapes2[k] for k in shapes1 if not k.startswith("head."))
    obs2 = torch.randn(b, v, c, h, w, generator=gen2)
    goals2 = torch.randn(b, m, v, c, h, w, generator=gen2)
    with torch.no_grad():
        z2_t = enc2(obs2)
        z2_g = enc2.encode_milestones(goals2)
        zs2_t, zs2_g = enc2.encode_siamese(obs2, goals2)
        z2_swapped = enc2(obs2.flip(1))                                   # views in the other order
        z2_single = enc2(obs2[1:2])                                       # B = 1
    assert z2_t.shape == (b, d) and z2_g.shape == (b, m, d) and zs2_t.shape == (b, d) and zs2_g.shape == (b, m, d)
    ref2_t = _ref_vit_forward(enc2, obs2)
    assert _close(z2_t, ref2_t) and _close(zs2_t, ref2_t) and _close(z2_single, ref2_t[1:2])
    for bi in range(b):
        for mi in range(m):
            ref_g = _ref_vit_forward(enc2, goals2[bi, mi].unsqueeze(0))[0]
            assert _close(z2_g[bi, mi], ref_g) and _close(zs2_g[bi, mi], ref_g)
    # the head sees the views in camera order: swapping them changes z
    assert not torch.allclose(z2_swapped, z2_t, rtol=1e-3, atol=1e-4)
    assert _close(z2_swapped, _ref_vit_forward(enc2, obs2.flip(1)))
    # a view axis is required with V >= 2 and rejected with V = 1; the view count must match
    _expect_value_error(lambda: enc2(obs))                                          # no view axis
    _expect_value_error(lambda: enc2(torch.randn(b, v + 1, c, h, w)))               # wrong V
    _expect_value_error(lambda: enc2(torch.randn(b, v, c, h + 8, w)))               # wrong frame size
    _expect_value_error(lambda: enc2.encode_milestones(goals))                      # [B, M, C, H, W]
    _expect_value_error(lambda: enc2.encode_siamese(obs2, goals))                   # single-view milestones
    _expect_value_error(lambda: enc2.encode_siamese(obs, goals2))                   # single-view observation
    _expect_value_error(lambda: enc(obs2))                                          # V = 1 rejects a view axis
    _expect_value_error(lambda: enc.encode_milestones(goals2))
    _expect_value_error(lambda: VisionEncoder(image_size=(h, w), patch_size=8, num_views=0))
    # E_psi_bar of a multi-view encoder: the same map with the EMA weights; gradients reach every view
    target2 = MomentumEncoder(enc2, momentum=mu)
    assert _close(target2(obs2), _ref_vit_forward(target2.encoder, obs2))
    obs2_g = obs2.clone().requires_grad_(True)
    enc2.zero_grad(set_to_none=True)
    (enc2(obs2_g) - target2(obs2_g.flip(1))).pow(2).sum().backward()               # non-zero residual
    assert obs2_g.grad is not None and all(bool(obs2_g.grad[:, vi].abs().sum() > 0) for vi in range(v))
    assert all(p.grad is not None for p in enc2.parameters())
    assert all(p.grad is None for p in target2.parameters())
    print("perception.py self-test passed")


if __name__ == "__main__":
    _self_test()