跳转至

Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

会议: NeurIPS 2026
arXiv: 2610.10857
论文: https://arxiv.org/abs/2610.10857
代码: https://keyframe-mnemonics.github.io
领域: 自监督/表示学习
关键词: 行为克隆, 自监督关键帧发现, 视界不变性, 简单性偏差, 非马尔可夫决策

一句话总结

提出无需人工标注的自监督关键帧记忆框架 Keyframe Mnemonics,利用神经网络简单性偏差与掩码采样从历史中筛选动作决策关键帧,赋予行为克隆策略跨越数个数量级超长时域的视界不变泛化能力。

研究背景与动机

在实际物理环境与机器人控制任务中,决策系统往往面临严重的部分可观测性与非马尔可夫动态,智能体必须在时间跨度上保留关键历史线索以明确当前最优动作。现有的行为克隆(BC)框架大多依赖近期的短期观测窗口(通常仅几秒),一旦面对延迟决策或长时序记忆任务便极易崩溃。为了赋予策略长程记忆能力,主流方案通常采用循环神经网络(通过 BPTT 训练)或 Transformer 注意力机制。然而,循环网络在长时间反向传播时极易发生梯度不稳和隐状态坍缩(Hidden-State Collapse),截断 BPTT 又会丢失远期线索;而 Transformer 则受到二次计算复杂度与有限上下文窗口的硬性约束,简单扩展窗口在长时域下不仅性能严重退化,还极易过拟合视觉伪相关。现有结合外部记忆矩阵或循环摘要的混合记忆架构,在超出训练时域时同样难逃状态重写和记忆模糊的厄运。

上述困境的深层矛盾在于:策略记忆机制通常对整条观测序列进行无差别编码与循环更新,导致与决策无关的冗余背景帧不断冲刷并模糊了决定性线索;而理想的决策上下文应当是稀疏且具有因果不变性的。先前尝试引入关键帧(Keyframe)条件的研究,无一例外依赖人工逐帧真值标注、启发式规则或昂贵的大模型推理提示词,缺乏通用的统计学习驱动机制。

本文另辟蹊径,将长时程部分可观测决策问题解耦为两个独立子任务:一是完全自监督地从演示中发现决定动作分布的关键帧,二是以所提关键帧为上下文条件训练轻量级行为克隆策略。核心 idea:利用深度神经网络固有的简单性偏差与组合掩码采样,构建出偏好关键帧的代理预测损失作为内在奖励,进而通过强化学习训练马尔可夫选择器策略抽取关键帧助记符,使行为克隆在有限工作记忆下获得无限时域不变泛化保证。

方法详解

整体框架

Keyframe Mnemonics(KM)的核心机制建立在两项结构假设之上:轨迹观测可划分为对置信状态具有近似充分性的稀疏关键帧子集 \(\mathcal{X}_{\mathrm{key}}\)(大小 \(C \ll T\))与冗余的背景帧 \(\mathcal{X}_{\mathrm{bg}}\),且关键帧对动作分布的互信息显著高于背景帧(\(I(a_t; \mathcal{X}_{\mathrm{key},t}) > I(a_t; \mathcal{X}_{\mathrm{bg},t})\))。基于此,KM 将学习过程解耦为三阶段串行流水线: 1. 阶段 1:掩码采样代理策略训练。对历史观测序列施加随机二值掩码,将抽样得到的受限上下文输入代理网络预测专家动作,利用组合爆炸效应放大关键帧的互信息优势; 2. 阶段 2:选择器策略强化学习优化。冻结代理预测损失作为内在奖励函数,在确定性离散 MDP 下使用 PPO 训练一个仅依赖单步观测的选择器策略,为观测帧赋予优先级并动态维护有限优先队列; 3. 阶段 3:关键帧增强的行为克隆策略训练。利用训练好的选择器遍历演示轨迹提取关键帧助记符,构建增强数据集,训练最终以关键帧和当前观测为联合条件的 BC 策略。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["离线专家演示轨迹 D<br/>包含非马尔可夫观测历史与动作"] --> B["掩码采样训练代理策略<br/>随机二值掩码抽样组合观测并预测专家动作"]
    B --> C["冻结代理损失目标 J_theta*<br/>构建偏好高互信息关键帧的内在奖励函数"]
    C --> D["PPO 优化马尔可夫选择器策略<br/>单步观测输出优先级并维护训练端 FIFO 队列"]
    D --> E["LIFO 优先队列生成关键帧助记符 B_t<br/>离线构建关键帧增强数据集 D_dagger"]
    E --> F["助记符条件行为克隆策略<br/>联合当前帧与助记符实现超长时域稳健决策"]

