CrystalBox: Future-Based Explanations for Input-Driven Deep RL Systems

Deep Reinforcement Learning (DRL) outperforms manual heuristics in a broad range of input-driven tasks, including congestion control (Jay et al. 2019), adaptive bitrate streaming (Mao, Netravali, and Alizadeh 2017), and network traffic optimization (Chen et al. 2018), among others. However, due to the difficulty in interpreting, debugging, and trusting (Meng et al. 2020), network operators are often hesitant to deploy them outside lab settings. The domain of explainable reinforcement learning (XRL) seeks to bridge this gap.

The field of Explainable Reinforcement Learning, which facilitates human understanding of a model’s decision-making process (Burkart and Huber 2021), can be categorized into two main areas: feature-based and future-based. The feature-based approach applies established techniques from Explainable AI (XAI) in supervised learning to DRL, adapting commonly used post-hoc explainers like saliency maps (Zahavy, Ben-Zrihem, and Mannor 2016; Iyer et al. 2018; Greydanus et al. 2018; Puri et al. 2019) and model distillation (Bastani, Pu, and Solar-Lezama 2018; Verma et al. 2018; Zhang, Zhou, and Lin 2020) techniques. However, these methods can fail to provide a comprehensive view or explain specific behaviors, because they do not provide insight into the controller’s view of the future.

More recently, there is a growing interest in future-based explainers, which generate explanations based on a controller’s forward-looking view. These explainers detail the effects of a controller’s decisions using either future rewards (Juozapaitis et al. 2019) or goals (van der Waa et al. 2018; Yau, Russell, and Hadfield 2020; Cruz et al. 2021). However, they necessitate either extensive alterations to the agent (Juozapaitis et al. 2019) or precise environment modeling (van der Waa et al. 2018; Yau, Russell, and Hadfield 2020; Cruz et al. 2021), both of which pose significant challenges in input-driven environments. Controllers in this domain cannot tolerate performance degradation and are deployed in highly variable real-world conditions.

We present CrystalBox, a model-agnostic, posthoc explainability framework that generates high-fidelity future-based explanations in input-driven environments. CrystalBox does not require any modifications to the controller; hence, it is easily deployable and does not sacrifice controller performance. CrystalBox generates succinct explanations by decomposing the controller’s reward function into individual components and using them as the basis for explanations. Reward functions in input-driven environments are naturally decomposable, where the reward components represent key performance metrics. Hence, explanations generated by CrystalBox are meaningful to operators. For example, in DRL-based congestion control (Jay et al. 2019), there are three reward components: throughput, loss, and latency.

We formulate the explainability problem as generating decomposed future returns, given a state and an action. We propose a streamlined model-agnostic supervised learning approach to estimate these returns using an on-policy action-value function. Our solution generalizes across both discrete and continuous control environments. We evaluate our framework on multiple real-world input-driven environments, showing its ability to generate succinct, high-fidelity explanations efficiently. We demonstrate the usefulness of CrystalBox’s explanations by providing insights when feature-based explainers find it challenging, and powering network observability and guided reward design.

In summary, we make the following key contributions:

• •

  We propose a new notion of future-based explanations in input-driven environments.

• •

  We exploit the decomposable and dense reward functions in input-driven environments and formulate a supervised learning problem to estimate future-based explanations.

• •

  We demonstrate the practical use of our ideas in several systems applications across input-driven environments.

We use Adaptive Bitrate Streaming (ABR) as a representative input-driven application to understand explainability. In ABR, two entities communicate: a client streaming a video over the Internet, and a server delivering it. The video is usually divided into small, seconds-long chunks, encoded in advance at various bitrates. The ABR controller’s goal is to maximize the Quality of Experience (QoE) for the client by selecting the optimal bitrate for subsequent video chunks based on network conditions. QoE is often defined as a linear combination rewarding higher quality and penalizing both quality changes and stalling (Mok, Chan, and Chang 2011).

