| """Small latent action model that doubles as a world model (LAPO-style). |
| |
| encoder: (frame t, frame t+k) -> 4 tokens of 3 numbers -> FSQ, 5 levels each |
| decoder: frame t + quantised latent -> predicted frame t+k |
| |
| The latent is tiny (125 codes per token, about 28 bits per frame pair), so it cannot carry |
| the image. The decoder gets the scene from frame t through skip connections, and the latent |
| only has room for what changed. Quantisation is FSQ (Mentzer et al. 2023), which has no |
| codebook that could collapse. An earlier version with an EMA VQ codebook and zero-initialised |
| FiLM layers did collapse: the decoder ignored the latent and the encoder became constant. |
| |
| References: LAPO, Learning to Act without Actions (Schmidt & Jiang, ICLR 2024); LAPA, Latent |
| Action Pretraining from Videos (Ye et al., ICLR 2025), which scales the idea to VLA |
| pretraining. |
| """ |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| def conv_block(cin: int, cout: int, stride: int) -> nn.Sequential: |
| return nn.Sequential( |
| nn.Conv2d(cin, cout, 3, stride, 1), nn.GroupNorm(8, cout), nn.SiLU(), |
| nn.Conv2d(cout, cout, 3, 1, 1), nn.GroupNorm(8, cout), nn.SiLU(), |
| ) |
|
|
|
|
| class Encoder(nn.Module): |
| """(B, cin, S, S) -> features at S/2, S/4, S/8, S/16.""" |
|
|
| def __init__(self, cin: int, widths=(32, 64, 128, 256)): |
| super().__init__() |
| chans = [cin, *widths] |
| self.levels = nn.ModuleList(conv_block(chans[i], chans[i + 1], 2) for i in range(len(widths))) |
|
|
| def forward(self, x): |
| feats = [] |
| for level in self.levels: |
| x = level(x) |
| feats.append(x) |
| return feats |
|
|
|
|
| class FiLM(nn.Module): |
| """Condition a feature map on the latent action: h * (1 + gamma) + beta. Default (non-zero) |
| initialisation, so the latent shapes the prediction from the first step.""" |
|
|
| def __init__(self, zdim: int, ch: int): |
| super().__init__() |
| self.lin = nn.Linear(zdim, 2 * ch) |
|
|
| def forward(self, h, z): |
| gamma, beta = self.lin(z).chunk(2, dim=-1) |
| return h * (1 + gamma[..., None, None]) + beta[..., None, None] |
|
|
|
|
| class Decoder(nn.Module): |
| """U-Net decoder over the context frame's features, FiLM-conditioned on the latent at |
| every scale; outputs the change to add to the context frame (tanh, in [-1, 1]).""" |
|
|
| def __init__(self, zdim: int, widths=(32, 64, 128, 256)): |
| super().__init__() |
| w0, w1, w2, w3 = widths |
| self.mid, self.film3 = conv_block(w3, w3, 1), FiLM(zdim, w3) |
| self.up2, self.film2 = conv_block(w3 + w2, w2, 1), FiLM(zdim, w2) |
| self.up1, self.film1 = conv_block(w2 + w1, w1, 1), FiLM(zdim, w1) |
| self.up0, self.film0 = conv_block(w1 + w0, w0, 1), FiLM(zdim, w0) |
| self.out = nn.Conv2d(w0, 3, 3, 1, 1) |
|
|
| def forward(self, feats, z): |
| f0, f1, f2, f3 = feats |
| h = self.film3(self.mid(f3), z) |
| for skip, block, film in ((f2, self.up2, self.film2), (f1, self.up1, self.film1), (f0, self.up0, self.film0)): |
| h = F.interpolate(h, scale_factor=2, mode="nearest") |
| h = film(block(torch.cat([h, skip], dim=1)), z) |
| h = F.interpolate(h, scale_factor=2, mode="nearest") |
| return torch.tanh(self.out(h)) |
|
|
|
|
| class FSQ(nn.Module): |
| """Finite scalar quantisation with odd level counts: z -> tanh -> round to `levels` values, |
| straight-through gradient. Returns values in [-1, 1] and an integer code per token.""" |
|
|
| def __init__(self, levels=(5, 5, 5)): |
| super().__init__() |
| assert all(n % 2 == 1 for n in levels), "use odd level counts" |
| self.n_codes = math.prod(levels) |
| |
| self.register_buffer("half_range", (torch.tensor(levels, dtype=torch.float32) - 1) / 2) |
| basis = [1] |
| for n in levels[:-1]: |
| basis.append(basis[-1] * n) |
| self.register_buffer("basis", torch.tensor(basis, dtype=torch.long)) |
| self.register_buffer("levels_t", torch.tensor(levels, dtype=torch.long)) |
|
|
| def forward(self, z): |
| bounded = torch.tanh(z) * self.half_range |
| rounded = torch.round(bounded) |
| q = bounded + (rounded - bounded).detach() |
| idx = ((rounded + self.half_range).long() * self.basis).sum(-1) |
| return q / self.half_range, idx |
|
|
| def codes_to_values(self, idx): |
| digits = (idx[..., None] // self.basis) % self.levels_t |
| return (digits.float() - self.half_range) / self.half_range |
|
|
|
|
| class LatentActionModel(nn.Module): |
| def __init__(self, img: int = 96, n_tokens: int = 4, levels=(5, 5, 5)): |
| super().__init__() |
| assert img % 16 == 0 |
| levels = tuple(levels) |
| self.cfg = {"img": img, "n_tokens": n_tokens, "levels": list(levels)} |
| self.n_tokens, self.token_dim = n_tokens, len(levels) |
| self.action_enc = Encoder(6) |
| self.to_z = nn.Sequential( |
| nn.Flatten(), nn.Linear(256 * (img // 16) ** 2, 512), nn.SiLU(), nn.Linear(512, n_tokens * len(levels)) |
| ) |
| self.fsq = FSQ(levels) |
| self.ctx_enc = Encoder(3) |
| self.dec = Decoder(n_tokens * len(levels)) |
|
|
| @property |
| def n_codes(self) -> int: |
| return self.fsq.n_codes |
|
|
| def encode(self, x0, x1): |
| """Continuous latent (B, T, d), quantised latent (B, T, d), code per token (B, T).""" |
| z = self.to_z(self.action_enc(torch.cat([x0, x1], dim=1))[-1]).view(-1, self.n_tokens, self.token_dim) |
| q, idx = self.fsq(z) |
| return z, q, idx |
|
|
| def decode(self, x0, q): |
| """World-model step: predicted next frame from the current frame and a latent action.""" |
| return x0 + self.dec(self.ctx_enc(x0), q.flatten(1)) |
|
|
| def codes_to_latent(self, idx): |
| """Codes (B, T) -> quantised latent (B, T, d), e.g. to roll out chosen latent actions.""" |
| return self.fsq.codes_to_values(idx) |
|
|
| def forward(self, x0, x1): |
| z, q, idx = self.encode(x0, x1) |
| return self.decode(x0, q), idx |
|
|
