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

01 Motivation

用 RL 微调预训练机器人策略，通常会继承 behavioral cloning (BC) 预训练留下的瓶颈：BC 学到的条件动作分布 p(a|c) 在示范稀疏或分布偏移的区域会“坍缩”，对最优动作赋予接近零的概率。于是在线 rollout 拿不到有效奖励，RL 也就学不动。Fine-tuning a pre-trained robot policy with RL usually inherits the bottleneck left behind by behavioral cloning (BC) pre-training: the conditional action distribution p(a|c) learned by BC “collapses” in regions where demonstrations are sparse or the distribution shifts, assigning near-zero probability to optimal actions. Online rollouts therefore receive no useful reward, and RL cannot make progress.

“In sparsely covered regions and under distribution shift, BC overfits to observed data: the conditional support collapses and BC can assign near-zero probability to optimal actions... how can we pre-train policies that adaptively broaden their action distribution during RL?”

作者的核心洞察是：训练一个能在条件分布 p(a|c) 与边缘分布 p(a) 之间插值的策略。p(a|c) 在熟悉的 context 下给出精确行为，p(a) 提供探索所需的宽泛覆盖；在两者之间插值，就能根据 context 的新颖程度动态平衡探索与利用。The authors' core insight is to train a policy that can interpolate between the conditional distribution p(a|c) and the marginal distribution p(a). p(a|c) yields precise behavior in familiar contexts, while p(a) provides the broad coverage that exploration requires; interpolating between the two dynamically balances exploration against exploitation according to how novel the context is.

TMRL 总览图 — Fig. 1：TMRL 通过在 context 上做“平滑”来连接 BC 预训练与 RL 微调。预训练阶段 Context-Smoothed Pre-training (CSP) 向 context `c` 注入噪声，借助扩散噪声尺度 `σ` 在“尖锐模仿” `p(a|c)` 与“更宽泛的边缘分布” `p(a)` 之间构造连续谱系（下方一行）；相近的 context（如左列的下游任务与 task 1）在更小的 `σ` 下就会重叠，而不相似的 context（task 2）需要更大的 `σ` 才重叠。RL 微调阶段，TMRL 学习动态调节这个条件强度，在 context-conditioned 行为与探索性行为之间插值，以获得更好的探索和适应能力。Fig. 1: TMRL bridges BC pre-training and RL fine-tuning by “smoothing” the context. During pre-training, Context-Smoothed Pre-training (CSP) injects noise into the context `c`, using the diffusion noise scale `σ` to build a continuous spectrum between “sharp imitation” `p(a|c)` and “a broader marginal distribution” `p(a)` (bottom row); similar contexts (such as the downstream task and task 1 in the left column) already overlap at a smaller `σ`, whereas dissimilar contexts (task 2) need a larger `σ` before they overlap. During RL fine-tuning, TMRL learns to modulate this conditioning strength dynamically, interpolating between context-conditioned and exploratory behavior for better exploration and adaptation.

101%state-based OGBench 任务上相对最优 baseline 的整体提升overall improvement over the best baseline on state-based OGBench tasks

200%cube-single 任务上最终性能相对最优 baseline 的提升improvement in final performance over the best baseline on cube-single

2.5×点云输入的灵巧抓取任务上相对 DSRL 的最终成功率final success rate relative to DSRL on the point-cloud dexterous grasping task

<1 hour真实机器人复杂操作任务的 RL 微调用时RL fine-tuning time on a complex real-robot manipulation task

02 Method

方法分两步：先用 Context-Smoothed Pre-training (CSP) 预训练出一个可按噪声等级查询的可控策略；再用 Timestep-Modulated RL (TMRL) 训练一个高层 RL 策略，在微调时动态选择这个噪声等级作为探索-利用的旋钮。The method has two stages: first, Context-Smoothed Pre-training (CSP) pre-trains a steerable policy that can be queried at any noise level; then Timestep-Modulated RL (TMRL) trains a high-level RL policy that dynamically selects this noise level during fine-tuning as an exploration-exploitation dial.

TMRL 方法示意图 — “Timestep-modulated exploration via context smoothing. (Left) During pre-training, a steerable policy p_θ is trained across all noise levels σ by corrupting the context c via the kernel q_σ(c̃|c), producing a policy that can be queried at any conditioning strength during inference. (Right) During RL fine-tuning, TMRL exposes p_θ with a context-noise dial σ as an explicit control variable for the high-level policy π_HL.”

Context-Smoothed Pre-training (CSP)

