Title: Sample Efficient Reinforcement Learning in Mixed Systems through Augmented Samples and Its Applications to Queueing Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Mixed Systems

Queueing Network

Sample Efficient
Abstract: This paper considers a class of reinforcement learning problems, which involve systems with two types of states: stochastic and pseudo-stochastic. In such systems, stochastic states follow a stochastic transition kernel while the transitions of pseudo-stochastic states are deterministic {\em given} the stochastic states/transitions. We refer to such systems as mixed systems, which are widely used in various applications, including Manufacturing systems, communication networks, and queueing networks. We propose a sample-efficient RL method that accelerates learning by generating augmented data samples. The proposed algorithm is data-driven (model-free), but it learns the policy from data samples from both real and augmented samples. This method significantly improves learning by reducing the sample complexity such that the dataset only needs to have sufficient coverage of the stochastic states. We analyze the sample complexity of the proposed method under Fitted Q Iteration (FQI) and demonstrate that the optimality gap decreases as  $O\left(\sqrt{\frac{1}{n}}+\sqrt{\frac{1}{m}}\right),$ where $n$ represents the number of real samples, and $m$ is the number of augmented samples per real sample. It is important to note that without augmented samples, the optimality gap is $O(1)$ due to the insufficient data coverage of the pseudo-stochastic states. Our experimental results on multiple queueing network applications confirm that the proposed method indeed significantly accelerates both deep Q-learning and deep policy gradient.

Sample Efficient Reinforcement Learning in Mixed Systems through Augmented Samples and Its Applications to Queueing Networks

Abstract

This paper considers a class of reinforcement learning problems, which involve systems with two types of states: stochastic and pseudo-stochastic. In such systems, stochastic states follow a stochastic transition kernel while the transitions of pseudostochastic states are deterministic given the stochastic states/transitions. We refer to such systems as mixed systems, which are widely used in various applications, including manufacturing systems, communication networks, and queueing networks. We propose a sample efficient RL method that accelerates learning by generating augmented data samples. The proposed algorithm is data-driven and learns the policy from data samples from both real and augmented samples. This method significantly improves learning by reducing the sample complexity such that the dataset only needs to have sufficient coverage of the stochastic states. We analyze the sample complexity of the proposed method under Fitted Q Iteration (FQI) and demonstrate that the optimality gap decreases as ˜ O ( √ 1 /n + √ 1 /m ) , where n is the number of real samples and m is the number of augmented samples per real sample. It is important to note that without augmented samples, the optimality gap is ˜ O (1) due to insufficient data coverage of the pseudo-stochastic states. Our experimental results on multiple queueing network applications confirm that the proposed method indeed significantly accelerates learning in both deep Q-learning and deep policy gradient.

1 Introduction

Reinforcement learning (RL) algorithms have recently achieved superhuman performance in gaming, such as AlphaGo (Silver et al., 2017), and AlphaStar (Vinyals et al., 2019), under the premise that vast amounts of training data can be collected. However, collecting data in real-world could be expensive and time-consuming applications such as clinical trials and autonomous driving, posing a significant challenge to extending the success of RL to broader applications. In this paper, we are interested in sample efficient algorithms for a class of problems in which environments (also called systems in this paper) include both stochastic and deterministic transitions, resulting in two types of states: stochastic states and pseudo-stochastic states. For example, in a queueing system, customers arrive and depart, following certain stochastic processes, but the evolution of the queues given arrivals and departures are deterministic. We call these kinds of systems mixed systems. Mixed systems are common in various RL applications, including data centers, ride-sharing systems, and communication networks,

which can be modeled as queueing networks so as mixed systems. The broad application of mixed systems makes it crucial to explore whether current RL approaches are efficient in learning optimal policies for such systems. However, existing RL approaches are often not sample efficient due to the curse of dimensionality. For instance, in a queueing system, the state space (i.e. the queue lengths) is unbounded, which means that a tremendous amount of data samples are required to adequately cover the state space in order to learn a near-optimal policy. To highlight the challenge, in a recent paper (Dai and Gluzman, 2022), a 96-core processor with 1 , 510 GB of RAM was used to train a queueing policy for a queueing network with only a few nodes.

This paper proposes a new approach to handling the curse of dimensionality based on augmented samples. Our approach is based on two observations:

When dealing with a mixed system that has a large state space, model-based approaches are incredibly challenging to use. For example, as of today, it remains impossible to analytically determine the optimal policy for a queueing network. On the other hand, using a deep neural network can be regarded as a numerical method for solving a large-scale optimization problem. Therefore, a data-driven, neural network-based solution could be a much more efficient approach.

It is well-known that training deep learning is data-hungry. While, in principle, model-free approaches can be directly used for mixed systems, they are likely to be very inefficient in sample complexity for mixed systems whose state space is large. In this paper, we will utilize the knowledge of deterministic transitions to generate new data samples by augmenting existing real samples. In other words, we can generalize real samples to a large (even infinite) number of samples that cover the unobserved pseudo-stochastic states. These samples are equally useful for training the neural network as the real samples.

Based on the two observations above, we consider a mixed system that includes two types of states, stochastic states and pseudo-stochastic states, where the transitions of the stochastic states are driven by a stochastic kernel, and the transitions of the pseudo-states are deterministic and conditioned on the current and next stochastic states. We comment that without conditioning the stochastic states, the pseudo-stochastic states become stochastic. The distributions of stochastic states and pseudostochastic states are correlated, which makes the problem different from MDPs with exogenous inputs.

With this state separation, we propose an augmented sample generator (ASG). The sample generator generates virtual samples from real ones while keeping stochastic states and augmenting the pseudostochastic states. Both the real samples and virtual samples are then used to train the deep neural networks, e.g., Q-networks or policy networks.

We analyze the sample complexity of the proposed approach for mixed systems under Fitted Q Iteration (FQI Ernst et al. (2005)) which is equivalent to DQN (Mnih et al., 2015a) for tabular setting. Specifically, we consider the scenario where the size of pseudo-stochastic state space is much larger than that of stochastic state space, and the set of available real data samples does not provide sufficient coverage of the joint state space. This is the situation where the proposed approach is expected to be particularly advantageous. Our analysis demonstrates that the proposed approach yields a significant improvement in the convergence rate for tabular settings. In particular, by generating m virtual samples for each real sample, the optimality gap between the learned policy and the optimal policy decreases as ˜ O ( √ 1 /n + √ 1 /m ) , where n is the number of real samples. Note that without augmented samples, the error is ˜ O (1) due to the lack of data coverage of the pseudo-stochastic states, which reduces to ˜ O ( √ 1 /n ) when we generate at least n augmented samples for each real sample. We also would like to emphasize that ω ( √ 1 /n ) is also the fundamental lower bound due to the coverage of the stochastic states.

1.1 Related Work

Deep reinforcement learning with augmented samples is not new. However, the approach in this paper has fundamental differences compared with existing works in the literature, which we will explain in detail.

Dyna and its variants: The first group of related methods is Dyna (Sutton, 1988) and its extensions. Dyna-type algorithms are an architecture that integrates learning and planning for speeding up learning or policy convergence for Q-learning. However, our proposed method differs fundamentally

from these approaches in several aspects: (1) Dyna-type algorithms are model-based methods and need to estimate the system model which will be used to sample transitions. On the contrary, ASG does not require such an estimation and instead leverages the underlying dynamics of the system to generate many augmented samples from one real sample. (2) Dyna is limited to making additional updates only for states that have been observed previously, whereas ASG has the potential to update all pseudo-stochastic states, including those that have not yet been explored. (3) Since Dyna is a model-based approach, the memory complexity is |S| 2 |X| 2 |A| , where S is the stochastic state space, X is the pseudo-stochastic state space and A is the action space. ASG employs a replay buffer only for real data samples, which typically requires far less memory space. Moreover, some Dyna-type approaches like Dyna-Q + (Sutton, 1988), Linear Dyna (Sutton et al., 2012), and Dyna-Q ( λ ) (Yao et al., 2009), require more computational resources to search the entire state space for updating the order of priorities.

Dyna-type approaches have been successfully used in model-based online reinforcement learning for policy optimization, including several state-of-the-art algorithms, including ME-TRPO (Kurutach et al., 2018), SLBO (Luo et al., 2019), MB-MPO (Clavera et al., 2018), MBPO (Janner et al., 2019), MOPO (Yu et al., 2020). As we already mentioned above, ASG is fundamentally different from these model-based approaches. The key differences between Dyna-type algorithms and ASG are summarized in Table 1.

