Sample Efficient Reinforcement Learning in Mixed Systems through Augmented Samples and Its Applications to Queueing Networks
Authors: Honghao Wei
Xin Liu
Weina Wang
Lei Ying
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Reinforcement Learning
Mixed Systems
Queueing Network
Sample Efficient
TL;DR: We propose a sample efficient RL method that accelerates learning in mixed systems.

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
Honghao Wei Washington State University honghao. wei sei@wsu.eee edu
Xin Liu ShanghaiTech University Muxxin77Shanghaittech edu. cn
Weina Wang Carnegie Mellon University weinawOcs. cmu. edu
Lei Ying
University Michigan, Ann Arbor leiying@umi.co ch. edu
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 stochastic transition kernel while the transitions of pscudo- stochastic states are deterministic given the stochastic states/transitions. We refer to such systems mixed systems, which widely used in various applications, in- cluding manufacturing systems, communication networks, and queueing networks. We propose sample efficient RL method that accelerates learning by generating augmented data samples. The proposed algorithm 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 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 1(/I/n + 1//m), where n is the number of real samples and is the number of augmented samples per real sample. It is important to note that without augmented samples, the optimality gap is (1) due to insufficient data coverage of the pseudo-stochasee 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.
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 significant challenge to extending the success of RL broader applications. In this paper, sample efficient algorithms for a class of problems in which envi 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,
interested in
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

which can be modeled as queueing networks so as mixed systems. The broad application of mixed systems makes crucial to explore whether current RL approaches are efficient in learning optimal policies for such systems. However, existing RI. 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) 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), 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 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 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. 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 large (even infinite) number of samples that cover the unobserved pseudo-stochaste 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 seeudo-stochssti states become stochastic. The distributions of stochastic states and pscudo- stochastic 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 pseudo- stochastic states. Both the real samples and virtual samples are then used to train the deep neural networks, e.g.. Q-networks 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-stochastee 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 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(VI/n + VI/m), where n is the number of real samples. Note that without augmented samples, the error Õ(1) due to the lack of data coverage of the pseudo-stochastic states, which reduces to (V/1n) when we generate at least n augmented samples for each real sample. We also would like to emphasize that w( /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 will explain in detail.
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, proposed method differs fundamentally
Dyna and its variants:
2

---

from these approaches in several aspects: (1) Dyna-type algorithms are model-basee methods and need estimate the system model which will be used 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 | 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( Q(A) (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-an 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
Table 1: Dyna-type v.s. ASG
Algorithm
estimate model?
update unseen states?
computational efficient?
store all transitions?
convergence analysis?
Dyna-type
X
ASG
X
Data augmentation: It has been shown that data augmentation can be used to efficiently train RL algorithms. Laskin 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 G- Convolution (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 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 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 oww-imensional representations in RL.
3

---

Factored MDP: A factored MDP (Kearns and Koller, 1999) is an MDP whose reward function and transition kemel exhibit some conditional independence structure. factored MDP, the state space is factored into a set of variables or features, and the action space is factored into 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 pseudd-stochasee states generally depend on the full state vector and are not factorized. However, these transitions deterministic,, which allows to generate augmented samples. Given the fundamental differences between the two, the analysis and the results 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: A recent paper considers MDP with exogenous inputs (Sinclair et al., 2022), where the system has endogenous states and exogenous states. The fundamental difference between exogenous sxogenous-state//ee state/endogenous state versus suochasiic-sttteppee state//seuudostee www.ssssttocasstii-tttete-- 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-stochaste state, the evolution of stochastic state, therefore, depends on the pscudo -tochosttcchastic state.
is that the
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 X', A. 7, no). space is S x X,, where 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 finite action space, R S x X' x A - [0, T'ax a deterministic and known cost function, and P S x X A (S x X)) is the transition kernel (where ((-) is the probability simplex).
whose state
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, at time can be represented
(1) We assume that this transition is independent of the pscudo-stochastic state X, (and everything else at or before time f). The transition of the pscudo-stochastic state is then deterministic, governed by a function g as follows:
(ay(a) P(S++1 = 5|S, = i, a).
X,+1 = g(,, X1, az, St+1).
Our system M is said to be 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
We focus discounted costs in this paper. Given a mixed system, the value function of policy is defined as
(*(
E
R(sh, Th, Th, an) (s0, = (8, a = n m(s,, mn)
h=
where 7 is the discount factor, and the expectation is taken over the transitions of stochastic states and randomness in the policy. Let Th = (e. [ (80, mo)], i.e., the expected cost when the initial distribution of the state is no. Let v" min, v".

---

The Q-value function of a policy is defined as
(=(s,za))== =
h R(sh, 2h on) (80, (s,z), a, on x (5h
(4)
Let Q be the optimal Q-value function. Then the Bellman equation can be written as
Q(s, = R[s, a) + P((sa)) min Q(s', gg(s, 2,,,
(S)
Note that the value function and the Q-value function both take values in [0, Vmmxx] for some finite Viax under any policy, due to the assumption that the cost [0, fmax].
2.1 Example: A Wireless Downlink Network
To better illustrate the structure of mixed systems. Consider the example of wireless downlink network shown in Fig. I with three mobile users. We model it as a discrete-time system such that A (i) the number of newly arrived data packets for mobile i at the beginning of time slot t, and (i) is the number of packets that can be transmitted to mobile during time slot if mobile i is scheduled. The base station can transmit to one and only one mobile during each time slot, and let A, E 21, 3} denote the mobile scheduled at time slot t. The length queue evolves
Flow
Flow
4++1 (1) = (1) + (i) 0,(0I(A,,
(6)
Flow
where I[() is the indicator function and (-)+ max
To minimize the total queue lengths, the problem can Figure 1: A Downlink Wireless Network be formulated as an RL problem such that the state is (A,, O.. qt), the action A,, and the cost is Σ; (i). We can see that in this problem, S; == (A,, O₂) is the stochastic state and qe is the pseudo-stochastic state. In general, the state space of the stochastic states is bounded, e.g., both A, and O, may be Bernoulli random variables, but the pseudo stochas-c state such as q: is large 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 given policy, will require orders of magnitude more samples. Therefore, vanilla versions of model-free approaches that do not distinguish between stochastic and sscudo-stochastie states require unnecessary amount of data samples and are not efficient.
Therefore, we propose 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 A,, O and A,, 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 input dataset D, an integer m, and a probability distribution 3-)) over the pscudo-stochastic states in X. For each sample (s, a, r, s', x²) D, we first sample secudo-stochastie state à from 8(-), and then construct virtual sample (s, a, s', H'), where i = R(s, f., a) and E' = g(s, â. a, 8'). Note that we are able to do this since 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 timeslot I, one of the real samples is
(At, Oz), X, qt. assume that some
(s, a, 8', = ((f(3, 4, 5), (1,2.0) (4.6 6), 1, 16, (1(2. 3, 4), (2,2,0)) (6. 10, 11)),
5

---

where the queue length qt+1 is calculated according to Eq. (6). Using ASG, we generate two pseudo-stochastic states (queue length) & = (1,2,3) x2 = (0, 2, 1),, then we are able to obtain two virtual samples
(,, 81 :- (((3,4,5,, (1,2, (1, (1,2,) ((,33)) (,,,,,4), (2,
(12,3, (s, 22, a, f2, ,2,2)) ((f3,4,5), (1,2,0)), 10,2, (1,,, 3,4), (2.2.0)), (2, 6)).
Note that all the virtual samples represent true transitions of the mixed system if the queue lengths were _1 and &2.
Algorithm 1: Augmented Sample Generator (ASG)
Input: A dataset D = {(s, x, a, r, s', ,²)),, a positive integer m, a distribution B()) on X; Initialize virtual sample dataset: D' 0;
3 for each sample (s, x, a, r, ,'',''' E D do
Sample virtual pseudo-stochastic states é1. B2...... I.m from B(-) independently;
Compute i; = R(8, L;, a) and i = g(s, di, a, 5') for i = -,,,.... mm; Add m virtual samples to virtual dataset: {(8, di, a, f., 8,, , 1, 2, ,,...,.... mm - D):
Output: D' D
3.2 Batch FQI with ASG
We now formally present algorithm, Batch FQI with ASG (Algorithm 2), for mixed systems. As in 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..dd and the distribution of each sample (s, a, r, x²) specified by a distribution E ASS x X' A) for (s, x. a), r R(s, x, a), s P(-|s, a). and unknown the agent.
x' - g x, a, s'). The distribution
Our algorithm follows the framework of a batch FQI algorithm. We consider class of candidate value-functions F (S X x A [0, Vmax]). We focus the scenario where the number of escudo-stochastee states, |x) is far more than that of stochastic states, (S], and the dataset D does not provide sufficient coverage of the pscudo-stochasto states (n < (.)).
The goal is to compute a near-optimal policy from the given dataset D, via finding an f E F that approximates the optimal Q-function Q. The algorithm runs in episodes. For each episode k. it first generates an augmented dataset Dj using ASG(D, m, BB2), where B& is some distribution of the pseudo-stochastic states satisfies that, for each typical pseudo-stochastic state, it can be generated with at least probability σT 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 3 for all the episodes, and in practice, we can simply adopt the estimation of the pseudo-stochastic states (i.c., with Gaussian distribution) as Br. The algorithm then computes fx as the minimizer that minimizes the squared loss regression objective over J.. The complete algorithm given 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 = {8, x, a, rr,'',, s', with n samples, w.p. at least 1 8, the oulput policy of Algorithm after k iterations, Tu satisfies
2
(m (m+ 11CC
(3
1 m
nK|FF1
35V2
Jk - 30 5
(1 7/2
(V
moi
72
V2
max log
8
+
max
VC max + x +*1 >)
where is a constant related to the data coverage assumption A.3, and Co, 01, 02 are constants defined in Assumption A.4 to ensure typical set of reasonable typical neudd-stochaster states.

---

Algorithm 2: Batch FQI with ASG
Input: A dataset D {(, R; 02 v, 8 x²) }. function class F, a positive integer m, distributions
Initialize approximate Q-function: Choose So randomly from F.;
3 for episode k = 1,2,.... ..., K do
Generate augmented dataset Dr = ASG (D m, Bc):
Define loss function as:
LD. (JJn-1) [Dl
1
Σ
(1(8,2,0) - - (dd',=))
fx arg mer L. (o min fs) f x) arg min (,)
Output: ffx
Theorem 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 3 that guarantees that 02 (l/n)) and σ / log ri). Therefore, for sufficiently large and sufficiently small 02 the first term is the dominating term. As m increases from to n, the first term of the bound decreases from (1) to ((/I/n). Remark that we cannot establish convergence result without using ASG (when m = since using the dataset alone doesn't have a sufficient data coverage guarantee, which is critical for batch RI. Due to the page limit, we only present the outline behind our analysis. The detailed proof deferred to the supplemental material.
Proof Outline:
1.
First we will show that using performance difference lemma (Kakade and Langford, 2002) it is sufficient by bounding f - Q" lla.s, 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 on the dis- tribution. Along with the augmented data generator under (.(()). Then later we show that the term
fll1 Q*112.5 /o where p&a, is the data distribution after augmenting virtual samples, The first two terms are related 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.
Il 22, can be bounded by O (Ilf T
3. Therefore, it is obvious that if we can have a bound on || - T ffs-1ll2,, which is independent of the episode k, we can prove the final result by expanding fl. - Q" ll2.ç iteratively for & times. We finally show that the term f fx-- I,,, can indeed be bounded by using the FQI minimizer at each episode and concentration bound on the offline dataset after augmenting virtual samples.
3.4 Extension to the case when X is infinite
Theorem assumes the number of the pseudo-stochaste stochastic states is finite. The result can be generalized to infinite pseudo-stochasto stochastic state space | if we make the following additional assumption:
Assumption 1. For the rypical set B of pseudo-stochasee in the supplemental material), for any 8,a S E
states (formally defined Assumption A.4
E B, then f(s a,,a) S Vmax
otherwise if x & B, we have [f(s, x, a) - Q (0,,,,00 S Vmax. Furthermore, for any given f E F, (8, E B, we huve [Vy(s', x') Vr(s", x²)l S Vmax; where x' =g(s, (,,,
-9,, u 8') x
The first part in assumption means that the function f F is always bounded when E 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, is also true for the optimal Q --uucctoo Q*. Therefore the difference between and Q" is assumed to be bounded, which is reasonable in practice. The second
7

---

part states that for any given E B, the difference between the value functions of any possible next transitions (s', =²) 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 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 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 I, 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 time. Server 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,. The service rate of a class i job is defined to be μ4s == 1/m,, and jobs of class 1 and class 3 arrive to the system following the Poisson processes with rates A and Az- To make sure the load can be supportable (i.e. queues can be stabilized), we assume Aym1 + X3ma < 1 and λ m2 < 1. In the simulation, consider the imbalanced medium traffic regime, where /11 = F3 = 2, 122 = 1.5, A = A 0.6.
first-n-fiirs--or (FIFO). Let ea(i)= E [-1,0,+1) denote the state of class i jobs at time f, where -1 means a departure, +1 means a arrival and 0 denotes no changes for job i. Then e, (e, e (3)) be the event happens at time t in the system, which is the stochastic state in this mixed system. Let q = (g((1.. 9: (2). g (3)) be the vector of queue lengths at time t, where (i) is the number of class jobs in the system. Obviously, 9: is the eseudo-stochastee state in the system, which can be derived from (q.-1, et-1) and action a,. We combine ASG with Q-learning and compare it with several baselines: (i) the vanilla Q-learning, (ii) proximal policy optimization (PPO) based algorithm proposed in Dai and Gluzman (2022) designed for queuing networks, (iii) random policy, and (iv) a priority policy such that the server always serves class 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 (4k 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 4k 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-stochaste 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.
All queues
4.2 Two-Phase Criss-Cross Network:
To further evaluate the power of ASG beyond the tabular setting, We combine ASG with Deep Q- network (Mnih al., 2015b) named DQN-ASG and evaluate it on more complicated variant of the
8

---

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 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). the simulation, all the parameters are set to be the same those in the previous case, and we use 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 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 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- -tochastt states (the queue lengths). However, the exogenous input in MDP with exogenous inputs has to be 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 A {2, 3], and the channel rates O to be {12, 12, 12). We use the Max- Max-We (Tassiulas and Ephremides, 1992) algorithm, a ptroughput-optinaee and heavy-traffic delay optimal algorithm, as the baseline. As shown in Fig. 2f, the learned policy outperforms Max- Reight. To the best of our knowledge, no efficient RL algorithms exist that can outperform Max- -eight.
class
class 3
antiaas
anivals
Flow
class arrivals
class 2 departwes
phase
plose
clus depertures
class artivils
Flow 2
Rg
Rg
Flow
1
class departures
cla
departures
(a) Criss-Cross network
(b) Phases Criss-Cross network
(c) Wireless network
2.5
7.5
20
DONAAGD
DON--AA ASG
2.4
7.0
38
Max Weighet
6.5
0.009
2.3
6.0
30
2.2
100
0 5.0
2.1
12
2.0g
4.90
200
320
3%
200
Training Spioode
Tralning Eclonde
Training
(d) Testing reward
(c) Testing reward
(f) Testing reward
Figure 2: Performance queuing systems
4.4 Additional Simulations on the criss-cross network
To further verify the performance of our algorithm on more scenarios, we first consider 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 jobs need multiple timeslots to be processed instead of 1. particular, when 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 special case where class 1, and 3 have fixed job size 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 obtain the optimal solution is very time-consuming. The priority policy already provides reasonable good baseline.
DON-AS5
DON-A5G
DCN-ASG
Total Cunse Lergh
0.5
orittty
1.5
Saseline
ty
Pifority Ramdom
6.0
Randam
1.4
Rancom
4.25
Total Conge
1.3
4.00
8.5
1.75
3.50
8.0
11
2.2
1.50
80
140
200
200
320
1%
80
140
200
250
320
3,00
BC
140
200
260
125
Tra nannnn Edioods
Trairing Esicodo
Training Spisade
(a)
(b)
(c)
Figure 3: Performance on Criss-Cross network with general job size
4.5 Additional Simulations on the Wireless Network
We also evaluate the performance of our algorithm the wireless network (Fig. 1) under different sets of arrival rates. The results of average queue length compared with those of Max-Weig -eight are summarized in Table 2. We 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 using ASG in policy gradient-type algorithms (i.c., TD3 (Fu et al., 10. 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.
Arrival Rates
Service Rates
Max- Weight
DQN-ASG
[2,3,
{12, 12, 1
10.979
10.403
7,2]
{12, 12, 12)
13.811
13.273
12, 12,2,6, 6}
[12, 12, 12}
13.774
13.137
[3,1,5]
[12, 12, 12]
10.895
10.091
Table 2: Best testing average queue length after training for 120k steps
Conclusions
In this work, considered RL problems for mixed systems which have two types of states: stochastic states and learning by augmenting virtual data samples. We theoretically quantified the (real) sample complexity gain by reducing the optimality gap from (11) to ((1/ 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.
states. We proposed a sample efficient RL, approach that accelerates
References
Agarwal, A., Kakade, S., Krishnamurthy, A., and Sun, (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, and Jiang, N. (2019). lnfornation-theereticc considerations in batch reinforcement learning. In Int. Conf. Machine Learning (ICML), pages 1042-1051, PMLR.
10

---

Clavera, L,, Rothfuss, J., Schulman, J., Fujita, Y., Asfour, T., and Abbeel. P. (2018). Model-basee reinforcement learning via meta-policy optimization. CoRR, abs/1809.05214.
Cohen, T. and Welling, M. (2016a), Group equivariant convolutional networks. Int. Conf. Machine Learning (ICML), pages 2990-29999 PMLR.
Cohen, S. and Welling, M. (2016b). Steerable enns, arXiv preprint arXiv:1612.08490
Dai, J. 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 with rich observations via latent state decoding. 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. 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. JEEE/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 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 (ICCAI), 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 1802. 1802.
Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, (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. 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, 51897540):59-533.

---

Mnih, V., Kavukcuoglu, K., Silver. D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A.. Riedmiller, M., Fidjeland, A. K., Ostrovski, G., al. (2015b). Human-level control through deep reinforcement learning. Nature, 518(7540))529-535
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:2 072005.5.
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 (ICMI.), pages 7424-7433. PMLR.
Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. (2017). Curiosity- 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 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-62))
Silver, D., Schritwieser,, 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-33))
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 preprinl arXiv: 1oXiv:/200 1206.3285.
Tassiulas, L. and Ephremides, A. (1992). Stability properties of constrained queueing systems and scheduling policies for maximum throughput in multihop radio networks. JEEE Trans. Autom. Control, 37:1936-1944
van der Pol, E., Worrall, D., van Hoof, H., Oliehoek, F., and Welling, M. (2020). MDP homomor- phic 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., al. (2019). Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)
Wang, D., Walters, R., and Platt, R. (2022). -qquivrrvva reinforcement learning. In iclr.
Weiler, M. and Cesa, G. (2019). General e(2)eequiaar 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), N9,, 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: basel offline policy optimization. Advances Neural Information Processing Systems (NeurIPS), 33:14129-14142.
12

---

Supplementary Material
A Proof of Theorem 1
In this section, 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(v) as distribution over states such that (s', ~ P(v) e x, a) & ~ a)) xx = g(s, a, 8'). In other words, it is the distribution of the next state if the state action pair follows v. For f : S x X' x A + R, E A(S x X' x A), where A(-) the probability simplex, and p > 1, define Il = (B(s.a)) /)/.,) Define T (5, == arg A. min (s. 2, a), f (s,
Recall that in Alg. 2, at iteration k
LD. (Ji(JiK
1 DDl
Σ
()) -V,, (6',
(s,7,,,,,s s"
where
f& : arg min D. (s: fk-1) and (8,x)) = min fa 28, 2.). JEF
aad
Assumptions:
Assumption A.1. (Realizability) For the optimal policy, Q' E
Assumption A.2. (Completeness) For the policy a to be evaluated, Vf F,Tf E F, where T : RXXXXA + x is the Bellman update operator, Vf:
(Tf)s, a) = R(s,2,,0 a) + P(-lea( (V' x' = g(s, x. a, s')].
We say a distribution / E A(S A) is admissible in MDP (S, X, A, 7). if there exist h 2 and policy such that v(s, a) = Lmex Pr(sæ- S, assumption is imposed to limit the distribution shift. Note that "admissible" is defined on the stochastic state and action. Later, admissible if v/ss, a)- = Σn v(s, 2, a) is admissible.
ag also, zo, T). The following
will also abuse the notation and call A(S X A) is
Assumption A.3. For a data distribution p. assume that there exists C < co such that for any admissible v and any (8, E ESXAXA
v(s,a) SC,
p(s,a)
where (ssaa) = - v(s, a) and a(s,a) -xrex M(s, x. 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 B of typical pseudo-stochastic states such that the distributions Su (x), k = -,...,, K used for augmenting virtual samples satisfy 3a (x) 2 aı, Vx E B. We also assume that the marginal distribution over using a reasonable policy म, that is, d (s, x)
(1 - 7) Ei=0>"" Pr(ss = 8, It x1) where Du the initial distribution. Furthermore, if we have for a distribution n of the states satisfying Eoes ExEx (1, x) < 02. then under any reasonable policy n, the marginal distribution
1o, π) satisfies Exes ΣxEx, d (s.x) x) S a2
n7 TT) satisfies EeEs TassExex Exex Th (s, I C0O2, Vh > 0. for some constant co 2 1. In particular, all the policies at each iteration, the optimal policy and their joint policies are assumed to be reasonable policies.
mi (s. 2) ;= Pr(sk- s, Ih z1(50,
We remark that in a queuing network, under any stable policy, the queue distribution has an exponential tail; in other words, large queue lengths oceur with small probability. In such case, we can
a uniform distribution for pseudo-stochastic states in set B to guarantee that o = θ (1o[(//;))
Therefore, if we choose 02 then 0 log
13

---

Auxiliary Lemmas
In the following lemma, we will show that when all admissible distributions are far away from the data distribution μ over stochastic state S and action A, we can have good coverage of S x X A by generating virtual samples.
For a given dataset of size [D] = n. and data distribution Jr, let aa denote the expected distribution of the the state action pair (s, x. a) in the combined datase after using Algorithm 2 with virtual sample distribution ((x).
Lemma A.1. Given virtual sample generating distribution ((x) of the ggeudo-stochane state, if ((x)) 2 01, Vo E B. Then given any admissible distribution v. then under Assumptions A.3, we have
for any (s, x. a) E X A. E B,
V(s 2,,
(m + 1)C
P1 (8,
moy
where
As(s, x, =,s,
Σ ((s,
(3)))
(8)
nm + n
nm + n
REX
Proof. Given the data distribution μt, we know that the real samples are drawn according to pls, x, a). Then
v(s, x. a)
Exexv(ss
Es x Esexu(s,e 13(5, 2, Exex ((,
(5(5,2,0)
(3(
(()
CBExM(s,,,
(s,0a)
a(s,a)
C1)
(3((,
(s
C
p(s,x,a)
m+1
m+1
(m 1)p(s,a) a)
<0
m Lxex (1s,, 2, ()x())
(m + 1) 1)p(s,)) a)
SC
Σ dEB M 2, ((x
17
SC
m
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 3(x) 2 ai, Va E B.
The next lemma transforms the norm in terms of distribution v to distribution já (Eq. (8))-
Lemma A.2. Let v be any admissible distribution, P. denote the new data distribution defined in Eq. (8) after generating virtual samples with B(x). If B(x) σ1, Vx E B, then under Assumption
A.3, for any function f S X x A + R, have |[
m+1 C.
01 [f where
1/2
llflew
Σ
((s,,2,))2vss,a
(
Proof. For any function have
1/2
||f
Σ
(xx1))SSxXXX.EB
14

---

1/2
(m + 1)C E ((0,,,,))
moi
(m + I)C ||fl
moi
where the first inequality is result of Lemma A.I.
Lemma A.3. Consider two functions f.f': S x & x A + R and define policy (s, == arg a min (J(s, x, a), ['(ss,(s,_))]] x, a)} Then we have Yu E ASS x x x A),
IV, - Vy 12l12.p(w = 11 Mlla..
Proof.
11, Vp ((2)
P(s'sss a) (min (6' min a'EA (66.))))
Σ
("EE
(s...))SxX xxx d'ES
$
Σ
((s,,,,,)
-f'ls' ((ss',(s,,,,,,5) 8'), g(s,,,.))
=Ils f'12,000xx
Lemma A.4. Under Assumptions A.3 and A.4, for any admissible distribution E A(S x X x A), and a data distribution Bg associated with virtual sample distribution B(x), define P(v) as distribution generated as s ~ P(v), then for any policy n, and f. : X A R, we have
115 < v
(m (m+ (mC
I)C ||5 (J-- + 7115" @'12..
(10)
ma
Proof.
[1 - 12 = -f 12
S(y)lS - T8 + -IT 12
(m + I)C
S(2)
|lf - TT + ITT Q2112
mo
(m + 1)C
1()
|lf Tf' + >VVy V*122,0)
mo
(m 1)C
|ff -f- + >1lf' Q' Q'
mo where inequality (1) holds because of triangle inequality, inequality (2) comes from lemma A.2, inequality (3) holds because
IIT 1²1 = 11T - TQQIE = Eia. [(T7))6.2.0)
-Ele.zaan ((TerrP((ea) (455
Sy a00 (/
= Ex
(5 ('(.,,
-> V1
and the last inequality holds due to Lemma A.3.
15

---

Lemma A.5. For given data sample (s, a. r, a,,,s, generated from data distribution μ.
such that (s, x2 a) p.s' a())s.a), x' g(s, a, s'), for any f.f E F, define Vy(s,z) mina' [18, v. a') ) then
E ( -N-
- 1 + an 22))
(11)
Proof.
E
- -V
=R - (1
=B
+
2E (1) (7))
(1)=0
-E - +
=IIʃ -T 12, + 7 () Var Vy. (s', ,z)))s,,,1))
where the equation (1) = 0 because that condition on (s, x, a), we have f and Vp are independent.
Lemma A.6. Under Algorithm 2, at iteration we have
Cas, (J; Lpm, (Tf:f') ((l=f')[I[-TT)) -15--
(12)
where Lpo, (f:5") = -J:J')=B((p (f:
Proof. Recall that DP = D U Dj and |Dx = nm + n. The expectation is w.r.t. the random draw of the dataset D and the random generation of dataset D2 with virtual sample distribution 8k. We know that
Cb. (f:f')
1
Σ
(f(s,I,a) -r_x,,))---Vyiie -V-V(6,))
Duul (s,a,a.,, ())))
Σ (( -
DD x")
Let M ((s.a,,sssa")) ,ax) denote the set of virtual samples that associated with the real sample (8, a, x') at iteration k. Then
Cpp, (fif) = B, (f:f')
+ (IIf - 1 + 2 Va x Is,
[T
(by using Lemma A.5)
E
Σ
Σ
nm + n
(s,2,0,r,s x E (s. (.M.R (aamaa)
(If - 1 12,uu + 2 " Var(V Varl 2'),,,
+ nim
Σ
E
Σ
- x')²
s, s s,
num1 + n
(s.r,a,r,a// (')
(3,2,0.888 a.
(1f-TTi (f''1
" (3, Var(V VarVa)) (s', 22)s,
nm + n
16

---

Σ
-
Tmm + n
(If - (If--f'll " VarlV)/
num + n
Σ (8, a)
-
nm + (s,I,))ESxXXXA
BEX
-TM
12, 2 B
TUM +
mn
Σ
nu + (8.8,0)E8xXxxA
(IS -T 2 (.. (ala:a/ww/aa((y(((s,,1///
nmm +
Tum
(115 - 1 +
a)])
(by using Lemma A.5)
Tmm +
where F = R(8, F. (), (,, = Exexm(s,d ((s)a)3.(x). Since we have P.g (8, 28 2,,0 a) m+1 (8, 0) m+1 ()) Therefore.
Cap (f:f)) Cpo, (Tf': -
The next lemma shows an upper bound on fl - T f
Lemma A.7. Given the MDP M (S, P. R,,), 7), we assume that the Q-function classes F satisfies vf E F,T E F. The dataset D is generated as: (s, x, a) ~ p,r R(s, x, a), P(-(s,a), x' g(s, x, x.,,U.., 8'), and the new dataset DE = DU Da is generated by following Alg. I with virtual sample generating distribution (x(x)) at kih iteration. Then with probubility at least 1 8. F, and k = ...... K we hvae
[|5++1 T < (=+++1) - v2 10g(nk(51118)
35V2
max
(13)
Proof. Using Lemma A.6 we know that
Il Tf Cpp, (f:f) - Cas, (Tf:f).
Then it is sufficient bound |lf T f' by bounding
Cas, f': Cpp, (T f'):)) LID, -
For any J, f", recall that
(f:f') Dul (emaa..n
1
Σ
CD(lf))
Σ
+ Dul (s..a..
(D.(110)
For any f' define
Y(::")=([(ssa) - (000:/00----//////)0/0/////--000
Then for each (s, s',x') D, get ii.. variables Y((fff') x((fff").......(.....)
17

---

We also define
(f -- (T
( (7
Rr. az)
where (st, Li, at Pr, s/ {.) is an augmented sample based on the ith real sample (st, Fi ai Ti, 8,). Denote the i..d virtual samples by X;, (f: f'), "..,,,,............... Xim f') Therefore
"
n m
nm + n Lb (f:f) - L.. (T):))
n
x Σ +)(fff)+
nmm x
X (f:f')
nm + nm
(14)
Taking the expectations on both sides, we obtain for any f, EF,
Las (f: - Liee (T Tem n
nu
E
Σ Ex(r)
THT + n
n1
We need to introduce 1. E f' and 1
X,,(f: E X1, (f: f') bound the above terms.
n For the first term, know that the variance Y can be bounded by:
m
(Yr: f)) {)))
-E -B (J(5, - -T -76000)
<4V2 max (f(s. -) Tf
= =V2 max Ilf - f
=V22 max (f; f')l.
(15)
where the last equality is true because
(|Y(::))) = E(LD(J:)))- B(D =IIf- - 1 (a. /ar
- -I[T -T (o (ar (Var( (s =IIʃ- 1
('',,'')ss,,, x,),
(using Lemma A.5)
Then by applying Bernstein's 's inequality, together with union bound over all f.f f F, we obtain with probability - 5 we have
2Vrr(Y (f: (og()))))(())))))
4V2 max log( log((=(3/8)
x((fff) sv
m
3n
8V2 E (f; f(] log( (12//0 max
4V2 max log 1og((=(3/8)
(16)
n
3n
For the second term, note that for any given ith sample (#,, 34, ajs 8/ x.) all the variables {X,,) are i..dd Then following similar argument, then for all f, f' E F, we have with probability at least
1 8/n,
(fx TG, as] M X (F:d")
8V2 E ((x,s; (s),, 0s) log (((1))
4V2 max log(n) F(2/5)
max
m
3m
Then it is easy to obtain that have for all f, EF,
f((
18

---

8V2
max
E (f:
lo
4V2 (ng(n/FF1//0))
max
+
8v2
max
(17)
71
m
3m
Combining (q.(1 and (q.(1 can obtain with probability least 1 o, for all S/ EF,
X((fif") - Σ
- Ex,0
+ x E (Y(:
Y((+f)++++
E
==
8V2 (f (4] [
4V mix
log (FF1/0)
+ x
Min
3n
nm
max 8V (f FV((
4V2
max
(ng(n) F12/))
5V
max
nm +
3m
n
(18)
Let fk-1, = fk. then according to Algorithm we know that f++1= = FR. = arg min L. (fisk)
JEF
According to Eq. (14), we have
L. fz: Lp (Tfa:
n
1
Σ
Ye(f:fa) +
THTRL
X
1 ΣΣ (i,(A
+
1=1
t=1
Then it is easy to observe that FR. for minimizes
Lb, (ti it also minimizes
+
nm
ΣΣ
(zf(1
TITR
THT
(=1j=1
because the two objectives only differ by a constant L fTi fi). Therefore under assumption A.2 know that Tík E F, are able obtain that
Y, (Tx.FFfaiie) +
X1, ffr.
7m +
mm + n
j=1
1
(Thxi =,
(19)
1=1
where the last equality holds due to the definitions of Y, and Xz,. Therefore plugging the result from (19) into (q.() we can obtain
E[ (fx+1i fell + +
m
E (fx,(fett:fe))
8V2 max
Eyy fa)) log (11188
4V2 max
(gg((F11//))
1
n
3n
1
8V2 E Xi(Jk+is fk)) log(n) 4V2 Log(n(FF1
5V2 max
max
+
nmm +
3m
By solving the quadratic formula, we get
Cpa, (fet1i fie) Cas, (Tfai 11++1 Tf
m + (fe nm + n j=l
(fx,(fa+2; fall
<5
-- m
2 log (n (8)
35V2
max
Finally, apply a union bound over all /. = 0 K, we conclude the proof.
19

---

A.1 Proof of Theorem 1
Now we are ready to show the main theorem. Given a dataset D. After generating virtual samples D₂ we get new combined dataset Di DU Dg at each iteration & with virtual sample generating distribution B& (2)) We first have
1
f
E
2/ (x.0) (T) (s,
1
(s
<
E
[Q" 28, [x(8, (T) + -
- (8,,
1
(a.r))-iio (a,0)
1
(1Q0* fuell fiell,dd1 the (8,0)x=Tx (s,,))
+ 11Q - -fxll 1,76 (o.x))x))
1 7
1 1-7 (IQº - - (s.0)x=%
+ II - fx 5 2Ie (s.m)xer+:
(20)
where the first equality follows from the performance difference lemma (Kakade and Langford, 2002). the first inequality holds because fT. E arg min, f& (s. x. a) and the last inequality true by using the fact that for any vector a = (aI,..., an) and valid distribution d = (di.... dn). Σ, d: =
||a|1add Σ lai ld; = Σ laslvdivll
Geuchy-Schwardd inequality)
Σ di x
Σ las1 Elasped = -11all2d.
v x
According to Assumption A.4, know that EvEs ExEX, dhi'' (s, S 02, which implies that xxxx,)) (d%
Tt (s x (8, a) S 02 and Eisma)e8xxxxxa5p (d (s,x) x
Tf. 28, 2,, a) S 02 Define { = (di (8, T (8, x, a). then we have
1/2
11/x Q7112s =
Σ
-((,,,,a)-(((60,000000 - ())
(SxN)))SxxxA
1/2
Σ
1/2
Σ
11(8,z,a) (x,,)))))(((,,1)
xxx.
f4
Q"112.5 + Vmax
(m 1)C ma1
l1 - T Ilg. + -1110--1-0"1ll.pxxxxxx + Vo2V max:
(21)
where the first incquality holds because va+b5 Va vb (a 0,6 0) and the last inequality comes from Lemma A.4.
By using LemmaA.7 we have with at least probability 1 - 8
1lfe T Ts--1
1.p.a I llfk - TT fk 12, S 61
(22)
where E 5 + 1) V2 yaxx log(n
+ 38v mus Therefore, we obtain
(m + 1)CEL
[I - Q'112 S I Q" 11'1111 P1)1
VozVmax.
(23)
mai
Now define g' P(E) x Tf -- Then based on Assumption A.4, it is easy to obtain EisxaeSSxxxAAES (',, a) S 0p Note that the distribution g" = PP(g') x fu-2QQ" still sat-
isfies Σ xxxx))))xxxx S" (s, a) < coo2 according to Assumption A.4, because Th 1 ft. and Q' are all assumed to be reasonable policies.
20

---

Repeating the expansion above for & times, we have
1-2 1 Q2 = 1-7 (Vmmttrr C)
114 -
+ vego2Vmax
+ 7 Vmaa.
moi
All the above analyses are still applied to the case when & = =(d70 (s,x) x J"}. Therefore, it is straightforward obtain
1mCa 1)Ca
c 1x
+ Vcoo2 Vmax 7 V
(24)
(1 7)2
mo1
Substituting 61 completes the proof.
A.2 Extension of Theorem
We repeat the assumption for extending our main results to the case when | can be infinite such that [s, x. a) may not be bounded by Vmax-
Repeat of Assumption 1: For the typical set B of pscudo-stochastic states (defined in Assump-
tion A.4), for any a E S A, f F, if æ E B, then f(s, x, a) < Vmmax otherwise if B, have [f(s x, a) (8, [Vy ..)' (s", x²") S maxx where x' g(s, I, fy, s'), g 8, x, f s").
S Vmax- Furthermore, for any given f E F, (s, x), E B, we have
There are two places we need to pay attention to: (1) bound on I Q" 2l, p())x= (2) :a bound on the variance of as shown in Eq. (15). For the first case, it automatically holds due to assumption 1. For the second term, first have
Var(Y((;:'')) < (J;
- -1 =E - 1116:000177(3330/111000777771000
2'1)
([(((s,,, - (8, ||
2r -2-22V766,
We also know that
- (1
<2( - TJ'
-
a) + (" 2/ (VV (s((s=z'')
<16V/ max
Therefore, we have Var(Y(ff:f))) - 16V2 max (f)
Then we can obtain similar result of the same order, which only differs for some constant č such that
(mm+ +1Cei + Vegor + x - 7)
2č
4¹ 5h
>)Vnxx Vmax
(25)
(1 7)²
mai
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
21

---

Algorithm 3: SAC-ASG
Input: Critic Networks: Qa Qo, Target Critie Networks: Qø; Qø;; Actor-Network: Ag, Empty sample reply buffer: D, Learning rate: A;
2
for each iteration do for each environment interaction do
Take action a, ~ Ag(a, (t, 1z), observe next state ($141, and reward ry Store the transition into replay buffer: D t D U {St, 2t, a,, Te, 84t1, #4+1); for each training step do of transitions from replay buffer D;
Sample mini-bated
for Each virtual training loop do
10
Obtain virtual dataset d' := ASG(d,m); Combine training dataset U d' == {8, x, a 8 x'};
11
12
à As(s', x²); y
(min Oo! x a) - a log( (o((As(()s(xx)))
13
(0,) (nm + n) -1 Σ (y Qa x, a) for E [1,2]
0; AV,, Jo,0; (0,) for E [1,2] // Update Critic networks
(+) (nm n) -n-1- log (a(a))) min Qn, (8)
16
é - Ø - (V)
// Update Actor network
17
0g 70, + (1-7)0)
// Update target network weights
Output: Actor Network As;
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. days).
7.0
SAC-A5G
Queue Length
6.5
Optimal
5.5
4.5
4.0
10
14
Training Fpisodes
Figure 4: Performance on the Two-phase Criss-Cross Network
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.
Setting
Arrival Rates
Service Rates
Job Size Range
(a]
[0.6, 0.6
(2, 1.5, 1.5)
2
(0.6, 0.6)
{7, 7}
(0.6,066)
[2.5, 4.5, 2.5)
Table 3: Detailed Environment Parameters
B.3 Two-phase Wireless Network
The two-phase wireless network a modified version of the downlink network in Section 2.1, which cannot be modeled as MDP with exogenous inputs. In this example, each packet has two phases. When packet from mobile is scheduled, it in phase 1, it leaves with probability 0.2 or moves from phase to phase 2; and if it is in phase 2, it leaves the system. We considered Poisson arrivals
22

---

with rates A (2,4.3), and the channel rates O = {12, 12, 12). As the result shown in Figure. ASG outperforms Max-weight in this environment well.
34
DQN-ASG
Max-Weight
Tota Oveus Length
32
30
28
220
60
100
140
180
220
260
300
Training Episode
Figure 5: Performance on the Two-phase Downlink Wireless Network
C Simulation Details
In our simulations, the reward function is the total queue length (or scaled version: divided by 20 the wireless networks). We used Gaussian distributions fitted with real data samples as BR, for each real sample, we generated augmented samples and repeated the process 10 times. We use 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 single NVIDIA GeForce RTX 2080 Super with AMD Ryzen 7 3700 8-Core Processor.
23


Reviews:
[{"review_id": "OBaMCBjAqW", "reviewer": "Reviewer_d68Z", "summary": "This 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. ", "questions": "1. 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\n", "limitations": "The 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. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- 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. ", "weaknesses": "- 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your time and thoughtful evaluation of our paper. We recap your comments and present our detailed response as follows.\n> (1) 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\nA number of important real-world applications can be modeled as mixed systems. In many of these systems, the separation of stochastic and pseudo-stochastic states is obvious. The paper discussed queueing systems extensively. Here are a few more examples: (i) Water resource management with dynamic water levels where weather and water conditions are stochastic, and the reservoir water level is pseudo-stochastic; and (ii) autonomous driving, where the speed, acceleration, and direction of a vehicle are pseudo-stochastic, and the surrounding environment such as traffic condition and weather of is stochastic. \n\nOf course, there are mixed systems, possibly many, for which it is nontrivial to identify stochastic and pseudo-stochastic states. The cost of separating the states varies from application to application. For these systems, we may take inspiration from the methods for factoring MDPs as suggested by the reviewer, despite factored MDPs and mixed systems have some fundamental differences (see our response to the next question). \n\nThis paper focuses on the design of sample-efficient RL algorithms assuming the separation is presented to us, which again is obvious in many applications, including queueing systems.\n\n> (2) 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\nWhile factored MDPs and mixed systems share some similarities, they are actually quite different. A factored MDP consists of subsets of states, which are weakly correlated. Let  $X_i$ be a state variable and vector $(X_1,\\dots,X_n)$ be the full state representation of an MDP. According to the definition in [R1], for a factored MDP, there is a set $Z_i$ associated with each state $i.$ Given $\\{Z_i\\}\\in\\{1,\\dots,n\\},$ the cost (similarly, reward) can be factored as \n\n$\n\\begin{equation}\ng(x)=\\sum_j g_j(X_{Z_j}),\n\\end{equation}\n$\n\nwhere $X_{Z_j}$ is a subset of the state variables defined by $Z_j$,  and the transition probabilities can be factorized (or localized) as \n\n$P_a(X_i(t+1)\\vert X(t) )=P_a(X_i(t+1)\\vert X_{Z_i}(t) ),\\forall i.$\n\nIn a mixed system, the cost/reward function is **not** assumed to be factored. Furthermore, although for 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. \n```\nFor example, in the example of a wireless downlink network in Figure 1, the queue lengths at the next time slot depend on the scheduling decision at the current time slot, which depends on the current values of all the queue lengths, the arrivals, and the channel states.\n```\n\nHowever, 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. Also, considering the batch off-line RL setting without enough coverage used in our paper, there are no existing approaches, including factored MDP, that guarantee convergence performance in this scenario.\n\nWe greatly appreciate the reviewer for bringing up factored MDPs. We will add the comparison to the related work in our revision to clarify the difference.\n> (3) 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\nThanks for the question! The insights could be useful for POMDPs. For example, the belief state in POMDPs is typically updated as follows: \n\n$b'(s') = \\frac{P(o'\\vert s',a)\\sum_s P(s'\\vert a,s)b(s) }{ \\sum_{s'}P(o'\\vert s',a)\\sum_s P(s'\\vert a,s)b(s) }, $ \n\nwhere $s$ denotes the state, $a$ is the action, $o$ stands for the observation state, $P(o|s,a)$ represents the observation probability, and $P(s'|s,a)$ is the transition probability. If the POMDP is also from a mixed system, we may generate augmented samples similarly by sampling according to the brief of stochastic states to improve sample efficiency.\n\n---\nWe hope that our response addresses your concerns about the applicability and connection to factored MDPs and are happy to answer any additional questions or concerns. We would also be grateful if the reviewer could consider reevaluating the review and rating based on our response.\n\n---------------------------references --------------------------------\n\n[R1] MIT 2.997 Decision-Making in Large-Scale Systems Lecture Note 19, https://ocw.mit.edu/courses/2-997-decision-making-in-large-scale-systems-spring-2004/c578943cb9b5ba63e3976e01de52c61e_lec_19_v1.pdf"}, {"review_id": "hBISkTipxO", "reviewer": "Reviewer_oHp5", "summary": "The 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.  ", "questions": "1. 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?  ", "limitations": "The paper does not extensively discuss limitations or potential negative societal impacts of the work. No further limitations need to be addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "1. 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", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We first would like to sincerely thank the reviewer for the willingness to reevaluate the rating. We recap your comments and present our detailed response as follows.\n\n> (1) Compared with MDPs with exogenous inputs.\n\n- Our discussion from lines 132-137 explains the fundamental difference. For the environments used in the paper, we would like to point out that the phases 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). Note that 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. \n\n- The reviewer is correct that the current wireless network environment can be modeled as an MDP with exogenous inputs. However, if we further consider phase-type service times, then it is not an MDP with exogenous inputs, for a reason similar to the phases Criss-Cross Network. In fact, scheduling in queue systems with nonexponential service times can often be modeled as mixed systems but are not MDPs with exogenous inputs.\n\nTo further address the reviewer's question, we considered a two-phase wireless downlink network, where 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.  \n\nThe simulation results are summarized in the following table:\n\n| Episodes| 20 | 60 |100|140|180|220|260|300|\n| ----------- | ----------- | ----------- |  ----------- | ----------- | ----------- | ----------- | ----------- | ----------- |\n| DQN-ASG | 1343.94 | 23.545 | 23.212 | 23.214 |  23.066 |  23.0125  | 23.104 | 22.940\n\nMax-Weight in this setting results in an average total queue length of **23.677**. Our approach still outperforms Max-weight.\n\nBesides queuing systems, we now further include a classical control problem **the Cart Pole problem** as an application of the mixed systems. Details can be found in the one-page pdf file we submitted. Note that the system is not an MDP with exogenous inputs because the stochastic states depend on the pseudo-stochastic states through actions (details in the uploaded PDF). \n\n> (2) 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\n- Thanks for the question! They are indeed CTMC. We use the standard uniformization method [R1,R2] to simulate them using DTMCs. We will clarify it in the revision.  \n\n> (3)  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\n- Thank you for pointing this out. In our simulations, augmented samples should be **valid** samples for which the action from the real sample remains to be feasible. For the example you gave, the augmented sample will be $\\{x, y, z\\}$ for $z\\geq 1,$ e.g. $(10, 0, 1),$ which ensures the action taken for the real sample is a valid action for the augmented sample as well. We will clarify this in the revision. \n\n---\nWe hope that our response addresses your concerns about the applicability and connection to MDPs with exogeous inputs. We are happy to answer any additional questions or concerns. We would also be grateful if the reviewer could consider reevaluating the review and rating based on our response. \n\n---------------------------references --------------------------------\n\n[R1] Serfozo, Richard F. \"An equivalence between continuous and discrete time Markov decision processes.\" Operations Research 27.3 (1979): 616-620.\n\n[R2] Puterman, Martin L. \"Markov decision processes: discrete stochastic dynamic programming.\" John Wiley \\& Sons, 2014."}, {"review_id": "mWk9Wog0LZ", "reviewer": "Reviewer_ciJb", "summary": "This 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.", "questions": "In 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.\"", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "Using 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.", "weaknesses": "In 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to sincerely thank the reviewer for the constructive comments. We will incorporate these great suggestions in our revision and we'd like to address the major concern in the following.\n\n> (1) Building 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\n- In a typical queueing system, 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 both a policy in F and the optimal policy can reduce the queue lengths exponentially fast (in expectation) and the difference in Q-functions therefore is bounded. The example in Section 2.1 therefore satisfies Assumption 1.  We will add further discussion in the revision. \n\n> (2) In line 93, the authors claim \"far less memory space.\" Could you provide a quantification, such as the complexity mentioned in line 91?\n\nSince our approach is model-free, the complexity of remembering the Q values for the tabular case is $SXA.$\n\n> (3) In 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\n- In our simulations, the reward function is the total queue length (or a scaled version).  We used Gaussian distributions fitted with real data samples as $\\beta_k.$  For each real sample, we generated $50$ augmented samples and repeated the process $10$ times. We will include more details about the algorithm in our revised manuscript.\n\n> (4) Regarding Figures 2 and 3, how is an episode defined in the context of the simulations?\n\nEach episode includes 5,000 steps.\n\n> (5)  In 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\nWe apologize for the confusion. Episode 0 in the label should be episode 20, which is the first testing episode. In our simulations, we test the performance of the learned policy every 20 episodes. DQN-ASG does not improve is because after 20 episodes, DQN-ASG already outperforms the baselines significantly and the performance only improves slowly after that. We include the new figures in the uploaded one-page pdf file.\n\n---\nWe hope the response above addresses the reviewer's concerns. We would also be happy to provide further clarifications if needed. We would also be grateful if the reviewer could consider reevaluating the review and rating based on our response."}]

Author rebuttal:
We appreciate the reviewers' insightful comments and suggestions. Your feedback has significantly contributed to enhancing the quality of our work. We would now like to address each reviewer's questions individually to ensure we provide comprehensive responses to your specific concerns. We include two additional simulations in the uploaded pdf. The first addresses comments by Reviewer *oHp5*, particularly concerning distinguishing mixed systems from MDPs with exogenous inputs. The second simulation explores the potential of our approach in other applications. 

Meta-review: The paper studies sample efficiency of RL in "mixed systems". The theoretical results and simulations are much appreciated and presented in a clean and well-organized manner. The proposed algorithm can make a big impact to queueing systems including wireless communication, resource management, autonomous driving, etc.

There are some initial concerns on the applicability of the mixed systems in practice. Through the author-reviewer discussions, these concerns are largely lifted. In addition, authors provide extended experimental results and explanations to core concept in the paper, such as the difference between the mixed systems and existing MDPs with exogenous inputs. I would encourage authors to carefully incorporate these discussions to the revision of the paper.
Decision: Accept (spotlight)