定义一个腐蚀核 q_σ(c̃|c)，以噪声尺度 σ 向 context 注入噪声，得到 context-smoothed 策略 p_θ,σ(a|c,z) = E_{c̃~q_σ}[p_θ(a|c̃,z)]。当 σ→0 时它退化为原始的条件策略；σ 增大时，被腐蚀的 context 信息量下降，诱导出的动作分布变成对相邻/被混叠 context 上行为的更宽混合——这不是随机噪声，而是从数据集里“借用”相邻 context 的连贯动作片段。具体实现上，作者用类似扩散模型前向过程的方式定义腐蚀核：q_{t_c}(c̃|c) = N(√ᾱ_{t_c}·c, (1-ᾱ_{t_c})I)，噪声等级 σ 由 diffusion timestep t_c 参数化——这也是 “timestep modulated” 名字的由来。训练时，策略被同时训练在所有噪声等级上（把 σ 作为策略的显式输入），这一步和标准模仿学习相比几乎没有增加复杂度，更像是一种数据增强。Define a corruption kernel q_σ(c̃|c) that injects noise into the context at noise scale σ, giving the context-smoothed policy p_θ,σ(a|c,z) = E_{c̃~q_σ}[p_θ(a|c̃,z)]. As σ→0 it degenerates into the original conditional policy; as σ grows, the corrupted context carries less information and the induced action distribution becomes a broader mixture over behaviors from neighboring or aliased contexts — this is not random noise, but coherent action segments “borrowed” from neighboring contexts in the dataset. Concretely, the authors define the corruption kernel in the manner of the forward process of a diffusion model: q_{t_c}(c̃|c) = N(√ᾱ_{t_c}·c, (1-ᾱ_{t_c})I), with the noise level σ parameterized by the diffusion timestep t_c — which is where the name “timestep modulated” comes from. During training the policy is trained on all noise levels at once (taking σ as an explicit policy input); compared with standard imitation learning this step adds almost no complexity and behaves more like a form of data augmentation.

Timestep-Modulated RL (TMRL)

在 steering 类算法的框架下（RL 只需选择一个隐变量 z 来控制预训练生成式策略的采样），TMRL 让高层策略 π_HL(z,σ|s) 同时输出 steering 隐变量 z 和 context-smoothing 噪声等级 σ：(z_t,σ_t)~π_HL(·|s_t)，用任意 off-policy actor-critic（如 SAC）优化期望回报。直觉上，σ 是一个探索-利用、动作覆盖的旋钮：大 σ 混叠更多 context（更强探索、更高覆盖），小 σ 收紧条件（找到进展后更强利用）。论文还给出理论分析：context smoothing 会提高不同 context 间动作分布的重叠（Theorem 1），且更近的 context 重叠更多，从而证明平滑后的策略相对基础 BC 策略能保证更高的 coverage 参数 κ（Corollary 1）。Within the framework of steering-style algorithms (RL only has to choose a latent variable z that controls sampling from the pre-trained generative policy), TMRL has the high-level policy π_HL(z,σ|s) output both the steering latent z and the context-smoothing noise level σ: (z_t,σ_t)~π_HL(·|s_t), optimizing expected return with any off-policy actor-critic (such as SAC). Intuitively, σ is a dial for exploration-exploitation and action coverage: a large σ aliases more contexts (stronger exploration, higher coverage), while a small σ tightens the conditioning (stronger exploitation once progress has been found). The paper also provides a theoretical analysis: context smoothing increases the overlap of action distributions across different contexts (Theorem 1), and closer contexts overlap more, proving that the smoothed policy guarantees a higher coverage parameter κ than the base BC policy (Corollary 1).

03 Experiments

作者在 8 个仿真与真实任务上评测 TMRL，覆盖 state-based OGBench 导航/操作、image-based π0 VLA (LIBERO)、点云输入的灵巧抓取（LEAP hand + Franka），以及 WidowX / Franka 上的真实世界 RL 微调，对比 baselines 包括 RLPD、SPiRL、DSRL、PostBC。The authors evaluate TMRL on 8 simulated and real tasks, covering state-based OGBench navigation/manipulation, image-based π0 VLA (LIBERO), point-cloud dexterous grasping (LEAP hand + Franka), and real-world RL fine-tuning on WidowX / Franka, with baselines including RLPD, SPiRL, DSRL and PostBC.

Setting	Baseline 表现Baseline result	TMRL 表现TMRL result	关键结论Key takeaway
OGBench pointmaze-giant / cube-single (state)	最优 baseline（DSRL 等）best baseline (DSRL and others)	整体提升 101%101% overall improvement	cube-single 上提升 200%（最终性能）200% improvement on cube-single (final performance)
LIBERO-90 (π0 / VLM embedding)	DSRL：仅学会重复抓同一物体，成功率接近 0DSRL: only learns to grasp the same object over and over, success rate near 0	TMRL 能学会任务TMRL learns the task	只有 TMRL 达到非平凡成功率only TMRL reaches a non-trivial success rate
灵巧抓取 (LEAP hand, pointcloud)dexterous grasping (LEAP hand, pointcloud)	DSRL	2.5× 最终成功率2.5× the final success rate	混叠点云 context 实现跨物体抓取策略共享aliasing point-cloud contexts lets grasping strategies be shared across objects
真实世界 (WidowX / Franka, π0)real world (WidowX / Franka, π0)	DSRL：三项任务均未学会DSRL: learns none of the three tasks	接近满分成功率near-perfect success rate	1 小时内完成微调fine-tuning completed within 1 hour

仿真任务 RL 成功率曲线 — “RL Success Rates for simulation tasks. TMRL attains near 100% success rate in both OGBench tasks, outperforming the best baselines by 14% in pointmaze-giant and 200% in cube-single at final performance. In libero-goal, TMRL and RLPD both reach 100% success. However, for the longer-horizon libero-90 task, only TMRL explores sufficiently to achieve non-trivial success rates.”

