Investigating Memory in Model-Free RL with POPGym Arcade¶

会议: ICML2026
arXiv: 2503.01450
代码: https://github.com/bolt-research/popgym-arcade
领域: 强化学习 / POMDP / 记忆模型
关键词: model-free RL, POMDP, 记忆模型, recurrent state, 价值涂抹

一句话总结¶

本文指出仅用回报来比较 RL 记忆模型并不可靠，作者构建了一个 GPU 加速的 MDP/POMDP "孪生"基准 POPGym Arcade，并提出 Observability Gap、Memory Bias、像素显著性和 Recall Density 四个工具，借此揭示了一种"价值涂抹（value smearing）"病理：记忆模型会把价值信用错误地分摊到无关的历史观测上，进而导致单个 OOD 观测就能通过 recurrent state 长期污染策略。

研究背景与动机¶

领域现状：在部分可观测决策（POMDP）场景下，主流做法是在策略前接一个记忆模型 \(f\)（RNN/GRU/LRU/Transformer/SSM 等），把历史轨迹 \(\mathbf{x}_t=(o_0,a_0,\dots,o_t)\) 压成一个固定大小的潜 Markov 状态 \(\hat{s}_t\)，再用其上的策略 \(\pi(\cdot\mid\hat{s}_t)\) 交互。评估记忆模型的事实标准就是：在若干 POMDP 任务上比较平均回报。

现有痛点：深度 RL 对模型规模、观测尺寸、任务难度、优化器、随机种子都极度敏感，而记忆模型本身又会引入额外的参数量、优化难度和正则效应。结果是：两种记忆模型在 POMDP 上的回报差异，根本无法判断到底来自"对部分可观测性的缓解"还是来自这些"无关混淆因素"。文献甚至出现了"在 MDP 上加记忆反而更好、在 POMDP 上加记忆反而更差"的悖论现象。

核心矛盾：回报这一标量同时承载了"策略能力"和"记忆能力"，把两者纠缠在一起。要诚实地评估记忆，必须能在保持其它变量不变的前提下，单独度量"加了记忆"和"换成部分观测"各自带来的影响——这就要求一组真正同源的 MDP/POMDP 孪生任务，且共享同一份观测/动作空间，以便复用同一个模型。

本文目标：(1) 构造一个共享 (Ω, A) 的 MDP/POMDP 孪生基准；(2) 给出能把回报"拆解"成可观测性差距与记忆偏差的度量；(3) 给出可视化和量化记忆使用模式的工具；(4) 用这些工具去诊断现有记忆模型实际学到了什么。

切入角度：作者注意到，只要给同一个底层 MDP 套上一个观测函数 \(O\)，就能得到与之配对的 POMDP；如果两者的状态/观测空间在像素层面完全一致，就可以用同一套网络分别训练，差值就自然分离出"部分可观测带来的难度"。在此基础上对 \(Q\) 值或策略对历史观测求梯度，则可量化"哪一帧历史在影响当前决策"。

核心 idea：用 MDP/POMDP 孪生任务把回报分解为 Observability Gap（部分可观测带来的损失）+ Memory Bias（引入记忆带来的副作用），再用 梯度型 Recall Density 衡量记忆实际"回看"了哪些时刻，发现并刻画"价值涂抹"病理。

方法详解¶

整体框架¶

POPGym Arcade 把记忆评估抽象为四件事配合使用：(1) 孪生环境：每个底层任务都同时提供 MDP 变体 \(\mathcal{M}\) 和 POMDP 变体 \(\mathcal{P}\)，二者共用 \(128{\times}128{\times}3\) 或 \(256{\times}256{\times}3\) 像素观测和 \(\{\uparrow,\downarrow,\leftarrow,\rightarrow,\times\}\) 五动作空间；(2) 两类对照度量——Observability Gap 与 Memory Bias，分别用配对 MDP/POMDP、配对"带记忆/不带记忆"策略来隔离不同因素；(3) 像素显著性可视化 \(\lVert\nabla_{o_t}Q\rVert\) 用于单条轨迹直观看"哪些像素被记住了"；(4) Recall Density 把梯度范数沿轨迹归一化并跨轨迹聚合，得到一个"时间→相对影响力"的分布函数，跨模型/任务可比。整套工具在 JAX 中以纯 GPU 流水线实现，吞吐量比 CPU 版 Atari 快约 \(10^4\) 倍。

