GitHub 地址：https://github.com/hzm8341/pi0.6
参考论文：MEM: Multi-Scale Embodied Memory for Vision Language Action Models，Physical Intelligence，2026年3月
一、为什么要做这个项目？
现有的机器人基础模型（如 π0、π0.5）能够执行"洗锅"、"叠衣服"、"做花生酱三明治"等单项任务，但面对连续多步骤的复杂任务------比如清洁整个厨房或从零开始做一顿饭------就会频繁"失忆"：
- 无历史记忆：每次推理只看当前帧，不知道自己做到哪一步了
- 无遮挡补偿：手臂挡住目标物体就找不到了
- 无错误恢复：同样的失败会无限重复，不会调整策略
- 任务时长限制：实验通常只在几分钟的短任务上演示
Physical Intelligence 在 2026 年 3 月发布的 **MEM（Multi-Scale Embodied Memory）**论文提出了双模态记忆架构，将机器人可处理的任务时长从约 2 分钟扩展到 15 分钟以上 ，平均任务完成率从约 30% 提升到 70%。
本项目基于开源的 openpi（π0.5 代码库），完整工程化实现了 π0.6-MEM 架构的所有核心组件，并将代码开源在 GitHub 上，希望帮助更多研究者和工程师快速上手。
二、MEM 核心思想：双模态分层记忆
MEM 的核心洞察是：机器人的记忆需求天然地分为两个时间尺度，需要用不同的模态来表达。
短期记忆（秒级） 长期记忆（分钟级）
───────────────────────────────── ──────────────────────────────────
• 需要精细空间信息 • 只需要语义级信息
• 处理遮挡（手臂挡住物体） • 跟踪任务进度（做到第几步了）
• 调整抓取姿势 • 记住已放好的食材
• 表示形式：图像序列（视频） • 表示形式：自然语言摘要
───────────────────────────────── ──────────────────────────────────关键数学分解（来自论文 Section III）：
把策略分解为高层策略 πHL 和低层策略 πLL：
π(a_{t:t+H}, l_{t+1}, m_{t+1} | o_{t-T:t}, m_t, g)
≈ πLL(a_{t:t+H} | o_{t-K:t}, l_{t+1}, g) ← 低层：生成动作
× πHL(l_{t+1}, m_{t+1} | o_t, m_t, g) ← 高层：子任务 + 更新记忆其中：
a= 动作序列，l= 子任务指令，m= 语言记忆，g= 任务目标K << T：低层策略只需要最近 K 帧，大幅降低计算量
三、项目整体架构
openpi-mem/
├── src/openpi/
│ ├── models/
│ │ ├── model.py ← ✅ 扩展 Observation（新增 MEM 字段）
│ │ ├── siglip.py ← ✅ 新增 SpaceTimeSeparableBlock 视频编码器
│ │ ├── pi0_config.py ← ✅ 新增 MEMConfig 配置类
│ │ ├── pi0.py ← ✅ embed_prefix 多模态记忆注入
│ │ ├── memory_manager.py ← ✅ 语言记忆标注生成与管理
│ │ └── high_level_policy.py← ✅ 高层策略（子任务 + 记忆更新）
│ ├── transforms.py ← ✅ VideoFrameStack / TokenizeMemory
│ ├── training/
│ │ └── data_loader.py ← ✅ MEMLeRobotDataset 多帧数据加载
│ └── policies/
│ └── policy.py ← ✅ MEMPolicy 推理接口（帧缓存）
└── scripts/
└── gen_memory_labels.py ← ✅ 离线语言记忆标注生成脚本改动统计：新增/修改 15 个文件，约 2668 行核心代码，以及 5 个测试文件（129 项测试全部通过）。
四、核心实现详解
4.1 短期视觉记忆：SpaceTimeSeparableBlock
这是整个架构最核心的创新------零参数量增加地让 SigLIP ViT 具备视频理解能力。
设计思路 ：在 ViT 每隔 4 层的位置，将原来的纯空间注意力替换为时空分离注意力------先做空间注意力（patch 之间），再对每个 patch 的不同时间步做因果时间注意力（历史帧之间）。
原始 Encoder1DBlock 的计算复杂度：O(n²)（n = patch 数）
时空联合注意力的复杂度：O(n²K²)（K = 帧数）⟶ 不可接受
SpaceTimeSeparableBlock：O(K·n² + n·K²) ⟶ 可接受！关键代码 （src/openpi/models/siglip.py）：
class SpaceTimeSeparableBlock(nn.Module):
"""
时空分离注意力 Block（MEM 核心）
关键特性：
- 不引入新参数（复用原 QKV 权重）
- K=1 时数值等价于原 Encoder1DBlock（零时间编码保证）
- 计算复杂度 O(K·n² + n·K²) vs 联合注意力 O(n²·K²)
"""
num_timesteps: int = 1 # K
@nn.compact
def __call__(self, x, deterministic=True):
K = self.num_timesteps
BK, N, D = x.shape
B = BK // K
# Step 1: 标准空间注意力（与原 Encoder1DBlock 完全相同）
x = self._spatial_attention(x, deterministic)
# Step 2: 因果时间注意力（MEM 核心）
if K > 1:
# 维度变换：(B*K, N, D) → (B*N, K, D)
x_t = x.reshape(B, K, N, D).transpose(0, 2, 1, 3).reshape(B * N, K, D)
# 时间位置编码（t=0 时值为 0，保证 K=1 等价性）
x_t = x_t + self._sinusoidal_time_emb(K, D)[None, :, :]
# 因果掩码：时刻 t 只能 attend 到 t' ≤ t
causal_mask = jnp.tril(jnp.ones((K, K), dtype=jnp.bool_))[None, :, :]
# 时间维度注意力
x_t = x_t + self._temporal_attn(x_t, causal_mask, deterministic)
# 维度还原：(B*N, K, D) → (B*K, N, D)
x = x_t.reshape(B, N, K, D).transpose(0, 2, 1, 3).reshape(B * K, N, D)
return x, {}
def _sinusoidal_time_emb(self, K, D):
"""时间位置编码，t=0 时恒为零（保证 K=1 数值等价）"""
positions = jnp.arange(K, dtype=jnp.float32)
half_d = D // 2
freqs = jnp.exp(-jnp.log(10000.0) * jnp.arange(half_d) / (half_d - 1))
args = positions[:, None] * freqs[None, :]
emb = jnp.concatenate([jnp.sin(args), jnp.cos(args)], axis=-1)
return emb - emb[0:1, :] # shift 使 t=0 → 零向量历史 Token 丢弃机制 ：在 ViT 上层（后 4 层），只保留当前帧的 token，输出形状从 (B*K, N, D) 还原为 (B, N, D)，推理 token 数量与单帧完全一致：
# Encoder.__call__ 中的 drop_history 逻辑
if use_video and lyr == drop_after:
BK, N, D = x.shape
B = BK // K
x = x.reshape(B, K, N, D)[:, -1, :, :] # 只取最后一帧（当前帧）
use_video = False # 后续层正常处理 B 帧4.2 长期语言记忆：MemoryDataGenerator
语言记忆是一段持续压缩更新的自然语言摘要，论文中最关键的发现是：
失败的子任务绝对不能更新记忆------否则会产生严重的训练-推理分布偏移。
# src/openpi/models/memory_manager.py
class MemoryDataGenerator:
"""离线生成语言记忆训练标注"""
def generate_labels_for_episode(self, episode_id, subtasks):
labels = []
current_memory = ""
for t, subtask in enumerate(subtasks):
# 只用成功步骤构建历史
successful = [
f"Step {i+1}: {s['instruction']}"
for i, s in enumerate(subtasks[:t])
if s.get("success", True) # ← 关键：只记录成功的步骤
]
new_memory = self._call_llm(successful) if successful else ""
labels.append(MemoryLabel(
memory_before=current_memory,
memory_after=new_memory,
...
))
# ← 核心：失败时不更新记忆！防止分布偏移
if subtask.get("success", True):
current_memory = new_memory
return labels为什么 Naive 方法（直接拼接历史指令）不行？因为训练数据里每个子任务基本只出现一次（人类演示接近最优），但推理时失败的子任务会反复出现，导致高层策略"懵了"。MEM 的压缩记忆机制完美规避了这个问题。
4.3 高层策略：HighLevelPolicy
高层策略以低频（每个子任务结束时）运行，使用 VLM 的语言生成能力，通过 chain-of-thought JSON 格式同时输出子任务指令和记忆更新：
# src/openpi/models/high_level_policy.py
_HL_PROMPT_TEMPLATE = """
Task goal: {task_goal}
Memory of completed steps:
{language_memory}
Based on current observation, output JSON:
{{
"subtask": "<next concrete robot action>",
"updated_memory": "<compressed past-tense summary>"
}}
"""
class HighLevelPolicy:
def update(self, observation_image, subtask_success=True):
raw = self._vlm(observation_image, prompt, max_tokens)
parsed = json.loads(raw.strip())
new_subtask = parsed["subtask"]
new_memory = parsed["updated_memory"]
# ← 核心约束：失败时保持原记忆不变
if subtask_success:
self._language_memory = new_memory
return new_subtask, self._language_memory4.4 模型配置：MEMConfig
所有 MEM 参数通过独立的 MEMConfig 管理，默认值全部关闭，与原 π0.5 完全向后兼容：
# src/openpi/models/pi0_config.py
@dataclasses.dataclass(frozen=True)
class MEMConfig:
# 短期视觉记忆
use_video_memory: bool = False # 默认关闭 → π0.5 行为
video_memory_frames: int = 6 # K 帧（预训练用 6，微调可到 18）
video_frame_stride_sec: float = 1.0
temporal_attn_every_n_layers: int = 4
drop_history_tokens_after_layer: int = -4
# 长期语言记忆
use_language_memory: bool = False
max_memory_tokens: int = 256
# 本体感觉历史嵌入
use_state_history: bool = False
state_history_frames: int = 6
# 启用 MEM 的配置示例：
config = Pi0Config(
pi05=True,
mem=MEMConfig(
use_video_memory=True,
video_memory_frames=6,
use_language_memory=True,
use_state_history=True,
)
)4.5 前向传播：embed_prefix 四段 Token
MEM 扩展了 embed_prefix，按以下顺序拼接 prefix token：
[语言记忆 tokens] → [图像 tokens × N_cams] → [语言 prompt tokens] → [状态历史 tokens]
↑ ↑ ↑ ↑
长期记忆上下文 当前帧/视频融合 任务描述 连续嵌入（非文本）
# src/openpi/models/pi0.py embed_prefix()
# (A) 长期语言记忆
if self.use_language_memory and obs.tokenized_memory is not None:
memory_tokens = self.PaliGemma.llm(obs.tokenized_memory, method="embed")
tokens.append(memory_tokens)
# (B) 图像 tokens（视频编码器融合历史帧）
for name in obs.images:
if self.use_video_memory and obs.image_history is not None:
image_tokens = self._encode_video_frames(
obs.image_history[name], # (B, K-1, H, W, C)
obs.images[name] # (B, H, W, C)
)
else:
image_tokens, _ = self.PaliGemma.img(obs.images[name], train=False)
tokens.append(image_tokens)
# (C) 语言 prompt（原有逻辑不变）
if obs.tokenized_prompt is not None:
tokens.append(self.PaliGemma.llm(obs.tokenized_prompt, method="embed"))
# (D) 本体感觉历史（连续线性投影，不用文本 token）
if self.use_state_history and obs.state_history is not None:
state_hist_tokens = self.state_history_proj(obs.state_history) # (B, K-1, D)
tokens.append(state_hist_tokens)4.6 推理接口：MEMPolicy
MEMPolicy 在推理阶段自动维护滑动帧缓存，完全向下兼容原 Policy 接口：
# src/openpi/policies/policy.py
class MEMPolicy(Policy):
def __init__(self, model, *, num_video_frames=6, ...):
# 滑动帧缓存（双端队列，maxlen = K-1）
self._frame_buf = {
cam: deque(maxlen=K-1) for cam in camera_keys
}
self._language_memory = ""
def reset_episode(self, task_goal=""):
"""新 episode 开始：清空所有缓存"""
for buf in self._frame_buf.values():
buf.clear()
self._language_memory = ""
def infer(self, obs, *, noise=None):
# 1. 从 deque 构建 image_history
obs_aug = self._build_obs_with_history(obs)
# 2. 注入语言记忆 token
if self._language_memory:
ids, mask = self._hl.tokenize_memory(self._language_memory)
obs_aug["tokenized_memory"] = ids
obs_aug["tokenized_memory_mask"] = mask
# 3. 调用底层推理
result = super().infer(obs_aug, noise=noise)
# 4. 推理完成后更新帧缓存（重要：推理后更新）
self._push_to_buffers(obs)
# 5. 按需触发高层策略更新
if self._hl and self._hl.should_update():
_, new_mem = self._hl.update(obs["base_0_rgb"])
self._language_memory = new_mem
return result五、快速上手
5.1 克隆项目
git clone https://github.com/hzm8341/pi0.6.git
cd pi0.6
pip install -e .5.2 启用 MEM 的最小配置
from openpi.models.pi0_config import Pi0Config, MEMConfig
# 只启用短期视觉记忆（不需要语言记忆标注）
config = Pi0Config(
pi05=True,
mem=MEMConfig(
use_video_memory=True,
video_memory_frames=6, # 6 帧历史（5秒@1Hz）
)
)
# 完整 MEM（需要预先生成语言记忆标注）
config_full = Pi0Config(
pi05=True,
mem=MEMConfig(
use_video_memory=True,
video_memory_frames=6,
use_language_memory=True,
max_memory_tokens=256,
use_state_history=True,
)
)5.3 推理（使用 MEMPolicy）
from openpi.policies.policy import MEMPolicy
from openpi.models.high_level_policy import HighLevelPolicy, HighLevelPolicyConfig
# （可选）创建高层策略
hl = HighLevelPolicy(
vlm_inference_fn=your_vlm_fn, # 你的 VLM 推理函数
tokenizer=your_tokenizer,
config=HighLevelPolicyConfig(subtask_trigger_steps=50),
)
# 创建 MEM 策略
policy = MEMPolicy(
model=model,
high_level_policy=hl,
num_video_frames=6,
)
# Episode 循环
policy.reset_episode(task_goal="Clean the kitchen")
for step in range(max_steps):
obs = env.get_observation()
result = policy.infer(obs) # 自动维护帧缓存和记忆状态
env.step(result["actions"])5.4 生成语言记忆训练标注
# 准备 subtask_annotations.jsonl（见文档格式说明）
python scripts/gen_memory_labels.py \
--dataset_path /path/to/your/dataset \
--output_path /path/to/memory_labels.jsonl \
--api_key YOUR_ANTHROPIC_API_KEY六、测试验证
项目包含 129 项测试，覆盖所有核心模块：
# 向后兼容性测试（无需 GPU）
pytest src/openpi/models/test_mem_observation.py -v
# 记忆管理器测试
pytest src/openpi/models/test_memory_manager.py -v
# 数据变换测试
pytest src/openpi/transforms_test_mem.py -v
# 帧缓存推理接口测试
pytest src/openpi/policies/test_policy_mem.py -v
# 视频编码器核心逻辑测试（需要 JAX/GPU）
pytest src/openpi/models/test_video_encoder.py -v主要测试覆盖点：
- K=1 时视频编码器数值等价于原 ViT（误差 < 1e-5）
- 时间信息确实从历史帧流向当前帧 token
- 因果掩码不泄露未来信息
- 失败子任务不更新语言记忆
- 帧缓存滑动窗口正确性（顺序、padding、maxlen）
七、关键实现细节与踩坑
7.1 K=1 数值等价性保证
最精妙的设计之一是时间位置编码的 shift 技巧：
def _sinusoidal_time_emb(self, K, D):
emb = ... # 标准正弦编码
emb = emb - emb[0:1, :] # ← 减去 t=0 的值
# 效果：emb[0] 恒为零向量
# K=1 时：时间注意力的 Q/K/V 都加了零偏移 → 退化为 self-attention → 等价于无时间注意力这保证了可以直接从 π0.5 的 checkpoint 热启动，无需重新初始化。
7.2 scan=False 的选择
π0.5 的 SigLIP ViT 使用 scan=True（JAX scan 加速）。视频编码器需要按层条件判断是否插入时间注意力，这与 scan 的"所有步骤相同函数"约束冲突，所以视频编码器强制 scan=False：
img = _siglip.Module(
scan=(K == 1), # K=1 用原始 scan，K>1 切换到逐层 non-scan
num_timesteps=K,
...
)7.3 论文结论复现的关键数字
| 指标 | 无记忆基线 | 仅视频记忆 | 仅语言记忆 | Naive 文本+视频 | MEM（完整） |
|---|---|---|---|---|---|
| 平均任务进度 | ~30% | ~40% | ~35% | ~33% | ~70% |
| Recipe Setup | 低 | 中 | 低 | 低 | 高 |
| Clean Kitchen | 极低 | 低 | 低 | 低 | 中高 |
Naive 文本+视频 （直接拼接历史指令，不做压缩）性能接近无记忆基线，这验证了记忆压缩机制（失败不更新）的关键性。
7.4 预训练必须包含视频记忆
论文实验：仅在微调阶段引入视频编码器 vs 预训练就加入视频编码器，前者比后者低约 20-30pp。因此训练数据混合中必须包含：
机器人遥操作数据（40%）
视频-语言数据，如视频字幕（20%） ← 关键：非机器人视频
带记忆标注的长任务数据（25%）
人工干预/纠错数据（10%）
单帧数据防遗忘（5%）八、与原 π0.5 的兼容性
所有新字段均为 Optional，默认 None，一行代码不改，行为与 π0.5 完全一致：
# 这段代码在新旧版本行为完全相同
config = Pi0Config(pi05=True) # mem 默认 MEMConfig()，所有 MEM 开关默认 False
model = config.create(rng)
obs = Observation.from_dict(old_data_dict) # 没有 image_history 字段 → 全 None
actions = model.sample_actions(rng, obs) # 完全走 π0.5 路径九、项目文件一览
| 文件 | 改动类型 | 核心内容 |
|---|---|---|
models/model.py |
修改 | 新增 5 个 MEM Optional 字段 |
models/siglip.py |
修改 | SpaceTimeSeparableBlock + 视频 Encoder |
models/pi0_config.py |
修改 | MEMConfig + Pi0Config 属性透传 |
models/pi0.py |
修改 | embed_prefix 4 段 token + _encode_video_frames |
models/memory_manager.py |
新增 | 语言记忆标注生成/管理 |
models/high_level_policy.py |
新增 | πHL：子任务 + 记忆更新 |
transforms.py |
修改 | VideoFrameStack + TokenizeMemory |
training/data_loader.py |
修改 | MEMLeRobotDataset 多帧加载 |
policies/policy.py |
修改 | MEMPolicy 帧缓存推理接口 |
scripts/gen_memory_labels.py |
新增 | 离线标注生成 CLI |
models/test_*.py（3 个） |
新增 | 核心模块测试 |
transforms_test_mem.py |
新增 | 变换测试 |
policies/test_policy_mem.py |
新增 | 推理接口测试 |
十、后续计划
-
\] **预训练权重发布**：基于开源机器人数据集训练 π0.6-MEM 基础权重
-
\] **记忆时间窗口扩展**：从 54 秒（18 帧）进一步扩展到小时级别
参考资料
- Physical Intelligence，MEM: Multi-Scale Embodied Memory for Vision Language Action Models ，2026，arxiv
- Physical Intelligence，π0: A Vision-Language-Action Flow Model for General Robot Control，2024
- Physical Intelligence，π0.5: A Vision-Language-Action Model with Open-World Generalization，2025
- Bertasius et al., Is Space-Time Attention All You Need for Video Understanding?，ICML 2021
- openpi 代码库：https://github.com/Physical-Intelligence/openpi
项目地址 ：https://github.com/hzm8341/pi0.6
欢迎 Star ⭐ / Issue / PR！如果对机器人具身智能、VLA 架构设计感兴趣，欢迎交流讨论。