CSP 预训练动作覆盖对比 — “CSP unlocks better action coverage before RL fine-tuning. We measure the Success@K for context-smoothed pre-training against standard BC and PostBC. CSP achieves greater success@K across all K on both tasks.” 在 cube 任务上，标准 BC 与 PostBC 在任意 K 下成功率都是 0，而 CSP 已经能取得非零 success@K，直接证明了它带来的 demonstrator action coverage 提升。On the cube task, standard BC and PostBC both have a success rate of 0 at every K, whereas CSP already attains non-zero success@K, which directly demonstrates the gain in demonstrator action coverage that it brings.

真实机器人上的 RL 微调结果 — “TMRL enables steering of π0 across three real-world tasks, while DSRL fails to learn any task.” 三个真实任务为 sausage-in-pot、shrimp-in-white-drawer、press-button；预训练好的 π0 起初无法完成任务（常抓错物体），TMRL 微调后成功率接近满分，DSRL 因无法超出基础策略的覆盖范围而表现很差。The three real-world tasks are sausage-in-pot, shrimp-in-white-drawer and press-button; the pre-trained π0 initially cannot complete them (it often grasps the wrong object), whereas after TMRL fine-tuning the success rate is near perfect, and DSRL performs poorly because it cannot reach beyond the coverage of the base policy.

Ablations

作者还做了 CFG-RL 消融：用 classifier-free guidance 的插值系数 w 替代 diffusion timestep 作为高层策略的控制变量（记为 TMRL-CFG）。结果“TMRL-CFG fails because the conditioning, which still relies on p(a|c), struggles to extrapolate to OOD contexts no matter the w”，而 TMRL 通过对 context 本身做腐蚀，能把 OOD context “拉回”分布内，在 OOD 场景下依然提供连贯的探索。此外，对比 TMRL 与 DSRL/RLPD 的探索行为可视化显示，RLPD 依赖纯随机动作噪声、DSRL 局限于基础策略条件分布附近，而 TMRL 的探索显著更宽但仍然连贯，是从相邻 context 混叠出来的动作序列而非随机噪声。论文还展示 TMRL 会随轨迹进程动态调节 timestep：例如在 “sausage-in-pot” 任务开始阶段用更大噪声帮助 π0 摆脱其默认偏向拿胡萝卜的行为，物体拿到后再降低噪声以精确放置。The authors also run a CFG-RL ablation: the interpolation coefficient w of classifier-free guidance replaces the diffusion timestep as the high-level policy's control variable (denoted TMRL-CFG). The result is that “TMRL-CFG fails because the conditioning, which still relies on p(a|c), struggles to extrapolate to OOD contexts no matter the w”, whereas TMRL, by corrupting the context itself, can pull an OOD context back into distribution and still provide coherent exploration in OOD situations. In addition, visualizations comparing the exploration behavior of TMRL with DSRL/RLPD show that RLPD relies on purely random action noise and DSRL is confined to the vicinity of the base policy's conditional distribution, while TMRL's exploration is markedly broader yet still coherent — action sequences aliased from neighboring contexts rather than random noise. The paper further shows that TMRL modulates the timestep dynamically as a trajectory unfolds: in the “sausage-in-pot” task, for example, it uses larger noise at the beginning to help π0 escape its default preference for picking up the carrot, then lowers the noise once the object has been picked up so that it can be placed precisely.

04 Limitations

Note: 以下两点均由作者在论文 “Discussion and Limitations” 章节中明确指出（stated），未做额外推断。Both points below are explicitly stated by the authors in the paper's “Discussion and Limitations” section; nothing further has been inferred here.

安全性：混叠动作可能导致不安全行为Safety: aliased actions may lead to unsafe behavior

“TMRL broadens the action distribution by aliasing actions across contexts, but this can lead to unsafe behavior. Mitigating these behaviors using a safety filter or a world model is necessary before TMRL can be used reliably.” 也就是说，为提升覆盖而做的 context 混叠，本身可能引入不安全的动作，作者认为在 TMRL 能被可靠使用之前，需要额外的安全过滤器或 world model 来约束。In other words, the context aliasing performed to increase coverage may itself introduce unsafe actions; the authors hold that an additional safety filter or world model is required to constrain the policy before TMRL can be used reliably.

样本效率在真实世界中仍不够实用Sample efficiency is still impractical in the real world

“While TMRL performs well in the real world, its sample efficiency remains impractical for many tasks of practical interest. We expect that further improvements to the steering algorithm and corruption kernel can increase TMRL's efficiency.” 即便真实机器人上一小时左右能完成部分任务的微调，作者仍认为对许多实际感兴趣的任务而言，当前的样本效率还不够，需要在 steering 算法与腐蚀核设计上进一步改进。Even though fine-tuning on a real robot can be completed in about an hour for some tasks, the authors still regard the current sample efficiency as insufficient for many tasks of practical interest, and see room for further improvement in the steering algorithm and the design of the corruption kernel.