mlsanityData to 14 Oct 2026

Policy-Based Trajectory Clustering in Offline Reinforcement Learning

arXiv 2506.09202 ↗First posted 10 Jun 2025cs.LG cs.AIcited by 2

On the map: Offline RL and imitation learning from data → Offline RL algorithms

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.

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.

Builds on

31 of 37 references.

Closest ancestors · 10

Same direction, newest first

  1. 2022Offline Reinforcement Learning with Closed-Form Policy Improvement Operators29 Nov
  2. 2022CORL: Research-oriented Deep Offline Reinforcement Learning Library13 Oct
  3. 2021Offline Reinforcement Learning with Implicit Q-Learning12 Oct
  4. 2021Offline RL Without Off-Policy Evaluation16 Jun
  5. 2021COMBO: Conservative Offline Model-Based Policy Optimization16 Feb
  6. 2020Conservative Q-Learning for Offline Reinforcement Learning8 Jun
  7. 2020MOReL : Model-Based Offline Reinforcement Learning12 May
  8. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems4 May
  9. 2020D4RL: Datasets for Deep Data-Driven Reinforcement Learning15 Apr
  10. 2019Behavior Regularized Offline Reinforcement Learning26 Nov

Other references · 21

  1. 2018World Modelscited by 1,513
  2. 2018Off-Policy Deep Reinforcement Learning without Explorationcited by 1,241
  3. 2019Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reductioncited by 840
  4. 2019When to Trust Your Model: Model-Based Policy Optimizationcited by 773
  5. 2015Unsupervised Deep Embedding for Clustering Analysiscited by 767
  6. 2020Deep Reinforcement Learning for Autonomous Driving: A Surveycited by 721
  7. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 644
  8. 2022Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learningcited by 562
  9. 2018Invariant Information Clustering for Unsupervised Image Classification and Segmentationcited by 455
  10. 2019Challenges of Real-World Reinforcement Learningcited by 291
  11. 2020Contrastive Clusteringcited by 265
  12. 2019DualDICE: Behavior-Agnostic Estimation of Discounted Stationary Distribution Correctionscited by 254
9 more
  1. 2020Accelerating Reinforcement Learning with Learned Skill Priorscited by 223
  2. 2019AlgaeDICE: Policy Gradient from Arbitrary Experiencecited by 188
  3. 2024Deep Reinforcement Learning for Robotics: A Survey of Real-World Successescited by 176
  4. 2020Structural Deep Clustering Networkcited by 145
  5. 2015Variational Information Maximisation for Intrinsically Motivated Reinforcement Learningcited by 144
  6. 2020GenDICE: Generalized Offline Estimation of Stationary Valuescited by 122
  7. 2017Deep Clustering via Joint Convolutional Autoencoder Embedding and Relative Entropy Minimizationcited by 113
  8. 2019BAIL: Best-Action Imitation Learning for Batch Deep Reinforcement Learningcited by 91
  9. 2021A Survey on Reinforcement Learning for Recommender Systemscited by 18

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026Less is More: Clustered Cross-Covariance Control for Offline RLcited by 1
  2. 2026ASALT: Adaptive State Alignment for Lateral Transfer in Multi-agent Reinforcement Learningcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024Dataset Clustering for Improved Offline Policy Learning
  2. 2024Discovering Behavioral Modes in Deep Reinforcement Learning Policies Using Trajectory Clustering in Latent Space
  3. 2025Harnessing Data from Clustered LQR Systems: Personalized and Collaborative Policy Optimization
  4. 2023Harnessing Mixed Offline Reinforcement Learning Datasets via Trajectory Weighting
  5. 2025Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning
  6. 2022Semi-Supervised Offline Reinforcement Learning with Action-Free Trajectories
  7. 2024POTEC: Off-Policy Learning for Large Action Spaces via Two-Stage Policy Decomposition
  8. 2023Beyond Uniform Sampling: Offline Reinforcement Learning with Imbalanced Datasets