Diffusion Timestep-Modulated Pre-training Enables Exploration for Efficient Policy Fine-tuning
BC pre-trains policies to imitate, but RL needs coverage to explore. We pre-train for that coverage.
Context-Smoothed Pre-training (CSP) trains a policy across all noise levels of a diffusion-style forward schedule, spanning the spectrum from precise imitation to broad coverage. Timestep-Modulated RL (TMRL) then trains a high-level policy to select the noise level per action chunk.
Outperforming baselines in simulation and real-world RL fine-tuning of π0 to near-perfect success on three tasks: under an hour on WidowX, under four on Franka.
The need for action distribution interpolation.
Robots trained by imitation copy demonstrations narrowly — in sparsely covered contexts the conditional support collapses to near-zero probability on optimal actions, so online rollouts yield no reward signal and RL fine-tuning stalls. We propose injecting forward-diffusion noise into policy inputs during pre-training, so the policy learns the full spectrum of distributions between the conditional p(a | c) and marginal p(a).
On a 2D pointmaze from OGBench, we
train two policies — one approximating p(a | c), the other
p(a) — where c is the goal position. Both train on
pointmaze-large and evaluate on the larger pointmaze-giant environment, with
an unseen goal location (agent in yellow, goal in pink).
Left, the conditional p(a | c) collapses and drifts the agent away from the goal. Right, the marginal p(a) covers broadly enough to reach it, but assigns goal-reaching behaviors only low likelihood.
This motivates a policy that can interpolate between the two — and a mechanism to learn where on that interpolation to operate.
Context smoothing & timestep modulation.
On the dexterous-grasping suite, the context c is the object pointcloud. Our
approach adds
noise to the policy's inputs — intuitively, blurring what it sees. As σ rises,
qσ
corrupts c until the out-of-distribution object aliases toward nearby in-distribution shapes from
training.
When we squint at an unfamiliar object, the details drop away and its rough shape resembles
objects we
already know how to grab. Similarily for the policy: a novel context blurs into familiar ones, so it
borrows a known-good
grasp instead of confidently guessing wrong.
The conditional ↔ marginal spectrum.
CSP trains one policy across all σ by corrupting the context c (here, the cube goal position) with qσ(c | c). The endpoints anchor the spectrum: σ = 0 recovers p(a | c), σ = T recovers p(a).
Drag the slider to see how the policy’s rollouts spread as the context (goal position) becomes more corrupted.
Coverage and RL sample efficiency.
Success@K — the fraction of contexts c (here, out-of-distribution
contexts)
where
at least
one of K base policy rollouts succeeds. On two tasks, CSP outperforms PostBC and
BC for all K, with the gap widest on cube-single where both baselines remain at
zero.
cube-single, BC and
PostBC stay near zero while CSP rises with K.We compare against several state-of-the-art baselines. TMRL outperforms or matches all baselines across four simulation tasks.
libero-90 task.
Context-Smoothed π0 in the real world.
sausage-in-pot, shrimp-in-drawer).
We fine-tune a context-smoothed π0 on BridgeData-v2 (WidowX 250) and DROID (Franka).
Base π0 versus context-smoothed π0, side by side. The learning curve below traces the full RL training run.
πHL treats the diffusion timestep σ as an adaptive exploration knob, choosing it per
action chunk. Below, the converged TMRL policy on shrimp-in-drawer learns to keep
σ low through the precision-critical pick and higher during reaching and placing.
“pick up the shrimp and put it into the white drawer”
@inproceedings{hong2026tmrl,
title = {TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning},
author = {Hong, Matthew M. and Zhang, Jesse and Nagabandi, Anusha and Gupta, Abhishek},
booktitle = {Robotics: Science and Systems (RSS)},
year = {2026}
}