Title: Enhancing Interpretability in Deep Reinforcement Learning through Semantic Clustering
Venue: NeurIPS 2025
Accepted: yes
Keywords: Interpretability

Reinforcement Learning

Clustering

Visualization

Neural Networks
Abstract: In this paper, we explore semantic clustering properties of deep reinforcement learning (DRL) to improve its interpretability and deepen our understanding of its internal semantic organization. In this context, semantic clustering refers to the ability of neural networks to cluster inputs based on their semantic similarity in the feature space. We propose a DRL architecture that incorporates a novel semantic clustering module that combines feature dimensionality reduction with online clustering. This module integrates seamlessly into the DRL training pipeline, addressing the instability of t-SNE and eliminating the need for extensive manual annotation inherent to prior semantic analysis methods. We experimentally validate the effectiveness of the proposed module and demonstrate its ability to reveal semantic clustering properties within DRL. Furthermore, we introduce new analytical methods based on these properties to provide insights into the hierarchical structure of policies and semantic organization within the feature space. Our code is available at https://github.com/ualiangzhang/semantic_rl.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a plug-in “semantic clustering module” for deep-reinforcement-learning (DRL) agents that pairs a lightweight 2-D feature-dimensionality-reduction (FDR) network with an online vector-quantiser. Jointly trained with any standard DRL algorithm (shown with PPO on the Procgen benchmark), the module yields a stable, human-interpretable partition of the latent state space into discrete clusters without hurting task performance. These clusters expose the structure of learned policies, enable video-level explanations such as “approaching the right-wall then jumping to a higher ledge,” and let practitioners inspect or segment roll-outs by semantic “skills.” Experiments across CoinRun, Ninja, Jumper and FruitBot demonstrate markedly clearer, seed-independent clusters than t-SNE visualisations, high agreement in a 15-person user study (Likert ≈ 4.2/5), and no significant score drop versus baseline PPO. Overall, the work shows that DRL already forms meaningful semantic groupings internally and that making them explicit dramatically improves transparency while remaining computationally light.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n* Novel approach to the important task of interpretable reinforcement learning\n* Strong and versatile analysis of the proposed method\n* Rigorous human study (though limited for making stronger behavioral claims)\n* Useful visualization tool\n\n**Weaknesses** (see detailed explanations in the questions):\n\n* Limited discussion on the benefits of adding vector codes to the state features\n* Inconsistency between Section 4.1 and Appendix C.2\n* Limited performance comparison and analysis of the sensitivity to the newly introduced hyper-parameter\n* Additional weaknesses noted in the discussion: choice of *k*, lack of automated description generation, etc.\n\nIn essence, this is a strong paper with several weaknesses that, however, do not undermine the overall contribution.\n\nQUESTIONS:\n* Limited discussion on the benefits of adding vector codes to the state features\n\nThe authors state that “Training the VQ code $k$ with a latent-conditioned policy $\\pi(a∣s,k)$ supports the extension to downstream tasks, such as macro action selection in hierarchical learning.” However, this appears to be the only justification provided for incorporating the VQ code into the state features. Given that the paper does not include any experiments in a hierarchical learning setting, nor an ablation study comparing performance with and without the VQ-code, this design choice feels insufficiently motivated. Since the main goal of the paper is to cluster state features—rather than to influence the policy—further discussion is warranted. An ablation study examining the effect of excluding the VQ-code could significantly enhance the clarity and rigor of the work.\n\n* Inconsistency between Section 4.1 and Appendix C.2\n\nIn Section 4.1, the authors state:\n“The t-SNE visualization of PPO (Figure 2a), spreads features across the space without forming clear clusters, limiting its utility for clustering analysis and requiring detailed manual examination of certain areas, as in previous studies.”\n\nHowever, in Appendix C.2, the “stop-gradient” experiment shows:\n“The results, as shown in Figures C.2 and C.3, demonstrated that states within the same semantic cluster continued to exhibit similar semantic interpretations, even without the influence of the semantic clustering module.”\n\nThis comparison seems somewhat inconsistent. If I understand correctly, the policy and the state representations in vanilla PPO and the stop-gradient method should be identical, assuming all other components remain the same. Figure C.2 in the appendix shows a t-SNE visualization of the stop-gradient method's feature space with cluster labels—yet these clusters should be equivalent to those from vanilla PPO, as the policy is unaffected. In that case, the two quoted statements are not strictly contradictory, but they do feel awkward when read together. The apparent discrepancy seems largely due to the absence of color labeling in Figure 2, which may lead to a misleading interpretation.\n\nI recommend revising Section 4.1 to explicitly incorporate and contrast the stop-gradient experiment, clarifying how your method differs from it and why joint training offers benefits. Additionally, it would be helpful to include an *unlabeled* version of your method's t-SNE plot, to mirror the presentation of vanilla PPO more symmetrically. This would provide a fairer and clearer comparison of visual cluster separation and further support your narrative.\n\n\nFurthermore, if the main difference between your method and the stop-gradient baseline lies in sharper cluster boundaries, then this distinction deserves richer discussion why is this important. Your analysis operates on the clusters themselves, not directly on the FDR or t-SNE visualizations (which shows sharper bounds). More evidence is needed to support the claim that sharper boundaries improve semantic grouping. For instance, examples of ambiguous or misclustered states at the boundaries would be valuable.\n\nLastly, the Semantic Formation experiment in Appendix F is not fully convincing. The statistics are averaged over 50,000 steps; it would be more informative to show episode-level metrics or frame-wise agreement scores. A comparison of Semantic Formation between your method and the stop-gradient baseline could help demonstrate the added value of your design. If your method indeed shows higher consistency or interpretability on these metrics, that would substantially strengthen your claims.\n\n* Limited performance comparison and sensitivity analysis of the newly introduced hyper-parameters\n\nPerformance comparisons are limited to a small set of Procgen environments. Evaluation on additional domains—such as Atari or other standard DRL benchmarks—would improve the generalizability of your claims.\n\nIt is commendable that your method preserves performance without requiring changes to the original PPO hyper-parameters and also demonstrates robustness to random seeds—both of which are important properties in the reinforcement learning domain and provides study of the effect of K. However, the method introduces new hyper-parameters (e.g., weights for the semantic loss components), and their sensitivity is not thoroughly explored. Understanding how these additional parameters affect performance and interpretability is crucial for assessing the stability and general usability of the approach. A more detailed sensitivity analysis would greatly strengthen the paper's practical value and reproducibility.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates methods for interpreting policies in deep reinforcement learning. In particular, they propose a method to cluster states used by the policy in a way that respects semantic differences in the policy behaviour. There are a few qualitative investigations in ProcGen, with the rationale that policies in procedural generation would face a generalization challenge where interpretability would be useful for understanding the learned policy.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n- The method appears to be generally applicable, allowing it to be added to a conventional RL algorithm without any change to its operation. (Although, I do have a question about whether it can change performance/learning, see below).\n\n- Qualitative analysis suggests that the proposed approach for clustering in drl is generally effective. The experiments seem to be well-designed, thought-out and provide compelling evidence for improvements over t-sne when it comes to interpreting the policy's behaviour semantically.\n\n### Weaknesses\n- The presentation of the experimental results could be improved (see below).\n\n- Unclear whether if there is an effect of the proposed feature extraction method on downstream performance. If so, then this is a weakness without an investigation to demonstrate that the performance impact is negligible.\n\nQUESTIONS:\n- Figure 1: Why is it that the VQ code is integrated into the feature, rather than the VQ embedding?\n- Section 3.1 (curse of dimensionality): Wouldn't the curse of dimensionality also apply to the problem of dimension reduction? I am not sure I understand this motivation if the VQ code needs to be learned.\n- Section 3.1 (distance relationships): Would the reduced features learned from preserving consistency across distance relationships also serve to represent relationships with respect to other distances, like bisimulation metrics?\n- Section 3.2 (FDR loss): It is not clear to me what the relationship is between p and q here, and why the loss takes the form that it does. By optimizing the FDR loss, are you ensuring that the similarities of induced by f and g \\circ f are similar?\n- Section 3.2 (vector magnitude in student T): I understand that using the same $\\alpha$ ensures the same distance relationship, but wouldn't the magnitudes of the vectors differ if f and g \\circ f are of different dimensions?\n- Section 3.2 (Modified VQ-VAE): without the third term, how does the encoder of the VQ-VAE learn to maintain the assigned embedding? I understand that the $e_k$ can cover the space as centroids of K means, but an important difference is that the underlying features are also changing as the encoder is learned.\n- Section 4 (Cluster separation and t-sne): I wonder what exactly the t-sne is capturing on top of the learned clusters of the embeddings. Does t-sne on the fused state feature essentially return the cluster index?\n- Section 4 (sensitivity to number of states/seed): While the overall structure of the clustering remains the same, I wonder if there is any way in which the identity of these clusters could be identified? For example, does the specific embedding (e_k) assigned to these clusters change?\n- Figure 3 and Table 1: The results presented here are a bit difficult to follow because the change in subsequent images is difficult to detect. For example, cluster 4 which is described as performing a jump has mostly images of the agent near the top of its jump rather than a scene closer to the ground at the beginning of the jump. I am also not sure if the mean images accurately depict the dynamic nature.\n- Table 2: While these results show the proposed method in a favourable light, they would be much stronger if the evaluation included clips from clusters of a baseline method for comparison.\n- Section 5: an important limitation that was left out is whether the method has any effect on performance. I understand that this not the primary concern of this paper, but training the policy with the proposed fused features could result in changes to performance.\n\n### Minor Comments\n- Equation 2: I suggest using another symbol for the \"control factor\", since $f$ is usually used to denote a function\n- Section 3.2 (Improved Clustering): the first paragraph feels out of place, and i think that it can be demonstrated in the experiments. It doesn't connect with the actual description of the method earlier in the section. THe second paragraph is interesting, but difficult to follow. Could the statement be made as a formal proposition or theorem?\n- Figure 2: the titles of the subfigures could be more informative or removed.\n- Figure 3: the mean images could benefit increased contrast to highlight the difference between the clusters.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors employ the following argument structure to motivate their paper:\n\n1. t-SNE has been used in DRL works to visualize the encodings of DRL agents\n2. t-SNE is unstable (sensitive to initialization and therefore has inconsistent output)\n3. t-SNE requires extensive manual cluster annotation\n4. It would be good to provide an alternative to t-SNE without the issues identified in 2 and 3\n\nTheir solution is to integrate (online) clustering into DRL training. Their solution is built on top of VQ-VAE.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- My own prior experiences using t-SNE to visualize DRL agent encodings left me desiring an alternative, so I am glad to see this paper propose one. \n- I think the authors successfully addressed the issue of t-SNE being unstable \n- Section 4.1 gives some evidence that their method yields more compact and separated clusters than t-SNE does.\n- Section 4.1 gives evidence that their method is not unstable like t-SNE is. \n- The authors provide a human evaluation. I think these are important for an XAI paper to be convincing. \n- I appreciate that the authors developed a visualization tool for analyzing the embedding space of DRL models. This is a nice addition that they didn't necessarily need to add. I liked the policy analysis at the end of section 4. I could see myself using it in the future if it is integrated into some standard python RL library. \n\n## Weaknesses\n\n- Nitpick: I would prefer if the authors gave a definition of \"semantic clustering\" that does not use the word \"semantic\" in it (unless the word \"semantic\" is also given a formal non-self-referential definition).\n\n- I would have appreciated if the **Online Clustering** paragraph in Section 3.1 did more to convince readers that previous work requires \"extensive manual annotation.\" The key word for me there is \"extensive.\" That should be substantiated. I view this as a significant omission. \n\n- I would have appreciated an explanation of Equation (4) rather than just showing the equation and ending the paragraph.\n\nQUESTIONS:\n- I found the sentence on line 107 beginning with \"We use the Student's t-distribution...\" unclear. What do you mean that it captures nonlinear structures?\n- What norm is being used in Equation (3) on line 110?\n- On Line 115-116 the authors say that \"the original distance relationship between features is maintained in the low-dimensional space.\" Is it a well-known or easy result that the student's t kernel is an isometry? If not, there should be a citation or proof for this in the appendix. \n- On line 127 did you mean to say that the clusters themselves are made more compact?\n- What are the qualifications of the human evaluators? I would have appreciated if this information was present in the main body of the paper near line 204-205. \n- How is the number of clusters controlled? Is it possible that there are overwhelmingly many for a human to interpret, and if so, what is to be done?\n\nLIMITATIONS:\nThe authors did not try to recover an inner product in their embedding space whose quantification of similarity corresponds to semantic similarity in the input space.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe work addresses the interpretability aspect of reinforcement learning, by developing a semantic clustering approach, which is trained end to end while learning a new policy. To do this the authors use a VQ-VAE inspired novel architecture quantaization to obtain low dimensional state embeddings. Through experiments on Procgen game environments, they demonstrate that DRL models naturally group states representing similar behaviors or \"skills\" together, such as \"jumping to higher platforms\" or \"approaching targets,\" without requiring external supervision.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \nThe empirical validation demonstrates compelling evidence for the method's effectiveness across multiple dimensions. The authors conduct experiments on four different Procgen environments and show through statistical comparisons that their approach outperforms traditional dimensionality reduction methods like UMAP, PCA, and t-SNE on standard clustering metrics.\n\nWeakness: \n1) No Ground Truth: They never establish what \"correct\" semantic clustering should look like\n2) The technical approach exhibits several design limitations that restrict its applicability. The number of clusters must be manually predetermined (they use K=8) with no adaptive mechanism to determine the optimal number for different environments or tasks.\n3) A critical practical limitation is the continued reliance on manual interpretation despite claims of enhancing interpretability. The method requires humans to examine clusters and manually assign semantic meanings, with no automated mechanism for generating natural language descriptions or understanding what behaviors each cluster represents.\n\nImprovements/suggestions: \n1) Introduction needs a bit more detail of the method the authors are proposing. The contribution section states that the paper proposes a novel architecture that addresses limitations. It would be better to directly mention what this method is here and what is the novelty. \n2) Missing related work: \n\n* patil et. al: contrastive abstraction for RL \n    * Uses contrastive learning to structure the space and uses hopfield network to obtain related clusters to a state. \n* Jan Robine et.al: smaller world models for reinforcement learning\n    * Use VQ-VAE to build world models. Related work for VQ-VAE and RL.\n\nQUESTIONS:\n1) “cluster inputs based on their semantic similarity in the internal space”: what is the internal space? \n2) Why element wise addition of the feature vector? why not concatenate it? \n3) If the goal is to analyse the DRL policy, then why make it conditional by adding to the policy term? What if the policy degrades? \n4) The results show that sematic similar states are clusterd together. But what if the features which procgen is changing (background), is ignored by the feature generator already (before the clustering module)? Implying the added VQ-VAE is not necessary.\n\nLIMITATIONS:\nThe limitation of the work, would be the set of environments the algorithm was tested on. It would be interesting to see it tested on partially observable environments like Minecraft etc.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 1, "ORIGINALITY": 4, "DATE": "2025-07-03"}