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

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

Matthew M. Hong, Jesse Zhang, Anusha Nagabandi, Abhishek Gupta · · 2026

rl imitation-learning manipulation
PDF Code (GitHub)

Framework

N/A

License

N/A

Stars

0

Summary

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.

Abstract Summary

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/.

Key Points

  • Introduces TMRL, advancing rl research.
  • Detailed evaluation demonstrates effectiveness on benchmarks and real-world scenarios.
  • Provides insights into methodology and architectural choices.

The Need for Action Distribution Interpolation

We begin with a simple experiment that illustrates why we need a pre-training procedure that can interpolate between the marginal action distribution $p(a)$ and a conditional distribution $p(a\mid c)$ when encountering unseen tasks. In a simple 2D point-maze from the simulated OGBench benchmark [ 35 ] , we train two policies to approximate $p(a)$ and $p(a\mid c)$ where $c$ is the goal position.

The Need for Action Distribution Interpolation — TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning
Fig. 2 : For an unseen goal position, the conditional action distribution $p\!\left(a\mid\mspace{-6mu}\raisebox{-0.2pt}{\includegraphics[height=8.32487pt]{figs/goal_marker.png}}\right)$ remains narrow, lacking sufficient support for the agent to reach the goal . In contrast, the marginal $p(a)$ provides the required coverage, though with too much action variance near the goal. In this setting, an optimal policy should more closely match $p(a)$ at the beginning of the trajectory and $p\!\left(a\mid\mspace{-6mu}\raisebox{-0.2pt}{\includegraphics[height=8.32487pt]{figs/goal_marker.png}}\right)$ at the end. Context-smoothed policy pre-training enables TMRL to interpolate between both distributions for more effective RL fine-tuning.

Timestep Modulated RL on Context-Smoothed Policies

Pre-training. We assume access to an offline dataset of near-optimal pre-training trajectories $\mathcal{D}=\{\tau_{j}\}$, where each trajectory $\tau_{j}=\{(c_{i},a_{i})\}_{i=1}^{T_{j}}$ contains context–action pairs. In this paper, we use context to abstractly define a variety of inputs to the policy, such as states, RGB images, point clouds, language instructions, and other task descriptors.

Timestep Modulated RL on Context-Smoothed Policies — TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning
Fig. 3 : Timestep-modulated exploration via context smoothing. ( Left ) During pre-training, a steerable policy $p_{\theta}$ is trained across all noise levels $\sigma$ by corrupting the context $c$ via the kernel $q_{\sigma}(\tilde{c}\mid c)$, producing a policy that can be queried at any conditioning strength during inference. ( Right ) During RL fine-tuning, TMRL exposes $p_{\theta}$ with a context-noise dial $\sigma$ as an explicit control variable for the high-level policy $\pi_{\text{HL}}$.

Theoretical Analysis: Benefits of Context Smoothing

In this section, we provide theoretical evidence and empirical analysis to better understand the gains from timestep-modulated, context-smoothed policies for RL fine-tuning. We first empirically validate that increasing $\sigma$ (or $t_{c}$) yields smooth, meaningful behavioral changes.

We show in Figure 4 that as corruption noise increases, the distribution of actions at a given context broadens, thereby increasing overlap with nearby contexts.

Theoretical Analysis: Benefits of Context Smoothing — TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning
Fig. 5 : Evaluation. We evaluate TMRL across 8 tasks spanning navigation and manipulation in simulation and the real world.

Experiments

Does context smoothing produce better action coverage over other pre-training approaches? Does TMRL effectively RL fine-tune policies trained across a variety of policy conditioning variables? Does TMRL enable real-world RL on VLAs? How does context-smoothing and TMRL compare to alternative coverage expansion mechanisms? How does TMRL learn to control context smoothing over time?

We first compare the effect of various pre-training procedures on action coverage by measuring success rates on unseen tasks.

Experiments — TMRL: Diffusion Timestep-Modulated Pretraining Enables Exploration for Efficient Policy Finetuning
Fig. 8 : Context Smoothed Pre-training for VLAs: we propose applying the corruption kernel ( Eq. 4 ) to VLM embeddings input to the action expert.

Sources and Demonstrations

Figures are reproduced from Hong et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.

Share

Related Papers

LeRobot: A Library for Real-World Robot Learning
NeurIPS 2024 Workshop

LeRobot: A Library for Real-World Robot Learning

Remi Cadene, Simon Alibert, Alexander Soare et al.

Hugging Face's LeRobot is an open-source PyTorch framework providing pretrained models, datasets, and training scripts for imitation and reinforcement learning on real robots, lowering the entry barrier to robot learning.

imitation-learning rl foundation-models manipulation
Code PDF Beginner
GitHub ★ 24,333 Updated May 2026
EgoForce: Forearm-Guided Camera-Space 3D Hand Pose from a Monocular Egocentric Camera
SIGGRAPH 2026

EgoForce: Forearm-Guided Camera-Space 3D Hand Pose from a Monocular Egocentric Camera

Christen Millerdurai, Shaoxiang Wang, Yaxu Xie et al.

Reconstructing the absolute 3D pose and shape of the hands from the user's viewpoint using a single head-mounted camera is crucial for practical egocentric interaction in AR/VR, telepresence, and hand-centric manipulation tasks, where sensing must remain compact and unobtrusive.

manipulation manipulation simulation
Code PDF Intermediate
Code ★ 0
Aligning Flow Map Policies with Optimal Q-Guidance
arXiv

Aligning Flow Map Policies with Optimal Q-Guidance

Christos Ziakas, Alessandra Russo, Avishek Joey Bose

Generative policies based on expressive model classes, such as diffusion-models and flow matching, are well-suited to complex control problems with highly multimodal action distributions.

rl diffusion-models manipulation
PDF Advanced