pi0.5

1
WebsocketClientPolicy.infer()  | 📂 packages/openpi-client/src/openpi_client/websocket_client_policy.py (仅跨进程调用时存在)
2
├── pack() (序列化观测数据)
3
├── 🌐 [跨进程通信: websocket send / recv]
4
└── WebsocketPolicyServer._handler()  | 📂 src/openpi/serving/websocket_policy_server.py
5
    ├── unpackb() (反序列化网络接收到的数据)
6
    └── Policy.infer()  | 📂 src/openpi/policies/policy.py (本地推断的默认起点)
7
        ├── _input_transform() (执行输入格式注入、数据清洗与归一化)
8
        ├── _sample_actions() (核心模型调用，被 JAX JIT 或 PyTorch eval 包装)
9
        │   └── Pi0.sample_actions()  | 📂 src/openpi/models/pi0.py
10
        │       ├── preprocess_observation()  | 📂 src/openpi/models/model.py
11
        │       ├── embed_prefix()  | 📂 src/openpi/models/pi0.py
12
        │       │   ├── PaliGemma.img() (通过 SigLIP 编码多视角图像特征)
13
        │       │   └── PaliGemma.llm() (通过 Gemma 编码文本指令 Prompt)
14
        │       ├── PaliGemma.llm() (Prefix 前向传播，计算并填充全注意力 KV Cache)
15
        │       └── jax.lax.while_loop() (Euler 降噪积分循环，默认迭代 10 步)
95 collapsed lines
16
        │           ├── embed_suffix() (组合当前动作向量与时间步特征)
17
        │           │   ├── action_in_proj() (将当前动作向量映射为特征 Token)
18
        │           │   ├── posemb_sincos() (计算时间步的 Sin-Cos 位置编码)
19
        │           │   └── time_mlp_in() / time_mlp_out() (生成供 π0.5 AdaRMS 使用的时间条件特征)
20
        │           ├── PaliGemma.llm() (结合 KV Cache 进行 Suffix 前向传播)
21
        │           ├── action_out_proj() (投影网络输出得到预测的速度向量 v_t)
22
        │           └── x_t + dt * v_t (利用 Euler 步长积分更新去噪动作 x_t)
23
        └── _output_transform() (执行结果后处理与反归一化)
24

25
torchrun ... scripts/train_pytorch.py / uv run scripts/train_pytorch.py  | 📂 /Users/julyfun/Documents/GitHub/openpi/scripts/train_pytorch.py (启动 PyTorch 训练)
26
├── [torchrun 模式] 启动 N 个 Python 训练进程 (跨进程：每 GPU 一个 rank)
27
main() (训练入口)
28
├── init_logging() (设置日志)
29
├── _config.cli()  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/training/config.py (tyro 解析 TrainConfig)
30
│   └── overridable_config_cli(...) (选择如 pi05_libero/debug，并应用 CLI override)
31
└── train_loop(config)  | 📂 /Users/julyfun/Documents/GitHub/openpi/scripts/train_pytorch.py (PyTorch 主训练流程)
32
    ├── setup_ddp() (初始化分布式)
33
    │   ├── 读取 WORLD_SIZE / LOCAL_RANK / RANK (判断是否 torchrun 多进程)
34
    │   ├── if use_ddp: torch.distributed.init_process_group(...) (跨进程：NCCL/Gloo 建立通信组)
35
    │   └── torch.cuda.set_device(device) (当前 rank 绑定 GPU)
36
    ├── set_seed(config.seed, local_rank) (设置 torch/numpy 随机种子)
37
    ├── 处理 checkpoint_dir / resume / overwrite (创建或复用实验目录)
38
    ├── if is_main: init_wandb(...) (只在主进程初始化 W&B)
39
    ├── build_datasets(config) (构造数据加载器)
40
    │   └── _data.create_data_loader(config, framework="pytorch", shuffle=True)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/training/data_loader.py
41
    │       ├── config.data.create(config.assets_dirs, config.model)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/training/config.py (生成 DataConfig)
