Behavioral Mode Discovery for Fine-tuning Multimodal Generative Policies

We address the problem of fine-tuning pre-trained generative policies with reinforcement learning (RL) while preserving the multimodality of their action distributions. Existing methods for RL fine-tuning of generative policies (e.g., diffusion policies) improve task performance but often collapse diverse behaviors into a single reward-maximizing mode. To mitigate this issue, we propose an unsupervised mode discovery framework that uncovers latent behavioral modes within generative policies. The discovered modes enable the use of mutual information as an intrinsic reward, regularizing RL fine-tuning to enhance task success while maintaining behavioral diversity. Experiments on robotic manipulation tasks demonstrate that our method consistently outperforms conventional fine-tuning approaches, achieving higher success rates and preserving richer multimodal action distributions.

Robots can often complete the same task in many different valid ways. For example, a robot may grasp an object from different sides or follow different paths to reach the same goal. Modern learning methods can capture this variety from demonstrations, but when the robot is later trained to improve its performance, it may forget many of these alternatives and settle on only one strategy. This paper introduces a method that helps a robot improve while still keeping multiple useful ways of acting. The method automatically identifies different behaviors already present in a trained robot policy and encourages the robot to preserve them during further learning. In robotic manipulation experiments, this leads to policies that perform tasks successfully while remaining more flexible and diverse. This is important for building robots that can adapt to new situations, recover from disturbances, and choose among several possible solutions rather than relying on a single fixed behavior.