TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning

📄 arXiv: 2605.12236v1 📥 PDF

作者: Matthew M. Hong, Jesse Zhang, Anusha Nagabandi, Abhishek Gupta

分类: cs.RO, cs.AI, cs.LG

发布日期: 2026-05-12

🔗 代码/项目: PROJECT_PAGE


💡 一句话要点

TMRL:扩散时间步调制预训练,提升探索能力,实现高效策略微调

🎯 匹配领域: 支柱一:机器人控制 (Robot Control) 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 强化学习 机器人策略 预训练 微调 扩散模型 探索 时间步调制

📋 核心要点

  1. 行为克隆预训练的策略动作分布窄,限制了强化学习微调阶段的探索能力,是现有方法的主要瓶颈。
  2. 论文提出上下文平滑预训练(CSP)和时间步调制强化学习(TMRL),通过注入噪声和动态调整扩散时间步来控制探索。
  3. 实验表明,TMRL显著提升了强化学习微调的样本效率,并在真实机器人操纵任务中实现了快速微调。

📝 摘要(中文)

本文提出了一种统一的框架,通过桥接行为克隆(BC)预训练和强化学习(RL)微调,实现了高效机器人策略微调所需的探索能力。该方法旨在解决BC预训练产生的狭窄动作分布,这种分布缺乏下游探索所需的覆盖范围。论文提出的上下文平滑预训练(CSP)方法,将前向扩散噪声注入策略输入,从而在精确模仿和广泛动作覆盖之间建立连续性。然后,通过时间步调制强化学习(TMRL)微调预训练策略,该方法训练智能体在微调期间动态调整这种条件,通过调制扩散时间步来显式控制探索。TMRL可以无缝集成到任意策略输入中,例如状态、3D点云或基于图像的VLA策略。实验结果表明,TMRL提高了RL微调的样本效率,并能够在不到一小时的时间内成功地在真实世界中微调复杂的操纵任务。

🔬 方法详解

问题定义:现有方法在预训练机器人策略时,通常采用行为克隆(BC)。然而,BC训练出的策略往往具有狭窄的动作分布,缺乏足够的探索能力,这严重限制了后续强化学习(RL)微调的效率和性能。因此,如何提升预训练策略的探索能力,使其能够更好地适应RL微调,是本文要解决的核心问题。

核心思路:本文的核心思路是通过在预训练阶段引入噪声,并允许在微调阶段动态调整噪声水平,从而实现对策略探索能力的精确控制。具体来说,通过扩散模型引入噪声,使得策略既能进行精确模仿,又能进行更广泛的探索。在微调阶段,通过强化学习来学习如何调整扩散过程的时间步,从而动态地控制探索的程度。

技术框架:该方法包含两个主要阶段:上下文平滑预训练(CSP)和时间步调制强化学习(TMRL)。在CSP阶段,将前向扩散噪声注入策略的输入,生成具有不同噪声水平的训练数据。在TMRL阶段,使用强化学习算法(如SAC)来微调预训练的策略,同时训练一个额外的网络来预测扩散时间步,该时间步用于控制策略的探索程度。整体流程是先使用CSP进行预训练,然后使用TMRL进行微调。

关键创新:最重要的技术创新点在于时间步调制强化学习(TMRL)的概念,它允许智能体在微调过程中动态地调整扩散时间步,从而显式地控制探索。与传统的探索方法(如ε-greedy或高斯噪声)相比,TMRL能够更精细地控制探索的程度,并根据任务的需要进行自适应调整。

关键设计:CSP阶段的关键设计在于如何选择合适的扩散模型和噪声水平。TMRL阶段的关键设计在于如何设计奖励函数和状态表示,以及如何训练时间步预测网络。论文中使用了SAC算法作为基础的强化学习算法,并设计了一个简单的神经网络来预测扩散时间步。损失函数包括标准的SAC损失和时间步预测损失。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,TMRL在多个机器人操纵任务上显著提高了RL微调的样本效率。例如,在真实世界的机器人操纵任务中,TMRL能够在不到一小时的时间内成功地微调策略,而传统的RL方法则需要更长的时间或者无法成功。与基线方法相比,TMRL在样本效率和最终性能方面均有显著提升。

🎯 应用场景

该研究成果可广泛应用于机器人操纵、自动驾驶、游戏AI等领域。通过预训练和微调相结合的方式,可以显著降低训练成本,提高策略的泛化能力和适应性。该方法尤其适用于需要在真实环境中进行微调的机器人任务,例如家庭服务机器人、工业机器人等。未来,该方法有望推动机器人技术的进一步发展,使其能够更好地服务于人类。

📄 摘要(原文)

Fine-tuning pre-trained robot policies with reinforcement learning (RL) often inherits the bottlenecks introduced by pre-training with behavioral cloning (BC), which produces narrow action distributions that lack the coverage necessary for downstream exploration. We present a unified framework that enables the exploration necessary to enable efficient robot policy finetuning by bridging BC pre-training and RL fine-tuning. Our pre-training method, Context-Smoothed Pre-training (CSP), injects forward-diffusion noise into policy inputs, creating a continuum between precise imitation and broad action coverage. We then fine-tune pre-trained policies via Timestep-Modulated Reinforcement Learning (TMRL), which trains the agent to dynamically adjust this conditioning during fine-tuning by modulating the diffusion timestep, granting explicit control over exploration. Integrating seamlessly with arbitrary policy inputs, e.g., states, 3D point clouds, or image-based VLA policies, we show that TMRL improves RL fine-tuning sample efficiency. Notably, TMRL enables successful real-world fine-tuning on complex manipulation tasks in under one hour. Videos and code available at https://weirdlabuw.github.io/tmrl/.