Prototype-Wrapper Networks for Interpretable RL

• The paper introduces PW-Nets, which wrap a pre-trained black-box RL policy with an interpretable prototype layer to explain action choices via exemplar comparisons.
• It employs an automated manifold-based prototype discovery method that selects representative latent states from rollout data, reducing manual tuning.
• Empirical results across diverse tasks show that PW-Nets maintain control performance while enhancing interpretability on both standard and out-of-distribution scenarios.

Prototype-Wrapper Networks (PW-Nets) are a prototype-based interpretability method for reinforcement learning (RL) in which a pre-trained black-box policy is wrapped by an interpretable prototype reasoning layer rather than replaced outright. In the formulation summarized in (Vamshi et al., 30 Mar 2026), PW-Nets preserve the original encoder and re-express action selection through similarities to reference prototypes, yielding exemplar-based explanations of why a policy chose a given action in a given state. Within RL, this design is motivated by the difficulty of transferring standard computer-vision and natural-language explanation methods to sequential decision-making, where explanations must remain meaningful over trajectories, policy-induced state distributions, and out-of-distribution behavior while avoiding degradation of control performance (Vamshi et al., 30 Mar 2026).

1. Problem setting and conceptual scope

As presented in (Vamshi et al., 30 Mar 2026), PW-Nets address a specific interpretability problem in RL: strong agents in games, robotics, autonomous driving, and language-model alignment are typically black boxes, yet deployment often requires understanding why an action was chosen, when the policy may fail, and how the agent behaves outside the training distribution. The RL setting complicates explanation because decisions are embedded in trajectories rather than isolated predictions, policy behavior depends on the induced state distribution, and many interpretable model classes underperform the original black-box controller.

PW-Nets are positioned as a compromise between these competing requirements. Rather than training a fully interpretable policy from scratch, they operate as a post-hoc wrapper around an already trained policy. This preserves the predictive or control competence of the underlying model while exposing decisions through human-inspectable reference states, or prototypes, associated with actions or behavioral tendencies. The resulting explanation style is case-based: the policy can be interpreted as choosing an action because the current state resembles representative past states.

A central limitation of earlier PW-Net usage, again as described in (Vamshi et al., 30 Mar 2026), is dependence on manually defined reference prototypes. That requirement introduces expert labor, environment-specific tuning, reduced reproducibility, and poor scalability, especially in high-dimensional or non-visual settings such as HumanoidStandup, where states are vector-valued and actions are continuous.

2. Architectural formulation

The black-box policy assumed by the PW-Net formulation has a final linear layer and is decomposed into an encoder and a linear readout:

$\pi_{\text{bb}}(s) = W f_{\text{enc}}(s) + b,$

where $s$ is the state, $z=f_{\text{enc}}(s)$ is the latent state, and $W,b$ are final-layer parameters (Vamshi et al., 30 Mar 2026). Prototype discovery or selection is performed over rollout data

$$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$

so the wrapper is grounded in the latent state space already used by the black-box policy.

The PW-Net wrapper computes action outputs through prototype similarities rather than direct latent linear scores. The action score for action $i$ is written as

$$a'_i = \sum_{j=1}^{N_i} W'_{i,j}\,\text{sim}(z_{i,j}, p_{i,j}),$$

where $p_{i,j}$ is the $j$-th prototype associated with action $i$, $s$0 is a wrapper weight, and $s$1 is the number of prototypes for that action (Vamshi et al., 30 Mar 2026). The similarity function is

$s$2

with $s$3 (Vamshi et al., 30 Mar 2026).

This formulation gives the wrapper a direct explanatory semantics. The prototypes act as representative examples; the similarity term measures resemblance between the current latent state and those examples; and the wrapper weights quantify how strongly each prototype supports an action. In that sense, PW-Nets translate a latent control policy into a prototype-conditioned argument over states previously encountered by the agent. The explanation is local because it addresses a specific decision, but it can also provide a limited global view by revealing what kinds of states characterize each action class (Vamshi et al., 30 Mar 2026).

A common misconception is that PW-Nets replace the black-box model with a prototype classifier. The formulation above indicates otherwise: the black-box encoder is retained, and the prototype layer is a wrapper over its representation. This makes the method post-hoc and modular rather than a full re-specification of the policy class.

3. Prototype acquisition: manual curation and automated selection

In earlier PW-Net practice, prototypes were manually specified. The 2026 extension in (Vamshi et al., 30 Mar 2026) treats that dependency as the main practical bottleneck and proposes a two-stage remedy, termed Principal Prototype Analysis on Manifold. The first stage performs automatic prototype discovery from rollout data, and the second stage trains the PW-Net wrapper using the discovered prototypes.