Data augmentation: It has been shown that data augmentation can be used to efficiently train RL algorithms. Laskin et al. (2020) showed that simple data augmentation mechanisms such as cropping, flipping, and rotating can significantly improve the performance of RL algorithms. Fan et al. (2021) used weak and strong image augmentations to learn a robust policy. The images are weakly or strongly distorted to make sure the learned representation is robust. However, these methods rely solely on image augmentation, and none of them consider the particular structure of the mixed system. The purpose of augmentation and distortion is to improve robustness and generalization. Our approach is to use virtual samples, which are as good as real samples to learn an optimal policy. The purpose of introducing the virtual samples is to improve the data coverage for learning the optimal policy.

Encoding symmetries into the designs of neural networks to enforce translation, rotation, and reflection equivariant convolutions of images have also been proposed in deep learning, like GConvolution (Cohen and Welling, 2016a), Steerable CNN (Cohen and Welling, 2016b), and E(2)Steerable CNNs (Weiler and Cesa, 2019). Mondal et al. (2020); Wang et al. (2022) investigated the use of Equivariant DQN to train RL agents. van der Pol et al. (2020) imposed symmetries to construct equivariant network layers, i.e., imposing physics into the neural network structures. Lin et al. (2020) used symmetry to generate feasible virtual trajectories for training robotic manipulation to accelerate learning. Our approach relies on the structure of mixed systems to generate virtual samples for learning an optimal policy and can incorporate much more general knowledge than symmetry.

In addition, existing works on low-dimensional state and action representations in RL are also related to our research focus on representing the MDP in a low-dimensional latent space to reduce the complexity of the problem. For example, Modi et al. (2021); Agarwal et al. (2020) studied provably representation learning for low-rank MDPs. Misra et al. (2020); Du et al. (2019) investigated the sample complexity and practical learning in Block MDPs. There are also practical algorithms (Ota et al., 2020; Sekar et al., 2020; Machado et al., 2020; Pathak et al., 2017) with non-linear function approximation from the deep reinforcement learning literature. Although the mixed model studied in this paper can be viewed as an MDP with a low-dimensional stochastic transition kernel, the state space of the MDP does not have a low-dimensional structure, which makes the problem fundamentally different from low-dimensional representations in RL.

Factored MDP: A factored MDP (Kearns and Koller, 1999) is an MDP whose reward function and transition kernel exhibit some conditional independence structure. In a factored MDP, the state space is factored into a set of variables or features, and the action space is factored into a set of actions or control variables. This factorization simplifies the representation and allows for more efficient computation and decision-making. While factored MDPs and mixed systems share some similarities, they are actually quite different. In a factored MDP, state variables are grouped into sets, allowing for the factoring of rewards that depend on specific state subsets, and localized (or factorized) transition probabilities. In a mixed system, the cost/reward function is not assumed to be factored. Furthermore, although the transition probabilities of stochastic states are localized, depending only on stochastic states and actions, the transitions of pseudo-stochastic states generally depend on the full state vector and are not factorized. However, these transitions are deterministic , which allows us to generate augmented samples. Given the fundamental differences between the two, the analysis and the results are quite different despite the high-level similarity. Additionally, for the batch offline RL setting without enough coverage used in our paper, there are no existing approaches, including factored MDPs, that guarantee performance in this scenario.

MDPs with exogenous inputs: Arecent paper considers MDP with exogenous inputs (Sinclair et al., 2022), where the system has endogenous states and exogenous states. The fundamental difference between exogenous-state/endogenous state versus stochastic-state/pseudo-stochastic-state is that the evolution of the exogenous state is independent of the endogenous state, while the evolution of the stochastic state depends on the action. Since the action is chosen based on the pseudo-stochastic state, the evolution of a stochastic state, therefore, depends on the pseudo-stochastic state.

The most significant difference between this paper and existing works is that we propose the concept of mixed systems and mixed system models. Based on the mixed system models, we develop principled approaches to generate virtual data samples, which are as informative for learning as real samples and enhance RL algorithms with much lower sample complexity. Equally significant, we provide sample complexity analysis that theoretically quantifies the sample complexity improvement under the proposed approach and explains the reason that principled data augmentation improves learning in RL.

2 Mixed Systems and Mixed System Models

We consider a discrete-time Markov decision process M = ( S × X , A , P, R, γ, η 0 ) , whose state space is S ×X , where S is the set of stochastic states and X is the set of pseudo-stochastic states, and both S and X are assumed to be finite. Further, A is a finite action space, R : S ×X ×A → [0 , r max ] is a deterministic and known cost function, and P : S × X × A → ∆( S × X ) is the transition kernel (where ∆( · ) is the probability simplex).

The transition kernel P is specified in the following way. The transitions of the stochastic state follow a stochastic transition kernel. In particular, the transition of the stochastic state S t at time t can be represented as

We assume that this transition is independent of the pseudo-stochastic state X t (and everything else at or before time t ). The transition of the pseudo-stochastic state is then deterministic, governed by a function g as follows:

Our system M is said to be a mixed system because it includes both stochastic and deterministic transitions. Our mixed system model then consists of the stochastic transition kernel and the deterministic transition function g .

We focus on discounted costs in this paper. Given a mixed system, the value function of a policy π is defined as

where γ is the discount factor, and the expectation is taken over the transitions of stochastic states and randomness in the policy. Let v π = E ( s 0 ,x 0 ) ∼ η 0 [ V π ( s 0 , x 0 )] , i.e., the expected cost when the initial distribution of the state is η 0 . Let v ∗ = min π v π .

The Q -value function of a policy π is defined as

Let Q be the optimal Q -value function. Then the Bellman equation can be written as

Note that the value function and the Q -value function both take values in [0 , V max ] for some finite V max under any policy, due to the assumption that the cost is in [0 , r max ] .

2.1 Example: A Wireless Downlink Network

To better illustrate the structure of mixed systems. Consider the example of a wireless downlink network shown in Fig. 1 with three mobile users. We model it as a discrete-time system such that Λ t ( i ) is the number of newly arrived data packets for mobile i at the beginning of time slot t, and O t ( i ) is the number of packets that can be transmitted to mobile i during time slot t if mobile i is scheduled.

The base station can transmit to one and only one mobile during each time slot, and let A t ∈ { 1 , 2 , 3 } denote the mobile scheduled at time slot t. The length of queue i evolves as

To minimize the total queue lengths, the problem can be formulated as an RL problem such that the state is ( Λ t , O t , q t ) , the action is A t , and the cost is ∑ i q t ( i ) .

Flow 1

Flow 2

Flow 3

We can see that in this problem, S t := ( Λ t , O t ) is the stochastic state and q t is the pseudo-stochastic state. In general, the state space of the stochastic states is bounded, e.g., both Λ t and O t may be Bernoulli random variables, but the pseudo-stochastic state such as q t is large or even unbounded. Therefore, while the distributions of the stochastic states can be learned with a limited number of samples, learning the distribution of the queue lengths, even under a given policy, will require orders of magnitude more samples. Therefore, vanilla versions of model-free approaches that do not distinguish between stochastic and pseudo-stochastic states require an unnecessary amount of data samples and are not efficient.

Therefore, we propose a sample-efficient, data-driven approach based on deep RL and augmented samples. Augmented samples guarantee enough coverage based on a limited number of real data samples. Note that given Λ t , O t and A t , we can generate one-slot queue evolution starting from any queue length.

3 Sample Efficient Algorithms for Mixed Systems

3.1 Augmented Sample Generator (ASG)

We first introduce the Augmented Sample Generator (ASG, Algorithm 1), which augments a dataset by generating virtual samples. In particular, ASG takes as input a dataset D , an integer m, and a probability distribution β ( · ) over the pseudo-stochastic states in X . For each sample ( s, x, a, r, s ′ , x ′ ) ∈ D, we first sample a pseudo-stochastic state ˆ x from β ( · ) , and then construct a virtual sample ( s, ˆ x, a, ˆ r, s ′ , ˆ x ′ ) , where ˆ r = R ( s, ˆ x, a ) and ˆ x ′ = g ( s, ˆ x, a, s ′ ) . Note that we are able to do this since we assume that functions R and g are given in our applications. For each sample in D , we repeat this procedure m times independently to generate m virtual samples. Taking the downlink wireless network as an example, where S t := ( Λ t , O t ) , X t := q t , assume that at some timeslot t, one of the real samples is

