REVIEW 4 major objections 5 minor 13 references
Policy-Based Trajectory Clustering in Offline Reinforcement Learning
T0 review · 4 major / 5 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read The paper proposes, formalizes, and tests policy-based trajectory clustering for offline reinforcement learning: group trajectories by the policy that generated them, and solve the resulting mixture-matching problem with two new algorithms.
desk verdict Worth reading for the problem framing and GridWorld results, but the advertised convergence proof doesn't hold for the implemented algorithm and the D4RL labels are heuristically constructed. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the policy centroid: each cluster $j$ is represented by a behavior-cloned policy $\pi_{\theta_j}$, and trajectories are assigned by the generation probability $P(\tau_i\mid \theta_j)$, which makes the clustering operate in policy space rather than raw trajectory space. For PG-Kmeans, the proof of finite-step convergence rests on the same monotonic-improvement argument used for standard K-means: the objective strictly improves until an assignment pattern repeats, and only finitely many patterns exist. For CAAE, the mechanism is a learnable Gaussian codebook $\{\mu_j\}_{j=1}^k$ with an attraction penalty $\min_j \|\mu_j - z_i\|_2$, plus a decoder that reconstructs actions conditioned on the latent code and observations. The K-coloring reduction supplies the ambiguity result: trajectories are nodes, conflicting state-action choices are edges, and a feasible clustering is a coloring, so multiple colorings can be equally valid.
What would settle it
Generate a synthetic two-policy dataset with a known conflict rate, then vary the conflict rate from near zero to one while holding everything else fixed; the paper's account predicts that Normalized Mutual Information should climb as conflict rate rises, so a flat or non-monotonic curve would show that conflict rate is not the limiting factor.
Extended reading notes
Core claim
The paper's central claim is that offline RL trajectories can be clustered by their generating policy, and that policy-level clusters are discoverable from state-action data without rewards. The authors formalize this as distribution matching: because the trajectory distribution is a mixture of policy-induced distributions, minimizing the KL divergence between the empirical data and the mixture gives the objective $\max_{\theta,W} \sum_{i,j} w_{i,j} \sum_t \log P(a_{i,t}\mid \theta_j, s_{i,t})$, and they solve it either by alternating assignment and behavior cloning (PG-Kmeans) or by training a codebook-regularized encoder-decoder (CAAE). They further claim the problem is fundamentally ambiguous: a reduction from K-coloring proves the general policy-based clustering problem NP-complete, and even in a two-state contextual bandit, different pairs of policies can generate identical trajectory distributions, so the optimal solution need not be unique. The experiments show both algorithms partition D4RL medium-expert and custom GridWorld datasets into meaningful clusters on most tasks, with the caveat that when the generating policies conflict on very few state-action pairs, no method separates them.
Load-bearing premise
The whole method depends on different policies leaving visibly different decision records: if two policies choose the same action at almost every state they visit, no clustering algorithm can tell their trajectories apart.
Editorial extensions
If this is right
- If the central claim is correct, offline datasets can first be decomposed into behavior modes, and downstream policy learning can train one policy per mode instead of fitting a single policy to conflicting behaviors.
- PG-Kmeans is guaranteed to converge in a finite number of iterations, so the policy-centroid loop is a terminating procedure, not an asymptotic one.
- The optimal clustering can be non-unique, which means applications should evaluate and use clusterings as one of several acceptable decompositions rather than a single ground-truth partition.
- CAAE's shared encoder-decoder gives it lower training cost and more stable optimization, while PG-Kmeans produces sharper cluster boundaries and direct one-step policies.
- The paper's motivations imply that after clustering, a small number of reward labels per cluster could be enough to guide downstream offline RL, supporting a semi-supervised training pipeline.
Reading between the lines
- Because the paper's failure cases track low conflict rates, conflict rate could be computed before clustering as a cheap screening predictor of whether a dataset is clusterable at all; the paper does not propose this step.
- The demonstrated non-uniqueness of policy decompositions suggests a more useful target than recovering one true partition: finding any decomposition whose center policies are simple or minimal, which the overparameterization-and-merge heuristic approximates.
- For trajectories that switch policies within one episode, CAAE's whole-trajectory embeddings are likely insufficient; a natural extension is a change-point-aware variant that segments trajectories first and clusters the segments.
- The same policy-based clustering objective could apply outside RL, wherever sequential decision records are generated by an unknown mixture of deterministic strategies, such as user behavior logs or control traces.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces policy-based trajectory clustering for offline RL, where trajectories are grouped according to the policy that generated them. Two algorithms are proposed: PG-Kmeans, which alternates between behavior-cloning policies for each cluster and assigning trajectories by policy likelihood, and CAAE, a VQ-VAE-style autoencoder with a learnable codebook. The authors claim finite-step convergence for PG-Kmeans, prove that the general clustering problem is NP-complete via a reduction to K-coloring, and report NMI scores on D4RL and custom GridWorld environments, arguing that both methods outperform existing baselines on most tasks.
Significance. If the claims hold, the paper would open a useful new task in offline RL data organization, with applications to data cleaning, semi-supervised policy learning, and interpretability. The KL-divergence derivation, the K-coloring reduction, and the extensive ablations on cluster count and regularization are valuable contributions. The paper is also careful in places, explicitly discussing the ambiguity of clustering solutions and reporting failures on low-conflict environments. However, the main theoretical guarantee is not supported for the implemented algorithm, the D4RL evaluation uses return-constructed labels rather than true policy labels, and the headline claim of effective clustering fails on two of eight reported tasks. These issues are load-bearing and require substantive revision.
major comments (4)
- [§5.1.2, Appendix A.1, Section 7] Theorem 5.1 and its proof assume that the M-step returns a global maximizer of J(W, θ), as stated in equation (5): max_θ J(W^{t-1}, θ) = J(W^{t-1}, θ^t). The implemented Algorithm 1 trains neural behavior-cloning policies with SGD, which provides no global optimality guarantee and does not even guarantee monotone increase of the objective. Without equation (5), the strict-increase argument collapses and the finite k^N termination bound no longer follows. Section 7 explicitly says 'the algorithms currently lack theoretical convergence guarantees', which directly contradicts the abstract and Theorem 5.1. The theorem should be restated under an exact-optimization assumption or removed, and the abstract's convergence claim should be qualified as empirical.
- [Appendix B.1, Table 1] The D4RL ground-truth labels used for NMI are not true generating-policy labels. As stated in B.1, episodes are labeled by maximizing the average return of the first and second halves, which is a return heuristic. The paper asserts that the resulting misclassification is 'negligible' without supporting evidence, yet the heuristic directly favors return-correlated methods: Return+Kmeans achieves NMI 0.97 on HalfCheetah and 0.86 on Hopper. Consequently, D4RL scores in Table 1 do not establish that the methods cluster by generating policy. The authors should either use a dataset with known policy labels or explicitly present the D4RL results as agreement with a return-based partition.
- [Table 1, Section 6.2] The central claim that both methods 'effectively partition trajectories into meaningful clusters' is not supported on Pathfollowing and Extra: PG-Kmeans obtains NMI 0.12 ± 0.11 and 0.02 ± 0.01, respectively, and CAAE obtains 0.15 ± 0.02 and 0.43 ± 0.22. On Pathfollowing, the DEC baseline (0.26 ± 0.10) outperforms both proposed methods. The explanation in Section 6.2—low conflict rates and inherent indistinguishability—is a scope condition, but it is not stated in the abstract or the introduction. The claims should be qualified to datasets where policies disagree sufficiently, and the failure cases should be presented as such rather than as a minor caveat.
- [Appendix B.3, Table 1] The SORL baseline is not the published SORL algorithm. Appendix B.3 states that the authors 'implement the clustering step only, and use hard clustering instead of soft clustering to represent SORL algorithm'. Since SORL's core design includes soft assignments and joint policy optimization, Table 1's 'SORL' column and the Section 6.2 claim that PG-Kmeans outperforms SORL are not supported by actual SORL. The baseline should be either run as originally published or renamed (e.g., 'SORL-clustering-step (hard)') and the comparison claims adjusted accordingly.
minor comments (5)
- [Appendix A.2, Theorem A.2] The reduction requires |S| ≥ 2d+1, not |S| ≥ 2d: for an edge (v_i, v_j), the union of trajectories can contain up to 2d states, so the smallest available index l may need to be 2d+1. Also, to call the problem NP-complete, the decision version needs a membership-in-NP argument, which is not provided.
- [Section 7 vs. Appendix A.3] Section 7 lists 'the uniqueness of clustering solutions remains an open question' as a limitation, but Section 4 and Appendix A.3 already give explicit examples of non-unique solutions. Please clarify whether the open question concerns conditions for uniqueness under additional identifiability assumptions, rather than uniqueness in general.
- [Appendix B.1] The claim that up to five mislabeled episodes have 'negligible impact' on NMI is asserted without a sensitivity analysis. A short experiment perturbing the label boundary would make this claim verifiable.
- [§6.2, Figure 3] Figure 3 is described as showing no significant correlation between NMI and loss and is used to motivate the overfitting discussion, but the same loss-based criterion is proposed as the internal metric for Best-of-N selection in Appendix B.4.1. The relationship between optimization loss and clustering quality should be clarified, since the two statements appear to conflict.
- [Various] There are several typos and formatting inconsistencies, including 'A alternative objective' in Section 3, 'Mathmatical' in the appendix title, inconsistent 'V AE' spacing, and equation (3) using h versus H in the summation index. These should be cleaned up.
Circularity Check
No construction-level circularity: the EM-style clustering derivation is self-contained, and the flagged concerns (return-derived D4RL labels and the unproven global-M-step convergence claim) are evaluation and proof-correctness issues rather than reductions of the derivation to its inputs.
full rationale
The paper's derivation chain is self-contained. The clustering objective (Eq. 3) is obtained by minimizing the KL divergence between the empirical trajectory distribution and a mixture of policy-induced distributions, with the environment dynamics canceled out to leave a per-step action-likelihood objective. The E-step (argmax assignment by policy likelihood) and M-step (behavior-cloning training on assigned trajectories) are standard EM-style coordinate ascent, and the finite-step argument mirrors the classical K-means convergence proof; this is an algorithmic self-consistency result, not a prediction extracted from fitted inputs. The central methods, PG-Kmeans and CAAE, do not receive ground-truth labels or returns as inputs, so their reported NMI is not forced by construction. The most salient weakness is evaluative rather than circular: for D4RL, the 'ground truth' labels are constructed from the same dataset by a return-maximization split (Appendix B.1: 'we partition the dataset by maximizing the average return of the first and second halves'), so NMI is measured against a return-derived proxy for policy identity rather than an independent label source, and Section 6.2 concedes near-chance results when policy conflicts are rare, such as Pathfollowing and Extra. This weakens the empirical claim but does not make the derivation equivalent to its input. The proof of Theorem 5.1 assumes an exact global maximizer in the M-step (Eq. 5: 'max_theta J(W^{t-1}, theta) = J(W^{t-1}, theta^t)'), which SGD-trained neural behavior cloning does not guarantee; this is a proof-correctness gap, and Section 7 itself states that 'the algorithms currently lack theoretical convergence guarantees.' That contradiction is a correctness risk, not a circularity. The only overlapping-author citation, SORL (Mao et al. 2024, co-authored by Hao Hu), is used as related work and is explicitly contrasted with PG-Kmeans rather than invoked to justify the paper's central premise. No step in the paper's own equations reduces to its own inputs by construction, so the circularity score is low.
Assumptions & free parameters
free parameters (4)
- CAAE codebook attraction weight alpha =
1
- CAAE codebook regularization coefficient 1/m^2 =
unspecified (m not defined)
- Initial cluster count k =
varies per environment (e.g., 4 on D4RL, overparameterized up to 3x k*)
- Best-of-N number of runs N =
not explicitly stated (used for PG-Kmeans, see Algorithm 3)
assumptions (5)
- domain assumption The dataset is a union of k balanced sub-datasets, each collected by a distinct deterministic behavior policy.
- standard math Environment dynamics are independent of the policy parameters theta_j and assignments w_ij, so trajectory likelihoods factor as policy action probabilities times a common dynamics term.
- domain assumption Each trajectory is generated by exactly one policy, captured by binary hard assignment weights w_ij.
- ad hoc to paper The M-step exactly maximizes J(W,theta), i.e., behavior cloning reaches the global optimum for each cluster.
- ad hoc to paper The return-based split of D4RL episodes into halves yields correct policy labels for NMI evaluation.
Cite this review
Pith. "Pith review of Policy-Based Trajectory Clustering in Offline Reinforcement Learning." pith.science (2026). https://pith.science/paper/T4R735ZM
@misc{pith2026250609202,
author = {Pith},
title = {Pith review of: Policy-Based Trajectory Clustering in Offline Reinforcement Learning},
year = {2026},
howpublished = {\url{https://pith.science/paper/T4R735ZM}},
note = {Machine review of arXiv:2506.09202}
}
read the original abstract
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.
Figures
Figures from the paper (2 more)
Reference graph
Works this paper leans on
-
[1]
Always move to the right, until the wall is reached
-
[2]
Always move to the down, until the wall is reached
-
[3]
Move to the right iff at the left-down half of the gridworld
-
[4]
Move to the right iff at the black grid, if we seem the gridworld as a chess board
-
[5]
Takeball There are four different rule-based policies, i-th policy will pick the i-th ball first
Move to the right iff at the white grid. Takeball There are four different rule-based policies, i-th policy will pick the i-th ball first. 16 Extra We use PPO algorithm to train three policies by giving positive/zero/negative rewards for the agent to reach the special grid. The rewards are set to +10, 0 or -10, respectively. When the agent reach the goal,...
-
[12]
5 times of distance to the polyline (−1, −1), (−1, 1), (1, 1)
-
[13]
5 times of distance to the polyline (−1, −1), (1, −1), (1, 1). These policies are corresponding to: No preference, prefer to go up first and prefer to go right first. B.3 Baseline Methods Deep Embedded Clustering (DEC). Deep Embedded Clustering (DEC) Xie et al. [2016] is a widely used deep clustering method that integrates representation learning with clu...
work page 2016
-
[2016]
URL https://arxiv.org/abs/1511.06335. Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. In M. Ranzato, A. Beygelzimer, Y . Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 28954–...
arXiv 2021
Show all 13 references
-
[2019]
Junyuan Xie, Ross Girshick, and Ali Farhadi
URL https://arxiv.org/abs/1911.11361. Junyuan Xie, Ross Girshick, and Ali Farhadi. Unsupervised deep embedding for clustering analysis,
1911 arXiv
-
[2020]
Yuanguo Lin, Yong Liu, Fan Lin, Lixin Zou, Pengcheng Wu, Wenhua Zeng, Huanhuan Chen, and Chunyan Miao
URL https://arxiv.org/abs/2009.09687. Yuanguo Lin, Yong Liu, Fan Lin, Lixin Zou, Pengcheng Wu, Wenhua Zeng, Huanhuan Chen, and Chunyan Miao. A survey on reinforcement learning for recommender systems. IEEE Transactions on Neural Networks and Learning Systems , 35(10):13164–131...
2009 arXiv
-
[2021]
Martin Klissarov, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup
URL https://arxiv.org/abs/2002.00444. Martin Klissarov, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Variational state encoding as intrinsic motivation in reinforcement learning. In Task-Agnostic Reinforcement Learning Workshop at Proceedings of the International Confere...
2002 arXiv
-
[2024]
Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou
URL https://arxiv.org/abs/2402.09550. Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/...
2023 arXiv
-
[2388]
URL http://dx.doi.org/10.1109/TNNLS.2023
doi: 10.1109/tnnls.2023.3280161. URL http://dx.doi.org/10.1109/TNNLS.2023. 3280161. Yihuan Mao, Chengjie Wu, Xi Chen, Hao Hu, Ji Jiang, Tianze Zhou, Tangjie Lv, Changjie Fan, Zhipeng Hu, Yi Wu, Yujing Hu, and Chongjie Zhang. Stylized offline reinforcement learning: Extracting ...
2023
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.