Plume: A Framework for High Performance Deep RL Network Controllers via Prioritized Trace Sampling

In Deep Reinforcement learning (DRL), an agent interacts with an environment. At each timestep, the agent receives the current system state $s_{t}$, and takes an action $a_{t}$, drawn from its policy, $\pi(a|s_{t})$. The environment plays the action out and transitions to the next state $s_{t+1}$, giving the agent a reward $r_{t}$ (Sutton and Barto, 2018; Achiam, 2018; Silver, 2015).

In network environments, non-deterministic network conditions are the primary sources of noise and uncertainty. These conditions determine the environment’s response to the controller’s chosen actions. For example, in congestion control, external traffic can dictate whether congestion will occur.

Formally, these conditions are called “inputs”, and input-driven environments form an Input-Driven Markov Decision Process (Mao et al., 2018), defined by the tuple $(S,A,Z,P_{s},P_{z},r,\gamma)$. Here, $S$ denotes the set of states, $A$ represents the set of actions, $Z$ is the set of training input traces, $P_{s}$ is the state transition function, $P_{z}$ is the input transition function, $r$ is the reward function, and $\gamma$ is the discount.

The state transition function $P_{s}(s_{t+1}|s_{t},a_{t},z_{t+1})$ defines the probability distribution of the next state $s_{t+1}$ given the current state $s_{t}$, action $a_{t}$, and upcoming input value $z_{t+1}$. Meanwhile, the input transition function $P_{z}(z_{t+1}|z_{t})$ defines the probability of the next input value based on the current one, leading to an effective transition function given by $P_{s}(s_{t+1}|s_{t},a_{t},z_{t+1})P_{z}(z_{t+1}|z_{t})$.

As depicted in Figure 1, the DRL learning process aims to guide the policy $\pi$ towards higher cumulative reward through a loop involving two steps: a policy evaluation step and a policy improvement step (Horgan et al., 2018). During policy evaluation, the agent assesses its current policy’s performance by gathering experience through acting in the environment and leveraging this experience in function learning. Here, it updates its neural network to learn a form of the value function $v^{\pi}(s)=\mathbb{E}_{\pi}[G|s_{0}=s]$, which is the expected return $G$ starting from state $s$, where $G$ is the discounted sum of rewards $G=\sum_{t=0}^{\infty}\gamma^{t}r_{t}$. Subsequently, in the policy improvement phase, the agent modifies policy $\pi$ to maximize $v^{\pi}$. Through this iterative process of estimating and maximizing the policy’s value function, the agent learns in the environment.

In this paper, we use adaptive bitrate streaming, congestion control, and load balancing as representative networking environments.

Input-driven DRL training environments used in networking settings suffer from several overarching challenges.

Next, we discuss commonly used ML techniques for handling skew and establish the need for prioritized trace sampling in input-driven environments.

While PER is effective in traditional DRL settings (Hessel et al., 2018; Horgan et al., 2018), it is limited in addressing the challenges presented by skew in input-driven environments. The reason is that while PER addresses the state skew in the function learning phase, the skew in input traces additionally affects the acting phase of the training loop (Fig. 1). The controller has limited opportunity to act in tail-end traces. Without modifying which traces are selected during acting phase, PER cannot increase the frequency of exploration in tail-end traces or ensure a comprehensive evaluation across the entire input trace distribution. Consequently, the importance sampling PER uses is not sufficient to handle the skew of input traces in the dataset.

To test our hypothesis, we experiment by enabling prioritization at two points in the DRL workflow: sampling transitions in the experience buffer at the function learning step (PER enabled vs. disabled) and sampling input traces in the acting step (Random sampling vs. 2-Class Equal Weighted). 2-Class Equal Weighted is a simple input trace prioritization scheme that divides the traces from the Puffer Platform into two classes, those with mean throughput higher/lower than the highest quality bitrate, 0.98 Mbps (Figure 2), and equally samples both classes. We evaluate the impact of each technique on a DQN variation of Gelato controller for ABR trained using the state-of-the-art algorithm Ape-X DQN (Horgan et al., 2018) (training settings detailed in § 5 and § 6.2).

In Figure 3, we observe that the simple 2-Class Equal Weighted gives the highest controller performance and training stability. By prioritizing the tail-end slow throughput traces, we achieved high performance in both all and slow network traces without compromising anything. Enabling PER does not significantly improve controller performance. Even though the replay buffer can store $2$ million transitions (over 5000 input traces), the controller performance falls short of the naive trace prioritization scheme. This highlights that the skew in the trace distribution cannot be easily overcome at the function learning step.