Let us consider two states, $S_{1}$ and $S_{2}$. We visualize them in Figure 1(a) showing past data like chunk sizes, transmission time, and buffer. The state histories of $S_{1}$ and $S_{2}$ are quite similar, with fluctuations in chunk sizes and transmission time, while the client’s buffer remains constant. Yet, the DRL controller picks medium-quality bitrate for $S_{1}$ and high quality bitrate for $S_{2}$. In this scenario, the network operator seeks to gain two insights to answer standard explainability questions (Doshi-Velez et al. 2017; van der Waa et al. 2018; Mittelstadt, Russell, and Wachter 2019; Miller 2019).

Answering $Q1$. Consider Figure 1(b) that shows results for $S_{1}$ and two actions: medium (left plot) and high (right plot) quality bitrates. LIME identifies the top features as the last few values of the chunk sizes, transmission times, and buffer. However, these features are the same for both actions. This leaves no way for the operator to gain an understanding of why the controller picks the top action in this state. The same observation also holds for $S_{2}$ (Figure 1(c)).

Answering $Q2$. We recall that in these states, the controller opts for medium quality and high quality bitrates, respectively. Looking at LIME’s top features in states $S_{1}$ (Figure 1(b)) and $S_{2}$ (Figure 1(c)) for the preferred actions, the same features appear in both explanations. Despite different preferred actions, LIME identifies similar key features for these decisions. This leaves us questioning why the controller opts for medium quality in one state and high quality in another. LIME falls short in explaining this discrepancy for $Q2$. We hypothesize that the same result holds for other feature-based explainers as they only use the state features.

For a comprehensive answer to these questions, it is necessary to look beyond the features and design explanations with details about the consequences of each action, as the controller inherently selects actions to maximize future returns. So, fully understanding the controller’s decision-making requires a forward-looking perspective.

Input-driven environments (Mao et al. 2018) represent a rich class of environments reflecting dynamics found in computer systems. Their dynamics are fundamentally different from traditional RL environments. We highlight two of their characteristics: their dependence on input traces and their naturally decomposable reward functions.

In input-driven environments, the system conditions are non-deterministic and constitute the main source of uncertainty. They determine how the environment behaves in response to the controller’s actions. For instance, in network traffic engineering, network demand influences whether congestion occurs on specific paths. These conditions are referred to as “inputs” (Mao et al. 2018), and the environments with inputs are called input-driven environments. Modeling these inputs remains an open research problem (Mao et al. 2019). In contrast, the presence of a stand-alone simulation environment is often assumed in traditional DRL problems (Brockman et al. 2016; Beattie et al. 2016).

A common theme across many of the DRL-based controllers for input-driven environments is a naturally decomposable reward function (Mao et al. 2019; Jay et al. 2019; Krishnan et al. 2018). This is because control in these environments involves optimization across multiple key objectives, which are often linearly combined, with each reward component representing a different objective.

Modeling the process behind these inputs can be challenging, as it can involve emulating environments such as the wide-area Internet. To circumvent this issue, state-of-the-art DRL solutions for input-driven environments do not model the complex input process. Instead, during training, they replay traces (or logged runs) gathered from real systems (Mao et al. 2019). In the beginning, the simulator selects a specific trace from the given dataset and generates the next state $s_{t+1}$ by simply looking up the subsequent logged trace value. Note that these traces are not accessible to the agent in advance and are only available during training.

The state transition function $P_{s}(s_{t+1}|s_{t},a_{t},z_{t+1})$ defines the probability distribution over the next state given the current state, action, and future input-value ($z_{t+1}$). The function $P_{z}(z_{t+1}|z_{t})$ defines the likelihood of the next input value given the current one. The combined transition function is $P_{s}(s_{t+1}|s_{t},a_{t},z_{t+1})P_{z}(z_{t+1}|z_{t})$.

Given $z_{t}$, obtaining $z_{t+1}$ through $P_{z}$ is an open problem (Mao et al. 2019), as mentioned above. During training, this problem is circumvented by selecting a specific $z=[..,z_{t},z_{t+1},..]$ from the set of traces $Z$ and accessing the subsequent value $z_{t+1}$ until the end of the trace. While this method works for training, it is not applicable for computing $P_{z}$ outside of training in algorithms such as Monte Carlo tree search (Browne et al. 2012) and model predictive control (Holkar and Waghmare 2010). This is because information about future network conditions, such as throughput values in our motivating example, is not available then.

