Ψ₀ (Psi-0) 模型架构#

1. 整体架构概览#

Ψ₀ 是 USC PSI Lab 面向人形机器人全身 loco-manipulation 的开源基础模型(RSS 2026,arXiv:2603.12263,CVPR 2026 第二届 3D-LLM/VLA Workshop 最佳论文)。它的核心主张不是某个新模块,而是一条数据配方:"scaling the right data in the right way" —— 先用大规模人类第一视角视频(EgoDex)让 VLM 学会任务语义和视觉表征,再用少量真机遥操作数据让动作专家学会本体动力学。声称 80 条轨迹即可微调出新的长程灵巧技能。

架构上是一个明确的三层系统分工:

关键点:VLM 与动作专家的接口只有一根线 —— VLM 最后一层的 hidden states 被当成"视觉 token 序列"(代码里叫 views)喂给动作专家(psi0.py:1625-1632)。既不做 pooling,也不额外抽 vision tower 特征,更没有中间层 Q-Former(默认路径下)。这让 System-2 和 System-1 可以完全解耦地分阶段训练。

2. 核心组件详解#

2.1 System-2:Qwen3-VL-2B-Instruct 主干#

Psi0Model.__init__(src/psi/models/psi0.py:1482)直接持有一个 Qwen3VLForConditionalGeneration。推理时把图像和指令拼成标准 chat 消息,过 AutoProcessor:

content = [{"type": "image", "image": img} for img in observation]
content.append({"type": "text", "text": instruction})
messages.append([{"role": "user", "content": content}])
texts = [self.vlm_processor.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in messages]
image_inputs, video_inputs = process_vision_info(messages, image_patch_size=16)

然后只取最后一层隐状态:

vlm_hidden_states = output.hidden_states[-1]        # (B, L, 2048)
vlm_hidden_states = vlm_hidden_states.unsqueeze(1)  # (B, 1, L, 2048) → 作为 "views"

注意 view_feature_dim 在所有 Ψ₀ 训练脚本里都写成 2048,正是 Qwen3-VL-2B 的 text_config.hidden_size。默认值 1920 是给外部 ResNet 特征用的历史遗留。

后训练与微调阶段 VLM 全程冻结(posttrain.py:126-128):

assert self.model_cfg.tune_vlm == False
for p in self.model.vlm_model.parameters():
    p.requires_grad = False

from_pretrained(psi0.py:1518)加载 checkpoint 时也很直接:按 vlm_model. / action_header. 前缀把 safetensors 劈成两半分别 load,并用 embed_tokens.weight 补上 tied 的 lm_head.weight。因为 FAST 预训练扩过词表,这里会按需 resize_token_embeddings(..., pad_to_multiple_of=192)。

2.2 System-1:VLATransformerBlock(SD3 风格 MMDiT)#

VLATransformerBlock(psi0.py:814)是整个动作专家的骨架,结构与 Stable Diffusion 3 的 JointTransformerBlock 同构,但把 SD3 的"图像流 / 文本流"换成了 动作流 A / 观测流 O:

JointVLAAttnProcessor(psi0.py:308)的核心就三行:

query = torch.cat([query, encoder_hidden_states_query_proj], dim=2)
key   = torch.cat([key,   encoder_hidden_states_key_proj],   dim=2)
value = torch.cat([value, encoder_hidden_states_value_proj], dim=2)
hidden_states = F.scaled_dot_product_attention(query, key, value, is_causal=False, attn_mask=attn_mask)
hidden_states, encoder_hidden_states = hidden_states[:, :residual.shape[1]], hidden_states[:, residual.shape[1]:]

即两条流各有独立的 q/k/v 投影和独立的 FFN,但共享同一次 attention。这与 π₀ 系列的"动作专家 cross-attend 到 VLM KV"不同 —— 这里观测流也会被动作流反过来更新,是真正的双向融合。attn_mask 由 VLM 的 attention_mask 扩展而来(padding token 被屏蔽),在 tokenize_obs 里追加 proprio token 时同步补 1(psi0.py:683)。

层数与维度:action_num_blocks=6,action_hidden_dim=1536,action_nheads=24,attention_head_dim=64(24 × 64 = 1536,严格整除)。

2.3 观测投影:ObservationProjection#

