# Policy-Based Trajectory Clustering in Offline Reinforcement Learning

arXiv: https://arxiv.org/abs/2506.09202 · HTML: https://mlsanity.com/paper/2506.09202/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 31 of its 37 references that are in the map.
- First posted on arXiv 10 Jun 2025.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 31)

- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Unsupervised Deep Embedding for Clustering Analysis](https://mlsanity.com/paper/1511.06335/) (2015)
- [Deep Reinforcement Learning for Autonomous Driving: A Survey](https://mlsanity.com/paper/2002.00444/) (2020)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Invariant Information Clustering for Unsupervised Image Classification and Segmentation](https://mlsanity.com/paper/1807.06653/) (2018)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Challenges of Real-World Reinforcement Learning](https://mlsanity.com/paper/1904.12901/) (2019)
- [Contrastive Clustering](https://mlsanity.com/paper/2009.09687/) (2020)
- [DualDICE: Behavior-Agnostic Estimation of Discounted Stationary Distribution Corrections](https://mlsanity.com/paper/1906.04733/) (2019)
- [Accelerating Reinforcement Learning with Learned Skill Priors](https://mlsanity.com/paper/2010.11944/) (2020)
- [AlgaeDICE: Policy Gradient from Arbitrary Experience](https://mlsanity.com/paper/1912.02074/) (2019)
- [Deep Reinforcement Learning for Robotics: A Survey of Real-World Successes](https://mlsanity.com/paper/2408.03539/) (2024)
- [Offline RL Without Off-Policy Evaluation](https://mlsanity.com/paper/2106.08909/) (2021)
- [Structural Deep Clustering Network](https://mlsanity.com/paper/2002.01633/) (2020)
- [Variational Information Maximisation for Intrinsically Motivated Reinforcement Learning](https://mlsanity.com/paper/1509.08731/) (2015)
- [GenDICE: Generalized Offline Estimation of Stationary Values](https://mlsanity.com/paper/2002.09072/) (2020)
- [CORL: Research-oriented Deep Offline Reinforcement Learning Library](https://mlsanity.com/paper/2210.07105/) (2022)
- [Deep Clustering via Joint Convolutional Autoencoder Embedding and Relative Entropy Minimization](https://mlsanity.com/paper/1704.06327/) (2017)
- [BAIL: Best-Action Imitation Learning for Batch Deep Reinforcement Learning](https://mlsanity.com/paper/1910.12179/) (2019)
- [A Survey on Reinforcement Learning for Recommender Systems](https://mlsanity.com/paper/2109.10665/) (2021)

## Built on by (most cited first, 2 of 2)

- [Less is More: Clustered Cross-Covariance Control for Offline RL](https://mlsanity.com/paper/2601.20765/) (2026)
- [ASALT: Adaptive State Alignment for Lateral Transfer in Multi-agent Reinforcement Learning](https://mlsanity.com/paper/2606.24601/) (2026)

## Abstract (arXiv)

We introduce a novel task of clustering trajectories from offline reinforcement learning (RL) datasets, where each cluster center represents the policy that generated its trajectories. By leveraging the connection between the KL-divergence of offline trajectory distributions and a mixture of policy-induced distributions, we formulate a natural clustering objective. To solve this, we propose Policy-Guided K-means (PG-Kmeans) and Centroid-Attracted Autoencoder (CAAE). PG-Kmeans iteratively trains behavior cloning (BC) policies and assigns trajectories based on policy generation probabilities, while CAAE resembles the VQ-VAE framework by guiding the latent representations of trajectories toward the vicinity of specific codebook entries to achieve clustering. Theoretically, we prove the finite-step convergence of PG-Kmeans and identify a key challenge in offline trajectory clustering: the inherent ambiguity of optimal solutions due to policy-induced conflicts, which can result in multiple equally valid but structurally distinct clusterings. Experimentally, we validate our methods on the widely used D4RL dataset and custom GridWorld environments. Our results show that both PG-Kmeans and CAAE effectively partition trajectories into meaningful clusters. They offer a promising framework for policy-based trajectory clustering, with broad applications in offline RL and beyond.