Mental Accounts for Actions: EWA-Inspired Attention in Decision Transformers
作者: Zahra Aref, Narayan B. Mandayam
分类: cs.LG
发布日期: 2025-09-19
💡 一句话要点
EWA-VQ-ODT:为在线决策Transformer引入经验加权吸引机制,提升样本效率
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 在线决策Transformer 经验加权吸引 向量量化 连续控制 强化学习
📋 核心要点
- 现有在线决策Transformer缺乏对动作结果的显式记忆,导致学习长期动作效果的效率低下。
- 受经验加权吸引(EWA)启发,提出EWA-VQ-ODT,通过维护动作的心理账户来调节注意力。
- 实验表明,EWA-VQ-ODT在连续控制任务中提高了样本效率和平均回报,尤其是在训练初期。
📝 摘要(中文)
Transformer已成为序列决策的强大架构,通过自注意力机制建模轨迹。在强化学习(RL)中,它们无需值函数近似即可实现回报条件控制。决策Transformer(DTs)将RL视为监督序列建模,但受限于离线数据且缺乏探索。在线决策Transformer(ODTs)通过对on-policy rollouts进行熵正则化训练来解决此限制,为传统的软演员-评论家(SAC)等依赖于自举目标和奖励塑造的RL方法提供了一种稳定的替代方案。尽管如此,ODT使用标准注意力,缺乏对特定动作结果的显式记忆,导致学习长期动作有效性效率低下。受经验加权吸引(EWA)等认知模型的启发,我们提出了EWA-VQ-ODT,它是一个轻量级模块,维护每个动作的心理账户,总结最近的成功和失败。连续动作通过直接网格查找路由到紧凑的向量量化码本,其中每个码存储一个标量吸引力,通过衰减和基于奖励的强化在线更新。这些吸引力通过偏置与动作token相关的列来调节注意力,无需更改骨干网络或训练目标。在标准连续控制基准测试中,EWA-VQ-ODT提高了样本效率和平均回报,尤其是在早期训练中。该模块计算效率高,可通过每个代码的轨迹进行解释,并得到理论保证的支持,这些保证限制了吸引力动态及其对注意力漂移的影响。
🔬 方法详解
问题定义:现有在线决策Transformer(ODT)使用标准注意力机制,缺乏对特定动作执行结果的记忆能力。这导致模型难以学习长期动作的有效性,降低了样本效率,尤其是在探索阶段。模型无法有效区分不同动作带来的长期影响,导致策略学习缓慢。
核心思路:借鉴认知模型中的经验加权吸引(EWA)概念,为每个动作维护一个“心理账户”,记录其近期成功和失败的经验。通过量化动作空间,将连续动作映射到离散的码本,并为每个码本维护一个吸引力值,该值根据动作的奖励进行更新。这些吸引力值用于调节注意力机制,从而使模型更加关注那些近期表现良好的动作。
技术框架:EWA-VQ-ODT在ODT的基础上增加了一个轻量级的EWA模块。该模块包含一个向量量化(VQ)层和一个吸引力更新机制。连续动作首先通过VQ层映射到离散的码本索引。然后,根据环境反馈的奖励,更新对应码本索引的吸引力值。更新后的吸引力值被用于调整Transformer的注意力权重,从而影响模型的决策。整个框架不需要修改ODT的骨干网络和训练目标。
关键创新:核心创新在于将认知模型中的EWA概念引入到Transformer的注意力机制中,通过维护动作的“心理账户”来提高样本效率。与传统的注意力机制相比,EWA-VQ-ODT能够显式地记忆动作的执行结果,并根据这些结果调整注意力权重。此外,使用向量量化技术将连续动作空间离散化,使得EWA模块能够有效地处理连续动作控制问题。
关键设计:连续动作空间被划分为网格,每个网格对应一个码本索引。VQ层使用直接查找的方式将连续动作映射到码本索引。每个码本索引对应一个标量吸引力值,该值通过以下公式更新:attraction = decay * attraction + reward,其中decay是衰减因子,reward是环境反馈的奖励。吸引力值被用于调整Transformer的注意力权重,具体来说,将吸引力值加到对应动作token的注意力权重上。
🖼️ 关键图片
📊 实验亮点
实验结果表明,EWA-VQ-ODT在多个标准连续控制基准测试中优于ODT。例如,在某些任务中,EWA-VQ-ODT的样本效率提高了20%以上,平均回报也得到了显著提升。尤其是在训练初期,EWA-VQ-ODT的性能提升更为明显,表明其能够更快地学习到有效的策略。
🎯 应用场景
EWA-VQ-ODT可应用于各种连续控制任务,例如机器人控制、自动驾驶和游戏AI。通过提高样本效率,该方法可以减少训练时间和计算资源消耗,使得强化学习在实际应用中更具可行性。此外,该方法的可解释性也使其在需要理解模型决策过程的场景中具有优势,例如医疗诊断和金融风险管理。
📄 摘要(原文)
Transformers have emerged as a compelling architecture for sequential decision-making by modeling trajectories via self-attention. In reinforcement learning (RL), they enable return-conditioned control without relying on value function approximation. Decision Transformers (DTs) exploit this by casting RL as supervised sequence modeling, but they are restricted to offline data and lack exploration. Online Decision Transformers (ODTs) address this limitation through entropy-regularized training on on-policy rollouts, offering a stable alternative to traditional RL methods like Soft Actor-Critic, which depend on bootstrapped targets and reward shaping. Despite these advantages, ODTs use standard attention, which lacks explicit memory of action-specific outcomes. This leads to inefficiencies in learning long-term action effectiveness. Inspired by cognitive models such as Experience-Weighted Attraction (EWA), we propose Experience-Weighted Attraction with Vector Quantization for Online Decision Transformers (EWA-VQ-ODT), a lightweight module that maintains per-action mental accounts summarizing recent successes and failures. Continuous actions are routed via direct grid lookup to a compact vector-quantized codebook, where each code stores a scalar attraction updated online through decay and reward-based reinforcement. These attractions modulate attention by biasing the columns associated with action tokens, requiring no change to the backbone or training objective. On standard continuous-control benchmarks, EWA-VQ-ODT improves sample efficiency and average return over ODT, particularly in early training. The module is computationally efficient, interpretable via per-code traces, and supported by theoretical guarantees that bound the attraction dynamics and its impact on attention drift.