where the queue length q t + 1 is calculated according to Eq. (6). Using ASG, we generate two pseudo-stochastic states (queue length) ˆ x 1 = (1 , 2 , 3) , ˆ x 2 = (0 , 2 , 1) , then we are able to obtain two virtual samples

Note that all the virtual samples represent true transitions of the mixed system if the queue lengths were ˆ x 1 and ˆ x 2 .

Algorithm 1: Augmented Sample Generator (ASG)

3.2 Batch FQI with ASG

We now formally present our algorithm, Batch FQI with ASG (Algorithm 2), for mixed systems. As in a typical setup for batch reinforcement learning (Chen and Jiang, 2019), we assume that we have a dataset D with n samples. The samples are i.i.d. and the distribution of each sample ( s, x, a, r, s ′ , x ′ ) is specified by a distribution µ ∈ ∆( S × X × A ) for ( s, x, a ) , r = R ( s, x, a ) , s ′ ∼ P ( ·| s, a ) , and x ′ = g ( s, x, a, s ′ ) . The distribution µ is unknown to the agent.

Our algorithm follows the framework of a batch FQI algorithm. We consider a class of candidate value-functions F ∈ ( S × X × A → [0 , V max ]) . We focus on the scenario where the number of pseudo-stochastic states, |X| , is far more than that of stochastic states, |S| , and the dataset D does not provide sufficient coverage of the pseudo-stochastic states ( n < |X| ).

The goal is to compute a near-optimal policy from the given dataset D, via finding an f ∈ F that approximates the optimal Q-function Q . The algorithm runs in episodes. For each episode k , it first generates an augmented dataset D k using ASG ( D,m,β k ) , where β k is some distribution of the pseudo-stochastic states satisfies that, for each typical pseudo-stochastic state, it can be generated with at least probability σ 1 . We assume that the pseudo-stochastic states will not transition to atypical states, such as extremely large values (possibly infinite) in the queuing example, under a reasonable policy. Details can be found in Assumption A.4 in the supplemental material. We remark that we can also use the same β for all the episodes, and in practice, we can simply adopt the estimation of the pseudo-stochastic states (i.e., with Gaussian distribution) as β k . The algorithm then computes f k as the minimizer that minimizes the squared loss regression objective over F . The complete algorithm is given in Algorithm 2.

3.3 Guarantee and Analysis

The guarantee of Algorithm 2 is summarized below.

Theorem 1. We assume the data coverage and completeness assumptions (details in the section A in the supplemental material). Given a dataset D = { s, x, a, r, s ′ , x ′ } with n samples, w.p. at least 1 -δ, the output policy of Algorithm 2 after k iterations, π f k , satisfies

where C is a constant related to the data coverage assumption A.3, and c 0 , σ 1 , σ 2 are constants defined in Assumption A.4 to ensure a typical set of reasonable typical pseudo-stochastic states.

1 Input : A dataset D = { ( s, x, a, r, s , x ) } , a function class F , a positive integer m, distributions { β } K ;

Algorithm 2: Batch FQI with ASG ′ ′ k k =1 2 Initialize approximate Q -function: Choose f 0 randomly from F . ; 3 for episode k = 1 , 2 , . . . , K do 4 Generate augmented dataset ˆ D k = ASG ( D,m,β k ); 5 Define loss function as: 1 k f k = arg min f ∈F L ˆ D k ( f ; f k -1 ) , V f k ( s, x ) = min a ∈A f k ( s, x, a ) ; 6 π f k ( s, x ) = arg min a ∈A f k ( s, x, a ) ; 7 Output: π f K ;

Theorem 1 shows that ASG significantly improves the (real) sample complexity, i.e., the number of real data samples needed. When the tail of pseudo-stochastic states decays fast, e.g. an exponential tails link in a typical queueing network, we can choose a sampling distribution β that guarantees that σ 2 = ˜ O (1 /n ) and σ 1 = Ω(1 / log n ) . Therefore, for sufficiently large k and sufficiently small σ 2 , the first term is the dominating term. As m increases from 1 to n, the first term of the bound decreases from ˜ O (1) to ˜ O ( √ 1 /n ) . Remark that we cannot establish a convergence result without using ASG (when m = 0 ) since using the dataset D alone doesn't have a sufficient data coverage guarantee, which is critical for batch RL. Due to the page limit, we only present the outline behind our analysis. The detailed proof is deferred to the supplemental material.

Proof Outline:

First we will show that using performance difference lemma (Kakade and Langford, 2002) it is sufficient by bounding ‖ f k -Q ∗ ‖ 2 ,ξ , where ξ is some distribution.

Then, by considering how rare the distribution of the pseudo-stochastic states is under ξ and then classify the pseudo-stochastic states into typical and atypical sets based on a threshold σ 2 on the distribution. Along with the augmented data generator under β k ( · ) . Then later we show that the term ‖ f k -Q ∗ ‖ 2 ,ξ can be bounded by O ( ‖ f k -T f k -1 ‖ 2 ,µ β k + γ ‖ f k +1 -Q ∗ ‖ 2 ,ξ ′ + √ σ 2 V max ) , where µ β k is the data distribution after augmenting virtual samples. The first two terms are related to the typical set given that a sufficient coverage of the data after data augmentation. The last term is easy to obtain by considering the atypical set.

Therefore, it is obvious that if we can have a bound on ‖ f k -T f k -1 ‖ 2 ,µ β k which is independent of the episode k, we can prove the final result by expanding ‖ f k -Q ∗ ‖ 2 ,ξ iteratively for k times. We finally show that the term ‖ f k -T f k -1 ‖ 2 ,µ β k can indeed be bounded by using the FQI minimizer at each episode and a concentration bound on the offline dataset after augmenting virtual samples.

3.4 Extension to the case when |X| is infinite

Theorem 1 assumes the number of the pseudo-stochastic states is finite. The result can be generalized to infinite pseudo-stochastic state space |X| if we make the following additional assumption:

Assumption 1. For the typical set B of pseudo-stochastic states (formally defined in Assumption A.4 in the supplemental material), for any s, a ∈ S × A , f ∈ F , if x ∈ B , then f ( s, x, a ) ≤ V max otherwise if x / ∈ B , we have | f ( s, x, a ) -Q ∗ ( s, x, a ) | ≤ V max . Furthermore, for any given f ∈ F , ( s, x ) , x ∈ B , we have | V f ( s ′ , x ′ ) -V f ( s ′′ , x ′′ ) | ≤ V max , where x ′ = g ( s, x, π f , s ′ ) , x ′′ = g ( s, x, π f , s ′′ ) .

The first part in assumption 1 means that the function f ∈ F is always bounded when x ∈ B . Otherwise, the difference between f and the optimal Q -function Q ∗ is bounded, which indicates that although f could be extremely large or even unbounded for the pseudo-stochastic states, including the infinite case, in the atypical set, it is also true for the optimal Q -function Q ∗ . Therefore the difference between f and Q ∗ is assumed to be bounded, which is reasonable in practice. The second

part states that for any given x ∈ B , the difference between the value functions of any possible next transitions ( s ′ , x ′′ ) and ( s ′′ , x ′′ ) is always bounded, which is also true in general especially in queuing literature since the changing of queue lengths between two timeslots is small.

Assumption 1 is easy to be satisfied in a typical queueing system (e.g. the wireless downlink network in Section 2.1). If the control/scheduling policy is a stabilizing policy, then the queues have exponential tails, i.e., the probability decays exponentially as queue length increases, and the policy can drive very large queues to bounded values exponentially fast. Therefore, if all functions in F are from stabilizing policies, Assumption 1 holds because any policy in F including the optimal policy can reduce the queue lengths exponentially fast (in expectation) and the difference in Q -functions therefore is bounded.

Under the additional Assumption 1, we show that the same order of the result is achievable. Details can be found in Section A.2 in the supplemental material.

4 Experiments

To evaluate the performance of our approach, we compared our results with several baselines in extensive queuing environments. In our simulations, when augmenting virtual samples, we only require that the augmented samples are valid samples for which the action from the real sample remains to be feasible.

