Behavioral Mode Discovery for Fine-tuning Multimodal Generative Policies

2026-05-11

AI Summary

This work addresses the collapse of multimodal action distributions during reinforcement learning fine-tuning of pretrained generative policies, a phenomenon often caused by over-optimization of extrinsic rewards that leads to loss of behavioral diversity. To mitigate this issue, the authors propose an unsupervised behavioral repertoire discovery framework that constructs intrinsic rewards via mutual information maximization, thereby guiding the policy to simultaneously improve task success rates and preserve multimodal behaviors. This approach uniquely integrates unsupervised representation learning with mutual information regularization, effectively counteracting diversity degradation during fine-tuning. Experimental results on robotic manipulation tasks demonstrate that the proposed method significantly outperforms conventional fine-tuning strategies, achieving both enhanced performance and richer behavioral repertoires.

Abstract

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.

Problem

• multimodal generative policies
• reinforcement learning fine-tuning
• behavioral diversity
• mode collapse
• action distribution

Innovation

• behavioral mode discovery
• multimodal generative policies
• reinforcement learning fine-tuning
• mutual information regularization
• diffusion policies