Input traces, which are time-dependent series of values that define complex external conditions, can be incredibly difficult to characterize and prioritize directly. Hence, the first step towards automated prioritization of traces is identifying the attributes $\Phi$ using critical feature identification.

To extract all features associated with the time series trace data, we rely on the popular feature extraction tool for the time series data, tsfresh (Christ et al., 2018). We extract a large set of features $[\phi_{1},\phi_{2},...\phi_{n}]$ broadly applicable to all input-driven DRL environments. However, because this large set of features may not be relevant to every application, we introduce an automated three-step process to narrow down to the critical ones, inspired by the idea of recursive feature elimination in supervised learning (sci, [n. d.]).

First, we start with the large set of features and apply clustering to create a fixed small number of clusters. This is denoted by $c=C([\phi_{1},\phi_{2},...\phi_{n}])$, where $c$ is the cluster labels, and $C$ is the clustering function.

Second, we obtain the features most relevant in producing this mapping. To do so, we use the cluster labels $c$ and train decision trees based on the features $[\phi_{1},\phi_{2},...\phi_{n}]$. With this training, we can compute the information gain $IG(c,\phi_{i})=H(c)-H(c|\phi_{i})$ for each feature $\phi_{i}$. Here, $H$ is the Shannon entropy of the cluster labels, which is a measure of the average level of “uncertainty”.

Third, we eliminate features with the lowest $IG$ values. We continue this cycle of clustering, classification, and feature elimination until we are left with only the features that have high information gain. As we eliminate less useful features, we increase the number of clusters to ensure that the final feature set is sufficiently expressive.

Note that the clustering at this stage is solely for feature selection and has no impact on the main clustering phase (§ 4.2).

The second stage involves clustering traces using the critical features identified in the previous stage. In this stage, we attempt to reduce the complexity of balancing the skew by obtaining their salient clusters.

To detect skew in the dataset, we can look at the attributes $\Phi$ of input traces. However, balancing the skew based solely on these attributes proves to be a complex task. This difficulty arises because the attributes, represented as $[\phi_{1},\phi_{2},\ldots,\phi_{n}]$, are continuous random variables that may not be independent. In other words, modifying the skew of one attribute could negatively affect the skew of another. To address this issue, we cluster the traces to obtain a single distribution to balance. We achieve this using a clustering algorithm $C$ to obtain the labels $c$ so that the mapping again becomes $c=C(\Phi)$. By doing so, we create a ranking function that allows us to instead prioritize a categorical distribution of input traces, where the cluster labels act as the categories. We represent this distribution as $y$, where $y_{i}$ is a category, or salient trace cluster within it.

To cluster the traces, we employ Gaussian Mixture Models (GMM) with Kmeans++ initialization (Pedregosa et al., 2011). Gaussian Mixture Models use a generalized Expectation Maximization algorithm (Exp, [n. d.]) and can effectively deal with the large variations found in input data. Thus, GMMs are a good fit for our real-world input trace datasets. However, GMMs can often converge to local optima and require us to know the number of clusters a priori. Hence, to produce an effective clustering automatically, we perform a two-stage search for random initializations used in GMMs and the number of clusters. First, for different cluster counts, we evaluate the GMM’s log-likelihood score for the trace features across a range of random initializations and identify the initialization that maximizes the log-likelihood score for each cluster count. Second, we determine the optimal number of clusters from the output of the previous stage based on the highest normalized Silhouette score (Sil, [n. d.]).

With critical feature identification and clustering stages complete, we have a categorical distribution of input traces $y$ that we can balance by prioritization.

So far, we have discussed balancing the distribution $y$. While this can be done in a number of ways, to ensure that the balancing leads to meaningful performance improvements, we introduce a target function to balance the distribution around: “reward-to-go”. Reward-to-go represents the additional rewards that a controller can still achieve. This can be formally defined by Equation 1:

In this equation, $y_{i}$ is a category (§ 4.2) in the input trace distribution $y$, $G=\sum_{t=0}^{\infty}\gamma^{t}r_{t}$ is the discounted return of the trace as described in Section 2.1, $G^{\pi^{*}}$ is the return under the optimal policy $\pi^{*}$, and $G^{\pi^{\theta}}$ is the return under the current policy. We aim to balance the input trace distribution based on how suboptimally the current policy performs, ensuring a uniform gap across all traces. In other words, we seek to ensure that target function $\Delta G_{y_{i}}=\Delta G_{y_{j}}$ for all categories $y_{i}$ and $y_{j}$. However, calculating reward-to-go is often not possible in real-world situations because it depends on variables such as the controller’s training parameters and state features, and can require solving an NP hard problem (Meskovic et al., 2015). In this work, we introduce two strategies to approximate this prioritization: Static and Dynamic.