Future-based explainers elucidate the controller’s decision-making process by giving a view into what the controller is optimizing: future performance. Given a state $s_{t}$ and an action $a_{t}$, a future-based explainer must provide a forward-looking view into the controller by capturing the consequences of that action in a concise yet expressive language.

We define the problem of generating future-based explanation as computing the decomposed future returns starting from a given state $s_{t}$ under action $a_{t}$ and policy $\pi$. Formally, decomposed future returns are the individual components within the $Q^{\pi}(s_{t},a_{t})=\sum_{c\in C}Q_{c}^{\pi}(s_{t},a_{t})$, where $C$ is the set of reward components in the environment. Here, $Q^{\pi}$ calculates the expected return for action $a_{t}$ in state $s_{t}$ and following policy $\pi$ thereafter. For the motivation example, the components in $C$ are quality, quality change, and stalling.

Now, we can define an explanation for a given state, action, and fixed policy as a tuple of return components:

Our explanation definition lets us clearly represent future action impact using return components that are typically the key performance metrics. It provides a future-centric view to analyze controller behavior over states and actions.

Let us return to the motivating example (§ 2). The operator seeks to explain two seemingly similar states, $S_{1}$ and $S_{2}$, where the controller chooses different actions (Fig. 1) and feature-based methods were not able to provide meaningful explanations. We answer these questions with future-based explanations generated by our framework, CrystalBox (§ 5).

Answering $Q1$. Consider Fig. 2(a) that provides an explanation for $S_{1}$ under both medium and high quality bitrates. The bars represent the positive or negative contributions of each performance metric to the total return. As we can see from the disaggregated performance metrics, ‘medium quality’ action is preferred as it leads to less stalling and has a higher overall return. Similarly, in $S_{2}$ (Fig. 2(b)), the disaggregated metrics suggest that the ‘high quality’ action was chosen due to better quality (blue bar) of the streamed video.

Answering $Q2$. The operator examines explanations for controller actions in both $S_{1}$ and $S_{2}$. They recognize that the expected stalling penalty in $S_{1}$ steers the preference towards medium quality, while the absence of such penalty in $S_{2}$ leads to the choice of high quality. Our example demonstrates that decomposed future returns serve as a useful technique for comparing both actions and states. They enable the operator to gain a comprehensive understanding of the agent’s decision-making process.

The next question is how to compute these decomposed future returns. To answer this question let us first formally define return components $Q^{\pi}_{c}(s_{t},a_{t})$, $\forall c\in C$:

where $r_{c}(s_{t},a_{t})$ is the reward value of component $c$ earned by the controller for taking action $a_{t}$ in state $s_{t}$. Therefore, to generate explanations, we need to compute this equation.

Unfortunately, decomposed future returns cannot be directly calculated because, as we see in Eq. 2, calculating the $Q^{\pi}_{c}$ function involves computing $P_{z}$ for the expected next state. However, $P_{z}$ cannot be computed outside of training, as explained in Section 3. There can of course be an approximation for $P_{z}$. But, obtaining low-bias and low-variance approximation is not currently possible (Mao et al. 2019).

To overcome this limitation, we propose to directly approximate these $Q^{\pi}_{c}$ components from examples, using training input traces $Z$. This approach enables us to bypass the need for $P_{z}$, making it feasible to generate explanations outside of training. In contrast to sparse-reward environments such as mazes, chess, and go, where such approximations may face considerable variance challenges (Sutton and Barto 2018), input-driven environments provide informative, consistent, non-zero rewards at each step. This feature makes our approximation practicable in input-driven contexts. Here, such approximation can be done either through a sampling-based or a learning-based procedure. We designed and evaluated several sampling-based approaches (§ 6.2) and found their performance unsatisfactory (§ 7).

