LeRobot 中的 WALL-OSS:基于 Qwen2.5-VL 与 Flow Matching 的跨具身 VLA 策略集成指南
【免费下载链接】lerobot🤗 LeRobot: Making AI for Robotics more accessible with end-to-end learning项目地址: https://gitcode.com/GitHub_Trending/le/lerobot
WALL-OSS 是由 XSquare Robot 团队于 2025 年提出的开源具身智能基础模型,现已以wall_x策略的形式完整集成进 Hugging Face LeRobot 生态。本文以 policy_walloss_README.md 与 walloss.mdx 为骨架,结合仓库内策略源码与测试,系统讲解 WALL-OSS 的模型架构、两种动作预测模式、安装方法、微调训练命令及全部核心参数,并深入到WallXConfig、ActionHead、MoE 路由等实现细节。读完本文,你将能够在 LeRobot 中独立完成 WALL-OSS 的安装、后训练(post-train)、评估与部署。
模型概述
WALL-OSS 是一个用于跨具身(cross-embodiment)机器人控制的视觉-语言-动作(Vision-Language-Action, VLA)模型。它捕获并压缩连续、高保真的物理交互数据流,通过在模型决策与实体“体感经验”之间建立直接反馈回路,驱动具备真正泛化能力的智能体产生——模型不仅要理解世界如何运作,还要知道如何在该世界中有效行动。
| 特性 | 说明 |
|---|---|
| 基础模型 | Qwen2.5-VL(视觉-语言模型) |
| 动作预测 | Flow Matching(扩散式)或 FAST(离散 token 式) |
| 架构 | 混合专家(Mixture of Experts, MoE),带动作专用路由 |
| 多模态输入 | 视觉(图像/视频)、语言、本体感觉(Proprioception) |
LeRobot 中的实现由官方 WallX,对应策略注册名为wall_x(见 modeling_wall_x.py 与WallXConfig的@PreTrainedConfig.register_subclass("wall_x")装饰器)。
核心技术路线
WALL-OSS 提出了一种紧耦合多模态架构(tightly-coupled MoE structure),将离散与连续两类动作建模策略统一在一个可微框架中,并通过两阶段训练(Inspiration → Integration)逐步统一语义推理与高频动作生成。其核心创新点包括:
- 具身感知增强的多模态预训练:在统一的视觉-语言-动作大规模数据上训练,强化空间、因果与操作理解能力。
- 统一跨层级思维链(Uni-CoT):单一可微框架统一高层指令推理、子任务分解与细粒度动作合成,形成从“理解”到“执行”的连续链路。
- MoE 动作头:根据任务阶段动态激活专家,在离散或连续空间中建模动作,保持稳定的 VLM 先验。
- 两阶段训练范式:
- Inspiration 阶段:注入离散动作先验,强化空间理解与语义-动作对齐;
- Integration 阶段:使用 Flow Matching 实现高频连续控制。
在源码层面,这两类动作建模分别对应WallXConfig.prediction_mode的diffusion与fast两种取值,具体实现位于 modeling_wall_x.py 的predict()方法中(见下文“两种动作预测模式”小节)。
环境安装
- 按 安装指南 安装 LeRobot 本体;
- 安装 WallX 额外依赖:
pip install -e ".[wallx]"wallx依赖组在 pyproject.toml 中定义,包含transformers、peft、scipy、torchdiffeq>=0.2.4,<0.3.0与qwen-vl-utils。其中:
torchdiffeq用于 diffusion 模式下动作采样的 ODE(常微分方程)积分;peft用于可选的 LoRA 低秩微调;qwen-vl-utils提供smart_resize等 Qwen2.5-VL 视觉预处理工具。
在代码中启用 WallX 策略只需指定策略类型:
policy.type=wall_x策略工厂会在make_policy_config中依据该值实例化WallXConfig(参见 tests/policies/wall_x/test_wallx.py 中的集成验证方式)。
使用 LeRobot 训练 WALL-OSS
训练 WALL-OSS 可直接使用 LeRobot 标准训练脚本与对应配置:
lerobot-train \ --dataset.repo_id=your_dataset \ --policy.type=wall_x \ --output_dir=./outputs/wallx_training \ --job_name=wallx_training \ --policy.repo_id=your_repo_id \ --policy.pretrained_name_or_path=x-square-robot/wall-oss-flow \ --policy.prediction_mode=diffusion \ --policy.attn_implementation=eager \ --steps=3000 \ --policy.device=cuda \ --batch_size=32训练参数说明
| 参数 | 说明 |
|---|---|
--dataset.repo_id | 训练数据集的 Hugging Face Hub 仓库 ID(例如lerobot/aloha_sim_insertion_human) |
--policy.type | 指定使用 WallX 策略架构 |
--output_dir | 训练 checkpoint 与日志的本地保存目录 |
--job_name | 本次训练运行的名称标识(用于日志/追踪) |
--policy.repo_id | 训练完成后模型将被推送的 Hugging Face Hub 仓库 ID |
--policy.pretrained_name_or_path | 用于初始化的预训练 WallX 权重路径(官方 WALL-OSS checkpoint,默认x-square-robot/wall-oss-flow) |
--policy.prediction_mode | 动作预测策略:diffusion使用迭代去噪生成动作,fast使用下一 token 预测 |
--policy.attn_implementation | 注意力实现后端——eager使用标准 PyTorch 注意力(WallX 目前仅支持eager) |
--steps | 训练总步数 |
--policy.device | 训练设备(cuda为 GPU,cpu为 CPU) |
--batch_size | 每批训练样本数 |
几点需要澄清的细节(依据源码):
- 文档参数表中的
--policy.pretrained_path在实现中对应配置字段为pretrained_name_or_path,其默认值为x-square-robot/wall-oss-flow,见 configuration_wall_x.py。 - 训练脚本中
--steps与--batch_size是lerobot-train顶层参数,而--policy.*前缀下的才是策略配置字段。 --policy.attn_implementation当前强制为eager:WallX 的双向动作 token“孤岛”需要显式注意力掩码,若配置为其他值会在WallXConfig.__post_init__中直接抛出ValueError(见 configuration_wall_x.py)。这与文档表格中提及的flash_attention_2/sdpa备选项存在差异,应以源码校验为准。
WallXConfig:核心配置字段深度解析
WallXConfig(configuration_wall_x.py)继承自PreTrainedConfig,是 WallX 策略的配置中枢。其关键字段与默认值如下:
输入 / 输出结构
| 字段 | 默认值 | 说明 |
|---|---|---|
n_obs_steps | 1 | 观察步数 |
chunk_size | 32 | 动作块长度(对应 wall-x 中的action_horizon),即每次模型调用预测的动作长度 |
n_action_steps | 32 | 每次实际执行的动作步数,必须满足n_action_steps <= chunk_size |
max_action_dim | 20 | 最大动作维度(WallX 采用 20 维统一动作表示,真实维度不足时补零) |
max_state_dim | 20 | 本体感觉状态最大维度 |
配置校验(__post_init__)会拦截以下非法组合:n_action_steps > chunk_size、prediction_mode不在["diffusion", "fast"]内、attn_implementation != "eager"、vision_attn_implementation不在{"auto", "sdpa", "varlen"}内。
动作预测相关
| 字段 | 默认值 | 说明 |
|---|---|---|
pretrained_name_or_path | x-square-robot/wall-oss-flow | 预训练权重来源 |
action_tokenizer_path | lerobot/fast-action-tokenizer | 动作 tokenizer 路径(仅 FAST 模式使用) |
prediction_mode | diffusion | diffusion或fast |
attn_implementation | eager | 文本/动作 token 注意力后端(固定eager) |
vision_attn_implementation | auto | 视觉注意力后端;auto在运行时支持时使用 PyTorch 打包变长注意力,否则回退到原生 SDPA |
__post_init__会根据prediction_mode自动设置use_fast_tokenizer:fast模式置True,diffusion模式置False并将action_tokenizer_path置为None(diffusion 模式不需要离散动作 tokenizer)。
优化器与调度器预设
| 字段 | 默认值 |
|---|---|
optimizer_lr | 2e-5 |
optimizer_betas | (0.9, 0.95) |
optimizer_eps | 1e-8 |
optimizer_weight_decay | 0.01 |
optimizer_grad_clip_norm | 1.0 |
scheduler_warmup_steps | 1000 |
scheduler_decay_steps | 100000 |
scheduler_decay_lr | 1e-6 |
它们分别通过get_optimizer_preset()生成AdamWConfig、通过get_scheduler_preset()生成CosineDecayWithWarmupSchedulerConfig,保证开箱即用的稳定训练配置。
特征校验与统一表示
validate_features()要求输入特征中至少有一个FeatureType.VISUAL视觉特征;若缺少observation.state或输出中缺少action,会按max_state_dim/max_action_dim自动补齐对应的PolicyFeature;同时校验真实状态/动作维度不超过上限。这一机制正是 WALL-OSS跨具身能力的体现——不同机器人自由度不同,通过统一维度填充(padding)与 DOF 掩码,同一套模型权重可以适配多种机械结构。
两种动作预测模式:diffusion 与 fast
WallXPolicy依据config.prediction_mode在predict_action_chunk中分派不同的推理路径(见 modeling_wall_x.py):
diffusion(默认):以max_action_dim作为action_dim、chunk_size作为pred_horizon,走流匹配(flow matching)路径。推理时从随机噪声出发,经torchdiffeq.odeint以euler方法沿num_inference_timesteps(默认 10)步时间网格积分 ODE 得到动作轨迹(见predict()中 diffusion 分支)。fast:以真实动作维度作为action_dim,通过自回归生成离散动作 token,再经action_processor.decode(time_horizon=pred_horizon, action_dim=action_dim)解码为连续动作。
ActionHead:Flow Matching 动作头实现
ActionHead(modeling_wall_x.py)是 continuous/flow 路径的核心模块,关键设计包括:
- Beta 分布噪声调度:
beta_alpha=1.5、beta_beta=1.0、s=0.999,训练时从 Beta 分布采样时间步t = (1 - sample) * s,并在 float32 下完成加噪noisy_action = (1-t) * noise + t * action与流目标flow = action - noise的计算,保证数值稳定性。 - 正弦时间嵌入:
SinusoidalPosEmb为每个时间步生成位置编码,与动作嵌入拼接后经三层线性投影(SiLU 激活)得到动作隐藏状态。 - 流匹配损失:
flow_loss在 float32 下计算动作预测与流目标间的 MSE,并支持用dof_mask掩蔽未激活的自由度通道,避免填充维度参与损失计算。 - 推理去噪:
step()实现单步去噪,供odeint反复调用;每一步将当前去噪动作嵌入替换到输入序列的<|action|>token 位置后过一遍 Transformer,取动作 token 位置的隐藏状态投影回动作空间。
FAST 模式:离散动作 token
fast模式下,模型将动作量化为一组特殊 token(<|action_token_i|>),define_action_token_id()依据 tokenizer 的action_token_vocab_size构建 token ID 集合,同时维护<|action|>与<|propri|>两个特殊 token ID(见 modeling_wall_x.py)。生成阶段使用max_new_tokens=100的自回归生成,随后过滤出动作 token 并解码为动作序列;当re_generate=True时使用temperature=0.7的采样以提高多样性。
多模态输入处理与 MoE 路由
图像预处理
常量定义于 constant.py:目标分辨率RESOLUTION = 256,先等比缩放至长边 256,再经smart_resize按IMAGE_FACTOR=28、MAX_PIXELS = 16384*28*28、MIN_PIXELS = 4*28*28量化到 Qwen2.5-VL 的 patch 网格。TOKENIZER_MAX_LENGTH = 768限制 tokenizer 输入长度。preprocess_inputs中对每路相机以 BCHW 批量张量在设备上完成双阶段 bicubic 缩放(见_resize_wall_x_image_batch),避免经 PIL 的 CPU 往返。
语言指令与 grounding 点
get_wallx_normal_text依据任务指令、chunk_size与frame_index构造模型输入文本模板,PRIORITY_ORDER(默认None)与GENERATE_SUBTASK_RATIO(默认0.0)分别控制相机优先级与子任务生成比例。process_grounding_points将 grounding 坐标从原始分辨率映射到缩放后分辨率。WallXTaskProcessor(processor_wall_x.py)保证任务描述以句号结尾,缺省时使用默认指令"Execute the robot action."。
本体感觉与动作的统一填充
preprocess_inputs(modeling_wall_x.py)中:
- 本体感觉
agent_pos与动作action中的 NaN 被置零,同时用~torch.isnan(...)构造agent_pos_mask/dof_mask; - 维度不足时按
max_state_dim/max_action_dim补零(mask 也同步补零);维度超限则报错; moe_token_types由<|action|>token 位置生成,用于 MoE 层的 token 级路由。
MoE 语言模型
src/lerobot/policies/wall_x/qwen_model/ 目录下包含移植的Qwen2_5_VLMoEModel(qwen2_5_vl_moe.py)与 vision_attention.py。训练前向(train_step_forward)将视觉、本体感觉、加噪动作三类嵌入分别masked_scatter到对应 token 位置,再统一过 MoE Transformer,最终总损失 = 语言建模交叉熵 +flow_loss_weight * flow_loss。to_bfloat16_for_selected_params将模型转 bf16 同时保留 LayerNorm 与action_preprocessor参数为 float32 以保数值稳定。
LoRA 微调支持
模型内置add_lora()(默认r=8, lora_alpha=32, lora_dropout=0.1,目标模块q_proj/v_proj),当配置中use_lora=True时自动应用,适合低资源场景下的参数高效微调。
处理器流水线与推理循环
make_wall_x_pre_post_processors(processor_wall_x.py)构建标准前后处理流水线:
- 预处理:
rename_observations→add_batch_dim→WallXTaskProcessor→normalize→to_device; - 后处理:
unnormalize→to_cpu。
推理时WallXPolicy维护动作队列(deque(maxlen=n_action_steps)):predict_action_chunk一次性预测chunk_size步动作,select_action每次弹出单步动作执行,实现平滑的滚动窗口控制(见 modeling_wall_x.py)。
仓库中的集成测试 tests/policies/wall_x/test_wallx.py 覆盖了 MoE 模型隐藏状态/注意力输出契约(test_moe_model_captures_requested_hidden_states_and_attentions),并在 CUDA + HF token 环境下验证策略实例化与完整数据流,可作为二次开发时的参考基线。
引用与许可证
如果使用本工作,请引用:
@article{zhai2025igniting, title = {Igniting VLMs Toward the Embodied Space}, author = {Zhai, Andy and Liu, Brae and Fang, Bruno and Cai, Chalse and Ma, Ellie and Yin, Ethan and Wang, Hao and Zhou, Hugo and Wang, James and Shi, Lights and Liang, Lucy and Wang, Make and Wang, Qian and Gan, Roy and Yu, Ryan and Li, Shalfun and Liu, Starrick and Chen, Sylas and Chen, Vincent and Xu, Zach}, journal = {arXiv preprint arXiv:2509.11766}, year = {2025} }WALL-OSS 采用Apache 2.0 许可证,与原 WallX 仓库保持一致。论文与官方代码等外部资源入口见 policy_walloss_README.md 的 Additional Resources 一节。
【免费下载链接】lerobot🤗 LeRobot: Making AI for Robotics more accessible with end-to-end learning项目地址: https://gitcode.com/GitHub_Trending/le/lerobot
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考