Diffusion Timestep-Modulated Pre-training Enables Exploration for Efficient Policy Fine-tuning
A unified framework bridging behavior-cloning pre-training and reinforcement-learning fine-tuning.
Context-Smoothed Pre-training (CSP) applies a diffusion-style forward-noising schedule to policy inputs, producing a continuous spectrum from precise imitation to broad coverage. Timestep-Modulated RL (TMRL) then trains a high-level policy πHL(z, σ | s) to select σ per action chunk.
Outperforming baselines in simulation and real-world 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. As σ rises, qσ progressively corrupts c: the pointcloud slowly loses its original structure, aliasing the out-of-distribution object with nearby in-distribution pointclouds seen during training. The policy can then borrow coherent grasps from related training contexts rather than failing on an unfamiliar shape.
Top: TMRL on a held-out drill (context c = pointcloud). Bottom: qσ removes structure as σ rises — aliasing the OOD object back to in-distribution shapes the policy has seen during training.
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 out-of-distribution states 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 coverage control, 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 Pre-training Enables Exploration for Efficient Policy Fine-tuning},
author = {Hong, Matthew M. and Zhang, Jesse and Nagabandi, Anusha and Gupta, Abhishek},
booktitle = {Robotics: Science and Systems (RSS)},
year = {2026}
}