【VLA与模仿学习】Wan2.1框架与训练源码导读

Felaim

17人浏览 · 2026-11-06 11:54:32

Felaim · 2026-11-06 11:54:32 发布

**先记输入、输出、真值：**Wan视频模型接收文字，I2V版本还接收起始图像，输出视频；训练数据是与文字配对的真实视频。真视频经VAE压缩后，与随机噪声一起构造速度监督。

DreamZero把这个已训练的视频模型继续训练为“未来视频+动作”的联合生成模型。Wan原始预训练、DreamZero机器人训练、Flash少步适配是不同阶段，训练目标与可训练模块不能混写。

1. 一张框架图，分别看训练与生成

在这里插入图片描述

**自绘教学图，按 Wan2.1 官方报告与公开仓库内嵌模块标注，不是原论文图。**图中的“视频latent GT”与“网络速度GT”是不同层次；推理时没有未来GT输入。

阶段	训练输入/条件	数据真值	训练什么
VAE训练	视频/图像	输入本身	压缩后还能还原内容
Wan DiT训练	带噪latent、噪声时间、文本，I2V另有首图条件	真实视频latent+噪声构成的速度	从噪声生成视频的更新方向
DreamZero训练	视频/动作噪声、真实历史、语言、机器人状态	未来视频与记录动作构成的联合速度	画面变化与动作的联系
DreamZero-Flash	视频偏高噪声、动作独立噪声	仍是联合速度	适应少步动作生成

2. 原始数据和条件怎么准备

一条教学训练数据：文字“机械臂拿起桌上的杯子”，视频包含靠近、抓住、提起的连续画面。输入条件决定“生成哪类内容”，而GT视频提供“这段内容实际怎样变化”的答案。

通用Wan预训练使用大规模图像/视频与描述，先筛选质量和运动等，再进行训练；不是要求每个互联网视频都有机器人关节标签。DreamZero用的机器人动作标签是后续适配新增的监督。

I2V训练可以从视频中选择起始图像作为条件，目标仍是视频。在代码里同一段GT视频同时用于取得首图条件与生成目标；模型允许看到首图，不等于允许看到全部未来干净答案。

3. Wan-VAE：为什么先把视频压缩

3.1 看一次真实形状换算

设B为batch、T为原始视频帧数、H/W为图像高宽。Wan2.1 VAE的输出形式：

$[B,3,T,H,W]\rightarrow[B,16,1+(T-1)/4,H/8,W/8],\quad T=4n+1.$

例如DreamZero拼好的33帧352×640：

RGB视频：[1,3,33,352,640]
VAE表示：[1,16,9,44,80]

前者有22,302,720个数，后者506,880个数，约44倍元素压缩。这个比值含通道由3增到16、首帧保留的影响，不能直接把4×8×8=256说成最终元素减少256倍；更不能由元素数直接推断实际GPU显存比例。

**33→9的具体分组：**第一幅单独编码，剩余32幅每4幅产生一个压缩时间位置，得到1+32/4=9；每个位置仍是16通道的空间特征网格，不是单个token。含逐组对应和编码循环的解释见源码导读§2.8。

3.2 本地实现怎样保证时间压缩与连续性

wan_video_vae.py:939中定义encoder/decoder与时间下采样。wan_video_vae.py:969的 encode：

建立卷积特征cache。
先编码第一帧。
之后每4帧送一次encoder，带着卷积历史特征。
汇总输出，经conv取 mu，按通道均值/标准差归一化后返回。

这里的“因果”是编码过去时不依赖未来帧。底层 CausalConv3d 将时间padding全部放在过去侧，若有 cache_x 就先拼历史特征再减去对应padding，见 wan_video_vae.py:34。VAE卷积的feature cache缓存特征，与DiT注意力的KV cache不同。

VAE内部有encoder输出均值/方差以及 sample 函数，不意味着DreamZero每次前向都随机采VAE后验：活动的 encode 返回标准化后的mu。其decoder在 wan_video_vae.py:998，高层封装在 wan_video_vae.py:1208。

