Training interpretable concept-based policies requires practitioners to manually select which human-understandable concepts an agent should reason with when making sequential decisions. This selection demands domain expertise, is time-consuming and costly, scales poorly with the number of candidates, and provides no performance guarantees. To overcome this limitation, we propose the first algorithms for principled automatic concept selection in sequential decision-making. Our key insight is that concept selection can be viewed through the lens of state abstraction: intuitively, a concept is decision-relevant if removing it would cause the agent to confuse states that require different actions. As a result, agents should rely on decision-relevant concepts; states with the same concept representation should share the same optimal action, which preserves the optimal decision structure of the original state space. This perspective leads to the Decision-Relevant Selection (DRS) algorithm, which selects a subset of concepts from a candidate set, along with performance bounds relating the selected concepts to the performance of the resulting policy. Empirically, DRS automatically recovers manually curated concept sets while matching or exceeding their performance, and improves the effectiveness of test-time concept interventions across reinforcement learning benchmarks and real-world healthcare environments.
Figures & tables
Figure 1: Standard pipeline for training concept-based policies. Practitioners select concepts for decision-making through a labor-intensive process of iteratively selecting candidate concepts, training concept-based policies, and evaluating their performance.
Figure 2: Concept-based models rely on a set of decision-relevant concepts that help distinguish between different states, yet currently these concepts are manually selected. In this work, we study how to identify and select decision-relevant concepts. Our key insight is that decision-relevant concepts best separate “different” states, where difference is defined by their decision consequences. We use this insight to develop algorithms for concept selection with performance guarantees.
Figure 3: Normalized reward of concept selection algorithms with perfect (top) and imperfect (bottom) concept predictors. Our algorithm, DRS, improves performance compared to the random, variance, and greedy baselines for four out of five environments in the perfect setting. DRS and DRS-log improve performance or are optimal in all environments in the imperfect setting.
Figure 4: We vary the number of timesteps that we train policies for in MiniGrid, while also varying the accuracy of concept predictors (left) or the number of concepts selected (right). Increasing the accuracy of concept predictors speeds up training, while increasing the number of concepts increases the maximum performance.
Figure 5: Impact of the number of concepts selected against the accuracy of the underlying concepts. Increasing either number or accuracy of concepts has a similar impact on performance, and that sufficiently accurate and many concepts are needed to ensure good performance.
Figure 6: We assess the impact of concept selection upon test-time intervention across four environments. Across all four environments, the DRS algorithm has the highest reward both before and after intervention, showing that well-selected concepts improve intervention.
Figure 7: The DRS algorithm can mimic manual concept selection in CUB. It matches performance (left) while selecting a subset of the manually selected concepts (right).
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: We assess the performance of a) not relaxing constraint P1c for DRS and b) varying ρ for CartPole and MiniGrid. We find that ρ=0.75 performs well across perfect and imperfect settings, while introducing P1c leads to lower performance in CartPole.
Figure 9: We compare concept selection algorithms while varying the number of concepts selected k with perfect (top) and imperfect (bottom) concept predictors. We find that for CartPole and MiniGrid, DRS outperforms baselines across values of k , while for environments such as Boxing and Glucose, a threshold number of concepts is needed before DRS outperforms baselines.
Concept Group
Avg. Selected
Total
Agent
Position X
1
5
Position Y
2
5
Direction
3
4
Action: Right
0
2
Action: Left
1
2
Appendix
Table 1: Average number of concepts selected per group by DRS in MiniGrid, with total concepts per group.
Figure 10: We compare the time to compute across methods and datasets, and find that all algorithms run in under ten minutes. While DRS and DRS-log take longer than variance and greedy, they still run quickly across environments.
Figure 11: We vary the number of timesteps that we train policies for in CartPole, while also varying the accuracy of concept predictors (left) or the number of concepts selected (right). When the number of concepts is small or the accuracy of the concepts is low, policies fail to train, with the efficiency of training increasing with both parameters.
Figure 12: We assess test-time intervention performance when varying the level of intervention α in CartPole. Across all α , the DRS algorithm performs best, with steadily improving performance as α is increased.
Figure 13: We vary the number of timesteps we use to train the policy π without concepts in MiniGrid. Typically, we have π=π∗ , but here, we demonstrate that the DRS and DRS-log algorithms perform well even when π is trained for only 100k timesteps.
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
Reinforcement learning policies are difficult to inspect, but interpreting them is a prerequisite for trustworthiness. Converting a trained policy into explicit decision-tree rules improves transparency and the resulting artifacts often remain too complex for human understanding. We present a pruning process that simplifies such rule-based policies while preserving task performance and making edits to the policy auditable. The process defines a small set of structural and usage-aware operators and evaluates candidate edits by re-executing the policy to measure return and interpretability proxies. This exposes an transformation process from complex to compact policy structures. We investigate this approach on classic control and MuJoCo benchmarks, where pruning traces reveal consistent interpretability improvements while maintaining high performance.
Mark Leon Ringer, Michel Tokic
Faculty of Mathematics, Informatics and Statistics, Ludwig-Maximilians-University Munich, Munich, Germany · Siemens AG, Data & Artificial Intelligence, Otto-Hahn-Ring 6, 81739 Munich, Germany
In-context learning (ICL) allows large language models to use demonstrations without updating their parameters. For interactive agents, demonstrations are trajectories of reasoning, actions, and observations that illustrate how to complete a task. Reusing demonstrations in new situations is challenging: directly inputting full trajectories consumes limited prompt space and constrains the overall effectiveness, while natural-language summaries require the agent to translate general guidance into specific action choices. Therefore, we study whether expert action choices can directly supervise a reusable scoring rule. We propose CB-Agent, which extends concept-based ICL to sequential action selection. We extract environment-action decisions from expert trajectories and fit a concept vector by maximum likelihood to model the conditional distribution over candidate actions given the context. This model combines the concept with frozen language-model features and priors. During ReAct interaction, the agent uses this distribution to select environment actions without including demonstrations in the action-scoring context. Across TextWorld, ALFWorld, and ScienceWorld, CB-Agent achieves the highest cross-benchmark average among the evaluated methods at all three tested model sizes and ranks first on TextWorld and ScienceWorld at each size. ALFWorld experiments support the learned concept's contribution beyond prior-only scoring and show further gains when combined with Reflexion. Code is available at https://anonymous.4open.science/r/CB-Agent-DFF9/.
Qingyue Zhang, Huaze Tang, Minyu Zhang +3
Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen, China
Neuro-symbolic Reinforcement Learning (NeSy-RL) combines symbolic reasoning with gradient-based optimization to achieve interpretable and generalizable policies. Relational concepts, such as "left of" or "close by", serve as foundational building blocks that structure how agents perceive and act. However, conventional approaches require human experts to manually define these concepts, limiting adaptability since concept semantics vary across environments. We propose GRAIL (Grounding Relational Agents through Interactive Learning), a framework that autonomously grounds relational concepts through environmental interaction. GRAIL leverages large language models (LLMs) to provide generic concept representations as weak supervision, then refines them to capture environment-specific semantics. This approach addresses both sparse reward signals and concept misalignment prevalent in underdetermined environments. Experiments on the Atari games Kangaroo, Seaquest, and Skiing demonstrate that GRAIL matches or outperforms agents with manually crafted concepts in simplified settings, and reveals informative trade-offs between reward maximization and high-level goal completion in the full environment.
Hikaru Shindo, Henri Rößler, Quentin Delfosse +1
Technical University of Darmstadt · Google Intrinsic · German Research Center for Artificial Intelligence (DFKI) +1