4.1 The Criss-Cross Network

The criss-cross network is shown in Fig. 2a, which consists of two servers and three classes of jobs. Each job class is queued in a separate buffer if not served immediately. Each server can only serve one job at a time. Server 1 can process jobs of both class 1 and class 3; server 2 can only process jobs of class 2 . Class 3 jobs become class 2 jobs after being processed, and class 2 and class 3 jobs leave the system after being processed. The service time of a class i job is exponentially distributed with mean m i . The service rate of a class i job is defined to be µ i := 1 /m i , and jobs of class 1 and class 3 arrive to the system following the Poisson processes with rates λ 1 and λ 3 . To make sure the load can be supportable (i.e. queues can be stabilized), we assume λ 1 m 1 + λ 3 m 3 < 1 and λ 1 m 2 < 1 . In the simulation, we consider the imbalanced medium traffic regime, where µ 1 = µ 3 = 2 , µ 2 = 1 . 5 , λ 1 = λ 3 = 0 . 6 .

All queues are first-in-first-out (FIFO). Let e t ( i ) ∈ {-1 , 0 , +1 } denote the state of class i jobs at time t, where -1 means a departure, +1 means a arrival and 0 denotes no changes for job i. Then e t = ( e t (1) , e t (2) , e t (3)) be the event happens at time t in the system, which is the stochastic state in this mixed system. Let q t = ( q t (1) , q t (2) , q t (3)) be the vector of queue lengths at time t, where q t ( i ) is the number of class i jobs in the system. Obviously, q t is the pseudo-stochastic state in the system, which can be derived from ( q t -1 , e t -1 ) and action a t . We combine ASG with Q -learning and compare it with several baselines: (i) the vanilla Q-learning, (ii) a proximal policy optimization (PPO) based algorithm proposed in Dai and Gluzman (2022) designed for queuing networks, (iii) a random policy, and (iv) a priority policy such that the server always serves class 1 jobs when its queue length is not empty, otherwise it serves class 3 jobs. Simulation results in Fig. 2d demonstrate that using ASG significantly improves the performance, and achieves the performance of the optimal policy, which was obtained by solving a large-scale MDP problem, after only 4 training episode ( 4 k training steps). All other policies are far away from the optimal policy, and barely learn anything, as we mentioned that the state space is quite large and is not even sufficiently explored after 4 k training steps. To further demonstrate the fundamental difference between ASG and Dyna-type algorithms, we also compared ASQ with the model-based approach: Q-learning with Dyna. We can observe that Dyna also fails to learn a good policy since, as we mentioned before, ASG is built for improving sample efficiency with augmented samples that represent true transitions for possible all all the pseudo-stochastic states, but dyna-type algorithms can only generate addition samples from an estimated model, which are limited by the samples that have been seen before.

4.2 Two-Phase Criss-Cross Network:

To further evaluate the power of ASG beyond the tabular setting, We combine ASG with Deep Qnetwork (Mnih et al., 2015b) named DQN-ASG and evaluate it on a more complicated variant of the

Criss-Cross network presented in Fig. 2b, where class 3 jobs have two phases. Class 3 jobs at phase 1 become phase 2 after being processed with probability 1 -p, leave the system with probability p ; and class 3 jobs at phase 2 leave the system after completed processing. Results in Fig. 2e indicate that combining ASG with DQN also has a significant improvement on the performance, and our method only needs 100 episodes to achieve a near-optimal policy. We also remark here that the training time for DQN-ASG is three times faster than the baseline ( 4 hours v.s. 12 hours). In the simulation, all the parameters are set to be the same as those in the previous case, and we use p = 0 . 8 . Both the criss-cross and two-phase criss-cross networks are continuous-time Markov chains (CTMCs). We use the standard uniformization (Puterman, 2014; Serfozo, 1979) approach to simulate them using discrete-time Markov chains (DTMCs).

We remark here that the two-phase Criss-Cross network cannot be modeled as MDPs with exogenous inputs. Note that the phase a job is in is a stochastic state, which depends on the action (whether a job is selected for service). Since the scheduling action depends on the queue lengths, the evolution of stochastic states depends on the pseudo-stochastic states (the queue lengths). However, the exogenous input in an MDP with exogenous inputs has to be an observable random process independent of the action. Therefore, the current phase of a job of the phases Criss-Cross network is not an exogenous input and the system is not an MDP with exogenous inputs.

4.3 Wireless Networks

In addition to the criss-cross networks, we also evaluate our approach on the downlink network described in Section 2.1 with three mobiles. We let the arrivals are Poisson with rates Λ = { 2 , 4 , 3 } , and the channel rates O to be { 12 , 12 , 12 } . We use the Max-Weight (Tassiulas and Ephremides, 1992) algorithm, a throughput-optimal and heavy-traffic delay optimal algorithm, as the baseline. As shown in Fig. 2f, the learned policy outperforms Max-Weight. To the best of our knowledge, no efficient RL algorithms exist that can outperform Max-Weight.

𝑆

!

𝑆

"

𝐵

"

𝐵

!

𝐵

#

class 1

arrivals

class 3

arrivals

class 2

departures

class 3

departures

(a) Criss-Cross network

𝑆

!

𝑆

"

𝐵

"

𝐵

!

𝐵

#

class 1

arrivals

class 3

arrivals

class 2

departures

phase 1

phase 2

𝑝

1 - 𝑝

class 3

departures

(b) Phases Criss-Cross network

Flow 1

Flow 2

Flow 3

(c) Wireless network

(d) Testing reward

(e) Testing reward

(f) Testing reward

一

7.5

20

DQN-ASG

DQN-ASG

7.0

Baseline

TotalQueueLength

18

Random

Priority

16

.5

14

5.0

12

4.5

4.90

80

140

200

260

320

10

40

80120160200240280

Training Episode

Training Episode

口

口

2.5

Q-learning-ASG

Baseline

2.4

Q-learning

Dyna

Random

Priority

2.3

Optimal

2.2

Total

2.1

2.0

1

2

3

4

5

6

Training Episode

4.4 Additional Simulations on the criss-cross network

To further verify the performance of our algorithm on more scenarios, we first consider a more general criss-cross network where the size of both class 1 and class 3 jobs can vary over time. Then all class 1 and class 3 jobs need multiple timeslots to be processed instead of 1 . In particular, when a new job of either class 1 or 2 joins the system, the job size of such job is chosen uniformly. The environment in Fig. 2b can be seen as a special case where class 1 , and 3 have a fixed job size 1 and 2 , respectively. The performances of different cases are presented in Fig. 3. Our approach still obtains the best result.

The details of the parameters can be found in section B.2. We do not include the optimal solution since, for the general case, the state space is quite large, and using value iteration to obtain the optimal solution is very time-consuming. The priority policy already provides a reasonable good baseline.

(a)

(b)

(c)

10.0

1.6

5.00

DQN-ASG

DQN-ASG

DQN-ASG

Total Queue Length

9.5

Baseline

Priority

Total QueueLength

1.5

Baseline

Priority

Lengt

4.50

4.75

Baseline

Priority

9.0

Random

1.4

Random

4.25

Random

Queue

1.3

4.00

8.5

1.2

Total

3.75

3.50

8.0

1.1

3.25

7.20

1.20

3.090

80

140

200

260

320

80

140

200

260

320

80

140

200

260

320

Training Episode

Training Episode

Training Episode

4.5 Additional Simulations on the Wireless Network

We also evaluate the performance of our algorithm on the wireless network (Fig. 1) under different sets of arrival rates. The results of average queue length compared with those of Max-Weight are summarized in Table 2. We can observe that ASG performs better than Max-Weight in all the cases. Simulation results on a more complicated two-phase wireless network can be found in the supplemental material (Section B.3).

4.6 Combining ASG with Policy Gradient-type Algorithm

We also investigate the possibilities of using ASG in policy gradient-type algorithms (i.e., TD3 (Fu et al., 2018), SAC (Haarnoja et al., 2018)) and compare our approach with one of the state-of-art algorithms MBPO (Janner et al., 2019). Our approach still achieves the best performance. Due to page limit, the detailed simulation and algorithm are deferred to Section B in the supplemental material.

5 Conclusions