**中间通道要分清：**本地 self.encoder 子网络输出32通道；conv1(...).chunk(2,dim=1)后才保留16通道mu。首帧与带历史cache的后续4帧均产生1个压缩时间位置。详见单帧/四帧encoder输出形状与meta核验。

3.3 VAE本身怎么训练（官方报告，非本地训练入口）

Wan报告§4.1.2：先训练2D图像VAE，再扩展为3D因果VAE，在低分辨率短视频上训练，最后用更高质量、多分辨率/时长视频微调；损失包括系数3的L1、3e−6的KL、3的LPIPS，最后阶段加入3D判别器GAN损失。原文

用教学公式理解各项职责：

$L_{VAE}=3L_{L1}+3L_{LPIPS}+3\times10^{-6}L_{KL}\quad\text{（最后阶段另有GAN项）}.$

L1：逐像素还原误差，颜色/位置有没有还原对。
LPIPS：比较感知特征，让还原画面的结构与视觉观感更接近原视频。
KL：约束压缩表示分布，避免编码器随意使用不受约束的表示；常见高斯后验对标准正态的教学公式为 0.5 Σ(μ²+σ²−logσ²−1)，不是本地重新实现过这套训练。
GAN：判别器区分真实与还原视频，促使细节更自然；这不保证物理正确。

本地边界：这里只带VAE网络与权重加载，不含复现上述完整VAE训练的data/optimizer/判别器入口。DreamZero冻结它，当前loss没有把KL、LPIPS、GAN一起继续优化。wan_flow_matching_action_tf.py:361。

3.4 为什么latent通道设为16

2026-11-04核对官方报告§4.1.1/4.1.4与本地VAE初始化。16是作者选定的latent特征宽度（超参数），不是从RGB的3通道、图像尺寸或时间压缩率算出来的。 官方报告明确采用16通道，但本次所查部分没有证明“16必然优于8或32”的专门通道消融，不能编造这个最优性结论。官方报告

每个压缩时空格用16个数来保存图像/运动信息。RGB的3通道有红绿蓝的固定含义；latent的16通道是网络学出的内部表示，不能指定成“第1通道是颜色、第2通道是深度”。

设计直觉是表示容量与压缩成本之间的取舍：通道更少，表示更紧凑，但可能更难保留细节；通道更多，允许保存更多信息，但latent存储和相关投影/解码开销增加。这里是一般原理，具体重建/生成质量还要靠训练实验，不能仅凭通道更多就保证更好。

代码直接将它设成默认参数：

class VideoVAE_(nn.Module):
    def __init__(self, dim=96, z_dim=16, ...):
        self.z_dim = z_dim
        self.encoder = Encoder3d(dim, z_dim * 2, ...)
        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
# encode中：
mu, _ = self.conv1(out).chunk(2, dim=1)

以上为教学节选，省略号不是完整可运行初始化。encoder先输出2×16=32通道，随后拆分并保留16通道mu；32的来源是先选定16再乘2，不能反过来说16由某个自然定律决定。

DreamZero继承Wan2.1预训练VAE及DiT，所以必须匹配已训练的16通道接口。改成32不只是改一个配置数：VAE编码/解码权重、通道标准化统计、DiT输入输出层以及训练数据表示都需适配。在保持同一种I2V条件设计时，输入通道为2×z_dim+4：16时36通道；32时68通道。

若时间/空间压缩率与patch大小不变，单纯改latent通道不会改变每时间位置880个token；变的是patch里的输入数及投影维度，DiT宽度5120也不必自动改变。不能说16→32会让整个14B网络计算量刚好翻倍。

记忆句：16是每个压缩时空格的特征宽度；Wan选定并训练好了这套16通道表示，DreamZero沿用它。

4. T5与CLIP：条件从哪里来

T5把指令编码成512个4096维token；注意力mask用于处理有效文本长度，见 wan_video_text_encoder.py:253、wan_flow_matching_action_tf.py:542。它给DiT“生成什么”的条件，不是让一个聊天模型逐字输出视频。

I2V还保留首图的两种条件：