42
    │       │   ├── LeRobotLiberoDataConfig.create(...) (LIBERO 路径；配置 repack/data/model transforms)
43
    │       │   │   ├── create_base_config(...) (加载 norm_stats，设置 repo_id/asset_id)
44
    │       │   │   ├── RepackTransform(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/transforms.py (LeRobot 字段重映射)
45
    │       │   │   ├── LiberoInputs(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/policies/libero_policy.py (转成模型输入格式)
46
    │       │   │   └── ModelTransformFactory.__call__()  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/training/config.py (resize/tokenize/pad)
47
    │       │   └── RLDSDroidDataConfig.create(...) (DROID RLDS 配置；但 PyTorch 下 RLDS loader 目前会 NotImplemented)
48
    │       └── create_torch_data_loader(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/training/data_loader.py (PyTorch 训练实际数据路径)
49
    │           ├── create_torch_dataset(...) (创建 LeRobotDataset / FakeDataset)
50
    │           ├── transform_dataset(...) (串联数据变换)
51
    │           │   └── CompositeTransform.__call__()  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/transforms.py
52
    │           │       ├── RepackTransform.__call__() (重排字段)
53
    │           │       ├── LiberoInputs.__call__()  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/policies/libero_policy.py (构造 image/state/actions/prompt)
54
    │           │       ├── Normalize.__call__()  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/transforms.py (归一化 state/actions)
55
    │           │       ├── ResizeImages.__call__() (图像 resize 到 224)
56
    │           │       ├── TokenizePrompt.__call__() / TokenizeFASTInputs.__call__() (prompt/token/action token)
57
    │           │       └── PadStatesAndActions.__call__() (补齐 action_dim)
58
    │           ├── if torch.distributed.is_initialized(): DistributedSampler(...) (跨进程：各 rank 分片取数据)
59
    │           └── TorchDataLoader.__iter__() (循环产出 batch；num_workers>0 时 spawn DataLoader worker 进程)
60
    │               └── DataLoaderImpl.__iter__() (dict -> Observation, actions)
61
    │                   └── Observation.from_dict(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models/model.py (uint8 图像转 [-1,1]，PyTorch 图像转 NCHW)
62
    ├── if is_main: 取 sample_batch 并 wandb.log(camera_views)  | 📂 /Users/julyfun/Documents/GitHub/openpi/scripts/train_pytorch.py
63
    ├── 构造 model_cfg (把 config.model 转成 Pi0Config，设置 dtype/action_dim/horizon)
64
    ├── PI0Pytorch(model_cfg).to(device)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/pi0_pytorch.py (创建 PyTorch pi0/pi05)
65
    │   ├── PaliGemmaWithExpertModel(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/gemma_pytorch.py (封装 PaliGemma VLM + Gemma action expert)
66
    │   ├── Linear(action_in_proj/action_out_proj/time_mlp/...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/pi0_pytorch.py
67
    │   └── 检查 transformers_replace 是否正确安装
68
    ├── model.gradient_checkpointing_enable() (开启显存优化)
69
    ├── if use_ddp: DistributedDataParallel(model, ...) (跨进程：反向传播时同步梯度)
70
    ├── if config.pytorch_weight_path: safetensors.torch.load_model(...) (加载已转换 PyTorch 权重)
71
    ├── torch.optim.AdamW(...) (构造优化器)
72
    ├── if resuming: load_checkpoint(...) (恢复 model/optimizer/global_step)
73
    ├── lr_schedule(step) (warmup + cosine decay)
74
    └── while global_step < config.num_train_steps (训练外循环)
75
        ├── if use_ddp and hasattr(loader, "set_epoch"): loader.set_epoch(...) (DDP shuffle epoch)
76
        └── for observation, actions in loader (训练 batch 循环)
77
            ├── observation/actions.to(device) (搬到当前 GPU)
78
            ├── 更新 optim.param_groups 的 lr
79
            ├── losses = model(observation, actions)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/pi0_pytorch.py
80
            │   └── PI0Pytorch.forward(...) (完整前向并返回逐元素 MSE loss)
81
            │       ├── _preprocess_observation(observation, train=True)
82
            │       │   └── preprocess_observation_pytorch(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/preprocessing_pytorch.py (resize/增强/mask)
83
            │       ├── sample_noise(...) / sample_time(...) (采样 flow matching 噪声和时间)
84
            │       ├── 构造 x_t = t*noise + (1-t)*actions, u_t = noise-actions
85
            │       ├── embed_prefix(images, masks, lang_tokens, lang_masks)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/pi0_pytorch.py
86
            │       │   ├── PaliGemmaWithExpertModel.embed_image(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/gemma_pytorch.py (SigLIP 视觉特征)
87
            │       │   └── PaliGemmaWithExpertModel.embed_language_tokens(...) (语言 token embedding)
88
            │       ├── embed_suffix(state, x_t, time)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/pi0_pytorch.py (action/time/state token；pi05 用 adaRMS 条件)
89
            │       ├── make_att_2d_masks(...) (构造 attention mask)
90
            │       ├── _prepare_attention_masks_4d(...) (转成 transformer 4D mask)
91
            │       ├── PaliGemmaWithExpertModel.forward(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/gemma_pytorch.py
92
            │       │   ├── for layer_idx in range(num_layers) (逐层联合处理 prefix VLM 与 suffix action expert)
93
            │       │   │   └── compute_layer_complete(...) (QKV 拼接、RoPE、attention、MLP、残差；可 checkpoint)
94
            │       │   └── compute_final_norms(...) (分别归一化 prefix/suffix 输出)
95
            │       ├── action_out_proj(suffix_out)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/models_pytorch/pi0_pytorch.py (预测 v_t)
96
            │       └── F.mse_loss(u_t, v_t, reduction="none") (flow matching loss)
97
            ├── loss = losses.mean()
98
            ├── loss.backward() (反向传播；DDP 时跨进程 all-reduce 梯度)
99
            ├── clip_grad_norm_(model.parameters(), ...) (梯度裁剪)
100
            ├── optim.step(); optim.zero_grad(...) (参数更新)
101
            ├── if global_step % log_interval == 0: wandb.log(...) (主进程记录 loss/lr/grad_norm)
102
            ├── global_step += 1
103
            ├── save_checkpoint(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/scripts/train_pytorch.py (按 save_interval 保存)
104
            │   ├── if not is_main: return (非主进程不保存)
105
            │   ├── safetensors.torch.save_model(...) (保存 model.safetensors)
106
            │   ├── torch.save(optimizer.state_dict()) (保存 optimizer.pt)
107
            │   ├── torch.save(metadata) (保存 metadata.pt)
108
            │   ├── _normalize.save(...)  | 📂 /Users/julyfun/Documents/GitHub/openpi/src/openpi/shared/normalize.py (保存 norm_stats)
109
            │   └── tmp_dir.rename(final_ckpt_dir) (原子替换 checkpoint)
110
            └── pbar.update / set_postfix (更新进度条)

1
# π₀.₅: PaliGemma prefix + Action Expert suffix + flow matching
2
v = SigLIP(imgs)                                    # [B, Nv, D]
3
t = PaliGemma.embed(prompt_with_discrete_state)     # state 在文本里
4
prefix = cat([v, t])                                # 图文双向 (prefix-LM)
5
x_t, t = noise * t + actions * (1-t)                # flow 插值
6
suffix = Linear_in(x_t)                             # [B, Na, D]
7
cond = MLP(sincos(t))                               # adaRMS 条件
8
(h_pre, h_suf) = DualGemma([prefix, suffix], mask=prefix_lm, adarms=[None, cond])
9
loss = MSE(Linear_out(h_suf[:, -Na:]), noise - actions)
10
# infer: cache prefix → Euler: x += dt * Linear_out(h_suf), t -= dt