Towards Automated Semantic Interpretability in Reinforcement Learning via Vision-Language Models

TL;DR

iTRACE introduces a fully automated workflow to achieve semantic interpretability in reinforcement learning by leveraging Vision-Language Models to discover human-understandable features and training an Interpretable Control Tree policy. A lightweight Extractor Module replaces costly VLM feature extraction to enable efficient RL training, while zero-shot feature discovery and segmentation grounding ensure generalization to unseen environments. Across Atari Skiing, BabyAI-GoToRedBall, and OpenAI Gym Highway, iTRACE attains performance competitive with black-box policies while providing transparent, verifiable decision-making. The framework demonstrates fast, trainable, and interpretable pixels-to-actions control, highlighting a practical path toward scalable semantic interpretability in RL.

Abstract

Semantic interpretability in Reinforcement Learning (RL) enables transparency and verifiability of decision-making. Achieving semantic interpretability in reinforcement learning requires (1) a feature space composed of human-understandable concepts and (2) a policy that is interpretable and verifiable. However, constructing such a feature space has traditionally relied on manual human specification, which often fails to generalize to unseen environments. Moreover, even when interpretable features are available, most reinforcement learning algorithms employ black-box models as policies, thereby hindering transparency. We introduce interpretable Tree-based Reinforcement learning via Automated Concept Extraction (iTRACE), an automated framework that leverages pre-trained vision-language models (VLM) for semantic feature extraction and train a interpretable tree-based model via RL. To address the impracticality of running VLMs in RL loops, we distill their outputs into a lightweight model. By leveraging Vision-Language Models (VLMs) to automate tree-based reinforcement learning, iTRACE loosens the reliance the need for human annotation that is traditionally required by interpretable models. In addition, it addresses key limitations of VLMs alone, such as their lack of grounding in action spaces and their inability to directly optimize policies. We evaluate iTRACE across three domains: Atari games, grid-world navigation, and driving. The results show that iTRACE outperforms other interpretable policy baselines and matches the performance of black-box policies on the same interpretable feature space.

Figures (5)

• Figure 1: High-level overview of iTRACE.
• Figure 2: Overview of iTRACE, our proposed framework. A. Feature Extraction follows three steps: (1) the VLM identifies key semantic features from domain knowledge and observations; (2) it generates code to extract features from SAM-based segmentations, extending to multi-frame inputs if needed; and (3) an exploration policy generates feature–image pairs, which are used to train a lightweight Extractor Module to replace costly SAM queries during RL. B. Tree-based RL Training then leverages the Extractor Module to obtain interpretable features, which serve as inputs to an Interpretable Control Tree for policy learning.
• Figure 3: An example of how raw image observations are processed in our pipeline. (a) Raw image observations are segmented using SAM. (b) A VLM assigns semantic labels to segments in a small set of example images. (c) A VLM then generates code that extracts target features from segmentation masks, applicable to all images. Step (b) is only performed on the example set to guide code generation.
• Figure 4: Randomly selected segmentation examples we provide to the VLM for domain Skiing. Example 1 is good while 2 and 3 are imperfect. Example 2 identifies the skier as a partial skier and a noisy segment as skier. Example 3 identifies both the actual skier on the right and decoration elements at the bottom as partial skier.
• Figure 5: Results across three domains. Higher scores are better.