VAE首图latent+mask：保留空间布局与已知帧信息。
CLIP首图特征：作为视觉上下文送cross-attention。

本地 wan_flow_matching_action_tf.py:563 把首图和后续零帧编码得到16通道条件，再拼4通道mask，得到20通道y；与16通道噪声视频拼接成为36通道。CLIP输出257个1280维token，经MLP投到5120维。两条路不是重复命名同一张量。

I2V分量	通道/形状	作用
noisy video	16通道	待生成的视频内容
condition latent	16通道	从首图+零帧编码得到的条件
known-frame mask	4通道	指示已知/待生成位置
合计DiT输入	36通道	channel轴拼接，不是时间轴拼接
CLIP条件	`[B,257,1280]`	另走图像cross-attention

官方Wan I2V机制见报告§5.1、官方image2video.py。本地mask组装为首latent位置置1的实现，不能不加核验就泛化到首尾帧插值等其他模式。

术语白话：I2V条件是什么，patchify是什么？含首图双分支与实际代码。

5. DiT：真正做生成的大网络

5.1 Patchify → 时空token

卷积核与步长 (1,2,2)：时间维不再缩小，latent空间每2×2变一个token。wan_video_dit_action_casual_chunk.py:1386。

$L=F\times(H_z/2)\times(W_z/2).$

示例F=9、Hz=44、Wz=80，得到L=9×22×40=7920，每token5120维。三维位置编码让不同时间/空间位置可区分。

看逐步图解：latent时间位置是什么，7920与5120分别怎么来。包含144→5120投影与flatten/transpose代码。

5.2 每个block怎样读取信息

基本顺序为：归一化/噪声时间调制→self-attention→残差→cross-attention→残差→FFN→门控残差。wan_video_dit_action_casual_chunk.py:1151。

self-attention让视频token相互读取，建模同帧空间关系和帧间变化。
cross-attention让当前生成内容读取文本/首图条件。
timestep embedding告诉网络当前是高噪声还是接近完成，以调整更新方式。
FFN对每个token进一步加工，不是单独执行一次机器人规划。

通用Wan的视频生成窗口可以做双向时空注意力；DreamZero为分块闭环改成因果历史访问并加入动作/状态token。不能把DreamZero的所有mask和KV规则当成原始Wan的默认结构。

活动类边界：本地DreamZero用 wan_video_dit_action_casual_chunk.CausalWanModel；普通Wan类在 wan2_1_submodule.WanModel等文件里存在，但不是当前训练配置的 _target_。I2V cross-attention实际由 wan2_1_submodule.py:308 提供。

5.3 输出是什么

输出head把token映射回patch内16通道的数值，再unpatchify回video latent形状。这一轮输出的是更新方向，不是最终视频RGB，也不是一定已经干净的latent。 scheduler用它推进当前样本；生成完成后才由VAE decoder还原RGB。

6. Wan的Flow Matching训练：一条样本如何变成loss

6.1 两种时间记法必须对应

Wan报告正文：x_t=t x_data+(1−t)ε，t从噪声0到数据1，target为 x_data−ε。本地scheduler则使用σ：

$x_\sigma=(1-\sigma)z+\sigma\epsilon,\quad v^*=\epsilon-z,\quad t=1-\sigma.$

本地 flow_match_scheduler.py:73、flow_match_scheduler.py:83。不是两个互相矛盾的目标，而是沿同一直线采用相反坐标方向。

6.2 六步训练流程

# 教学伪代码；本地实际返回的是视频/动作两分支
z = frozen_vae.encode(real_video)       # 干净压缩视频GT
c = frozen_text_encoder(instruction)   # 条件
noise = randn_like(z)
sigma = sample_schedule()
z_noisy = (1-sigma)*z + sigma*noise
pred_velocity = dit(z_noisy, sigma, c, image_condition)
target_velocity = noise-z
loss = weighted_mse(pred_velocity, target_velocity)
loss.backward()                       # 梯度更新DiT
optimizer.step()