关键设计

1. 掩码采样代理训练与简单性偏差放大:将关键帧发现转化为自监督目标优化 为了在没有关键帧标注的情况下量化哪些观测对动作具有决定性影响,本文将关键帧发现形式化为受限容量内存的缓冲子集优化问题。在阶段 1 中,代理策略 \(f_\theta: \mathcal{B} \to \mathcal{A}\) 接收容量至多为 \(k\) 的观测缓冲 \(\mathcal{B}_t\) 并预测专家动作 \(a_t\)。为了构建训练缓冲,系统对历史序列 \(h_t\) 施加二值掩码 \(\Psi_k\),其中采样大小服从 \(n \sim \min(\mathrm{Binomial}(t, 0.5), k)\) 的无放回均匀采样。根据曲棍球棒恒等式(Hockey-stick Identity),全数据集中诱导出的所有可能样本对总数达到组合级规模: $\(M = \sum_{i=1}^{N} \left( T_i + \sum_{j=1}^{k} \binom{T_i + 1}{j + 1} \right)\)$ 基于 Kolmogorov 复杂度理论,以背景帧为条件的假设需要拟合巨大的样本支持集,其程序描述长度随着支持集规模 \(M\) 线性剧增;相比之下,仅依赖因果关键帧的最小描述长度假设则保持紧凑。深度神经网络在随机梯度下降(SGD)下天生具备偏向更低复杂度的简单性偏差(Simplicity Bias)。当支持集规模 \(M\) 足够大时,SGD 收敛到基于关键帧的不变假设类 \(\mathcal{H}_{\mathrm{in}}\) 的概率呈指数级逼近 1: $\(P_{\mathrm{SGD}}(\mathcal{H}_{\mathrm{in}} \mid \mathcal{D}) \approx \frac{1}{1 + 2^{- [M (I(A; \mathcal{X}_{\mathrm{key}}) - I(A; \mathcal{X}_{\mathrm{bg}})) - \mathcal{O}(1)]}}\)$ 因此,训练完成后冻结的代理预测负对数似然损失 \(J_{\theta^*}(\mathcal{B}) = \mathbb{E}[\mathcal{L}(f_{\theta^*}(\mathcal{B}), A)]\) 便成为天然的关键帧评分器:缓冲内包含的关键帧越多,预测损失越小。

2. 基于确定性 MDP 与双阶段队列的选择器策略学习:实现高效信用分配与无限保持 由于评估候选子集的空间呈组合爆炸,直接在离散子集上进行搜索无法适用于连续在线决策。本文将关键帧抽取建模为一个在演示轨迹上转移的确定性离散 MDP,状态为当前帧 \(o_t\),动作是由轻量级选择器网络 \(\pi_{\mathrm{s}}(p_t \mid o_t)\) 输出的标量优先级 \(p_t \in [0, 1]\)。该设计具有两大精妙机制: 首先是惊奇加权的对数内在奖励机制,环境根据更新后的缓冲输入冻结代理计算损失并返回奖励: $\(R(\mathcal{B}_t) = T_i \cdot \frac{e^{J_{\theta^*}^{j=t} / \kappa}}{\sum_{j=1}^{T_i} e^{J_{\theta^*}^j / \kappa}} \cdot \frac{\log(J_{\theta^*}(\mathcal{B}_t) + \lambda)}{\log(\lambda)}\)$ 归一化对数变换在目标值接近 0 时提供尖锐梯度,而基于时步期望损失的惊奇系数(Surprise Weighting)则在决策高度不确定、严重依赖记忆的历史时刻放大奖励权重,引导强化学习集中攻克关键时步。 其次是训练与推理非对称的优先队列淘汰机制。在阶段 2 训练选择器时,缓冲采用先进先出(FIFO)结合拒绝阈值 \(\eta_{\mathrm{r}} = 0.5\) 的淘汰规则:若优先级打平,优先逐出最早入队的旧帧。这强制选择器不能投机取巧地输出全 1 饱和优先级,必须对真正关键的帧输出极致区分度的高分。而在阶段 3 及在线部署时,缓冲转为后进先出(LIFO)并引入邻近优先级去重阈值 \(\eta_{\mathrm{n}} = 0.05\):一旦初始关键帧以高分占据队列,后续平分的背景帧只能在末端自我淘汰,从而在数学上严格保证了关键帧在任意无限长时域下永不被逐出。