psi0.py:423。默认配置(n_conditions=0、token_fusion="concat"、use_film=False)下路径其实很短:

_obs_strat = "add_token" 是硬编码的(psi0.py:589)—— 本体状态永远作为一个独立 token 拼在观测序列末尾,而不是逐 token 广播拼接。

代码里还留着三条未启用的分支,值得知道它们存在但默认不走:n_conditions > 0 时会用 ResNet-18 编码 2D 轨迹条件图(traj2ds),token_fusion 可选 cross(轨迹 token cross-attend 视觉)或 perceiver(256 个 latent query 分别聚合视觉/轨迹),use_film 会插一层 FiLM 调制。所有 Ψ₀ 训练脚本都设 --model.n_conditions=0 --model.no-use_film。

2.4 动作投影与时间条件#

入口 ActionProjectionIn(psi0.py:746):

self.ac_proj = nn.Sequential(nn.Linear(action_dim, action_dim), nn.GELU("tanh"), nn.Linear(action_dim, output_dim))
self.dec_pos = nn.Parameter(torch.empty(action_pred_horizon, output_dim))  # xavier_uniform 初始化
...
action_hidden_states = self.ac_proj(noise_acs) + self.dec_pos.unsqueeze(0)

动作位置编码是可学习的(不是正弦),而观测侧是固定正弦 —— 一个小但明确的不对称设计。

