How to?

Robot Transformer Basics

Jan 1, 2025
2501
4 Minutes
709 Words

“使用 Image Encoder 以及 Text Encoder 并使用 FiLM 进行 Fusion 后用 Transformer 处理的模型” 用高度凝练的几行 pytorch 伪代码告诉我这里的写法。尤其是 film 和如何后续采用 transformer 处理

FiLM

来自 RT1

使用 MLP(txt token) 的输出缩放 img token(为 txt token 唯一用处)

1
# img: [B, C, H, W], txt_ids: [B, L]
2
v = ImageEncoder(img) # [B, Nv, D] (视觉 token)
3
t = TextEncoder(txt_ids).mean(1) # [B, Dt] (文本全局语义)
4
5
gamma, beta = Linear(t).chunk(2, dim=-1) # 各 [B, D]
6
v_film = v * (1 + gamma[:, None, :]) + beta[:, None, :] # FiLM: feature-wise affine
7
8
x = torch.cat([CLS.expand(B,1,D), v_film], dim=1) # [B, 1+Nv, D]
9
h = TransformerEncoder(x) # [B, 1+Nv, D]
10
y = Head(h[:, 0]) # 用 CLS 做下游预测

更现代的(把现成 LLM Transformer 当作 fusion 主干)

1
v = VisionEncoder(img) # [B, Nv, Dv]
2
v_tok = VisionProjector(v) # [B, Nv, D] 对齐到 LLM 维度
3
4
t_tok = LLM.embed_tokens(text_ids) # [B, Nt, D]
5
x = torch.cat([BOV, v_tok, EOV, t_tok], dim=1) # 视觉+语言统一序列(Fusion in LLM)
6
7
h = LLM.transformer(x, causal_mask=True) # 统一 Transformer 融合
8
a_logits = ActionHead(h[:, -Na:, :]) # 取动作位置 hidden state
9
a = sample_or_argmax(a_logits) # 自回归/并行输出离散动作 token
10
11
如果动作是连续量,也常见:
12
a = MLP(h[:, -1, :]) # 直接回归 7DoF/控制量

预测 attn map

来自 RT1

就是你能想到的最简单的自主加权. (每个 token 自己算自己对 m 个结果的权重,权重和为 1)

1
# x: [B, N, D] (N个输入token)
2
attn_logits = MLP(x) # [B, N, M] -> 每个token对M个map的打
3
4
attn = torch.softmax(attn_logits, dim=1) # 在token维归一化: 每张map覆盖N个token
5
6
# 用每张attention map对token做加权汇聚
7
# x_out[b,m,d] = sum_n attn[b,n,m] * x[b,n,d]
8
x_out = torch.einsum('bnm,bnd->bmd', attn, x) # [B, M, D]

RT-2 离散动作 token

1
# RT-2: 连续动作 → 词表尾部 K 个 token,同一 VLM 自回归 CE
2
# V 是词表大小. 若 x 和 yaw 量化到同一档,会得到同一个 token id,这是设计如此,不是 bug,要靠序列位置区分语义
3
act_ids = V - digitize(clip(a, a_min, a_max), K_bins) # [B, A]
4
seq = cat([vision_tok, text_tok, act_ids[:, :-1]], dim=1) # 都是离散 id
5
h = VLM_Transformer(seq, images=img)
6
loss = CE(LM_Head(h)[action_mask], act_ids) # 仅动作位算 loss
7
8
# infer: 自回归 A 步,logits 尾部子集 argmax → bin_centers 反量化
9
context = cat([vision, text])
10
for _ in range(A):
11
next_id = argmax(LM_Head(VLM(context)[:, -1:])[:, :, V-K:], -1)
12
context = cat([context, next_id], dim=1)
13
a = bin_centers[V - stack(context[:, -A:]) - 1]

pi0.5

1
# π₀.₅: PaliGemma prefix + Action Expert suffix + flow matching
2
v = SigLIP(imgs) # [B, Nv, D=2048]. 已经是连续 embedding
3
t = PaliGemma.embed(prompt_with_discrete_state) # state 在文本里
4
prefix = cat([v, t]) # 图文双向 (prefix-LM)
5
x_t, t = noise * t + actions * (1-t) # flow 插值
6
# [B, Na, D=1024]. 编码到 token-embedding-space(和上面 space 完全不同). 代码中叫做 action_in_proj
7
suffix = Linear_in(x_t)
8
cond = MLP(sincos(t)) # adaRMS 条件
9
# Gemma 内跑 suffix 的那条路是 Action Expert (expert 1,Gemma-300M)
10
(h_pre, h_suf) = DualGemma([prefix, suffix], mask=prefix_lm, adarms=[None, cond]) # 输出 hidden states
11
loss = MSE(Linear_out(h_suf[:, -Na:]), noise - actions)
12
# infer: cache prefix → Euler: x += dt * Linear_out(h_suf), t -= dt # 解码回动作空间 (flow 的速度场)
Article title:Robot Transformer Basics
Article author:Julyfun
Release time:Jan 1, 2025
Copyright 2026
Sitemap