用扩散模型精修决策变换器,提升低质量轨迹下的策略性能。
DRDT3: Diffusion-Refined Decision Test-Time Training Model
- 结合TTT层与自注意力机制,捕捉上下文信息进行粗略动作预测。
- 通过迭代扩散生成,逐步优化动作预测,使结果更接近最优策略。
- 适合在数据质量差的离线强化学习任务中使用,尤其对轨迹不完美场景有效。
决策变换器(DT)在多个经典控制任务中表现优异,但难以从次优且带奖励标签的轨迹中学习最优策略。本文探索利用条件生成建模实现高质量轨迹拼接。受近年来递归神经网络(RNN)线性复杂度和序列建模能力的启发,我们引入测试时训练(TTT)层,在测试阶段更新隐藏状态以建模轨迹。提出统一框架DRDT3,超越传统DT模型。核心为决策TTT(DT3)模块,融合自注意力与TTT层优势,捕捉近期上下文并生成粗略动作预测。随后,通过生成式扩散模型迭代精修动作,逐步逼近最优动作。采用统一优化目标联合集成DT3与扩散模型。在D4RL基准的多任务实验中,无扩散精修的DT3已优于标准DT,而完整版DRDT3进一步超越当前最先进的基于DT及离线强化学习方法。
原文摘要 · Abstract (English)
Decision Transformer (DT), a trajectory modelling method, has shown competitive performance compared to traditional offline reinforcement learning (RL) approaches on various classic control tasks. However, it struggles to learn optimal policies from suboptimal, reward-labelled trajectories. In this study, we explore the use of conditional generative modelling to facilitate trajectory stitching given its high-quality data generation ability. Additionally, recent advancements in Recurrent Neural Networks (RNNs) have shown their linear complexity and competitive sequence modelling performance over Transformers. We leverage the Test-Time Training (TTT) layer, an RNN that updates hidden states during testing, to model trajectories in the form of DT. We introduce a unified framework, called Diffusion-Refined Decision TTT (DRDT3), to achieve performance beyond DT models. Specifically, we propose the Decision TTT (DT3) module, which harnesses the sequence modelling strengths of both self-attention and the TTT layer to capture recent contextual information and make coarse action predictions. DRDT3 iteratively refines the coarse action predictions through the generative diffusion model, progressively moving closer to the optimal actions. We further integrate DT3 with the diffusion model using a unified optimization objective. With experiments on multiple tasks in the D4RL benchmark, our DT3 model without diffusion refinement demonstrates improved performance over standard DT, while DRDT3 further achieves superior results compared to state-of-the-art DT-based and offline RL methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。