In this work, we considered RL problems for mixed systems which have two types of states: stochastic states and pseudo-stochastic states. We proposed a sample efficient RL approach that accelerates learning by augmenting virtual data samples. We theoretically quantified the (real) sample complexity gain by reducing the optimality gap from ˜ O (1) to ˜ O (1 / √ n ) . Experimental studies further confirmed the effectiveness of our approach on several queuing systems. The possible limitation of using ASG is that it introduces more computation during each learning iteration when generating virtual samples, though it is not obvious in our experiments.

References

Agarwal, A., Kakade, S., Krishnamurthy, A., and Sun, W. (2020). Flambe: Structural complexity and representation learning of low rank mdps. In Advances Neural Information Processing Systems (NeurIPS) , page 20095-20107. Curran Associates Inc.

Chen, J. and Jiang, N. (2019). Information-theoretic considerations in batch reinforcement learning. In Int. Conf. Machine Learning (ICML) , pages 1042-1051. PMLR.

Clavera, I., Rothfuss, J., Schulman, J., Fujita, Y ., Asfour, T., and Abbeel, P. (2018). Model-based reinforcement learning via meta-policy optimization. CoRR , abs/1809.05214.

Cohen, T. and Welling, M. (2016a). Group equivariant convolutional networks. In Int. Conf. Machine Learning (ICML) , pages 2990-2999. PMLR.

Cohen, T. S. and Welling, M. (2016b). Steerable cnns. arXiv preprint arXiv:1612.08498 .

Dai, J. G. and Gluzman, M. (2022). Queueing network controls via deep reinforcement learning. Stochastic Systems , 12(1):30-67.

Du, S., Krishnamurthy, A., Jiang, N., Agarwal, A., Dudik, M., and Langford, J. (2019). Provably efficient rl with rich observations via latent state decoding. In Int. Conf. Machine Learning (ICML) , pages 1665-1674. PMLR.

Ernst, D., Geurts, P., and Wehenkel, L. (2005). Tree-based batch mode reinforcement learning. Journal of Machine Learning Research , 6:503-556.

Fan, L., Wang, G., Huang, D.-A., Yu, Z., Fei-Fei, L., Zhu, Y., and Anandkumar, A. (2021). SECANT: self-expert cloning for zero-shot generalization of visual policies. In Int. Conf. Machine Learning (ICML) .

Fu, L., Fu, X., Zhang, Z., Xu, Z., Wu, X., Wang, X., and Lu, S. (2018). Joint optimization of multicast energy in delay-constrained mobile wireless networks. IEEE/ACM Transactions on Networking , 26(1):633-646.

Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. (2018). Soft Actor-Critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Int. Conf. Machine Learning (ICML) , pages 1861-1870.

Janner, M., Fu, J., Zhang, M., and Levine, S. (2019). When to trust your model: Model-based policy optimization. In Advances Neural Information Processing Systems (NeurIPS) , volume 32.

Kakade, S. and Langford, J. (2002). Approximately optimal approximate reinforcement learning. In Int. Conf. Machine Learning (ICML) , pages 267-274.

Kearns, M. and Koller, D. (1999). Efficient reinforcement learning in factored mdps. In Proc. Int. Joint Conf. Artif. Intell. Org (IJCAI) , volume 16, pages 740-747.

Kurutach, T., Clavera, I., Duan, Y., Tamar, A., and Abbeel, P. (2018). Model-ensemble trust-region policy optimization. arXiv preprint arXiv:1802.10592 .

Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A. (2020). Reinforcement learning with augmented data. In Advances Neural Information Processing Systems (NeurIPS) . Curran Associates Inc.

Lin, Y., Huang, J., Zimmer, M., Guan, Y., Rojas, J., and Weng, P. (2020). Invariant transform experience replay: Data augmentation for deep reinforcement learning. IEEE Robotics and Automation Letters , 5(4):6615-6622.

Luo, Y., Xu, H., Li, Y., Tian, Y., Darrell, T., and Ma, T. (2019). Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. In Int. Conf. on Learning Representations (ICLR) .

Machado, M. C., Bellemare, M. G., and Bowling, M. (2020). Count-based exploration with the successor representation. In AAAI Conf. Artificial Intelligence , pages 5125-5133.

Misra, D., Henaff, M., Krishnamurthy, A., and Langford, J. (2020). Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In Int. Conf. Machine Learning (ICML) , pages 6961-6971. PMLR.

Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., and Ostrovski, G. (2015a). Human-level control through deep reinforcement learning. Nature , 518(7540):529-533.

Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015b). Human-level control through deep reinforcement learning. Nature , 518(7540):529-533.

Modi, A., Chen, J., Krishnamurthy, A., Jiang, N., and Agarwal, A. (2021). Model-free representation learning and exploration in low-rank mdps. arXiv preprint arXiv:2102.07035 .

Mondal, A. K., Nair, P., and Siddiqi, K. (2020). Group equivariant deep reinforcement learning. ICML 2020 Worshop on Inductive Biases, Invariances and Generalization in RL .

Ota, K., Oiki, T., Jha, D., Mariyama, T., and Nikovski, D. (2020). Can increasing input dimensionality improve deep reinforcement learning? In Int. Conf. Machine Learning (ICML) , pages 7424-7433. PMLR.

Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. (2017). Curiosity-driven exploration by self-supervised prediction. In Int. Conf. Machine Learning (ICML) , pages 2778-2787. PMLR.

Puterman, M. L. (2014). Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons.

Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, D. (2020). Planning to explore via self-supervised world models. In Int. Conf. Machine Learning (ICML) , pages 8583-8592. PMLR.

Serfozo, R. F. (1979). An equivalence between continuous and discrete time markov decision processes. Operations Research , 27(3):616-620.

Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. (2017). Mastering the game of go without human knowledge. nature , 550(7676):354-359.

Sinclair, S. R., Frujeri, F., Cheng, C.-A., and Swaminathan, A. (2022). Hindsight learning for mdps with exogenous inputs. arXiv preprint arXiv:2207.06272 .

Sutton, R. S. (1988). Learning to predict by the methods of temporal differences. Machine learning , 3(1):9-44.

Sutton, R. S., Szepesvári, C., Geramifard, A., and Bowling, M. P. (2012). Dyna-style planning with linear function approximation and prioritized sweeping. arXiv preprint arXiv:1206.3285 .

Tassiulas, L. and Ephremides, A. (1992). Stability properties of constrained queueing systems and scheduling policies for maximum throughput in multihop radio networks. IEEE Trans. Autom. Control , 37:1936-1948.

van der Pol, E., Worrall, D., van Hoof, H., Oliehoek, F., and Welling, M. (2020). MDP homomorphic networks: Group symmetries in reinforcement learning. In Advances Neural Information Processing Systems (NeurIPS) , volume 33, pages 4199-4210.

Vinyals, O., Babuschkin, I., Czarnecki, W. M., Mathieu, M., Dudzik, A., Chung, J., Choi, D. H., Powell, R., Ewalds, T., Georgiev, P., et al. (2019). Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature , 575(7782):350-354.

Wang, D., Walters, R., and Platt, R. (2022). SO(2) -equivariant reinforcement learning. In iclr .

Weiler, M. and Cesa, G. (2019). General e(2)-equivariant steerable cnns. In Advances Neural Information Processing Systems (NeurIPS) , volume 32, page 14357-14368. Curran Associates, Inc.

Yao, H., Bhatnagar, S., and Diao, D. (2009). Multi-step linear dyna-style planning. In Advances Neural Information Processing Systems (NeurIPS) , NIPS'09, page 2187-2195, Red Hook, NY, USA. Curran Associates Inc.

Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma, T. (2020). Mopo: Model-based offline policy optimization. Advances Neural Information Processing Systems (NeurIPS) , 33:14129-14142.

A Proof of Theorem 1

In this section, we present the proof of Theorem 1. We first introduce and recall some necessary notations and assumptions. Then, we present some auxiliary lemmas and their proofs. Finally, we combine the lemmas to prove the main result.

Notations: Define P ( ν ) as distribution over states such that ( s ′ , x ′ ) ∼ P ( ν ) ⇔ ( s, x, a ) ∼ ν, s ′ ∼ P ( s ′ | s, a ) , x ′ = g ( s, a, s ′ ) . In other words, it is the distribution of the next state if the state action pair follows ν. For f : S × X × A → R , ν ∈ ∆( S × X × A ) , where ∆( · ) is the probability simplex, and p > 1 , define ‖ f ‖ p,ν = ( E ( s,x,a ) ∼ ν [ | f ( s, x, a ) | p ]) 1 /p . Define π f,f ′ ( s, x ) := arg min a ∈A min { f ( s, x, a ) , f ′ ( s, x, a ) } .