3. 助记符增强的行为克隆策略训练与闭环执行:消除因果混淆并保障时域不变性 在阶段 3 中,冻结的选择器策略以 LIFO 规则在整个演示轨迹集上离线运行,将每个时间步生成的稳定关键帧集合 \(\mathcal{B}_t\) 与当前观测 \(o_t\) 及专家动作 \(a_t\) 拼装,构建出助记符增强数据集 \(\mathcal{D}^\dagger = \{(\mathcal{B}_t, o_t, a_t)\}\)。最终的策略模型 \(\pi_{\mathrm{BC}}(a_t \mid o_t, \mathcal{B}_t)\) 既能通过紧凑的 \(\mathcal{B}_t\) 彻底消除因果混淆、获得足以判别动作模态的置信状态,又通过保留当前观测 \(o_t\) 捕获环境局部连续变化的动力学细节。部署时,选择器和 BC 策略完全保持马尔可夫特性(仅依赖单步输入),避免了任何 BPTT 递归或全局注意力缓存,使得单步推理时间与内存开销恒定为 \(\mathcal{O}(1)\),彻底与执行时域长度解耦。

损失函数 / 训练策略

  • 阶段 1 代理策略:在图像任务中使用 ResNet-18 视觉主干与 Transformer 编码器,以在线掩码抽样训练 60 个 epoch,优化器为 AdamW,学习率 \(1 \times 10^{-4}\),对于离散动作采用负对数似然(NLL)分类损失,连续控制任务则对动作块(Chunking 长度 \(H=8\))使用 L1 回归损失。
  • 阶段 2 选择器策略:采用 ResNet-18 与轻量两层 MLP(隐藏层维度 500),通过 Stable-Baselines3 的 PPO 算法与 32 个并行轨迹环境进行 \(5 \times 10^6\) 步交互训练。折扣因子设为 \(\gamma=0.99\),熵正则系数 0.01,裁剪范围 0.2,PPO epoch 为 5。
  • 阶段 3 BC 策略:默认采用 4 层单步回归 Transformer(SST),参数量约 62.13M,训练 15,000 步;高频动态抓取任务(如 Intercept 系列)则采用 4 层流匹配交叉注意力 Transformer(FM),步长 \(H=8\),采用欧拉离散流匹配调度器训练 100,000 步,推理仅需 5 步离散欧拉反向积分。训练过程对视觉缓冲注入包含亮度、对比度、仿射旋转在内的强烈数据增强以抵御因果混淆。

实验关键数据

主实验

在涵盖空间记忆、序列记忆、目标记忆的机器人操作基准 MIKASA-Robo(23 个非零任务,每任务 3 个种子各 100 次评测)与合成序列/网格基准上,KM 展现出压倒性优势。