为什么不需要标签员标注“噪声去哪里”？ 因为z来自真实视频，noise是程序自己抽的，两者都已知，差值就能算出来。这与机器人动作GT不同：动作GT仍需演示记录，原始Wan不会凭视频自动知道某台机器人的关节数值。

为什么训练只随机取一个噪声程度，生成却要多步？ 训练在许多样本、许多随机噪声时间上学习局部速度；推理时拿这个学到的速度场从噪声走到数据，需要数值求解器分若干步更新。训练不是每条样本都先完整生成一段视频再算loss。

6.3 一次手算验证

假设一维干净latent z=2、噪声ε=10、σ=0.75：

输入：xσ=0.25×2+0.75×10=8。
GT速度：ε−z=8。
若预测速度7，未加权平方误差 (7−8)²=1。
若预测正确，步长σ从0.75走到0.5：新值 8+(0.5−0.75)×8=6。
再走到0：6+(0−0.5)×8=2。

实际是高维张量、非恒定网络速度，多步求解并不保证精确恢复任意GT；算例只帮助看清符号与方向。

6.4 大规模预训练的阶段

官方报告§4.2.2：先低分辨率图像预训练，再图像/视频联合训练，逐步增加分辨率与时长；使用Flow Matching，后训练进一步使用高质量视频。原始时间采样描述为logit-normal，与DreamZero本地按shift时间表抽索引不同。原文

因此“Wan怎么训练”包含两层：先训练可用的VAE，再用压缩视频训练DiT。DreamZero接收的是这些已学好的权重，继续做机器人适配，而不是从随机参数重新完成互联网规模的视频预训练。

7. Wan推理：从噪声到视频

文字/首图 → 条件编码
随机视频latent噪声
  → DiT给更新方向 → scheduler推进一次
  → DiT给更新方向 → scheduler再推进
  → … → 视频latent → VAE解码 → RGB视频

本地简单scheduler的更新式：x_next=x+(sigma_next−sigma)*pred_velocity，见 flow_match_scheduler.py:43。实际DreamZero推理调用FlowUniPC多步求解器，不能把这个Euler式当其每个分支逐行执行的完整算法。

CFG（无分类器引导）按有指令和空/负条件两次预测形成 v_uncond+s*(v_cond−v_uncond)，强化条件影响；不是分类器输出概率。DreamZero对视频使用该组合，动作读取条件分支，见 wan_flow_matching_action_tf.py:1274。

8. 从Wan到DreamZero：哪些保留，哪些改变

部分	保留/改变	源码入口
T5、CLIP、VAE	载入并冻结	wan_flow_matching_action_tf.py:174、wan_flow_matching_action_tf.py:361
视频DiT参数	从Wan初始化，再继续训练	wan_flow_matching_action_tf.py:264
状态/动作接口	新增编码和解码	wan_video_dit_action_casual_chunk.py:1367
注意力	改为分块历史条件，视频与当前动作互相读取	wan_video_dit_action_casual_chunk.py:816
GT与loss	视频速度以外，加动作速度MSE	wan_flow_matching_action_tf.py:760
推理记忆	已发生真实视频进入KV	wan_flow_matching_action_tf.py:1143
Flash	训练时视频和动作噪声采样解耦	wan_flow_matching_action_tf.py:679

详细调用、真实shape、mask、loss缩放与配置陷阱见 2026-11-04 - DreamZero源码全链路与Wan2.1训练实现。

9. 代码阅读顺序与证据边界

wan_flow_matching_action_tf.py:603：看完整训练前向，先知道各模块串在哪里。
wan_video_vae.py:969：看时间压缩和mu返回；不必一开始读完全部卷积层。
wan_flow_matching_action_tf.py:563：看I2V首图的双分支。
wan_video_dit_action_casual_chunk.py:1998：看patch、action/state拼接、teacher-forcing和输出。
wan2_1_submodule.py:308、wan_video_dit_action_casual_chunk.py:1151：看cross-attention和一个block。
flow_match_scheduler.py:73、wan_flow_matching_action_tf.py:784：把target和loss与前向对应。
wan_flow_matching_action_tf.py:973：最后看生成采样和KV，而不是从TensorRT/服务端优化开始。

