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

TL;DR

This paper introduces a novel framework combining diffusion-based pretraining and timestep-modulated reinforcement learning to enhance exploration and sample efficiency in robot policy fine-tuning, achieving real-world success in under an hour.

Contribution

It proposes CSP for pretraining with diffusion noise and TMRL for dynamic exploration control during RL fine-tuning, improving efficiency and applicability across various policy inputs.

Findings

TMRL improves RL fine-tuning sample efficiency.

TMRL enables real-world manipulation tasks in under one hour.

The framework seamlessly integrates with different policy input modalities.

Abstract

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