openpi π₀ 训练与推理 矩阵维度详解

1. 关键维度参数┌─────────────────────────────────────────────────────────────────────────────────┐ │ 关键维度参数 │ ├─────────────────────────────────────────────────────────────────────────────────┤ │ │ │ 【输入相关】 │ │ B = batch_size = 32 # 批次大小 │ │ H = W = 224 # 图像分辨率 │ │ C = 3 # RGB 通道 │ │ num_images = 3 # 相机数量 (base, left_wrist, right_wrist)│ │ L = max_token_len = 200 # 语言最大 token 数 │ │ state_dim = 8 # 状态维度 (DROID: 7关节+1夹爪) │ │ │ │ 【动作相关】 │ │ action_horizon = 15 # 预测未来步数 │ │ action_dim = 32 # 每步动作维度 │ │ │ │ 【SigLIP 视觉编码器】 │ │ patch_size = 14 # patch 大小 │ │ num_patches = (224/14)² = 256 # 每张图的 patch 数 │ │ vision_dim = 1152 # SigLIP 隐藏维度 │ │ │ │ 【PaliGemma (Gemma 2B)】 │ │ pali_width = 2048 # 隐藏维度 │ │ pali_depth = 18 # Transformer 层数 │ │ pali_heads = 8 # 注意力头数 │ │ pali_head_dim = 256 # 每头维度 │ │ pali_mlp_dim = 16384 # FFN 维度 │ │ │ │ 【Action Expert (Gemma 300M)】 │ │ expert_width = 1024 # 隐藏维度 │ │ expert_depth = 18 # Transformer 层数 │ │ expert_heads = 8 # 注意力头数 │ │ expert_head_dim = 256 # 每头维度 │ │ expert_mlp_dim = 4096 # FFN 维度 │ │ │ └─────────────────────────────────────────────────────────────────────────────────┘2. 训练过程 - 矩阵维度流转STEP 1: 输入数据输入 batch: ├── images: │ ├── base_0_rgb: [B, H, W, C] = [32, 224, 224, 3] float32, 值域[-1,1] │ ├── left_wrist_0_rgb: [B, H, W, C] = [32, 224, 224, 3] │ └── right_wrist_0_rgb: [B, H, W, C] = [32, 224, 224, 3] ├── image_masks: │ ├── base_0_rgb: [B] = [32] bool │ ├── left_wrist_0_rgb: [B] = [32] │ └── right_wrist_0_rgb: [B] = [32] ├── state: [B, state_dim] = [32, 8] float32 ├── tokenized_prompt: [B, L] = [32, 200] int32 ├── tokenized_prompt_mask: [B, L] = [32, 200] bool └── actions (真实动作): [B, action_horizon, action_dim] = [32, 15, 32] float32STEP 2: 采样噪声和时间noise = random.normal(shape=actions.shape) = [B, action_horizon, action_dim] = [32, 15, 32] float32 time = random.beta(a=1.5, b=1, shape=[B]) = [32] float32, 值域(0,1) time_expanded = time[..., None, None] = [32, 1, 1] 用于广播STEP 3: 计算插值点 x_t 和目标向量场 u_tx_t = time_expanded * noise + (1 - time_expanded) * actions