图中的shape来自明确尺寸的教学记录和本地代码算术；CPU/meta核验不含整网VAE/DiT权重。VAE训练阶段据官方报告，不冒充公开仓库包含原始全量预训练配方。Wan2.2或其他后续版本的48通道VAE不能直接套用本文16/36通道图。

来源与交叉链接

Wan2.1官方仓库、官方model.py、官方image2video.py。
Wan报告：§4.1 VAE、§4.2 DiT训练、§5.1 I2V。
Wan2.1概念卡、2026-11-04 - DreamZero论文精读。

10. Wan2.1 与 Wan2.2：先分清具体型号

2026-11-04 核对官方Wan2.2仓库/配置，以及本地DreamZero的两套YAML。以下比较基础T2V/I2V路线，不覆盖该系列所有语音/动画扩展型号。

核心区别：Wan2.2的大模型A14B路线引入高/低噪声双专家；TI2V-5B路线则是较小的稠密模型，配更高空间压缩的VAE。不能把“Wan2.2”整体等同于MoE或5B。

项目	Wan2.1-I2V-14B	Wan2.2-I2V-A14B	Wan2.2-TI2V-5B
基本职责	图像/文字条件生成视频	图像/文字条件生成视频	同一模型支持文生视频与图生视频
DiT结构	单个稠密主干	高噪声/低噪声双专家	单个稠密主干
参数口径	约14B	官方约27B总参数，每步约14B活跃	约5B
VAE	Wan2.1 VAE	仍用Wan2.1 VAE	Wan2.2 VAE
时间×高×宽压缩	4×8×8	4×8×8	4×16×16
latent通道	16	16	48
每个DiT的宽/层/头	5120 / 40 / 40	5120 / 40 / 40	3072 / 30 / 24
与当前DreamZero关系	原论文采用的基座	本次所查Wan22配置不是这条路线	公开仓库另有适配配置

来源：Wan2.2官方介绍、A14B官方配置、TI2V-5B官方配置。

10.1 MoE是什么意思：按去噪阶段分工

MoE是Mixture of Experts（混合专家），这里可理解为两套分别擅长不同噪声阶段的生成网络。

随机噪声 → 高噪声专家：先形成总体布局与运动
         → 到达切换阈值
         → 低噪声专家：进一步完善细节
         → 视频latent → 解码

Wan2.1主要由同一个主干处理各个噪声阶段；Wan2.2 A14B把阶段分给两位专家。这里是按去噪时间切换整套专家，不能直接套用某些语言MoE“每个token选择几个FFN专家”的细节。A14B的A可以按active参数口径理解；总参数与单步参与计算的参数不同。两位专家的权重仍需存储/装载，不能据单步14B断言存储开销与单模型完全一样。

10.2 5B为什么可能更省计算：空间token更少

用同样33帧、256×256视频做教学比较，T=33均变成9个latent时间位置：

量	Wan2.1 VAE	Wan2.2 TI2V-5B VAE
latent形状（B=1）	[1,16,9,32,32]	[1,48,9,16,16]
latent元素数	147456	110592
patch=(1,2,2)后每时间位置token	16×16=256	8×8=64
总视频token数	2304	576

高度宽度各多压缩2倍，所以空间token少至1/4；但通道由16变48，latent总元素仅变为3/4。不要把token减少4倍说成全部显存或推理必然快4倍。网络宽度/层数、注意力、采样步数、硬件和分辨率都会影响成本。

10.3 本地DreamZero实际换了什么

源码：Wan2.1配置、Wan2.2配置。

本地适配参数	Wan2.1主线	Wan2.2-TI2V-5B适配
DiT宽/层/头	5120/40/40	3072/30/24
VAE通道/空间步长	16 / 8	48 / 16
输入/输出通道	36 / 16	48 / 48
concat_first_frame_latent	true	false
当前代码输入视频尺寸	三路176×320拼成352×640	在head内统一到160×320
每latent时间位置视频token	880	50