For a pre-trained policy, rollout data are collected in latent space. If the original action space is continuous, actions are discretized by taking absolute values of action components, applying a sigmoid transformation, and assigning the class index of the maximum activation. The stated objective is not exact reconstruction of the continuous action space, but grouping by dominant action dimension (Vamshi et al., 30 Mar 2026).

Prototype discovery uses a lightweight transformation network $s$4 with a Linear layer from latent size to prototype size, followed by InstanceNorm1d and ReLU; the prototype size is fixed to 50 in all environments (Vamshi et al., 30 Mar 2026). The method assumes that encoded states lie near a lower-dimensional manifold and therefore constructs local piecewise-linear manifolds. Around anchor points, neighborhoods are expanded while preserving PCA reconstruction quality above threshold $s$5, and the local submanifold dimension is set to $s$6 (Vamshi et al., 30 Mar 2026).

Training combines metric discrimination and geometry preservation through

$s$7

where $s$8 is a modified Proxy Anchor loss and

$s$9

aligns transformed Euclidean distances with a manifold similarity $z=f_{\text{enc}}(s)$0 (Vamshi et al., 30 Mar 2026). The similarity is built from orthogonal distance to a local manifold and projected distance along that manifold, with the stated intent that off-manifold deviation be penalized more sharply than along-manifold separation.

Two proxy sets are maintained: a backpropagated proxy $z=f_{\text{enc}}(s)$1 and a momentum proxy $z=f_{\text{enc}}(s)$2 updated by

$z=f_{\text{enc}}(s)$3

After training, each momentum proxy is mapped to the nearest encoded training instance, ensuring that the final prototypes are real data examples rather than abstract latent vectors (Vamshi et al., 30 Mar 2026). This is a significant design choice because the explanatory object remains an inspectable state sampled from the agent’s experience.

The reported training parameters for prototype discovery are 200 epochs, Adam, learning rate $z=f_{\text{enc}}(s)$4, scheduler decay $z=f_{\text{enc}}(s)$5, batch size 128, threshold $z=f_{\text{enc}}(s)$6, $z=f_{\text{enc}}(s)$7, and submanifold dimension $z=f_{\text{enc}}(s)$8 (Vamshi et al., 30 Mar 2026). The second-stage PW-Net training uses Adam, learning rate $z=f_{\text{enc}}(s)$9, scheduler decay $W,b$0, 10 epochs, 5 independent iterations, and 30 simulation runs per evaluation iteration (Vamshi et al., 30 Mar 2026). Non-parallel prototype discovery is reported at about 7 hours and 640 MB, while a parallelized manifold-similarity implementation is reported at about 2 hours with 4700 MB peak memory (Vamshi et al., 30 Mar 2026).

4. Empirical behavior in reinforcement learning

The evaluation summarized in (Vamshi et al., 30 Mar 2026) spans CarRacing, BipedalWalker, HumanoidStandup, Pong, LunarLander, and Acrobot, with black-box teachers including PPO, Double Dueling DQN, TD3, Actor-Critic, and PPO with an MLP policy from Stable-Baselines3. Comparisons include the automated prototype-selection method, PW-Net, PW-Net*, VIPER, k-means prototype selection, and the original black-box.

On CarRacing, PW-Net achieves $W,b$1 and the black-box achieves $W,b$2, while the automated extension reaches $W,b$3 (Vamshi et al., 30 Mar 2026). On BipedalWalker the corresponding values are $W,b$4, $W,b$5, and $W,b$6; on HumanoidStandup, $W,b$7, $W,b$8, and $W,b$9; on Pong, $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$0, $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$1, and $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$2; on LunarLander, $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$3, $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$4, and $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$5; and on Acrobot, $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$6, $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$7, and $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$8 (Vamshi et al., 30 Mar 2026). The direct conclusion drawn in that study is that automated prototype selection matches or exceeds baseline PW-Net across all tested environments while remaining competitive with, and sometimes slightly above, the original black-box policy.

Interpretability was assessed directly in a CarRacing user study. Two groups of 25 participants were shown 20 scenarios: 10 in-distribution safe-driving cases and 10 out-of-distribution failure cases from a modified CarRacing environment with new road types and red obstacles. Participants rated expected safety on a five-point Likert scale (Vamshi et al., 30 Mar 2026). The reported finding is that both groups behaved similarly on in-distribution cases, but prototype-based explanations improved anticipation of failures in out-of-distribution cases; the black-box group tended to overestimate safety under failure scenarios (Vamshi et al., 30 Mar 2026).