基准 / 任务 评估指标 KM (本文) 最强基线 (ELMUR / TrXL) 提升幅度
MIKASA-Robo (23 任务全集均值) 绝对成功率 (SR ↑) 56.6% 42.7% (ELMUR) +13.9% (17/23 项胜出)
MIKASA Shell Game Push 成功率 (SR ↑) 0.99 ± 0.01 0.87 ± 0.04 (ELMUR) +12.0%
MIKASA Intercept Grab Fast 成功率 (SR ↑) 0.37 ± 0.02 0.20 ± 0.01 (ELMUR) +17.0%
MIKASA Rotate Lenient Pos Neg 成功率 (SR ↑) 0.71 ± 0.06 0.26 ± 0.00 (ELMUR) +45.0%
MIKASA Remember Color3 成功率 (SR ↑) 0.99 ± 0.01 0.89 ± 0.07 (ELMUR) +10.0%
MIKASA Remember Shape Color3x2 成功率 (SR ↑) 0.52 ± 0.07 0.19 ± 0.02 (ELMUR) +33.0%
Add (超长时域 \(h=10^4\)) 最终均方误差 (MSE ↓) 3.7e-5 ± 4.0e-5 0.494 ± 0.320 (LSTM-B) 误差降低 4 个数量级
Scattered-Copy (\(h=10^4\)) 复制成功率 (SR ↑) 1.000 ± 0.000 0.040 ± 0.089 (LSTM-T) +96.0% (基线近乎全崩)
T-Maze (\(h=10^4\)) 决策成功率 (SR ↑) 1.000 ± 0.000 0.840 ± 0.261 (LSTM-B) +16.0% (方差降至 0)
MiniGrid-Memory (\(h=1001\)) 终点成功率 (SR ↑) 1.000 ± 0.000 0.500 ± 0.173 (DLSTM) +50.0%
LTMB-Hallway (\(h=1000\)) 终点成功率 (SR ↑) 0.933 ± 0.058 0.000 ± 0.000 (所有基线) +93.3% (全基线归零)

在极限时域测试中,将 Add、Scattered-Copy 与 T-Maze 的序列长度外推至 \(h=10^5\) 与 \(h=10^6\)(超出训练时域 5 个数量级),KM 仍保持 100% 成功率与 \(10^{-5}\) 级超低 MSE。在真机 RememberColor2 实验中,面对高达 20 倍延迟时间窗(延时从 5 秒拉长至 100 秒),扩散策略(DP)与 PTP 均彻底失效,而 KM 仍稳健维持 80% 的实物抓取成功率。

消融实验

在核心控制任务与延迟保持基准 RememberColor3-\(\Delta T\) 上对核心组件进行系统剥离(训练时线索至动作延迟 \(\Delta T=15\),测试时推至 \(\Delta T=250\))。

配置 / 变体 \(\Delta T=15\) 成功率 (%) \(\Delta T=250\) 成功率 (%) 机制分析说明
Full KM (完整方法) 72.3 ± 5.8 65.3 ± 7.6 表现出完美的视界不变性,长程衰减极微弱
随机选择器 (\(\pi_{\mathrm{rnd}}\)) 替代学习选择器 31.3 ± 3.1 26.7 ± 1.2 缺少因果关键帧,动作分布无法去歧义,全面溃退
贪婪单步选择优化 (\(\gamma=0\)) 100.0 ± 0.0 (仅active T-Maze) 0.0 ± 0.0 (延迟任务) 无法处理远期动作奖励延迟分配,仅在即时任务生效
代理阶段强行输入当前观测 \(o_t\) 18.7 ± 5.9 12.3 ± 4.1 破坏简单性偏差假设,代理过拟合局部捷径,选择器失效
缓冲队列改用 LIFO \(\to\) FIFO 75.0 ± 1.6 57.7 ± 3.8 推理阶段易被后期背景噪点逐出,长时程保持能力受损
缓冲容量减半 (\(0.5\times k\)) 94.7 ± 2.5 40.0 ± 23.2 (T-Maze) 缓冲无法完整容纳任务必需的关键帧,信息存在瓶颈