Recall that in Alg. 2, at iteration k,

where

Assumptions:

Assumption A.1. (Realizability) For the optimal policy, Q ∗ ∈ F .

Assumption A.2. (Completeness) For the policy π to be evaluated, ∀ f ∈ F , T f ∈ F , where T : R S×X×A → R S×X×A is the Bellman update operator, ∀ f :

We say a distribution ν ∈ ∆( S × A ) is admissible in MDP ( S , X , A , P, R, γ ) , if there exist h ≥ 0 and a policy π such that ν ( s, a ) = ∑ x ∈X Pr( s h = s, x h = x, a h = a | s 0 , x 0 , π ) . The following assumption is imposed to limit the distribution shift. Note that 'admissible' is defined on the stochastic state and action. Later, we will also abuse the notation and call ν ∈ ∆( S × X × A ) is admissible if ν ( s, a ) = ∑ x ν ( s, x, a ) is admissible.

Assumption A.3. For a data distribution µ, we assume that there exists C < ∞ such that for any admissible v and any ( s, a ) ∈ S × A ,

Assumptions A.2-A.3 are standard assumptions in batch reinforcement learning (Chen and Jiang, 2019). However, in assumption A.3, we only require the data coverage of stochastic states which is the fundamental difference.

Assumption A.4. We assume that there exists a set B of typical pseudo-stochastic states such that the distributions β k ( x ) , k = 1 , . . . , K used for augmenting virtual samples satisfy β k ( x ) ≥ σ 1 , ∀ x ∈ B . We also assume that the marginal distribution over using a reasonable policy π, that is, d π η 0 ( s, x ) := (1 -γ ) ∑ ∞ t =0 γ t -1 Pr( s t = s, x t = x | ( s 0 , x 0 ) ∼ η 0 , π ) satisfies ∑ s ∈S ∑ x ∈X ,x/ ∈B d π η 0 ( s, x ) ≤ σ 2 , where η 0 is the initial distribution. Furthermore, if we have for a distribution η of the states satisfying ∑ s ∈S ∑ x ∈X ,x/ ∈B η ( s, x ) ≤ σ 2 , then under any reasonable policy π, the marginal distribution η π h ( s, x ) := Pr( s h = s, x h = x | ( s 0 , x 0 ) ∼ η, π ) satisfies ∑ s ∈S ∑ x ∈X ,x/ ∈B η π h ( s, x ) ≤ c 0 σ 2 , ∀ h > 0 , for some constant c 0 ≥ 1 . In particular, all the policies at each iteration, the optimal policy and their joint policies are assumed to be reasonable policies.

Weremark that in a queuing network, under any stable policy, the queue distribution has an exponential tail; in other words, large queue lengths occur with a small probability. In such a case, we can use a uniform distribution for pseudo-stochastic states in set B to guarantee that σ 1 = Θ ( 1 log(1 /σ 2 ) ) .

Therefore, if we choose σ 2 = 1 n , then σ 1 = 1 log n .

Supplementary Material

Auxiliary Lemmas

In the following lemma, we will show that when all admissible distributions are not far away from the data distribution µ over stochastic state S and action A , we can have a good coverage of S × X × A by generating virtual samples.

For a given dataset D of size | D | = n and data distribution µ, let ¯ µ β denote the expected distribution of the the state action pair ( s, x, a ) in the combined datase after using Algorithm 2 with a virtual sample distribution β ( x ) .

Lemma A.1. Given a virtual sample generating distribution β ( x ) of the pseudo-stochastic state, if β ( x ) ≥ σ 1 , ∀ x ∈ B . Then given any admissible distribution ν, then under Assumptions A.3, we have for any ( s, x, a ) ∈ S × X × A , x ∈ B ,

where

Proof. Given the data distribution µ, we know that the real samples are drawn according to µ ( s, x, a ) . Then

where the inequality (1) holds because of Assumption A.3, the equality (2) holds by substituting equation (8) and the last inequality is true because the fact that β ( x ) ≥ σ 1 , ∀ x ∈ B .

The next lemma transforms the norm in terms of distribution ν to distribution ¯ µ β (Eq. (8)).

Lemma A.2. Let ν be any admissible distribution, ¯ µ β denote the new data distribution defined in Eq. (8) after generating virtual samples with β ( x ) . If β ( x ) ≥ σ 1 , ∀ x ∈ B , then under Assumption A.3, for any function f : S × X × A → R , we have ‖ f ‖ 2 ,ν ≤ √ m +1 m C σ 1 ‖ f ‖ 2 , ¯ µ β , where

Proof. For any function f, we have

Proof.

Lemma A.4. Under Assumptions A.3 and A.4, for any admissible distribution ν ∈ ∆( S × X × A ) , and a data distribution ¯ µ β associated with a virtual sample distribution β ( x ) , define P ( ν ) as a distribution generated as s ′ i ∼ P ( ν ) , then for any policy π, and f, f ′ : S × X × A → R , we have

Proof.

where inequality (1) holds because of triangle inequality, inequality (2) comes from lemma A.2, inequality (3) holds because

and the last inequality holds due to Lemma A.3.

where the first inequality is a result of Lemma A.1.

Lemma A.3. Consider two functions f, f ′ : S × X × A → R and define a policy π f,f ′ ( s, x ) := arg min a ∈A min { f ( s, x, a ) , f ′ ( s, x, a ) } . Then we have ∀ ν ∈ ∆( S × X × A ) ,

Lemma A.5. For a given data sample ( s, x, a, r, s ′ , a ′ ) generated from a data distribution µ, such that ( s, x, a ) ∼ µ, s ′ ∼ P ( ·| s, a ) , x ′ = g ( s, x, a, s ′ ) , for any f, f ′ ∈ F , define V f ( s, x ) = min a ′ f ( s, x, a ′ ) , then

Proof.

where the equation (1) = 0 because that condition on ( s, x, a ) , we have f and V f ′ are independent.

Lemma A.6. Under Algorithm 2, at iteration k, we have

β k D k

Proof. Recall that ˆ D k = D ∪ D k and | ˆ D k | = nm + n. The expectation is w.r.t. the random draw of the dataset D and the random generation of dataset D k with virtual sample distribution β k . We know that

Let M ( s,x,a,s ′ ,x ′ ) k denote the set of virtual samples that are associated with the real sample ( s, x, a, s ′ , x ′ ) at iteration k. Then

where ¯ r = R ( s, ¯ x, a ) , µ k ( s, x, a ) = ∑ x ′ ∈X µ ( s, x ′ , a ) β k ( x ) = µ ( s, a ) β k ( x ) . Since we have ¯ µ β k ( s, x, a ) = 1 m +1 µ ( s, x, a ) + m m +1 µ ( s, a ) β k ( x ) .

Therefore,

The next lemma shows an upper bound on ‖ f k +1 -T f k ‖ 2 2 , ¯ µ β k .

Lemma A.7. Given the MDP M = ( S , X , P, R, γ ) , we assume that the Q -function classes F satisfies ∀ f ∈ F , T f ∈ F . The dataset D is generated as: ( s, x, a ) ∼ µ, r = R ( s, x, a ) , s ′ ∼ P ( ·| s, a ) , x ′ = g ( s, x, a, s ′ ) , and the new dataset ˆ D k = D ∪ D k is generated by following Alg. 1 with virtual sample generating distribution β k ( x ) at k th iteration. Then with probability at least 1 -δ, ∀ f ∈ F , and k = 0 , . . . , K we hvae

Proof. Using Lemma A.6 we know that

Then it is sufficient to bound ‖ f -T f ′ ‖ 2 2 , ¯ µ β k by bounding

For any f, f ′ , recall that

For any f, f ′ define

Then for each ( s, x, a, s ′ , x ′ ) ∈ D, we get i.i.d. variables Y 1 ( f ; f ′ ) , . . . , Y n ( f ; f ′ ) .

We also define