We define a new supervised learning problem to learn decomposed future returns $Q^{\pi}_{c}$ in input-driven environments. We exploit the key insight that $Q^{\pi}_{c}(s_{t},a_{t})$ form a function that can be directly parameterized and learned. Formulating a supervised learning problem requires us to design (i) how we obtain the samples, (ii) the function approximator architecture, and (iii) the loss function. Fig. 3(a) shows the high-level flow of our approach and Fig. 3(b) focuses on the details of the learning procedure.

First, we propose a procedure to obtain samples of $Q^{\pi}_{c}$. From Eq. 2, we see that $Q^{\pi}_{c}$ depends on the combined transition function of the environment. This combined transition function forms a power distribution involving the variables $P_{z}$ and $P_{s}$. $P_{z}$ cannot be approximated with low bias or variance, leading to high bias estimations of $Q^{\pi}_{c}$. To solve this problem, we circumvent approximating $P_{z}$ by instead using its unbiased samples. These unbiased samples are exactly the traces from the training set of traces $Z$. We replay these traces with a given policy $\pi$ to obtain estimates of $Q^{\pi}_{c}$ (See Fig. 3(b) ‘Data Collection’ block) using Monte-Carlo rollouts (Sutton and Barto 2018). Second, we define our function approximator architecture with parameters $\theta$. We use a simple split fully connected layers architecture (Appendix A.5¹¹1This is the extended version of the paper (Patel, Abdu Jyothi, and Narodytska 2024)) to jointly learn the return components. Finally, we define a loss for each individual component $c$ and take a weighted sum as the overall loss. We use standard Mean Squared Error for our individual components.

While our formulation is unbiased, it can still have high variance due to its dependence on $P_{z}$. As a result, $Q^{\pi}_{c}$ can have high variance in realistic settings, requiring a large number of samples to achieve sufficient coverage. We employ two generic strategies to overcome this. First, we notice that reusing the embeddings $\phi(s_{t})$ from the policy instead of the states $s_{t}$ directly helps to reduce the variance. The embeddings $\phi(s_{t})$ capture important concepts about states, reducing the variance at the early stages of predictor learning. Second, we collect two kinds of rollouts: on-policy and exploratory (see Fig. 3(b)). On-policy rollouts strictly follow the policy, whereas exploratory rollouts start with an explorative action. This approach enables CrystalBox to achieve higher coverage of $Q^{\pi}_{c}$ for all actions.

Overall, CrystalBox will converge to the true $Q^{\pi}$ function, reflecting the policy’s performance. This stems from the fact that our formulation is a special instance of the Monte Carlo Policy Evaluation algorithm (Silver 2015; Sutton and Barto 2018) for estimating $Q^{\pi}_{\theta}$. Here, $Q^{\pi}_{\theta}$ is split into smaller parts $Q^{\pi}_{\theta,c}$, summable to the original. Consequently, the Monte Carlo Policy Evaluation’s correctness proof applies.

We highlight CrystalBox applies to any fixed policy $\pi$, even if $\pi$ is non-deterministic or has continuous action space. All that is required is access to a simulator with training traces, which is publicly available for most input-driven RL environments (Mao et al. 2019).

In this section, we give an overview of metrics and baselines that we use for evaluating CrystalBox.

We present an evaluation of CrystalBox ²²2https://github.com/sagar-pa/crystalbox. We aim to answer the following questions: Does CrystalBox produce high-fidelity explanations? Is CrystalBox efficient? Does joint training help? What applications can CrystalBox enable?

Prior work (Juozapaitis et al. 2019; Anderson et al. 2019) introduced a method to generate explanations based on a decomposed deep Q-learning technique for game environments and evaluated this idea on grid-world games. This approach is designed for Q-learning policies, and Q-learning cannot be applied to continuous control input-driven environments. Moreover, they assume access to and the ability to modify the training procedure. In contrast, our approach works for both discrete and continuous control systems without altering the training procedure.

Investigating the use of feature-based techniques for creating interpretable future return predictors represents a compelling research direction. Another potential avenue is to explore whether we can employ future return predictors during policy learning for human-in-the-loop frameworks.