While there exists no analytical way to compute $\Delta G_{y_{i}}$, in some cases, we can show that static prioritization effectively balances the skew. First, consider that under random trace sampling, the imbalance can be arbitrarily large:

Given these propositions, it is evident that under static prioritization, irrespective of the initial input trace distribution, the relative reward-to-go ratio $\frac{\Delta G^{\prime}{y_{i}}}{\Delta G^{\prime}{y_{j}}}$ is close to one, while this ratio can take arbitrarily large values under random trace sampling. For both propositions, an underlying assumption is that the ratio $\frac{\Delta G_{y_{i}}}{\Delta G_{y_{j}}}$ is approximately equal to the inverse of the ratio of probability densities for the relevant categories. This assumption mirrors a real scenario: as the sampling frequency of a category increases, the controller also becomes better at handling traces from that category. Consequently, the reward-to-go gap decreases with increasing sampling probability for the category.

As the optimal return cannot be calculated, we replace the return $G^{\pi^{*}}$ with the expected return of the trace computed based on the corpus of seen traces, $\hat{G}(\Phi)$. Note that $\hat{G}(\Phi)$ is a function approximator that is trained continuously, in parallel with the controller. In a broad sense, this allows us to measure the improvement the controller can still achieve on a given trace. Nonetheless, this approximation is vulnerable to bias, especially in input traces where the controller’s performance is poor. In such poor conditions, the estimate may become overly pessimistic and might not accurately capture the reward-to-go. To address this concern, we introduce the second term, $-\mathbb{E}_{y_{i}}[G^{{\pi}^{\theta}}]$, which gives priority to traces that have low returns.

The dynamic weights are proportional to the normalized sum of components of $\Delta G_{y_{i}}$ (Eq. 3). Note that our prioritization stage is outside the DRL algorithm’s training loop in the Trace Selection module (Fig. 4).

We now turn to detail our implementation of all the experiments performed in this paper. We implement Plume as a Python library compatible with all major DRL frameworks.

A straightforward implementation of Plume can directly interfere with the various distributed training paradigms used in many DRL algorithms (Mnih et al., 2016; Horgan et al., 2018). To this degree, we implement our prioritization strategy using the distributed shared object-store paradigm in Ray (Moritz et al., 2018). This allows us to share the sampling weights across distributed RL processes without interfering with any DRL workflows.

With our implementation, the overhead for Plume is minimal. The Critical Feature Identification and clustering stages are completed once before training, with runtimes in the order of minutes. In Plume-Dynamic, we train a neural network to map the attributes $\Phi$ of an input trace to the return $G^{\pi^{\theta}}$ in that trace parallel to the training. We maintain a short bounded history of the trace-return pairs for each category and use this history to compute the two components of our prioritization function. To compute the first term in our approximation, we take the ground-truth samples of trace feature-return pairs, measure the mean absolute error of the neural network for these samples, and average them across each category. To calculate the compensation term, we take the negative of the mean return found in each category. We do this prioritization process continuously, adjusting the weights to the controller’s current needs. This dynamic prioritization calculation adds a computational overhead on the order of milliseconds per iteration. This added prioritization computation is handled in parallel to the DRL training and does not slow it down.

In this section, we present the settings used in our experiments. We present our results as averages over $4$ instances ($4$ controllers trained using the same scheme with different initial random seeds). This is consistent with the standard reporting practice in the RL community (Horgan et al., 2018; Kapturowski et al., 2018; Mnih et al., 2016). For testing on the Puffer platform, we select the best of these four seeds for benchmarking. For more details on these settings, see Appendix 6 and D.

We compare Gelato-Random and Gelato-Plume-Static with the performance of the Buffer-based controller BBA (Huang et al., 2014), the in-situ continuous training controller Fugu’s February version, Fugu-Feb (Yan et al., 2020), and CausalSim (Alomar et al., 2023), a version of Bola (Spiteri et al., 2020) tuned by trace-driven causal simulation. We note that on the Puffer platform, Gelato-Random is called by its code name “unagi”, Gelato-Plume-Static is called “maguro”. We additionally compare Gelato with the original version of Fugu over the period 07 March 2022 - 05 October 2022 in Figure 14 in Appendix 6.

In this section, we present the results of our experiments evaluating Plume in ABR, CC and LB. We aim to answer the following questions: How does the performance of Plume compare with random sampling? How do controllers trained with Plume perform in the real world?