The same study compared automated and human-curated prototypes. For acceleration, automated and human-curated prototypes were reported as comparable, whereas for some other action classes human-curated prototypes were slightly preferred (Vamshi et al., 30 Mar 2026). This indicates that the explanatory quality of prototypes is not exhausted by control performance: state representativeness remains a separate concern, particularly in visual domains.

Ablations in Pong and BipedalWalker further show that local manifold dimension, proxy momentum, the decay exponents governing off-manifold and along-manifold similarity, reconstruction threshold $$D \leftarrow \{(z_i,\pi_{\text{bb}}(s_i))\}_{i=1}^n,$$9, and number of prototypes all matter (Vamshi et al., 30 Mar 2026). Higher $i$0 and higher momentum $i$1 help; larger $i$2 helps while larger $i$3 hurts; more prototypes help up to a plateau in BipedalWalker but eventually degrade performance in Pong, suggesting that symbolic states tolerate larger prototype sets more readily than pixel encodings.

5. Relation to the broader prototype-model literature

Several earlier prototype-based models illuminate architectural motifs that later become central to PW-Nets, although they do not use the PW-Net term. A foundational precursor is the prototype-based neural layer framework that shows how distance-to-prototype computations can be realized using standard matrix multiplication, bias addition, and convolution, and how softmax over negative distances can serve as a classifier head (Saralajew et al., 2018). This suggests that the PW-Net wrapper is not an isolated construct but a task-specific instance of a broader family of differentiable prototype layers.

Other lines of work reveal wrapper-like behaviors around prototypes. One-way Prototypical Networks replace explicit negative supports with a null class centered at zero, yielding a prototype-plus-rejection mechanism for membership decisions (Kruspe, 2019). AutoProtoNet adds a decoder to a Prototypical Network so that latent prototypes can be reconstructed as images and manually refined through latent interpolation, showing that a prototype learner can be augmented with a human-facing interpretability layer and editable class representations (Sandoval-Segura et al., 2022). Concept-level debugging of ProtoPNets goes further by allowing a supervisor to mark prototypes to forget or keep and then fine-tuning the model with forgetting and remembering losses, indicating that prototype systems can support targeted human intervention at the concept level rather than only at the level of raw inputs (Bontempelli et al., 2022).

Two additional strands are especially suggestive for PW-Net design. Winner-take-all networks interpreted through positive and negative prototype pairs show that a discriminative neuron can be viewed as comparing distance to a supportive prototype against distance to a confusing prototype, rather than matching a single prototype alone (Sabzevar et al., 2020). Robust Classification-by-Components models with interpretable RBF foundations argue that deep prototype-based networks are closely related to RBF classifiers and that classifier-head constraints, class-wise component priors, and probabilistic negative reasoning are essential if prototype explanations are to remain semantically coherent (Saralajew et al., 2024). A plausible implication is that PW-Nets can be understood not merely as storing prototypes, but as combining prototype comparison with an explicit decision wrapper whose design materially affects faithfulness, robustness, and explanatory content.

6. Limitations, misconceptions, and open directions

PW-Nets should not be conflated with fully interpretable RL policies. Their interpretability is wrapper-level and exemplar-based, not a complete semantic account of the original representation. The black-box encoder remains in place, so the method explains decisions through similarity to prototypes without fully explaining why the encoder arranged latent states as it did (Vamshi et al., 30 Mar 2026). This aligns with a broader observation from deep prototype-based models: once prototypes operate in a learned latent space on top of a deep backbone, interpretability is at best partial rather than inherent (Saralajew et al., 2024).

A second limitation is operational. Earlier PW-Net use depends on manual prototype selection, which requires expert knowledge and scales poorly. The automated manifold-based extension addresses this, but introduces its own assumptions: encoded states must admit useful locally linear manifold structure; continuous actions must be discretized into dominant action classes; and prototype discovery remains post-hoc and policy-specific rather than integrated into RL training (Vamshi et al., 30 Mar 2026). The same source identifies additional open issues, including scalability to very large action spaces, sensitivity of pixel domains to excessive prototype counts, and the fact that human-curated prototypes can still be slightly preferred in subjective evaluations.

Future directions proposed in (Vamshi et al., 30 Mar 2026) include hierarchical or clustering-based prototype strategies for larger output spaces, geometry-aware prototype discovery during policy training as a soft regularizer, multi-task or transfer settings with shared behavioral primitives, and evaluation in high-stakes domains such as autonomous driving and robotic manipulation. Taken together, these directions indicate that PW-Nets are best understood not as a fixed architecture, but as an emerging prototype-wrapper paradigm for RL in which performance preservation, exemplar faithfulness, and human interpretability must be negotiated simultaneously.