{"id":"d8b843c9-edfb-4c35-a37f-e711cc5f2918","arxiv_id":"2506.09202","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":4,"one_line_summary":"Proposes Policy-Guided K-means and Centroid-Attracted Autoencoder for clustering offline RL trajectories by their generating policy, with a finite-step convergence result and an NP-completeness connection to graph coloring.","lead":"This paper frames the task of grouping trajectories in offline reinforcement learning datasets by the policy that produced them, and introduces two clustering algorithms, PG-Kmeans and CAAE. The work also proves a convergence result for one algorithm and links the task to graph coloring, though experimental performance is near chance on two of eight tested environments.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Theorem 5.1's finite-step convergence guarantee is unproven for the implemented algorithm: the proof assumes an exact global M-step optimum that neural BC training does not provide, and Section 7 concedes no convergence guarantees.","rationale":"I read the paper's central claim as having two components: the empirical claim that PG-Kmeans and CAAE effectively partition trajectories into meaningful clusters, and the theoretical claim of finite-step convergence for PG-Kmeans. The reader's weakest assumption (detectability of policies from action conflicts) is real but explicitly acknowledged in Section 6.2 as an 'inherent indistinguishability' limitation; it is a property of the data, not a flaw in the algorithm's argument. The convergence proof is a direct internal inconsistency: Section 5.1.2's Theorem 5.1 asserts a formal guarantee, Appendix A.1 proves it using an exact global M-step optimum, yet Algorithm 1 implements an approximate neural M-step and Section 7 concedes that no convergence guarantees exist. This is a concrete, checkable correctness risk that goes beyond the data limitation. I also considered the D4RL label construction (Appendix B.1) as a candidate; it is a genuine evaluation-validity concern, but the GridWorld experiments with true ground-truth policies partially mitigate it, and the paper's methods are action-based rather than return-based. The convergence theorem, by contrast, is unsupported as proven for the actual algorithm. The reader's CONDITIONAL verdict remains appropriate: the empirical contributions are substantial, but the theoretical guarantee and the D4RL labeling protocol need revision. Hence UNCHANGED.","tokens_in":18494,"tokens_out":12441,"duration_ms":126505,"concrete_test":"Instrument Algorithm 1 to record J(W^t, θ^t) and J(W^{t+1}, θ^{t+1}) on the GridWorld datasets (e.g., Takeball, Diagonal). Before termination, the proof requires strict increase of max_θ J(W^{t-1}, θ); at minimum, any run that visits the same assignment matrix twice, or in which J fails to strictly increase after an E-step, falsifies the guarantee for the implemented method. To isolate the cause, also run PG-Kmeans with an exact tabular maximum-likelihood M-step (replace neural BC with counts) and verify the proof's chain holds; if it does, the flaw is specifically the approximate neural M-step, confirming that Theorem 5.1 as stated does not cover Algorithm 1.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The proof of Theorem 5.1 (Appendix A.1) builds a strictly increasing sequence {max_θ J(W^{t-1}, θ)}. The key step is equation (5), max_θ J(W^{t-1}, θ) = J(W^{t-1}, θ^t), which requires the M-step of Algorithm 1 to return a global maximizer of J over policy parameters. In the implemented algorithm, each M-step trains a neural-network behavior-cloning model with SGD; it has no guarantee of reaching the global maximum, and the training objective may decrease during optimization. Without (5), the chain (5)–(8) collapses, and the finite-termination argument (bounded by k^N assignment matrices) no longer applies; the algorithm could in principle cycle among assignments. Section 7 explicitly states 'the algorithms currently lack theoretical convergence guarantees', directly contradicting Theorem 5.1 and the abstract's claim. The empirical observation that PG-Kmeans converged within 20 iterations on the tested datasets is not a surrogate for a proof and does not hold for the approximate M-step. This is load-bearing because the convergence guarantee is presented as a headline theoretical contribution; if the theorem is unsupported, the paper's theoretical claims must be retracted or reduced to 'empirical convergence only'.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":18730,"tokens_out":4404,"duration_ms":47765,"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":[{"comment":"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.","section":"§5.1.2, Appendix A.1, Section 7"},{"comment":"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.","section":"Appendix B.1, Table 1"},{"comment":"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.","section":"Table 1, Section 6.2"},{"comment":"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.","section":"Appendix B.3, Table 1"}],"minor_comments":[{"comment":"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":"Appendix A.2, Theorem A.2"},{"comment":"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.","section":"Section 7 vs. Appendix A.3"},{"comment":"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.","section":"Appendix B.1"},{"comment":"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.","section":"§6.2, Figure 3"},{"comment":"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.","section":"Various"}],"recommendation":"major_revision","confidential_remarks":"The paper has a genuine tension between its advertised theoretical guarantee and its own limitation section; the D4RL label construction is a substantial evaluation concern. The SORL comparison should also be corrected before this is suitable for publication. None of these issues appear to be intentional misrepresentation, but they are load-bearing and need to be fixed in the manuscript."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nThe paper is a solid first cut at a real problem—clustering offline RL trajectories by the policy that generated them. The two algorithms (PG-Kmeans, CAAE) are sensible, the KL-divergence derivation is clean, and the GridWorld experiments with known ground-truth policies show both methods can work when policies actually conflict. The K-coloring reduction is a nice way to expose the inherent non-uniqueness of the clustering solution. If you work on offline RL or dataset decomposition, it's worth a read.\n\nThe soft spots are real but not all equal. The biggest is the convergence claim. Theorem 5.1 says PG-Kmeans converges in finite steps; the proof in Appendix A.1 depends on an exact global M-step optimum. The implemented algorithm trains a neural policy with SGD, which does not provide that. Section 7 then says the algorithms 'currently lack theoretical convergence guarantees'—a direct contradiction. This is load-bearing because convergence is advertised as a main contribution. The authors should either prove convergence under an approximate M-step with a sufficient-decrease condition, or drop the theorem and present empirical convergence only. As written, the theorem is unsupported.\n\nThe D4RL evaluation has a second issue: the 'ground truth' labels are not the true generating policies but a return-based split of the first and second halves of the dataset. The authors acknowledge an error of up to five episodes; that is probably minor for NMI, but it is not the same as knowing the policies. The GridWorld results, where labels are known, are the stronger evidence and should be given more weight.\n\nThe two failures (Pathfollowing NMI 0.12, Extra 0.02 for PG-Kmeans) are honestly explained as low conflict rates and overfitting. That's a legitimate limitation, but it means the abstract's claim of 'effectively partition trajectories into meaningful clusters' is too strong as stated. It should be qualified.\n\nMinor: the paper says no prior algorithm was designed for policy-based clustering, while its own related work cites SORL and Wang et al. 2024, which do policy-level clustering. That inconsistency should be fixed.\n\nOverall: the empirical contribution, especially on GridWorld, is credible, and the problem framing is useful. The theory section needs repair. I'd send it to peer review with a clear request to address the convergence contradiction and recalibrate the claims. It's not a desk reject.","headline":"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.","tokens_in":19272,"tokens_out":4050,"would_cite":false,"duration_ms":36504,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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.","keywords":["offline reinforcement learning","trajectory clustering","policy mixture","behavior cloning","policy-guided K-means","centroid-attracted autoencoder","normalized mutual information"],"falsifier":"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.","tokens_in":1612,"feed_emoji":"🧭","tokens_out":1840,"duration_ms":103391,"temperature":0.7,"pith_summary":"The paper proposes a new offline reinforcement-learning task: take a fixed dataset of trajectories and group them by the policy that generated each one, so each cluster center is itself a policy. It argues this is worth doing because offline datasets often mix several behavior policies, and training one policy on the whole mixture causes interference; clustering first could improve data use, policy evaluation, and semi-supervised RL. To attack the task, the paper derives a clustering objective from the KL divergence between the empirical trajectory distribution and a mixture of policy-induced distributions, and introduces two algorithms: Policy-Guided K-means (PG-Kmeans), which alternates behavior cloning and cluster assignment, and a Centroid-Attracted Autoencoder (CAAE), which pulls trajectory embeddings toward codebook entries. On D4RL and GridWorld benchmarks, both methods recover the generating policies with high Normalized Mutual Information in most settings, though performance collapses when the policies rarely disagree, a limitation the paper ties to the datasets' inherent indistinguishability. A reduction from K-coloring shows the general problem is NP-complete and that equally valid clusterings can be structurally different.","feed_headline":"Two algorithms sort offline RL trajectories by their generating policy","feed_subtitle":"Two new methods sort mixed-policy trajectories on D4RL and GridWorld, failing only when the policies rarely disagree.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Supplies the K-means convergence argument used as the template for PG-Kmeans's finite-step convergence theorem.","marker":"Bottou and Bengio [1995]"},{"why":"Provides the D4RL medium-expert datasets used for the continuous-control experiments.","marker":"Fu et al. [2020]"},{"why":"Defines DEC, the deep clustering baseline that PG-Kmeans and CAAE must outperform.","marker":"Xie et al. [2016]"},{"why":"Introduces SORL, the EM-style policy clustering algorithm used as the primary point of comparison.","marker":"Mao et al. [2024]"},{"why":"Defines the Normalized Mutual Information metric used for all clustering scores.","marker":"Strehl and Ghosh [2002]"},{"why":"The variational autoencoder formulation underlying CAAE's architecture and the VAE+Kmeans baseline.","marker":"Kingma and Welling [2013]"}],"fun_headline_variants":["Cluster offline RL trajectories by their generating policy","Offline RL: identify which policy produced each trajectory","Sort mixed trajectories into policy groups without rewards","New methods cluster trajectories by hidden policy, no rewards needed","Policy-based clustering uncovers trajectory sources in offline RL"],"cache_read_input_tokens":21376,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["Cluster offline RL trajectories by their generating policy","Offline RL: identify which policy produced each trajectory","Sort mixed trajectories into policy groups without rewards","New methods cluster trajectories by hidden policy, no rewards needed","Policy-based clustering uncovers trajectory sources in offline RL"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000295,"raw_usage":{"total_tokens":1740,"prompt_tokens":999,"completion_tokens":741,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":615,"completion_tokens_details":{"reasoning_tokens":668}},"tokens_in":615,"tokens_out":741,"duration_ms":7950,"temperature":1.0,"reasoning_tokens":668,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T04:54:55.464743+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[],"review_version":1}