关键设计¶

Observability Gap 与 Memory Bias（双因子分解）:
- 功能：把"加了记忆的策略在 POMDP 上拿到 \(X\) 分回报"这一标量拆解为两个量纲相同、可直接比较的诊断信号。
- 核心思路：固定记忆模型 \(f\) 和策略 \(\pi\)，分别在孪生 MDP 上和 POMDP 上跑出回报 \(J(f,\pi,\mathcal{M})\) 与 \(J(f,\pi,\mathcal{P})\)，二者之差即 \(\text{Gap}(f,\pi,\mathcal{M},\mathcal{P})=J(f,\pi,\mathcal{M})-J(f,\pi,\mathcal{P})\)，刻画"\(f\) 没能完全把轨迹还原成 Markov 状态"造成的损失；再固定 POMDP 的底层 MDP，比较"带记忆"和"不带记忆"两个策略在 MDP 上的回报差 \(\text{Bias}(f,\pi,\mathcal{M})=J(f,\pi,\mathcal{M})-J(\pi,\mathcal{M})\)，捕捉参数量、优化难度、隐式正则等"与可观测性无关"的副作用。两者与回报同尺度，可直接对比量级。
- 设计动机：以往只看回报时，无法判断"GRU 比 Transformer 好 5 分"到底是因为 GRU 更会推断状态、还是优化更容易、还是参数量更合适。Gap/Bias 同时给出"该模型对付部分观测的能力"和"该模型自带的附加代价"两条独立曲线，文中实验中 MinGRU 与 GRU 之间的 Bias 差异（0.05）就和 Gap 差异（0.05）量级相当，足以颠倒回报排名。
像素显著性 + Recall Density（梯度型记忆使用度量）:
- 功能：在轨迹层面回答"当前决策到底用了哪几帧历史"。
- 核心思路：给定轨迹 \(\mathbf{x}_n\)，按 Eq.1 推出 \(\hat{s}_0,\dots,\hat{s}_n\)，然后对历史观测求 \(Q\) 或 \(\pi\) 的输入梯度，\(\sum_{a_n}\lVert\nabla_{o_t}Q(\hat{s}_n,a_n)\rVert_2^2=\sum_{a_n}\lVert\frac{\partial Q}{\partial \hat{s}_n}\frac{\partial \hat{s}_n}{\partial o_t}\rVert_2^2\)，在像素层面叠成热力图（链式法则同时穿过 CNN 和记忆模型）。为避免"挑樱桃"，再对单条轨迹的 \(L_1\) 梯度范数做轨迹内归一化得到经验密度 \(\delta_Q(\mathbf{x}_n,t)\)，把绝对时刻 \(t\) 映射到归一化时间 \(\tau=t/n\in[0,1]\)，跨多条轨迹取均值，得到 Recall Density \(\mathbb{E}_{\pi,f}[\delta_Q(\mathbf{x},\tau)]\)，并提供 \(\pi\) 梯度版本以兼容策略梯度方法。
- 设计动机：仅看回报或单帧显著性都不足以判断"模型有没有把记忆用对地方"；Recall Density 提供了一个跨轨迹长度可比、跨模型可比的量化曲线，可直接对照"MDP 下理论上密度应集中在 \(\tau\to 1\)"这一 ground truth 做异常检测，是后续发现"价值涂抹"病理的关键测量手段。
POPGym Arcade 孪生环境与 JAX 流水线:
- 功能：为前两类工具提供"硬件加速且形式上可比"的实验底座。
- 核心思路：将每个任务的状态拆为低维隐 Markov 状态 \(\tilde{s}\in\tilde{S}\)（如 MineSweeper 的雷位）和像素 Markov 状态 \(s\in S\)（如带数字提示的棋盘像素），二者都满足 Markov 性，但只要套上观测函数 \(O:\tilde{S}\mapsto\Delta(\Omega)\) 即可生成 POMDP 孪生；所有任务统一到同一 \(S=\Omega\) 与同一动作空间，由此可以单一网络跨任务复用，甚至可以"训练中途从 POMDP 切回 MDP"以做对照实验。10 个基础环境 × 12 种难度/观测组合给出 120 个任务，并标注 Reward Memory Length 是 \(O(k)\) 还是 \(O(n)\)，方便区分"窗口帧叠加就能解"与"必须真正记忆"两类难度。
- 设计动机：要让 Gap/Bias 度量在统计上有意义，必须既保证 MDP 与 POMDP 共享相同的最优可达回报上限，又能高吞吐地跑多种子多模型；纯 JAX + GPU 让作者得以在 7 种记忆模型 × 5 种子 × 120 配置上完成完整 sweep，这是该研究"统计置信"的基础设施前提。