In Figure 6, we present our results comparing Plume with random sampling and other controllers. We present our observations below.

Over this 1-year period, the algorithms we analyze streamed over $58.9$ stream-years of videos to over $200,000$ viewers across the Internet (puf, [n. d.]). Over this duration, Gelato-Plume-Static achieves $75\%$, $78\%$ and $81\%$ stall reduction compared to CausalSim, Fugu and BBA respectively (Fig. 6(d)). Gelato-Plume-Static additionally achieves SSIM improvements of $0.28$, $0.12$ and $0.15$ dB over CausalSim, Fugu and BBA respectively. Note that this quality improvement over BBA is more than $5\times$ that of Fugu, which only managed a $0.03$ dB improvement over BBA. CausalSim did not provide an SSIM improvement over BBA over this period. Gelato-Plume-Static has an average SSIM variation of $0.77$ dB, compared to $0.67$, $0.53$ and $0.78$ dB of CausalSim, Fugu and BBA respectively. Moreover, we find that Gelato-Random is a strong baseline, achieving $0.27$ dB SSIM improvement and $45\%$ stall reduction over CausalSim. We make similar observations when comparing Gelato with the original version of Fugu in Figure 14 in Appendix 6.

To evaluate Plume and various prioritization strategies, we introduce a controlled ABR environment for microbenchmarking prioritization: TraceBench. TraceBench makes two key changes to the standard ABR environment. First, it simplifies the quality-of-experience measurement to include only two terms, quality and stalling. Second, it parameterizes the traces by key characteristics of real-world traces: mean and variance of network throughput. With these modifications, we can reliably and thoroughly evaluate the controller across various network conditions. While this design is a simplification of the real-world environment, it gives a good approximation of a wide range of realistic settings. We believe that the development of such a framework is important for the community as real-world datasets do not allow microbenchmarking of DRL controllers or prioritization strategies. We envision that it can be used by other controller designers. We note that parameterized trace generation is a part of TraceBench, used to create a variety of scenarios to evaluate controllers on. It is not part of any of the prioritization strategies.

In generating Traces for TraceBench, we focus on traces with two levels of mean throughput, slow and fast, and two levels of variance of the throughput, high variance and low variance. In this benchmark, we generate three sets of datasets with different proportions of these traces: Majority Fast, Balanced, and Majority Slow. See Figure 7 for a visualization of the trace distributions, and see Figure 13 in Appendix B for a visualization of example traces. Note that the training and testing datasets remain disjoint: a controller trained on the Majority Fast dataset does not have access to the testing version of it.

We use the off-policy RL algorithm Ape-X DQN (Horgan et al., 2018). To evaluate prioritization in isolation, we use the same DRL hyper-parameters for all agents and present results averaged over $4$ instances. For details of our training parameters, see Appendix B.

Our experiments investigate two important questions. First, we evaluate how the versions of Plume’s prioritization, Plume-Static and Plume-Dynamic, compare to random trace sampling, the standard trace sampling technique. We additionally evaluate the impact of PER (Schaul et al., 2015), the state prioritization technique described in Section 2. Second, we investigate how sensitive these methods are to network conditions distribution shifts. We would like to emphasize here that these experiments are possible in real-world settings.

In Figure 9, we analyze the performance of Plume across various training and testing trace distributions. Particularly, we analyze the following scenarios:

• •

  Scenario 1: The training distribution is similar to the real world but the testing is adversarially different, i.e., we train on the Majority Fast but test on the Majority Slow dataset.

• •

  Scenario 2: Both training and testing have a balanced set of traces, i.e., we train and test on the Balanced dataset.

• •

  Scenario 3: The training distribution largely consists of the tail end of the testing distribution, i.e., we train on the Majority Slow but test on the Majority Fast dataset.

Below, we summarize the findings of our experiments with ABR, CC and LB presented in Section 6, and the analysis of our extensive Plume benchmarking presented in this section.

• •

  Plume is a generalized solution for DRL training in input-driven environments that automatically balances the trace distribution, and offers significant improvement in performance over random sampling in ABR, CC and LB, in simulation and in real-world testing, over both on-policy and off-policy algorithms.

• •

  Plume’s prioritization strategies work across trace distributions, providing controllers with greater performance and robustness in all.

• •

  Gelato trained with Plume offers the best performance when compared to prior ABR controllers on the real-world Puffer platform. It achieves $75\%$ and $78\%$ reduction in stalls over CausalSim (Alomar et al., 2023) and Fugu (Yan et al., 2020) respectively. It also achieves a statistically significant SSIM improvement of $0.28$ dB over CausalSim and $0.12$ dB over Fugu.