关键发现

  • 代理训练绝对不可包含当前帧:在 MiniGrid-Memory 与 RememberColor3 中,如果在阶段 1 代理输入中加入了当前观测 \(o_t\),性能分别从 100% 和 99% 暴跌至 33.3% 和 18.7%。这是由于当前帧构成极易拟合的低复杂度局部捷径,破坏了模型对不变因果线索的学习。
  • PPO 折扣因子揭示了关键帧抽取的长程信用分配本质:当 \(\gamma=0\) 时,选择器在绝大多数存在时延的基准上性能直接归零(例如 Add 的 MSE 恶化至 0.184,Scattered-Copy 成功率跌至 0%);只有将 \(\gamma\) 提升至 0.99,才能跨越数十步时滞将关键帧线索与决策终点的预测收益联结起来。
  • 与马尔可夫基础任务零性能倒退(No Regression):在无需记忆的 Robomimic 标准任务(Lift、Can、Square)上,引入 KM 框架的策略与基线 Diffusion Policy 成功率完全持平,证实增加关键帧机制不会损害标准环境中的底层控制质量。

亮点与洞察

  • 将数据采样作为正则化杠杆:在时序维度实施掩码采样子集,通过曲棍球棒恒等式制造组合级超大样本空间,借由神经网络原生简单性偏差自然“逼出”因果关键帧,完全摆脱了对预训练大模型或手工标注的依赖。
  • FIFO 训练与 LIFO 推理的双队列设计:训练时用先进先出规则对饱和高优先级实施“惩罚性溢出”,促使选择器精细调整优先级尺度;推理时切换为后进先出,以极简的规则在无需重调权重的前提下锁死首发关键帧,达成数学上的无限时域保持保证。
  • 将非马尔可夫决策转化为马尔可夫化执行:选择器单步计算 \(p_t\),BC 策略单步输出 \(a_t\),整个推理系统完全无状态、无循环反向传播、无显存平方爆炸,却能完美掌控长程历史上下文。

局限与展望

  • 单帧可分性假设(Assumption 1)的严苛限制:当前框架要求关键帧能够在单个观测中被清晰识别。若任务依赖多帧连续动作或时空光流聚合才能构成因果线索,当前的单帧马尔可夫选择器将无法有效表征。
  • 掩码采样在超多关键帧下的组合概率衰减:当一项任务需要同时记住 5 个以上分散的关键帧时,随机掩码在单次采样中同时命中所有关键帧的概率呈指数下降,导致代理策略在极高容量记忆任务(如 ChainOfColors 系列)中训练收敛极慢。
  • GPU 原生强化学习工程瓶颈:由于阶段 2 选择器训练依赖 CPU-GPU 频繁通信的离散环境交互,训练耗时占据了全流程的 50% 以上,后续需迁移至 GPU 原生并行强化学习库(如 Isaac Gym / CleanRL)。

相关工作与启发

  • vs ELMUR / RATE 等显式记忆网络:ELMUR 依靠外挂可重写记忆槽或 Transformer 键值缓存,在时域变长时不可避免地反复重写和稀释记忆表征,一旦测试时延迟拉长便发生灾难性遗忘;KM 物理上独立锁存离散关键帧,从根源上规避了时域伸缩带来的特征漂移。
  • vs PTP (Past-Token Prediction):PTP 采用启发式历史子采样方案作为视觉扩散策略的上下文,当延迟超出采样回溯窗口时策略立即失效;KM 采用自监督强化学习显式挖掘信息极大化的关键帧,保持能力超脱于物理采样步长之外。
  • vs 基于 VLM 提示词的关键帧提取:现有结合大视觉模型的方法需为每项机器人任务手工编写 Prompts 并在推理时产生高昂延迟与 API 成本;KM 则是纯粹从离线示教中涌现特征,轻量紧凑,推理开销低至毫秒级。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从简单性偏差与信息论角度严格论证并实现了完全自监督的关键帧发现框架,构思精巧。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 5 大合成/网格时序基准、23 项 MIKASA 真实仿真任务、真机闭环与 Robomimic 对照,横跨 6 个数量级时域外推。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论证明完备,从引理、定理到算法伪代码逻辑链条极其严密清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为长时序具身智能与长视界行为克隆提供了彻底摆脱记忆坍缩与算力瓶颈的全新范式。