损失函数 / 训练策略¶

主算法采用片上 TD(\(\lambda\)) 的 Q-learning 实现 PQN（Gallici et al., 2024），刻意避开 target network、replay buffer、共享主干等常见混淆因素；同时在附录中用 PPO 与 DQN 复现关键结论以排除算法效应。所有记忆模型都加了一条绕过记忆的 skip connection，让策略在 MDP 上有"忽略记忆"的能力——这一设计也使得后面观察到的"记忆仍涂抹历史"更具说服力。共测试 7 种记忆模型：Transformer、Recurrent Linear Transformer、Linear TTT、Gated DeltaNet、MinGRU、GRU、LRU SSM。

实验关键数据¶

主实验¶

评测维度	工具	关键发现
跨 7 种记忆模型在所有任务上聚合	Return / Gap / Bias 三联图（Fig. 5）	Memory Bias 在不同模型间差异显著且全部为负；MinGRU 与 GRU 的 Bias 差（0.05）与 Gap 差（0.05）量级相同——回报排名完全可以被 Bias 颠倒
BattleShip / MineSweeper 上 sweep 层数 \(L\) 与隐藏维 \(H\)	Gap–Bias Pareto 前沿（Fig. 6）	层数 \(L\uparrow\) 通常恶化 Bias；隐藏维 \(H\uparrow\) 通常改善 Gap；两条曲线构成 Pareto 前沿，可用于挑容量
MDP 上的像素显著性 + Recall Density	Fig. 3、Fig. 7	理论上 MDP 的 \(V^*(s_t)\) 与 \(s_{t-k},\dots,s_{t-1}\) 无关，密度应集中在 \(\tau\in[0.66,1)\)；实测却在 \(\tau<0.66\) 区段获得显著权重——所有模型、所有任务均如此，即"价值涂抹"
OOD 注入实验	单帧噪声（Fig. 9） + 轨迹前缀打乱（Fig. 10）	仅注入一帧 OOD 观测就能让 LRU 策略的相对 \(Q\) 值与贪心动作发生显著扰动；打乱轨迹前缀（排除 CNN 混淆）后，效应在 BattleShip/MineSweeper（LRU）和 CartPole（Transformer）上依然成立——recurrent state 被 OOD 污染并将策略扰动延伸至远期

消融实验¶

配置	关键指标	说明
Full GRU / LRU 完整模型	训练曲线低方差收敛（Fig. 8）	排除"价值涂抹是优化不稳定的伪影"这一备选解释
加 skip connection（策略可绕过记忆）	仍出现 smear	说明策略并未在 MDP 下选择"忽略记忆"，记忆模型确实学到了把信用涂抹到无关过去的解
Transformer（无 recurrent state）	在轨迹前缀打乱时仍受影响	表明 OOD 污染并非 RNN 特有，是部分可观测下记忆-策略联合解的共性问题
同实验在 PPO/DQN 复现（附录）	现象一致	排除算法（on-policy 价值法）特异性

