Discrete tokenizers: VQ-VAE and FSQ

@torch.no_grad()
def restart_dead_codes(
    quantizer: VectorQuantizer, z: torch.Tensor, dead: torch.Tensor
) -> int:
    """Moves each dead code onto a random encoder output; returns how many moved.

    A code nobody chooses gets no gradient and never recovers on its own. Placing it
    where the data actually is gives it a chance to win some inputs back.
    """
    count = int(dead.sum())
    if count:
        picks = torch.randint(0, len(z), (count,))
        quantizer.codebook.weight[dead] = z[picks]
    return count


def codebook_usage(indices: torch.Tensor, num_codes: int) -> tuple[float, float]:
    """Share of codes used, and perplexity exp(entropy) of their frequencies.

    Perplexity equals ``num_codes`` when every code is used equally and 1 when one
    code takes everything: the effective vocabulary size.
    """
    counts = torch.bincount(indices, minlength=num_codes).float()
    probs = counts / counts.sum()
    nonzero = probs[probs > 0]
    entropy = -(nonzero * nonzero.log()).sum()
    return (counts > 0).float().mean().item(), entropy.exp().item()

def bin_tokens(
    chunk: torch.Tensor, bins: int, low: float = -1.0, high: float = 1.0
) -> tuple[torch.Tensor, torch.Tensor]:
    """Per-step, per-dimension binning (as in RT-2 and OpenVLA): one token per number.

    Returns:
        Integer tokens with the chunk's shape, and the de-tokenised chunk (bin
        centres).
    """
    width = (high - low) / bins
    tokens = ((chunk.clamp(low, high - 1e-6) - low) / width).floor().long()
    return tokens, low + (tokens.float() + 0.5) * width


def dct_matrix(n: int) -> torch.Tensor:
    """Orthonormal DCT-II matrix: row k is the k-th cosine basis over n samples."""
    k = torch.arange(n)[:, None].float()
    i = torch.arange(n)[None, :].float()
    matrix = torch.cos(torch.pi / n * (i + 0.5) * k) * (2.0 / n) ** 0.5
    matrix[0] /= 2**0.5
    return matrix


def dct_tokens(chunk: torch.Tensor, scale: float) -> tuple[torch.Tensor, torch.Tensor]:
    """FAST-style tokens: DCT along time, scale, round to integers.

    Smooth motion puts almost all its energy in the low-frequency coefficients, so
    after rounding most high-frequency coefficients are 0. FAST then compresses the
    integer sequence further with byte-pair encoding; here we just count non-zeros.

    Args:
        chunk: Actions, shape (time, dims).
        scale: Multiplier before rounding; larger keeps more detail and more tokens.

    Returns:
        The integer coefficients (time, dims) and the reconstructed chunk.
    """
    basis = dct_matrix(chunk.shape[0])
    coefficients = (basis @ chunk * scale).round()
    return coefficients.long(), basis.T @ (coefficients / scale)


def smooth_chunk(steps: int = 50, dims: int = 7, seed: int = 0) -> torch.Tensor:
    """A smooth multi-joint action chunk: a few low-frequency sines per dimension."""
    generator = torch.Generator().manual_seed(seed)
    t = torch.linspace(0, 1, steps)[:, None]
    freqs = torch.rand(3, dims, generator=generator) * 2 + 0.2
    phases = torch.rand(3, dims, generator=generator) * 6.28
    amps = torch.rand(3, dims, generator=generator) * 0.3
    return (
        amps[None] * torch.sin(6.28 * freqs[None] * t[..., None] + phases[None])
    ).sum(1)

class VectorQuantizer(nn.Module):
    """Snaps each vector to its nearest entry in a learned codebook."""

    def __init__(self, num_codes: int, dim: int, beta: float = 0.25) -> None:
        """Creates the codebook.

        Args:
            num_codes: Number of codebook entries (the vocabulary size).
            dim: Size of each entry.
            beta: Weight of the commitment loss that keeps the encoder near its code.
        """
        super().__init__()
        self.num_codes = num_codes
        self.codebook = nn.Embedding(num_codes, dim)
        self.codebook.weight.data.uniform_(-1.0, 1.0)
        self.beta = beta

    def forward(
        self, z: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Quantizes ``z`` of shape (batch, dim).

        Returns:
            The quantized vectors (with a straight-through gradient to ``z``), the
            integer code indices (batch,), and the VQ loss.
        """
        indices = torch.cdist(z, self.codebook.weight).argmin(dim=-1)
        chosen = self.codebook(indices)
        # Move codebook entries toward the encoder outputs, and the encoder toward
        # its chosen entries (the commitment term).
        loss = ((chosen - z.detach()) ** 2).mean()
        loss = loss + self.beta * ((z - chosen.detach()) ** 2).mean()
        # Straight-through estimator: the forward pass uses the code, the backward
        # pass pretends quantization was the identity.
        return z + (chosen - z).detach(), indices, loss

class FiniteScalarQuantizer(nn.Module):
    """Rounds each dimension to a few fixed levels; no codebook is learned."""

    def __init__(self, levels: list[int]) -> None:
        """Creates the quantizer.

        Args:
            levels: Number of levels per dimension; every entry must be odd. The
                vocabulary size is the product of the levels.

        Raises:
            ValueError: If any level count is even or smaller than 3.
        """
        super().__init__()
        if any(level < 3 or level % 2 == 0 for level in levels):
            raise ValueError(f"levels must be odd and at least 3, got {levels}")
        self.register_buffer("levels", torch.tensor(levels))

    @property
    def num_codes(self) -> int:
        """Size of the implicit codebook."""
        return int(self.levels.prod().item())

    def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Quantizes ``z`` of shape (batch, len(levels)).

        Returns:
            Quantized values in [-1, 1] (straight-through gradient) and one integer
            code per row in [0, num_codes).
        """
        half = (self.levels - 1) / 2
        bounded = torch.tanh(z) * half
        rounded = bounded + (bounded.round() - bounded).detach()
        digits = (rounded + half).long()
        radix = torch.cumprod(
            torch.cat([torch.ones(1, dtype=torch.long), self.levels[:-1]]), dim=0
        )
        return rounded / half, (digits * radix).sum(dim=-1)

def failed_runs(
    make_quantizer: Callable[[], nn.Module],
    seeds: range,
    steps: int = 1200,
    threshold: float = 0.2,
    *,
    restart: bool = False,
) -> int:
    """Counts training runs whose final reconstruction error exceeds ``threshold``.

    A bad run usually means a codebook entry was never used (a "dead code"), so a
    whole cluster is reconstructed as another one. The quantizer is created after
    seeding, so a seed fixes its initialisation too.
    """
    failures = 0
    for seed in seeds:
        torch.manual_seed(100 + seed)
        _, _, _, error = train_autoencoder(
            make_quantizer(), steps, seed, restart=restart
        )
        failures += int(error > threshold)
    return failures