怎么给 SmolVLA 打这个补丁

state dropout 的全部机制:state 在模型里走哪条路、补丁挂在哪一行、置零之后模型实际收到的是什么、退火曲线怎么算、以及怎么确认它真的挂上了。看完应该能自己判断这个补丁对不对,而不是相信我说它对。

一、state 在 SmolVLA 里走哪条路

先弄清楚要动的是什么。SmolVLA 把观测拼成一串 token 喂给 VLM,整个本体状态只占其中一个 token:

图像 token …… ─┐
语言 token …… ─┼─→ 前缀序列 ─→ VLM ─→ 动作专家 ─→ 50 步动作块
state token ×1 ─┘

那一个 token 的来路(modeling_smolvla.py):

观测 17 维
   ↓  归一化处理器(在策略【外面】,是 preprocessor 的一步)
   ↓  (x − mean) / (std + 1e-8)
   ↓
SmolVLAPolicy.prepare_state(batch)          # :410
   state = batch[OBS_STATE]
   state = pad_vector(state, 32)            # max_state_dim,补零到 32
   ↓
VLAFlowMatching.embed_prefix()              # :620
   state_emb = self.state_proj(state)       # nn.Linear(32 → 960)
   embs.append(state_emb[:, None, :])       # 变成序列里的【一个】token
   att_masks += [1]                         # 图像/语言【不】回看 state

二、补丁挂在 prepare_state,不是别的地方

prepare_state 是唯一同时被训练和推理两条路调用的地方:

挂在这里的好处是一个函数管住两条路,坏处是必须自己区分 —— 所以补丁第一件事就是查 self.training:

def prepare_state(self, batch):
    state = orig(self, batch)                 # 先让原函数做 padding
    if not self.training:                     # ★ 推理原样返回 ★
        return state                          #   真机上永远拿到真实 state
    ...                                       # 只有训练才置零

三、置零之后,模型实际收到的是什么

这一点值得说清楚,因为「置零」听起来像是喂了一个无意义的 0。实际不是:

归一化【之后】的 state 置零
   ↓  归一化是 (x − mean)/std,所以 0 对应的原始值就是 mean
   ↓  = 「训练集里最平均的那个姿态」
   ↓
state_proj(0) = W·0 + b = b                   # 只剩偏置
   ↓
一个【固定的、可学习的】向量

四、逐样本,不是整批

keep = (torch.rand(state.shape[0], device=state.device) >= p).to(state.dtype)
out  = state * keep.view(-1, *([1] * (state.ndim - 1)))

torch.rand 出的是 batch_size 个独立随机数,所以同一批里有的样本看得到 state、有的看不到。

五、退火曲线

整个调度就是一个纯函数,好测也好读:

def p_at(step, total):
    vf, af = SD_VISION_FRAC, SD_ANNEAL_FRAC   # 0.3, 0.4
    p_end  = SD_P_END                          # 0.2
    frac = step / total
    if frac < vf:                    return 1.0                    # 阶段 A
    if frac >= vf + af:              return p_end                  # 阶段 C
    return 1.0 + (p_end - 1.0) * ((frac - vf) / af)                # 阶段 B 线性

六、怎么注入(不改 lerobot 源码)

训练入口是 lerobot-train 这个 console 脚本。要在它建模型之前打补丁,要么改 lerobot 源码(每次升级都要重打,而且改动在 git 之外),要么在同一个进程里先 import 再交棒。选后者:

# train_state_dropout.py
total = _steps_from_argv(sys.argv[1:])        # 从 --steps=N 里读总步数,用来算退火进度
state_dropout.install(total_steps=total)      # ← 打补丁

from lerobot.scripts.lerobot_train import main as train_main
train_main()                                  # 交棒,参数原样透传

install() 做的事:

orig = SmolVLAPolicy.prepare_state             # 存原函数
def prepare_state(self, batch): ...            # 包一层
SmolVLAPolicy.prepare_state = prepare_state    # 替换类方法

替换的是类方法,所以之后 lerobot-train 自己 from_pretrained 造出来的实例自动带补丁,不用把 policy 对象传来传去。install() 自带重复调用保护。

七、两个坑

步数计数器 vs 梯度累积

补丁自己数前向次数,而一次优化器更新 = ACCUM 次前向。当前配方 MICRO_BATCH=16, B_EFF=16 所以 ACCUM=1,两者相等。如果以后改成 ACCUM > 1,退火进度会按前向次数走,而 --steps 传的也是前向次数(脚本里 STEPS = UPDATES × ACCUM),所以仍然对得上 —— 但这是巧合般的一致,改配方时要重新确认。SD_TOTAL_STEPS 可以显式覆盖。

总步数读不到

命令行里没有 --steps 又没给 SD_TOTAL_STEPS 时,total=0,p_at 会全程返回 p_end —— 阶段 A 直接没了,而且不报错。所以 train_state_dropout.py 在这种情况下会打一条明确的警告到 stderr。

八、怎么确认它真的挂上了

这一步不能省。补丁没挂上的表现和挂上了一模一样:训练照跑,loss 照降。

python state_dropout.py

BENCH=1 SD_DEBUG=1 ./train_smolvla_state_dropout.sh

python 03-software/scripts/vision_sweep_phase.py <新ckpt> <数据集> "<task>" right 12

九、这个补丁可能怎么坏