聚合掩码自编码在离线强化学习中的未来预测与策略优化

本文提出聚合掩码未来预测（AMFP）框架，通过整合加权聚合与掩码自编码（MAE），利用全局融合令牌增强轨迹重建与状态-动作耦合表示，有效抑制随机动作噪声，提升离线强化学习（RL）策略的鲁棒性与性能。实验在D4RL基准验证其竞争力。

　　

亮点

• (1)
  我们提出聚合掩码未来预测（AMFP），一种基于序列建模的新型自监督框架，用于离线强化学习（RL）。
• (2)
  我们设计了一种新的预训练任务，整合加权聚合与掩码自编码（MAE）以重建状态轨迹，增强表征中的核心信息有效性，并降低随机动作引发的错误编码风险。
• (3)
  我们引入全局融合令牌，促进聚合与掩码自编码的融合，同时编码状态序列表征和隐式状态-动作耦合。
• (4)
  所提出的AMFP在D4RL基准上进行了评估，与现有离线RL基线相比展现出高度竞争力。

自监督学习中的掩码技术

近年来，自监督学习已成为深度学习的一个重要研究方向。在各种自监督学习技术中，掩码（Masking）被广泛采用为一种有效方法。在自然语言处理领域，BERT采用掩码语言建模方案，随机掩盖某些令牌，并要求模型预测它们。BERT及其改进版本均展现出卓越性能。

预训练

离线RL方法旨在从预先收集的数据集中学习策略，无需在训练期间与真实环境交互。在实际应用中，这些数据集通常通过多种策略获得，例如专家演示、随机采样或两者结合。当离线RL被构建为条件序列建模问题时，数据集中长度为T的轨迹样本τ通常表示为τ = (C₁, s₁, a₁, C₂, s₂, a₂, …, C_t, s_t, a_t, …, C_T, s_T, a_T)，其中s_t ∈

方法论

在我们的工作中，我们采用名为AMFP的两阶段训练框架。第4.1节介绍预训练阶段——聚合掩码自编码，旨在学习核心表征。第4.2节描述用于下游策略学习的微调阶段。

实验

在本节中，我们探讨了AMFP在解决离线决策任务中的有效性。实验在11个D4RL数据集上进行，涵盖三个运动域场景和一个Franka Kitchen域场景。所有实验均在配备Intel Core i7-12700KF处理器和NVIDIA RTX 4090 GPU的计算机上使用PyTorch框架实现。我们的代码公开在https://github.com/monolithycq/AMFP。

结论

在这项工作中，我们提出了聚合掩码未来预测（AMFP），一种用于离线RL的新型自监督学习框架。AMFP引入了一种新的预训练任务，通过全局融合令牌整合加权聚合与掩码自编码。该设计使模型能够选择性地强调可靠的状态转换并抑制由次优动作引起的误导性转换，同时保留隐式的动作-状态依赖关系。D4RL基准上的实验结果证明了该方法的有效性。