File size: 6,109 Bytes
c1ebdfc 434b56b c1ebdfc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | # 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"]))
@torch.no_grad()
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)
|