v0.2.1: model.py finds instructions.json from the Hugging Face cache; link to the LeRobot version
434b56b verified | # SPDX-FileCopyrightText: 2026 Adaptive Machine Patterns Ltd | |
| # SPDX-License-Identifier: CC-BY-NC-4.0 | |
| """AMP Tadpole: a small vision-language-action model for the AMP Toy | |
| Robotics Benchmark. | |
| in images uint8 (3, 3, 512, 512): the top, side and iso cameras | |
| state (7,): 6 joint angles (rad) + gripper aperture (0 open, 1 closed) | |
| clause one instruction clause, in canonical wording | |
| out (8, 7): the next 8 joint states + gripper aperture | |
| images -> a CNN shared by the cameras -> 16 x 16 tokens per camera | |
| clause -> token ids (instructions.json) -> word embeddings | |
| state -> one token | |
| all tokens + 8 learned action tokens -> 6-layer transformer | |
| each action token -> one step's change; the changes are added up | |
| This file is a model package (see the benchmark's README): build() makes | |
| the module, predict() runs it on one observation. It needs only PyTorch. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| HERE = Path(__file__).absolute().parent # not resolve(): in the Hugging Face cache the files are links to blobs | |
| GRID = 16 # tokens per camera side | |
| CHUNK = 8 # steps per answer | |
| STATE_DIM = 7 | |
| TOKEN_LEN = 48 # instruction slots | |
| STATE_FEATURES = 31 | |
| INC_SCALE = [0.04] * 6 + [0.125] # one step's change, per output unit | |
| JOINT_LIMITS = [(0.1, 0.8 * math.pi), (math.pi / 2, 3 * math.pi / 2), | |
| (-0.7 * math.pi, 0.7 * math.pi), (-0.6 * math.pi, 0.6 * math.pi), | |
| (-math.pi, math.pi), (-math.pi, math.pi)] | |
| class Tadpole(nn.Module): | |
| def __init__(self, vocab: int, d: int = 256, layers: int = 6, heads: int = 8): | |
| super().__init__() | |
| def block(i, o): | |
| return nn.Sequential( | |
| nn.Conv2d(i, o, 3, 2, 1), nn.GroupNorm(8, o), nn.GELU(), | |
| nn.Conv2d(o, o, 3, 1, 1), nn.GroupNorm(8, o), nn.GELU()) | |
| self.cnn = nn.ModuleList([ # 512 -> 16 | |
| nn.Sequential(nn.Conv2d(5, 48, 7, 4, 3), nn.GroupNorm(8, 48), nn.GELU()), | |
| block(48, 96), block(96, 192), block(192, d)]) | |
| self.pos = nn.Parameter(torch.zeros(1, GRID * GRID, d)) | |
| self.view = nn.Parameter(torch.zeros(3, 1, d)) | |
| self.word = nn.Embedding(vocab, d) | |
| self.word_pos = nn.Parameter(torch.zeros(1, TOKEN_LEN, d)) | |
| self.state = nn.Linear(STATE_FEATURES, d) | |
| self.query = nn.Parameter(torch.zeros(1, CHUNK, d)) | |
| layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.0, | |
| batch_first=True, norm_first=True) | |
| self.fuse = nn.TransformerEncoder(layer, layers, enable_nested_tensor=False) | |
| self.head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, STATE_DIM)) | |
| self.register_buffer("lo", torch.tensor([a for a, _ in JOINT_LIMITS] + [0.0]), persistent=False) | |
| self.register_buffer("hi", torch.tensor([b for _, b in JOINT_LIMITS] + [1.0]), persistent=False) | |
| self.register_buffer("inc", torch.tensor(INC_SCALE), persistent=False) | |
| def state_features(self, state): | |
| """The 7 numbers scaled to -1..1, plus sines and cosines of the joint | |
| angles, so that small differences in pose are not lost.""" | |
| q = state[:, :6] | |
| return torch.cat([2 * (state - self.lo) / (self.hi - self.lo) - 1, | |
| torch.sin(2 * q), torch.cos(2 * q), | |
| torch.sin(8 * q), torch.cos(8 * q)], dim=1) | |
| def forward(self, images, tokens, token_mask, state): | |
| """images uint8 (B, 3, 3, H, W); tokens (B, 48) ids from | |
| instructions.json; token_mask (B, 48) bool, True for real tokens; | |
| state (B, 7). Returns (B, 8, 7): the next 8 states.""" | |
| B, H = images.shape[0], images.shape[-1] | |
| x = images.float() / 255.0 | |
| w = self.word(tokens) + self.word_pos | |
| ax = torch.linspace(-1, 1, H, device=images.device) | |
| yy, xx = torch.meshgrid(ax, ax, indexing="ij") | |
| f = torch.cat([x.flatten(0, 1) * 2 - 1, | |
| xx.expand(B * 3, 1, H, H), yy.expand(B * 3, 1, H, H)], dim=1) | |
| for stage in self.cnn: | |
| f = stage(f) | |
| f = f.flatten(2).transpose(1, 2) + self.pos | |
| f = (f.view(B, 3, GRID * GRID, -1) + self.view).flatten(1, 2) | |
| s = self.state(self.state_features(state))[:, None] | |
| q = self.query.expand(B, -1, -1) | |
| h = torch.cat([f, w, s, q], dim=1) | |
| keep = torch.ones(B, f.shape[1], dtype=torch.bool, device=h.device) | |
| tail = torch.ones(B, 1 + CHUNK, dtype=torch.bool, device=h.device) | |
| h = self.fuse(h, src_key_padding_mask=~torch.cat([keep, token_mask, tail], dim=1)) | |
| steps = self.head(h[:, -CHUNK:]) * self.inc | |
| return steps.cumsum(dim=1) + state[:, None] | |
| _TABLE = None | |
| def instructions() -> dict: | |
| """The token table: every canonical clause -> its token ids.""" | |
| global _TABLE | |
| if _TABLE is None: | |
| _TABLE = json.loads((HERE / "instructions.json").read_text()) | |
| return _TABLE | |
| def encode(clause: str): | |
| """(ids, mask) for one canonical clause, each of length 48.""" | |
| t = instructions() | |
| try: | |
| ids = t["instructions"][clause] | |
| except KeyError: | |
| raise ValueError(f"AMP Tadpole only reads canonical clauses " | |
| f"(e.g. 'Pick up the lime star'); got {clause!r}") from None | |
| pad = t["pad"] | |
| return ids + [pad] * (TOKEN_LEN - len(ids)), [True] * len(ids) + [False] * (TOKEN_LEN - len(ids)) | |
| def build() -> Tadpole: | |
| return Tadpole(vocab=len(instructions()["vocab"])) | |
| def predict(module: Tadpole, images, state, clause: str) -> np.ndarray: | |
| dev = next(module.parameters()).device | |
| ids, mask = encode(clause) | |
| out = module(torch.as_tensor(np.asarray(images)[None], device=dev), | |
| torch.tensor([ids], device=dev), torch.tensor([mask], device=dev), | |
| torch.as_tensor(np.asarray(state, dtype=np.float32)[None], device=dev)) | |
| return out[0].cpu().numpy().astype(np.float64) | |