出口 ActionProjectionOut(psi0.py:784)是 DiT 的 final layer,用 temb 做 adaLN 调制后线性投到 action_dim。注意实现里写的是 x = x * scale + shift(没有先过 norm_final,代码注释 # TODO be careful! 保留在原地)。

时间嵌入 _TimeNetwork(time_dim=256, out_dim=1536)(psi0.py:60)是经典的对数频率正弦嵌入 + 两层 MLP。它有一处专门为 RTC 做的改造:

def forward(self, x):
    # RTC: support 2D timesteps (B,Tp)
    x = x[..., None] * self.w
    x = torch.cat((torch.cos(x), torch.sin(x)), dim=-1)
    return self.out_net(x)

输入可以是 (B,) 也可以是 (B, Tp)。配套地,AdaLayerNormZero.forward(psi0.py:217)同时接受 (B, D) 和 (B, T, D) 的 emb;VLATransformerBlock 在给观测流做调制时用 temb[:,-1] if len(temb.shape) > 2 else temb(psi0.py:940/947)—— 因为观测流没有"每个动作步的噪声等级"这个概念,只能取最后一个。这三处改动合起来,就是 RTC 的全部模型侧支持。

2.5 备选动作头:DiTActionTransformerModel#

psi0.py:1211,由 --model.use-dit 开启,默认关闭。它把观测 token 先过一个 Q-Former(256 个可学习 query + CrossAttentionBlock)压成定长,再送进一个标准 DiT(depth 24,16 heads,learn_sigma=False):

query = self.query_tokens.unsqueeze(0).expand(B, -1, -1)   # (B, 256, D)
queried = self.cross_attention(query, condition_tokens)     # (B, 256, D)
action_output = self.DiT(noisy_action, timestep, queried)

代价是丢掉了 per-token timestep 能力(DiT 只接受标量 timestep),所以这条路径与 RTC 不兼容。

3. 三阶段训练流水线#

3.1 阶段 1:FAST 离散动作预训练(PretrainTrainer)#

src/psi/trainers/pretrain.py。这一阶段没有扩散、没有动作专家 —— 动作被 FastActionTokenizer(DCT + BPE,physical-intelligence/fast 系)离散成 2048 个 bin,追加到 Qwen 的词表里,然后当作普通语言 token 做自回归下一 token 预测:

self.vlm.resize_token_embeddings(len(self.tokenizer) + self.model_cfg.action_tokenizer.bins, ...)
...
loss = output.loss    # 标准 CE

监控指标 action_l1_loss 会把预测 token 解码回连续动作再算 L1(pretrain.py:224-254),只用于观察,不回传。

发布的主干 checkpoint 是二者串联的结果:psi0/pre.fast.1by1.2601091803.ckpt.ego200k.he30k。

3.2 阶段 2:Flow Matching 后训练动作专家(PostTrainTrainer)#

src/psi/trainers/posttrain.py:525。VLM 冻结,只训 ~500M 的动作头。用的是标准 rectified flow:

sigmas = torch.rand((bsz,), device=...)                       # σ ~ U(0,1)
timesteps = sigmas * 1000
noisy_actions = (1 - σ) * actions + σ * noise
target_action = noise - actions                                # 速度场
loss_action = F.mse_loss(pred, target, reduction="none")
loss_action = (loss_action * actions_mask).sum(1)              # 沿 horizon 求和(不是平均)
loss_action = (loss_action.mean(0) * self.loss_w).sum()        # 按维度加权

两个细节值得注意:

沿 horizon sum(1) 而非 mean(1)。 代码注释写明理由:"to keep gradient consistent when training with different action chunks" —— 这样从 chunk=16 换到 chunk=30 时,单步梯度尺度不会因分母变化而漂移。

逐维度损失权重 loss_w。 默认 [0.1, 0.2, 0.1] 对应 xyz / rpy / gripper,展开后强制校验 sum == 1.0(posttrain.py:64)。当 action_dim != 7 时退化为均匀的 1/Da。

优化器按前缀分组(posttrain.py:134):action_header.* 用 train.learning_rate,vlm_model.* 用 lang_backbone_lr=1e-5(冻结时该组为空)。

3.3 阶段 3:任务微调 + RTC 训练(FinetuneTrainer / SonicTrainer)#

src/psi/trainers/finetune.py:539。这一阶段打开 --model.rtc --model.max-delay=8,损失函数多出一段:

delay = torch.randint(low=0, high=self.model_cfg.max_delay, size=(bsz,))
prefix_mask = torch.arange(Tp)[None, :] < delay[:, None]        # (B, Tp)
sigmas = torch.where(prefix_mask, 0.0, sigmas[:, None])         # 前 delay 个动作步噪声等级归零
timesteps = sigmas * 1000                                        # → (B, Tp) 逐 token timestep
...
mask = mask * (~prefix_mask)[:, :, None].float()                # prefix 不计入 loss

含义:每个训练样本随机抽一个延迟 d ∈ [0, max_delay),把 chunk 的前 d 步当作已知的干净动作(σ=0),其余步正常加噪,并且只在后半段算 loss。模型因此学会"给定任意长度的干净前缀,补全剩下的动作"。这正是推理时 RTC 需要的能力 —— 上一块动作还在执行,新一块必须与之无缝衔接。

全局 batch size 在所有真机与仿真实验里统一为 128(device bs × GPU 数 × 梯度累积)。

4. 推理流水线#

4.1 基础路径:predict_action#

psi0.py:1650。一次 VLM 前向 + N 步动作去噪(eval_diffusion_steps 默认 10):

action_samples = torch.randn(bsz, self.action_horizon, self.action_dim)
self.noise_scheduler.set_timesteps(num_inference_steps)
for timestep in self.noise_scheduler.timesteps:
    model_pred = self.action_header(timestep=timestep.expand(bsz), joint_attention_kwargs=dict(
        action_hidden_embeds=action_samples, views=vlm_hidden_states, obs=states, traj2ds=traj2ds))
    action_samples = self.noise_scheduler.step(model_pred, timestep, action_samples).prev_sample

VLM 只跑一次,循环里只有动作头 —— 6 层 × 1536 维,推理开销可以忽略。

4.2 RTC:四个变体#

Real-Time Chunking 要解决的是推理延迟导致的动作断裂:上一 chunk 已经执行了 d 步,新 chunk 必须在前 d 步与旧的对齐。仓库里实现了四条路线:

训练一致的 RTC(部署实际使用的路径,psi0.py:1770-1846):

prefix_mask = torch.arange(H)[None, :] < d
for timestep in self.noise_scheduler.timesteps:
    batched_timestep_masked = torch.where(prefix_mask, 0, timestep)          # 前缀 timestep 恒为 0
    action_samples = torch.where(prefix_mask[:, :, None], prev_actions, action_samples)  # 前缀替换为真实旧动作
    model_pred = self.action_header(timestep=batched_timestep_masked, ...)
    action_samples = self.noise_scheduler.step(model_pred, timestep, action_samples).prev_sample

因为训练时见过完全相同的分布,这里不需要任何引导项 —— 直接钉住前缀即可。

伪逆引导 RTC(psi0.py:1978-2043)则是训练无关的方案,把"前缀要贴近 prev_actions"写成一个引导项加到速度场上:

pred_x0 = x_t - τ * v                                   # 外推干净动作
error = (prev_actions - pred_x0.detach()) * W           # W 为软掩码
pinv_correction = torch.autograd.grad(pred_x0, x_t, grad_outputs=error)[0]   # VJP
inv_r2 = (τ**2 + (1 - τ)**2) / τ**2
c = torch.nan_to_num(τ / (1 - τ), posinf=guidance_weight)
g = torch.clamp(c * inv_r2, max=guidance_weight)        # guidance_weight 默认 5.0
v = v - g * pinv_correction

软掩码 W 由 _create_soft_mask(psi0.py:2377)按 RTC 论文 Eq.5 生成,三段式:

$$ W_i = \begin{cases} 1 & i < d \quad\text{(冻结区,已执行)}\[2pt] c_i \cdot \dfrac{e^{c_i}-1}{e-1},\ \ c_i = \dfrac{H-s-i}{H-s-d+1} & d \le i < H-s \quad\text{(过渡区)}\[2pt] 0 & i \ge H-s \quad\text{(自由区)} \end{cases} $$

可选 exponential(默认)/ linear / hard / simple 四种调度。约束 d ≤ s ≤ H - d 在入口处强校验(psi0.py:1897)。

4.3 部署约束#

src/psi/deploy/psi0_serve_real_sonic.py:72-81:

self.Ta = action_exec_horizon or launch_config.model.action_exec_horizon
assert self.Ta <= self.Tp, "action_exec_horizon is too big"
if enable_rtc:
    assert launch_config.model.rtc, "rtc is not supported for this model"
    self.rtc_max_delay = launch_config.model.max_delay
    assert self.Tp - self.Ta <= self.rtc_max_delay

即推理延迟被定义为 d = Tp - Ta(预测长度减执行长度),必须落在训练时见过的 [0, max_delay) 区间内。服务端把上一次预测的剩余部分拼成 prev_actions 传给模型,形成滚动衔接。

服务与客户端脚本:

4.4 System-0:SONIC 全身控制器#

Ψ₀ 输出的是全身动作指令,下肢部分由 SONIC(NVlabs GR00T-WholeBodyControl)的 RL 跟踪控制器执行。接入 SONIC 后动作维度从 36 涨到 78、状态维度从 36 到 43,数据经 scripts/data/raw_sonic_to_psi_lerobot.py 转成 Ψ₀ 的 LeRobot 格式,训练走 SonicTrainer(src/psi/trainers/sonic.py)。SONIC 侧细节见 GR00T-WholeBodyControl 总览。

SIMPLE 仿真评测按数据来源分两个入口:遥操作采的任务(*Teleop)用 eval_decoupled_wbc.py + psi0_decoupled_wbc agent,CuRobo 运动规划生成的任务(*MP)用 eval.py + psi0 agent。

5. 关键超参数表#

模型结构(Psi0ModelConfig,src/psi/config/model_psi0.py)#

Flow Matching#

RTC#

学习率与优化#

6. 关键源文件表#

7. 数据与评测#

数据集#

EgoDex 的动作预处理复用了 H-RDT 的流程(src/h_rdt/datasets/pretrain/)。

评测环境:SIMPLE#

SIMPLE 是本项目配套的人形基准仿真器,MuJoCo 负责物理、Isaac Sim 负责渲染。提供 level-0/1/2 三档域随机化。因为 IsaacSim 用同步渲染 API,单个 episode 评测需要 6~10 分钟。

仓库内基线#

Ψ₀ 仓库把 8 个对比方法都收进来共享同一套 src/(通过 uv 分组管理独立虚拟环境):

发布的 Checkpoint#

消融组的命名本身就说明了论文的核心实验:拆开验证"人类视频预训练"这一步到底贡献了多少。

8. 与其他 VLA 的对照#

Ψ₀ 最有辨识度的两点:一是把 FAST 离散预训练 + flow matching 后训练拆成两个完全独立的阶段(而不是像 π₀-FAST 那样二选一),二是把 RTC 从推理技巧变成训练目标 —— 逐 token timestep 这个改动只有三处几行代码,但让实时衔接从"引导近似"变成了"分布内推理"。