SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Honghao Wei
honghao.wei@wsu.edu
Weina Wang
weinaw@cs.cmu.edu
liuxin7@shanghaitech.edu.cn
Lei Ying
University of Michigan, Ann Arbor
leiying@umich.edu
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  $\tilde{\mathcal{O}} (\sqrt{1 / n} +\sqrt{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  $\tilde{\mathcal{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.
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  $\tilde{\mathcal{O}}(\sqrt{1/n} + \sqrt{1/m})$ , where  $n$  is the number of real samples. Note that without augmented samples, the error is  $\tilde{\mathcal{O}}(1)$  due to the lack of data coverage of the pseudo-stochastic states, which reduces to  $\tilde{\mathcal{O}}(\sqrt{1/n})$  when we generate at least  $n$  augmented samples for each real sample. We also would like to emphasize that  $\omega(\sqrt{1/n})$  is also the fundamental lower bound due to the coverage of the stochastic states.
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 |\mathcal{X}|^2 |\mathcal{A}|$ , where  $S$  is the stochastic state space,  $\mathcal{X}$  is the pseudo-stochastic state space and  $\mathcal{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-  $\mathrm{Q}^{+}$  (Sutton, 1988), Linear Dyna (Sutton et al., 2012), and Dyna-  $\mathrm{Q}(\lambda)$  (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.
Table 1: Dyna-type v.s. ASG
Algorithm	estimate model?	update unseen states?	computational efficient?	store all transitions?	convergence analysis?
Dyna-type	✓	X	X	X	X
ASG	X	✓	✓	✓	✓
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 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 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: 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-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 = (\mathcal{S} \times \mathcal{X}, \mathcal{A}, P, R, \gamma, \eta_0)$ , whose state space is  $\mathcal{S} \times \mathcal{X}$ , where  $\mathcal{S}$  is the set of stochastic states and  $\mathcal{X}$  is the set of pseudo-stochastic states, and both  $\mathcal{S}$  and  $\mathcal{X}$  are assumed to be finite. Further,  $\mathcal{A}$  is a finite action space,  $R: \mathcal{S} \times \mathcal{X} \times \mathcal{A} \to [0, r_{\max}]$  is a deterministic and known cost function, and  $P: \mathcal{S} \times \mathcal{X} \times \mathcal{A} \to \Delta(\mathcal{S} \times \mathcal{X})$  is the transition kernel (where  $\Delta(\cdot)$  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
$$
p _ {i j} (a) = P \left(S _ {t + 1} = j \mid S _ {t} = i, a _ {t} = a\right). \tag {1}
$$
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:
$$
X _ {t + 1} = g \left(S _ {t}, X _ {t}, a _ {t}, S _ {t + 1}\right). \tag {2}
$$
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  $\pi$  is defined as
$$
V ^ {\pi} (s, x) := \mathbb {E} \left[ \sum_ {h = 0} ^ {\infty} \gamma^ {h} R \left(s _ {h}, x _ {h}, a _ {h}\right) \mid \left(s _ {0}, x _ {0}\right) = (s, x), a _ {h} = \pi \left(s _ {h}, x _ {h}\right) \right], \tag {3}
$$
where  $\gamma$  is the discount factor, and the expectation is taken over the transitions of stochastic states and randomness in the policy. Let  $v^{\pi} = \mathbb{E}_{(s_0,x_0)\sim \eta_0}[V^{\pi}(s_0,x_0)]$ , i.e., the expected cost when the initial distribution of the state is  $\eta_0$ . Let  $v^{*} = \min_{\pi} v^{\pi}$ .
The  $Q$ -value function of a policy  $\pi$  is defined as
$$
Q ^ {\pi} (s, x, a) := \mathbb {E} \left[ \sum_ {h = 0} ^ {\infty} \gamma^ {h} R \left(s _ {h}, x _ {h}, a _ {h}\right) \mid \left(s _ {0}, x _ {0}\right) = (s, x), a _ {0} = a, a _ {h} = \pi \left(s _ {h}, x _ {h}\right) \right]. \tag {4}
$$
Let  $Q$  be the optimal  $Q$ -value function. Then the Bellman equation can be written as
$$
Q (s, x, a) = R (s, x, a) + \gamma \mathbb {E} _ {s ^ {\prime} \sim P (\cdot | s, a)} \left[ \min  _ {a ^ {\prime}} Q \left(s ^ {\prime}, g (s, x, a, s ^ {\prime}), a ^ {\prime}\right) \right]. \tag {5}
$$
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  $\Lambda_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} \in \{1,2,3\}$  denote the mobile scheduled at time slot  $t$ . The length of queue  $i$  evolves as
$$
q _ {t + 1} (i) = \left(q _ {t} (i) + \Lambda_ {t} (i) - O _ {t} (i) \mathbb {I} \left(A _ {t} = i\right)\right) ^ {+}, \tag {6}
$$
where  $\mathbb{I}(\cdot)$  is the indicator function and  $(\cdot)^{+} = \max \{\cdot ,0\}$ .
To minimize the total queue lengths, the problem can be formulated as an RL problem such that the state is  $(\mathbf{\Lambda}_t,\mathbf{O}_t,\mathbf{q}_t)$ , the action is  $A_{t}$ , and the cost is  $\sum_{i}q_{t}(i)$ .
We can see that in this problem,  $S_{t} \coloneqq (\mathbf{\Lambda}_{t}, \mathbf{O}_{t})$  is the stochastic state and  $\mathbf{q}_t$  is the pseudo-stochastic state. In general, the state space of the stochastic states is bounded, e.g., both  $\mathbf{\Lambda}_t$  and  $\mathbf{O}_t$  may be Bernoulli random variables, but the pseudo-stochastic state such as  $\mathbf{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  $\Lambda_t$ ,  $\mathbf{O}_t$  and  $A_t$ , we can generate one-slot queue evolution starting from any queue length.
Figure 1: A Downlink Wireless Network
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  $\beta(\cdot)$  over the pseudo-stochastic states in  $\mathcal{X}$ . For each sample  $(s, x, a, r, s', x') \in D$ , we first sample a pseudo-stochastic state  $\hat{x}$  from  $\beta(\cdot)$ , and then construct a virtual sample  $(s, \hat{x}, a, \hat{r}, s', \hat{x}')$ , where  $\hat{r} = R(s, \hat{x}, a)$  and  $\hat{x}' = g(s, \hat{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} := (\Lambda_{t}, \mathbf{O}_{t}), X_{t} := \mathbf{q}_{t}$ , assume that at some timeslot  $t$ , one of the real samples is
$$
(s, x, a, r, s ^ {\prime}, x ^ {\prime}) := \left(\left(\left\{\left(3, 4, 5 \right\rbrace , \{1, 2, 0 \}\right), \left\{4, 6, 6 \right\}, 1, 1 6, \left(\left\{\left(2, 3, 4 \right\}, \left\{2, 2, 0 \right\}\right)\right), \left\{6, 1 0, 1 1 \right\}\right), \right.
$$
where the queue length  $\mathbf{q}_{\mathbf{t} + 1}$  is calculated according to Eq. (6). Using ASG, we generate two pseudo-stochastic states (queue length)  $\hat{\mathbf{x}}_1 = (1,2,3),\hat{\mathbf{x}}_2 = (0,2,1)$ , then we are able to obtain two virtual samples
$$
(s, \hat {x} _ {1}, a, \hat {r} _ {1}, s ^ {\prime}, \hat {x} _ {1} ^ {\prime}) := ((\{3, 4, 5 \}, \{1, 2, 0 \}), \{1, 2, 3 \}, 1, 6, (\{2, 3, 4 \}, \{2, 2, 0 \}), \{3, 6, 8 \}),
$$
$$
(s, \hat {x} _ {2}, a, \hat {r} _ {2}, s ^ {\prime}, \hat {x} _ {2} ^ {\prime}) := ((\{3, 4, 5 \}, \{1, 2, 0 \}), \{0, 2, 1 \}, 1, 3, (\{2, 3, 4 \}, \{2, 2, 0 \}), \{2, 6, 6 \}).
$$
Note that all the virtual samples represent true transitions of the mixed system if the queue lengths were  $\hat{\mathbf{x}}_1$  and  $\hat{\mathbf{x}}_2$ .
Algorithm 1: Augmented Sample Generator (ASG)
1 Input: A dataset  $D = \{(s,x,a,r,s',x')\}$ , a positive integer  $m$ , a distribution  $\beta (\cdot)$  on  $\mathcal{X}$ ;
2 Initialize virtual sample dataset:  $D^{\prime} = \emptyset$
3 for each sample  $(s,x,a,r,s',x')\in D$  do
4 Sample  $m$  virtual pseudo-stochastic states  $\hat{x}_1, \hat{x}_2, \dots, \hat{x}_m$  from  $\beta(\cdot)$  independently;
5 Compute  $\hat{r}_i = R(s,\hat{x}_i,a)$  and  $\hat{x}_i' = g(s,\hat{x}_i,a,s')$  for  $i = 1,2,\ldots,m$ ;
6 Add  $m$  virtual samples to virtual dataset:  $\{(s,\hat{x}_i,a,\hat{r}_i,s',\hat{x}_i'),i = 1,2,\ldots ,m\} \to D'$
7 Output:  $D^{\prime}\cup D$
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  $\mu \in \Delta(S \times \mathcal{X} \times \mathcal{A})$  for  $(s,x,a)$ ,  $r = R(s,x,a)$ ,  $s' \sim P(\cdot | s,a)$ , and  $x' = g(s,x,a,s')$ . The distribution  $\mu$  is unknown to the agent.
Our algorithm follows the framework of a batch FQI algorithm. We consider a class of candidate value-functions  $\mathcal{F} \in (\mathcal{S} \times \mathcal{X} \times \mathcal{A} \to [0, V_{\max}])$ . We focus on the scenario where the number of pseudo-stochastic states,  $|\mathcal{X}|$ , is far more than that of stochastic states,  $|\mathcal{S}|$ , and the dataset  $D$  does not provide sufficient coverage of the pseudo-stochastic states ( $n < |\mathcal{X}|$ ).