where ( s i , ˆ x i , a i , ˆ r i , s ′ i , ˆ x ′ i ) is an augmented sample based on the i th real sample ( s i , x i , a i , r i , s ′ i ) . Denote the m i.i.d virtual samples by X i 1 ( f ; f ′ ) , . . . , X i m ( f ; f ′ ) . Therefore

Taking the expectations on both sides, we obtain for any f, f ′ ∈ F ,

We need to introduce 1 n ∑ n i =1 Y i ( f ; f ′ ) and 1 m ∑ n i =1 ∑ m j =1 X i j ( f ; f ′ ) to bound the above terms. For the first term, we know that the variance of Y can be bounded by:

where the last equality is true because

Then by applying Bernstein's inequality, together with a union bound over all f, f ′ ∈ F , we obtain with probability 1 -δ we have

For the second term, note that for any given i th sample ( s i , x i , a i , s ′ i , x ′ i ) all the variables { X i j } are i.i.d. Then following a similar argument, then for all f, f ′ ∈ F , we have with probability at least 1 -δ/n,

Then it is easy to obtain that we have for all f, f ′ ∈ F ,

Combining Eq.(17) and Eq.(16) we can obtain with probability at least 1 -δ, for all f, f ′ ∈ F ,

Let f = f k +1 , f ′ = f k , then according to Algorithm 2 we know that

According to Eq. (14), we have

Then it is easy to observe that ˆ T k, F f k minimizes L ˆ D k ( · ; f k ) , it also minimizes

because the two objectives only differ by a constant L ˆ D k ( T f k ; f k ) . Therefore under assumption A.2 we know that T f k ∈ F , we are able to obtain that

where the last equality holds due to the definitions of Y i and X i j . Therefore plugging the result from Eq.(19) into Eq.(18), we can obtain

By solving the quadratic formula, we get

Finally, apply a union bound over all t = 0 . . . K, we conclude the proof.

A.1 Proof of Theorem 1

Now we are ready to show the main theorem. Given a dataset D. After generating virtual samples D k we get a new combined dataset ˆ D k = D ∪ D k at each iteration k with virtual sample generating distribution β k ( x ) . We first have

where the first equality follows from the performance difference lemma (Kakade and Langford, 2002), the first inequality holds because π f k ∈ arg min a f k ( s, x, a ) and the last inequality is true by using the fact that for any vector a = ( a 1 , . . . , a n ) and a valid distribution d = ( d 1 , . . . , d n ) , ∑ i d i = 1

According to Assumption A.4, we know that ∑ s ∈S ∑ x ∈X ,x/ ∈B d π f k η 0 ( s, x ) ≤ σ 2 , which implies that ∑ ( s,x,a ) ∈S×X×A ,x/ ∈B { d π f k η 0 ( s, x ) × π ∗ } ( s, x, a ) ≤ σ 2 , and ∑ ( s,x,a ) ∈S×X×A ,x/ ∈B { d π f k η 0 ( s, x ) × π f k } ( s, x, a ) ≤ σ 2 . Define ξ = { d π f k η 0 ( s, x ) × π f k } ( s, x, a ) , then we have

where the first inequality holds because √ a + b ≤ √ a + √ b ( a ≥ 0 , b ≥ 0 ) and the last inequality comes from Lemma A.4.

By using LemmaA.7 we have with at least probability 1 -δ

where glyph[epsilon1] 1 = 5 ( 1 n + 1 m ) V 2 max log( nK |F| 2 /δ ) + 3 δV 2 max n . Therefore, we obtain

Now define ξ ′ = P ( ξ ) × π f k -1 ,Q ∗ . Then based on Assumption A.4, it is easy to obtain ∑ ( s,x,a ) ∈S×X×A ,x/ ∈B ξ ′ ( s, x, a ) ≤ c 0 σ 2 . Note that the distribution ξ ′′ = P ( ξ ′ ) × π f k -2 ,Q ∗ still satisfies ∑ ( s,x,a ) ∈S×X×A ,x/ ∈B ξ ′′ ( s, x, a ) ≤ c 0 σ 2 according to Assumption A.4, because π f k -1 , π f k -2 and Q ∗ are all assumed to be reasonable policies.

Repeating the expansion above for k times, we have

All the above analyses are still applied to the case when ξ = { d π f k η 0 ( s, x ) × π ∗ } . Therefore, it is straightforward to obtain

Substituting glyph[epsilon1] 1 completes the proof.

A.2 Extension of Theorem 1

We repeat the assumption for extending our main results to the case when |X| can be infinite such that f ( s, x, a ) may not be bounded by V max .

Repeat of Assumption 1: For the typical set B of pseudo-stochastic states (defined in Assumption A.4), for any s, a ∈ S × A , f ∈ F , if x ∈ B , then f ( s, x, a ) ≤ V max otherwise if x / ∈ B , we have | f ( s, x, a ) -Q ∗ ( s, x, a ) | ≤ V max . Furthermore, for any given f ∈ F , ( s, x ) , x ∈ B , we have | V f ( s ′ , x ′ ) -V f ( s ′′ , x ′′ ) | ≤ V max , where x ′ = g ( s, x, π f , s ′ ) , x ′′ = g ( s, x, π f , s ′′ ) .

There are two places we need to pay attention to: (1) : a bound on ‖ f 0 -Q ∗ ‖ 2 ,P ( ξ ) × π f 0 ,Q ∗ , (2) : a bound on the variance of Y as shown in Eq. (15). For the first case, it automatically holds due to assumption 1. For the second term, we first have

We also know that

Therefore, we have Var ( Y ( f ; f ′ )) ≤ 16 V 2 max E [ Y ( f ; f ′ )] .

Then we can obtain a similar result of the same order, which only differs for some constant ˜ c such that

B Additional Simulations

B.1 Combining PSG with Policy Gradient-type algorithms

In this section, we investigate the possibilities of using ASG in policy gradient-type algorithms. In particular, we use ASG in the phase of policy evaluation. An algorithm (SAC-ASG) that incorporates ASG into SAC is presented in Alg. 3. We also compare our algorithm SAC-ASG with state-of-art Dyna-type model-based approaches, i.e., MBPO (Janner et al., 2019) on the phases criss-cross network environment (Fig. 2b). The simulation results are shown in Fig.4. We can observe that

Algorithm 3: SAC-ASG

the performance of our approach is significantly better than the baselines'. We also would like to emphasize that the training time of our approach is much less than that of MBPO ( 4 hours v.s. 3 days).

14

7.5

SAC

7.0

SAC-ASG

IQueueLength

6.5

MBPO

Optimal

6.0

5.5

Total

5.0

4.5

4.0

0

2

4

6

8

10

12

TrainingEpisodes

B.2 Details of the Environment in Section 4.4

In this section, we summarize the detailed parameters used in section 4.4 in Table 3.

B.3 Two-phase Wireless Network

The two-phase wireless network is a modified version of the downlink network in Section 2.1, which cannot be modeled as an MDP with exogenous inputs. In this example, each packet has two phases. When a packet from a mobile is scheduled, if it is in phase 1 , it leaves with probability 0 . 2 or moves from phase 1 to phase 2; and if it is in phase 2 , it leaves the system. We considered Poisson arrivals

with rates Λ = { 2 , 4 , 3 } , and the channel rates O = { 12 , 12 , 12 } . As the result shown in Figure. 5 ASG outperforms Max-weight in this environment as well.

34

DQN-ASG

Length

32

30

Total Queue L

28

26

24

230

60

100

140

180

220

260

300

TrainingEpisode

C Simulation Details

In our simulations, the reward function is the total queue length (or a scaled version: divided by 20 i the wireless networks). We used Gaussian distributions fitted with real data samples as β k , for each real sample, we generated 50 augmented samples and repeated the process 10 times. We use a simple two layer neural network with hidden size 128 with learning rate 0 . 0003 . The batch size used in all the simulations is 256 .

C.1 Experimental settings

