smolvla
主要是代码理解和修改
概念理解
基础
- smolvla的核心训练范式始终是模仿学习。
- 推理的时候,smolvla是只用当前的信息推理,完全没有用之前的,但async
inference有融合一点点以前的action
- 不显式引入历史状态或动作序列
- 但是Action Chunking是有一点的
- 然后sa有一点:
We employ a causal attention mask for the SA layers, ensuring that each action token can only attend to past tokens within the chunk, preventing future action dependencies. Empirically, we find that interleaving CA and SA layers provides higher success rates and faster inference time.我们为 SA 层采用因果注意力掩码,确保每个作令牌只能关注块内过去的令牌,从而防止未来的作依赖。根据经验,我们发现交错 CA 和 SA 层提供了更高的成功率和更快的推理时间。
self-atten和cross-atten
看得稀里糊涂的
假设一个batch是一个车间的工人(假设64个),输入零件,输出组装的东西。工人利用vlm得到高级特征token,比如螺丝和螺母搭好算一个。 cross-attention:对比生成的半成品和图纸 self-attention:确保当前动作和历史动作连贯。仅关注自身序列的历史
state在整个结构里面处在什么位置 SmolVLA 由一个紧凑的预训练视觉语言模型组成,丢弃最后的 L − N 层(剪刀图标)。其余层嵌入了三个输入:(i)语言指令,(ii)RGB图像,以及(iii)机器人感觉运动状态。它们合并的标记为交替交叉注意力(金色)和自我注意力(浅黄色)块的动作专家提供信息,这些块通过流匹配进行训练,以输出 n 个低级动作块,在 . . . , at+n。SmolVLA 在公共社区数据集上进行预训练,并在低成本机器人上进行评估。
train
- 除了基础参数外,我需要加的
- 训练的时候用/不用深度图,需要在输入特征里面
- 训练的时候用哪一种language方法
- "":baseline
- mtask_relative
- mtask_grid_xxcm
- add_task_to_state好像得单独一个
主函数
def train(cfg: TrainPipelineConfig):
validate
对policy_path
从 CLI(命令行参数)里获取 --policy 参数对应的路径。会再解析命令行里的 policy 相关覆盖参数(cli_overrides)。用 PreTrainedConfig.from_pretrained(...) 加载 policy 的配置(只加载配置,不加载权重)。记录 self.policy.pretrained_path = policy_path。
也就是当train的cfg有policy的时候,cfg.policy=PreTrainedConfig.from_pretrained(policy_path...)创建了一个类,包含的通常是超参数、路径、模型结构信息等,然后 self.policy.pretrained_path = policy_path。
此时还没有加载模型的配置,那模型的输入特征是在这里决定的吗?
--policy.path=models/forsmolvla/smolvla_base这个是训练的时候的设置,所以只是加载预训练的smolvla_base的输入设置
我看了这个smolvla_base的input feature规定state是6维,所以我如果想fine-tune的话,得改这个base的设置🤔
dataset
dataset = make_dataset(cfg)
应该就是提取cfg里面dataset的部分吧?感觉训练的时候关于datset的处理(用不用深度/文本处理应该写在dataset里面,而不是写在train.py里面)
这个要改就得改lerobotdataset的处理,希望可以自定义一个配置类,先这么些
- lerobotdataset类的init新增
1
2
3
4
5# 下面是为了train的自定义内容
language_tip_mode: str="", # 当空的时候就等于baseline
add_location_to_state: bool = False, # 控制是将location加到state里面
exclude_features: list[str] = None, # 控制需要过滤的 key,最终的训练(主要是为了过滤掉depth)
obj_detector=None, # 可选的目标检测器,为了state加的 - 把 FilteredBatchLoader.add_location_to_state 移到 lerobotDataset里面:_add_location_to_state函数(dataset是无辜的要怪就怪模型结构),调用加在geitem函数里面
- filter加载getitem函数的最后,return 之前
- make_dataset传入参数
- train清理之前的config,传入参数到make_dataset里边
- 正向流程
- 整理train cfg的模板,增加,同样config/train.py里面也要修改
1
2language_tip_mode: str="", # 当空的时候就等于baseline
add_location_to_state: bool = False, # 控制是将location加到state里面 - train里面dataste=make_dataset(cfg)
- 调用lerobotdataset初始化
- getitem得到一帧对应的color,depth,state,force,action等
- 返回一帧的color,state,action
- 整理train cfg的模板,增加,同样config/train.py里面也要修改
- 测试:
- 能不能正常train baseline
- 能不能正常train language_mode 看task就行
- 能不能train state(显然不行) 看state
policy
- 用cfg.policy(就是validate的时候cfg.policy=PreTrainedConfig.from_pretrained(policy_path...)创建的类)和dataset.meta创建policy
1
2
3
4policy = make_policy(
cfg=cfg.policy,
ds_meta=dataset.meta,
) - make_policy用meta干什么,meta也要修改,但不能在lerobotdataset里面,而是make_policy的时候。这个policy是我最后要得到的policy吧。meta就规定了input
feature和output feature
- 过滤depth和force这两个feature
- 当且仅当state要求改的时候才改
- policy = policy_cls.from_pretrained(**kwargs)
- smolvla里面没有代码,应该调用的是Pretrainedpolicy这个基类的from_pretrained
- 能保证传入的kwargs是正确的。然后instance = cls(config,
**kwargs)是生成SmolVLAPolicy实例的,为什么这里的cls里面已经没有10那个信息了?
- cls 实际上是 SmolVLAPolicy 类的初始化,初始化里面有一个加载空flowmatching模型的,就是cls.model=VLAFlowMatching(config)。然后VLAFlowMatching的初始化里面有self.state_proj = nn.Linear( self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size )。
- 所以instance的结果里面只有两个包含state的内容
- SmolVLAPolicy( (normalize_inputs): Normalize( (buffer_observation_state): ParameterDict( (mean): Parameter containing: [torch.FloatTensor of size 6] (std): Parameter containing: [torch.FloatTensor of size 6] ) )这个是预训练模型的前6维的权重,后面还需要用
- (state_proj): Linear(in_features=32, out_features=960, bias=True)这个是投影到32维
- from_pretrained里面policy = cls._load_as_safetensor(instance,
model_file, config.device, strict)
- cls还是 SmolVLAPolicy,就是调用_load_as_safetensor这个class method
- 用的是smolvlapolicy里面的_load_as_safetensor
- safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)把权重加载到模型实例里
- return load_smolvla( model, model_file, device=map_location, checkpoint_keys_mapping="model._orig_mod.//model.", )调用特定的 loader 做额外处理
- load_smolvla函数,我需要在这里确保state的权重能1.前6位正常加,7-10位初始化一下
- 有没有办法识别input feature此时是6还是10?目前这个模型传到这里只剩下cls两个关于state的部分了,所以从外部引入
- 给7-10维重置,变成用 Xavier 均匀初始化覆盖,bias 清零
train
- policy.train()运行的是SmolVLMWithExpertModel的train,设置冻结哪一块
- 开始按步数训练 train_tracker, output_dict = update_policy(policy)
- update_policy函数
- 训练逻辑 问题可能发生在,state更新参数的时候7-10维变化比较大
1
2
3
4
5
6loss, output_dict = policy.forward(batch)
grad_scaler.scale(loss).backward()
grad_scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(policy.parameters(), ...)
grad_scaler.step(optimizer)
optimizer.zero_grad()
- 训练逻辑
- normalize归一化没有匹配维度
- modeling_smolvla里面SmolVLAPolicy的init有self.normalize_inputs = Normalize(config.input_features, config.normalization_mapping, dataset_stats)
- state前6维的mean,std是从dataset_stats里来的,dataset_stats是从SmolVLAPolicy的init传入的
- 之前make_policy传入的kwargs["dataset_stats"] = meta_for_policy.stats,相当于ds_meta.stats
- ds_meta是ds_meta=dataset.meta,所以还是本地传入的,看到episode_stats.jsonl文件里面有类似min,max,std的,get
- 怎么写入save_episode有一个write_episode_stats函数
- 怎么读 直接从meta里面读的,这个我好像改不了
- 怎么计算:ep_stats = compute_episode_stats(episode_buffer, self.features)
- 形式: "observation.state": {"min": [-47.14912414550781, -96.49805450439453, 9.134847640991211, -1.6986300945281982, 2.4175825119018555, 0.6323396563529968], "max": [8.552631378173828, 44.06614685058594, 99.51667785644531, 42.136985778808594, 52.91819381713867, 56.36856460571289], "mean": [-13.58438491821289, -31.481496810913086, 56.781070709228516, 22.7191219329834, 25.275249481201172, 15.44166374206543], "std": [19.134763717651367, 52.86572265625, 35.76657485961914, 15.830253601074219, 18.950584411621094, 18.50875473022461], "count": [181]},也就是一个episode里面所有state每个维度的统计。
- 然后变成全局的:aggregate_feature_stats
- 修改过程,使得可以在算出来
- 在正式normalize之前,需要确保能传入datset["observation.state"]以计算
- batch =
self.normalize_inputs(batch)这里用的是Normalize的forward,正好传入了batch
- 一个batch里面有什么?若干个epsiode?
- 怎么normalize的forward里面的mean和std又变成全局的了,全局还不好算?
- normalize类用create_stats_buffers这个函数把所有episode的mean和std算到一个全局的里面
VLAFlowMatching
init
self.vlm_with_expert = SmolVLMWithExpertModel(...)构建带专家头的多模态主干,这个对象里包含:
视觉编码器(如 SigLIP)+ 语言嵌入层;
一个 Transformer(“VLM”主体);
动作“专家(Expert)”分支以及处理器(processor / tokenizer)。
此时的self.config.vlm_model_name="models/forsmolvla/HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
所以vlm_with_expert加载的就是冻结的vlm?
1
2
3
4
5
6
7
8
9
10
11self.vlm_with_expert = SmolVLMWithExpertModel(
model_id=self.config.vlm_model_name,
freeze_vision_encoder=self.config.freeze_vision_encoder,
train_expert_only=self.config.train_expert_only,
load_vlm_weights=self.config.load_vlm_weights,
attention_mode=self.config.attention_mode,
num_expert_layers=self.config.num_expert_layers,
num_vlm_layers=self.config.num_vlm_layers,
self_attn_every_n_layers=self.config.self_attn_every_n_layers,
expert_width_multiplier=self.config.expert_width_multiplier,
)这个代码初始化了一个什么?
self.state_proj这个是把原来的32维的state投影到文本隐藏维(text_config.hidden_size),这个维度是从 self.vlm_with_expert里面来的
action_in_proj / action_out_proj:将 动作维 映射到专家隐藏维 D_exp,便于送入专家分支;将专家输出再投回 动作维,得到对目标速度场的预测
set_requires_grad,如果config.train_state_proj是true的话,就会设置state_proj里面的所有参数require_grad为true,所以这里应该是true
一些关于图像的token,不关注了
一个prefix_length,默认是 prefix_length: int = -1
embed_prefix
- 输入state
- 目的:把图像、(可选的)图像特殊 token、语言 token、状态拼成一段“前缀序列”,并生成对应的 pad mask 与 跨段注意力屏蔽标记。
- 图像嵌入
- 语言嵌入
- 状态嵌入 pad_mask和att_mask
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20# 状态嵌入,要改肯定改这里,如果要改的话还要加一个,self.config里面有吗?
# 可以用self.add_location_to_state(或者改成train_add_location_to_state)
# 投影到self.vlm_with_expert.config.text_config.hidden_size对应的维数,把状态维度投影到 和语言/图像 token 相同的 hidden_size
state_emb = self.state_proj(state)
# 如果只有2维?就加个维。这里也应该肯定是2d的吧
print("prefix_embed这里是1d的",state_emb.ndim == 2)
state_emb = state_emb[:, None, :] if state_emb.ndim == 2 else state_emb
# 追加到总的embs,embs从(image_end_token,lang_emb)变成(image_end_token,lang_emb,state_emb)
embs.append(state_emb)
# bsize就是batchsize
bsize = state_emb.shape[0]
device = state_emb.device
# states_seq_len就是状态token的个数
states_seq_len = state_emb.shape[1]
# 这段token没有padding,形状是((B, states_seq_len))
state_mask = torch.ones(bsize, states_seq_len, dtype=torch.bool, device=device)
pad_masks.append(state_mask)
# Set attention masks so that image and language inputs do not attend to state or actions
att_masks += [1] * (states_seq_len)
pad_masks(padding mask)
作用:标记哪些 token 是“真实有效的”,哪些是 padding(补齐长度)。
取值:布尔值 True/1 = 有效,False/0 = padding。
形状:(B, L_total),对应整个序列每个位置。
att_masks:attention 域掩码
作用:告诉 make_att_2d_masks 不同 token 属于哪个“域(domain)”,从而控制谁可以 attend 谁。
这个函数最后,pad_masks是表示哪个token是有效的(state这里应该都是有效),att_masks的话,图和task是0,state是1
embed_suffix
把 动作序列 与 时间 t 融合,作为“后缀序列”,只包含动作端的 token。跟state倒是没什么关系 att_masks 为长度 T_act 的 1 序列。 配合前缀里的 0/1,通常表示“后缀的动作 token 属于另一个注意力域”,从而避免图像/文本去看动作;也可允许状态/动作之间的相互可见性(具体由 make_att_2d_masks 决定)。
为什么action也要是1?大概懂了点,要分为条件输入和决策输入,emm主要是我不知道还能放在哪儿,理论上来说这个位置信息肯定是放在0比较好的
forward
目标:学习一个 速度场 v_θ(x_t, t),去逼近真速度 u_t = ε - x_0 的等价形式(此实现里是 ε - actions),属于 flow matching/噪声驱动的动作建模思想。
SmolVLMWithExpertModel
init
- self.vlm和self.procesor的区别
- self.vlm是AutoModelForImageTextToText.from_pretrained( model_id),g会加载一个 视觉-语言-文本(VLM)大模型
- self.processor = AutoProcessor.from_pretrained(model_id)看起来都是from_pretrained(model_id)
- self.vlm是模型权重 + 前向推理逻辑,类似 transformers 里常见的
AutoModelForCausalLM。这里是 Image+Text → Text
的多模态大模型(视觉-语言模型,里面应该包含
- model
- 是整个 VLM 的 主体 backbone,直接调用 self.vlm.model(...) 会自动执行 vision_model + connector,给你一个对齐到 text hidden_size 的 embedding
- .vision_model视觉编码器(把图片转成
embedding)纯粹的encoder,只负责「把图像转成 patch embedding」
- 和VLAFlowMatching里面的prefix_embed的区别是:embed_prefix 不自己做图像编码,它只是一个「拼接器」
- img_emb = self.vlm_with_expert.embed_image(img)这里调用的就是self.vlm_with_experts.vlm.model
- .text_model:文本编码器/解码器
- 同理,embed_prefix里面对state和language的处理也类似,关注state:
- .connector:模态对齐层(把 vision hidden state 映射到 text hidden space)
- forward
- generate
- model
- self.processor是输入预处理器,负责把原始数据(图像、文字字符串)转成
self.vlm 能接受的张量
- 对图像:resize、归一化、转 tensor
- 对文本:tokenize → token_id → attention_mask
- self.vlm.model.text_model.layers减少到前num_vlm_layers(16)层
- 构建更窄的 expert 部分lm_expert
- lm_expert_config = copy.deepcopy(config.text_config)为什么只用text_config?因为专家模型只管 text 部分(不管 vision),vision也被拼接到text里面了。先得到一个大概的config结构,和self.vlm.config.text_config是一样的大小。
- 更窄
- 原来mlp里面的结构是这样的:hidden_size → intermediate_size → hidden_size,也就是先把先把 hidden 向量投影到一个更宽的空间(intermediate_size),再投影回 hidden_size,这样模型有更强的表达能力。
- hudden_size (每个 token embedding 的维度,Transformer 层的输入/输出维度)
- intermediate_size FFN 里间层宽度,挺复杂的不用管这个经验公式
- num_hidden_layers 堆多少层 Transformer,因为 Expert 要“对应”VLM 的层结构,方便做跨注意力对齐。这样一层 Expert 对应一层 VLM,更容易设计 cross-attn。如果额外设置了一个num_expert_layers,vlm跟expert不是一层层对应,但是expert_layer要是16的因数,为了方便稀疏交互
- self.lm_expert = AutoModel.from_config(lm_expert_config)结构和 VLM.text_model 一样,但更小、更窄
- 大小确定了,值是需要自己训练的,也就是expert需要单独训练。
- atten_mode
- 因为 cross attention 模式下,Expert 不直接用自己的 K/V,而是用大 VLM 的 K/V 表示,改值,所以需要维度也匹配一下,不然vlm的值拿不过来:当该层做 cross attention 时 → Expert 的输入是 VLM 传下来的特征(比如 hidden_size=1024),必须经过这个替换过的 Linear,投影到 Expert 自己的宽度(比如 hidden_size=512)。。vlm不是冻结的吗,也有kv?
- 不是每一层 Expert 都用 cross attention,每 xx 层保留一次 纯自注意力,其他层就把 K/V 换成来自 VLM 的(投影过的)K/V,也就是cross attention和self-attention
- 专家不自己负责词嵌入,统一用 VLM 的词嵌入或上游传入的 inputs_embeds。self.lm_expert.embed_tokens = None
- set_requires_grad
- 如果self.freeze_vision_encoder设置为true:整个vlm冻结,expert可以训练,只更新expert的权重
- train_expert_only=True
- self.vlm,self.vlm.model,self.vlm.model.vision_encoder
forward
- batch_size
- 对每一层要么self-atten要么cross-atten,记录结果att_outputs, past_key_values
1 | att_outputs,past_key_values=self_or_cross-attention( |
att_outputs 和 past_key_values 在整个模型中的位置
这两个是 中间结果,在模型 forward 的流水线上扮演不同角色:
att_outputs
这是这一层(vlm + expert)的 attention 输出 hidden states。
会作为 下一层的输入,一路传到最后一层 → 接 decoder head / classifier head。
所以它处于 「主干计算图」 中,梯度会往回传。
相当于 transformer 层的 hidden_states。
past_key_values
保存每一层的 KV (key, value) 张量。
在训练时 可以不用(因为我们每次都全量计算),
在推理时(自回归生成),用它来避免重复计算旧 token 的 KV。
它处于 「缓存/优化分支」,通常不会参与梯度计算。
只存储 forward 的中间结果,不会影响训练更新
训练流程
- train
- smolvlapolicy
- init
dataset_stats和config.input_features里面的state维度要对应,改一下dataset_stats
- policy = make_policy(ds_meta)这个ds_meta就是dataset_stats吗?加上
- forward此时传入的batch的state是10维
- init
dataset_stats和config.input_features里面的state维度要对应,改一下dataset_stats
- vlaflowmatching
- self.vlm_with_expert = SmolVLMWithExpertModel
- self.state_proj = nn.Linear( self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size )
- (, suffix_out), = self.vlm_with_expert.forward( attention_mask=att_2d_masks, position_ids=position_ids, past_key_values=None, inputs_embeds=[prefix_embs, suffix_embs], use_cache=False, fill_kv_cache=False, )
- smolvlmwithexpert
- init
- forward,传入的
- inputs_embeds:外部准备好一个 [batch, seq_len, hidden_dim] 的张量,它都会进入 forward 流程
- position_ids:这只影响 RoPE(旋转位置编码),它决定每个 token 在序列里的位置。和你要训练的 state 的值域/维度没有直接关系,只是告诉注意力计算“第 i 个 token 的位置是多少”。
dataloader
之前写在train.py里面的,for key in batch:之后,为了保存数据到本地
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15states = batch["observation.state"].detach().cpu()
if states.dim() == 3: # e.g. [B, 1, 10] → squeeze掉
states = states.squeeze(1)
episode_indices = batch["episode_index"].detach().cpu().tolist()
frame_indices = batch["frame_index"].detach().cpu().tolist()
states_list = states.numpy().tolist()
with open("modified_states.jsonl", "a") as f: # 追加写入
for ep, fr, st in zip(episode_indices, frame_indices, states_list):
record = {
"episode_index": int(ep),
"frame_index": int(fr),
"state": st
}
f.write(json.dumps(record, ensure_ascii=False) + "\n")
evaluate
- 要新增的功能
- 服务器端推理次数统计:只要在policy_server端增加log就行,init里面搞个计数器
- 能测试baseline,修改task和修改state
- task各种形式
- state要proj和normalize
config
基本上处理都可以从pretrained_name_or_path里面看出来吧,只要在配置的yaml文件里面写清楚model的path就行,然后服务器根据名字选择出 mtask,mstate和baseline
- baseline
- mtask_
- relative
- grid_2cm
- grid_5cm
- mstate_
- relative
- 1m
policy_server
- 根据config得到控制
- 推理之前_predict_action_chunk里面处理state和task
- 推理state的时候是否还需要修改normalize?
- prepare_batch函数里面看到了,来改吧
- 之前forward函数里面有batch = self.normalize_inputs(batch,adding_state_stat)
- 需要确保self.add_location_to_state有对应的值
- self.add_location_to_state=config.add_location_to_state
- cfg是SmolVLAConfig,也就是模型初始化的时候要有add_location_to_state这个设置
- self.add_state_dim加载正确
- 推理的时候去掉side_depth这个多余的feature
robot_clinet
- 注释掉validate_robot_cameras_for_policy(lerobot_features, policy_image_features),因为本地不用这个
8.29重新整理代码,今晚收集第二个物体+整理代码
git remote add upstream https://github.com/huggingface/lerobot.git 之前更新的时候不小心丢掉了
- camera:realsense camera更改比较多,主要是新增保存双通道的深度图
- scripts/train.py
- configs/train.py 增加了3个命令行参数 # 自定义 #
当空的时候就等于baseline
language_tip_mode: str = "" # 改成模式,有 pure和grid # 控制是将location加到state里面 add_location_to_state: str = "" freeze_except_7_10: bool= False
factory.py 里面make_dataset