关键发现¶

价值涂抹是普遍现象：MDP 下 \(V\) 理应只依赖当前状态，但 7 种记忆模型在 10 个任务下的 Recall Density 都把大量权重摊到轨迹前 2/3，说明记忆–价值联合优化倾向于把"恰好出现在该轨迹的过去"当成解释变量，存在对当前策略下轨迹分布的过拟合。
回报不可信，必须看 Bias：如果只看回报，会得到"加记忆就更好/更差"的相互矛盾结论；Bias 直接揭示记忆模型在没有可观测性问题时仍带来净负效应，意味着部分既有 SOTA 比较存在隐藏混淆。
OOD 污染是涂抹的实践代价：因为价值被涂抹到无关历史，单个异常观测足以经由 recurrent state 长期改变策略，对真实世界部署/离线 RL 等场景构成实质风险。
可解释干预：Gap 大就加大隐藏维 \(H\)（缓解状态混淆），Bias 负得多就降低层数 \(L\)（缓解优化难度），给出了一个具体可执行的调参方向。

亮点与洞察¶

"孪生 + 量纲相同的两个差值" 是非常优雅的因果分解：把一个混杂量分解成 Gap 与 Bias 两条同尺度信号，方法学上把 RL 中"参数量/优化器/任务难度"的诸多混淆首次摆到台面上量化讨论，可推广到任何"加一个模块是否真有用"的评测场景。
用 MDP 当 oracle 去检验 POMDP 工具：先在 ground-truth 已知的 MDP 上验证 Recall Density 的预期形态（应集中在末段），再用其异常来定义病理，是非常稳的实验范式，对未来评估其它"记忆/上下文/检索"模块都适用。
价值涂抹病理可迁移到 LLM：作者明确指出，若现代 RLHF 后的 LLM 在长上下文 ICL 任务中也存在类似涂抹，可以解释为何模型对长上下文的"无关插入"高度敏感——这给把 Recall Density 类工具引入 LLM 长上下文诊断打开了想象空间。

局限与展望¶

实验集中在 像素 model-free RL 上，model-based RL（如世界模型）与 RL-finetuned LLM 是否存在同款涂抹尚未验证，作者建议下一步重点测量。
"最佳记忆模型"的结论严重依赖比较轴的选择（这里选了隐藏维 \(H\)），换成参数量或 wall-clock 可能完全改变排名，因此 LRU 的"总体最优"应谨慎理解。
POMDP 下没有 ground-truth credit 分布，目前只能用 MDP 替身论证涂抹现象的存在，未来需要可定量量化 POMDP 中涂抹强度的方法。
价值涂抹的根因（优化难度？过拟合轨迹分布？容量不足？）仍是猜想，需更系统的可控实验来证因果。
当前 Recall Density 以梯度范数为代理，对饱和激活或截断 BPTT 的模型可能低估真实信息流，未来可结合 attention rollout、积分梯度或行为干预实验做交叉验证。
POPGym Arcade 的动作空间统一为 5 个离散动作，对连续控制 POMDP（如部分观测 MuJoCo）的迁移性还需后续工作扩展。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 把 RL 记忆评估从"比回报"升级为"做因果分解 + 梯度可解释 + 病理诊断"，并首次刻画"价值涂抹"现象。
实验充分度: ⭐⭐⭐⭐⭐ 7 种记忆模型 × 10 任务 × 多难度 × 5 种子，并用 PQN/PPO/DQN 三套算法互相印证，结论稳健。
写作质量: ⭐⭐⭐⭐ 论点逻辑链清晰（基准 → 度量 → 病理 → 实践后果），定义形式化；个别图（Fig. 7）信息密度偏大略难读。
价值: ⭐⭐⭐⭐⭐ 重写了"记忆模型如何评估"的方法学，并直接撼动一批"加记忆就更好"的旧结论；JAX 化的孪生基准也是社区可直接复用的资产。