For all the simulations, We used a single NVIDIA GeForce RTX 2080 Super with AMD Ryzen 7 3700 8 -Core Processor.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on a specific class of Markov Decision Processes (MDPs) in which a state can be decomposed into a stochastic component and a pseudo-stochastic component. The transition dynamics of stochastic states depend solely on the current stochastic state and the action taken, while the transition dynamics of pseudo-stochastic states are determined deterministically by the current state, the next stochastic state, and the action. Such MDPs are suitable for modeling applications like queueing systems.\n\nThe authors initially consider a scenario with finite state spaces and batch reinforcement learning. Given a set of independent and identically distributed (i.i.d.) transition and reward data, they propose a data augmentation algorithm that generates m new virtual samples from each data point. This is achieved by replacing the pseudo-stochastic components with newly sampled ones. Subsequently, they propose the fitted Q iteration with an augmented dataset. Under certain coverage assumptions for stochastic states and with n initial samples, they establish an error bound on the Q function, which scales as 1/\\sqrt{n} + 1/\\sqrt{m}. The authors extend their results to scenarios with an infinite pseudo-stochastic state space, along with additional assumptions. Simulations are conducted to demonstrate the benefits of the proposed algorithm.\n\nSTRENGTHS:\nUsing reinforcement learning to determine optimal policies for systems with large state spaces is a well-known challenge. However, this paper presents a valuable and intuitive framework that effectively captures the inherent structure of environments, specifically queueing systems, through the introduction of pseudo-stochastic states. This innovative modeling approach is commendable. The concept of utilizing fixed transition functions to generate new data points is novel and has practical relevance for systems influenced by limited stochastic inputs. The theoretical results provided in the paper are sound and contribute to the overall strength of the work.\n\nWEAKNESSES:\nIn Section 2.1, the authors provide an example where the pseudo-stochastic state space is not finite, which goes beyond the scope of the model as indicated by the assumption of a finite state space in line 148. Although the authors later extend their results, it remains unclear how this example aligns with Assumption 1.\n\nBuilding upon point 1, Assumption 1 poses challenges in designing the function class F for the Q function. It seems difficult to define a class where the value of an atypical state under any function within the class is consistently close to the optimal Q value. The authors need to provide better justification for why this property can be easily achieved. One approach could be to establish a broad enough function class that satisfies Assumption 1 for a simple queueing system. Justifying the reasonableness of Assumption 1 is crucial, especially considering the paper's claim that the theory and method are applicable to queueing systems with unbounded state spaces.\n\nThe simulation section lacks sufficient explanation and some of the plots appear unreasonable. Further details and clarification are needed to address these issues.\n\nThere are several typos that need to be corrected throughout the paper.\n\nQUESTIONS:\nIn line 93, the authors claim \"far less memory space.\" Could you provide a quantification, such as the complexity mentioned in line 91?\n\nIn the simulations, how is the reward function determined? How are the state sampling distributions \\beta_k set? Additionally, what is the number of virtual samples used?\n\nRegarding Figures 2 and 3, how is an episode defined in the context of the simulations?\n\nIn Figure 3 (a, b, c), it appears that the performance of DQN-ASG does not show improvement after training with episodes. Furthermore, why is the performance in episode 0, where the algorithm should not have learned anything yet, already better than other algorithms? The same question applies to Figure 4 in the appendix.\n\nIn lines 24 and 25, from \"time-consuming\" to the period before \"In this paper,\" the text does not form a complete sentence.\n\nIn line 56, the last \"and\" is unnecessary.\n\nIn line 222, please change \"satisfies\" to \"satisfying.\"\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel setting called mixed systems, in which environments comprise both stochastic and deterministic transitions. This leads to two types of states: stochastic states and pseudo-stochastic states. Stochastic states follow a stochastic transition kernel, while pseudo-stochastic states have deterministic transitions given the stochastic states/transitions. Such systems are common in real-world reinforcement learning (RL) applications such as data centers, ride-sharing systems, and communication networks.  \n\nExisting RL approaches can suffer from sample inefficiency due to the curse of dimensionality. The paper proposes an augmented sample generator (ASG) to improve sample efficiency. ASG augments transitions by retaining stochastic states while sampling new pseudo-stochastic states.  \n\nThe paper provides a motivating example of a wireless downlink network, where stochastic states are the number of data packets arriving and the number of packets that can be transmitted at each time step. The pseudo-stochastic state is the length of the queue, which can be computed deterministically given the stochastic states.  \n\nThe authors offer a convergence guarantee in the form of an optimality gap, connecting the learned value and optimal value with the number of real samples and augmented samples per real sample.  \n\nExperiments are conducted on a criss-cross network and the motivating example. The results demonstrate that combining ASG with Q-learning or policy gradient-type algorithms significantly improves performance.\n\nSTRENGTHS:\n1. The paper presents a new setting that is widely observed in real applications.  \n2. A highly suitable algorithm (ASG) is proposed for this setting.  \n3. The effectiveness of ASG is demonstrated through solid theoretical proof and several experiments.\n\nWEAKNESSES:\nThe main weakness of the paper is that the difference between the new mixed systems setting and existing MDPs with exogenous inputs is not entirely clear. Although the analysis in Line 133 is sound, the environments used in this paper can be modeled with MDPs with exogenous inputs. For example, in the wireless downlink network, stochastic state is independent of actions taken. Seems that all environments used can be viewed as MDPs with exogenous inputs, and also a special case of mixed systems. However, the new setting is valuable, and more real-world application examples that can only be modeled with mixed systems are desirable. **If the authors address this issue, I recommend strongly accepting the paper**.  \n\nA minor weakness is that the algorithm assumes access to the reward function $R$ and pseudo-stochastic state transition function $g$, potentially limiting its scope. However, this assumption seems reasonable in most examples mentioned.\n\nQUESTIONS:\n1. In Sections 4.1 and 4.2, the system seems to be defined on a continuous time span. It is unclear how continuous time is discretized. I suggest adding details about this issue to the main text or supplementary material.  \n\n2. Regarding the wireless downlink network example in Section 2.1, let's consider a situation where the real pseudo-stochastic state has queue lengths of {0, 0, 10}, and a good policy would take action 3 to schedule mobile 3 at this time step. ASG could augment a pseudo-stochastic state as {10, 0, 0}, and in this case, the real action 3 taken seems inappropriate. Although this augmented data still follows the transition function g, it is information less to the policy learning. How will such augmented data affect policy learning and the convergence result?\n\nLIMITATIONS:\nThe paper does not extensively discuss limitations or potential negative societal impacts of the work. No further limitations need to be addressed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper consider the sample efficiency of RL in a setting involving \"mixed systems\" where the state space is factored into a stochastic component and another much larger component that is deterministic when conditioned on the stochastic component. The authors develop an algorithm for this setting relying on generating relevant synthetic states that have not been encountered before. This approach when applied to fitted q-learning is able to show sample efficiency bounds for which it is not necessary that the data includes all pseudo-stochastic states. The authors also demonstrate the practical utility of their approach in application to the deep function approximation settings based on queuing networks.\n\nSTRENGTHS:\n- The paper is well written and makes a focused contribution. \n\n- Not only is significant theoretical analysis provided, but also it is shown that this theory makes a meaningful impact for deep function approximation. \n\n- While the empirical domains are not very complex, they do represent real-world use cases, which makes up for this deficiency in my mind.\n\nWEAKNESSES:\n- I think the approach will make a big impact where it is applicable, but I worry that this setting is a bit niche. I don't think this is a big problem, but just am a bit less excited as a result. \n\n- It is unclear to what extent the approach can be extended to partially observable domains with more complex input spaces.\n\n- Given the similarities, I would appreciate a more clear connection with the factored MDPs literature.\n\nQUESTIONS:\n1. How realistic is it that the factorization of stochastic states and pseudo-stochastic states is provided to you? Can the cost of identifying these be bounded? \n\n2. Can you talk to the connection between your setting and work on factored MDPs, including existing sample efficiency bounds in that setting? To what degree are you using novel ideas in the bound and approach derivation and to what degree are the insights different from that literature? \n\n3. Do you think that any of the insights of this paper can be useful in partially observable versions of mixed systems, or is it restricted to fully observable settings?\n\nLIMITATIONS:\nThe authors do a good job for the most part of discussing assumptions and limitations. I agree that the assumption for the case of an infinite number of pseudo-stochastic states seems mild. Looking at appendix A, the data coverage and completeness assumptions also seem mild, although I think at least more intuition about this should be provided in the main text. One thing I think should be further highlighted and discussed though is that not only do these components of the state space need to exist but they also need to be provided to the user. In my mind, this knowledge about the state space is relatively hard to come by and limits the areas of application. Maybe taking inspiration from the literature on factored MDPs, the authors can account for the cost of the discovery process as well.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}