SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {L} _ {\hat {D} _ {k}} (f; f _ {k}) - \mathcal {L} _ {\hat {D} _ {k}} (\mathcal {T} f _ {k}; f _ {k}) = \frac {n}{n m + n} \times \frac {1}{n} \sum_ {i = 1} ^ {n} Y _ {i} (f; f _ {k}) + \frac {n m}{n m + n} \times \frac {1}{n m} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {m} X _ {i j} (f; f _ {k}).
$$
Then it is easy to observe that  $\widetilde{\mathcal{T}}_{k,\mathcal{F}}f_k$  minimizes  $\mathcal{L}_{\hat{D}_k}(\cdot ;f_k)$ , it also minimizes
$$
\frac {n}{n m + n} \times \frac {1}{n} \sum_ {i = 1} ^ {n} Y _ {i} (\cdot ; f _ {k}) + \frac {n m}{n m + n} \times \frac {1}{n m} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {m} X _ {i _ {j}} (\cdot ; f _ {k})
$$
because the two objectives only differ by a constant  $\mathcal{L}_{\hat{D}_k}(\mathcal{T}f_k;f_k)$ . Therefore under assumption A.2 we know that  $\mathcal{T}f_{k}\in \mathcal{F}$ , we are able to obtain that
$$
\begin{array}{l} \frac {1}{n m + n} \sum_ {i = 1} ^ {n} Y _ {i} (\hat {\mathcal {T}} _ {k, \mathcal {F}} f _ {k}; f _ {k}) + \frac {1}{m n + n} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {m} X _ {i j} (\hat {\mathcal {T}} _ {k, \mathcal {F}} f _ {k}; f _ {k}) \\ \leq \frac {1}{n m + n} \sum_ {i = 1} ^ {n} Y _ {i} \left(\mathcal {T} f _ {k}; f _ {k}\right) + \frac {1}{n m + n} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {m} X _ {i j} \left(\mathcal {T} f _ {k}; f _ {k}\right) = 0, \tag {19} \\ \end{array}
$$
where the last equality holds due to the definitions of  $Y_{i}$  and  $X_{i_j}$ . Therefore plugging the result from Eq.(19) into Eq.(18), we can obtain
$$
\begin{array}{l} \frac {1}{m + 1} \mathbb {E} \left[ Y \left(f _ {k + 1}; f _ {k}\right) \right] + \frac {m}{m n + n} \sum_ {i = 1} ^ {n} \mathbb {E} \left[ X _ {i} \left(f _ {k + 1}; f _ {k}\right) \right] \\ \leq \frac {1}{1 + m} \left(\sqrt {\frac {8 V _ {\operatorname* {m a x}} ^ {2} \mathbb {E} \left[ Y \left(f _ {k + 1} ; f _ {k}\right) \right] \log \left(| \mathcal {F} | ^ {2} / \delta\right)}{n}} + \frac {4 V _ {\operatorname* {m a x}} ^ {2} \log \left(| \mathcal {F} | ^ {2} / \delta\right)}{3 n}\right) \\ + \frac {n m}{n m + n} \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \left(\sqrt {\frac {8 V _ {\max } ^ {2} \mathbb {E} [ X _ {i} (f _ {k + 1} ; f _ {k}) ] \log (n | \mathcal {F} | ^ {2} / \delta)}{m}} + \frac {4 V _ {\max } ^ {2} \log (n | \mathcal {F} | ^ {2} / \delta)}{3 m}\right) + \frac {\delta V _ {\max } ^ {2}}{n}\right) \\ \end{array}
$$
By solving the quadratic formula, we get
$$
\mathcal {L} _ {\hat {\mu} _ {\beta_ {k}}} \left(f _ {k + 1}; f _ {k}\right) - \mathcal {L} _ {\hat {\mu} _ {\beta_ {k}}} \left(\mathcal {T} f _ {k}; f _ {k}\right) = \left\| f _ {k + 1} - \mathcal {T} f _ {k} \right\| _ {2, \bar {\mu} _ {\beta_ {k}}} ^ {2}
$$
$$
= \frac {1}{m + 1} \mathbb {E} [ Y (f _ {k + 1}; f _ {k}) ] + \frac {m}{n m + n} \sum_ {j = 1} ^ {n} \mathbb {E} \left[ X _ {i} \left(f _ {k + 1}; f _ {k}\right) \right]
$$
$$
\leq 5 \left(\frac {1}{n} + \frac {1}{m}\right) V _ {\max } ^ {2} \log (n | \mathcal {F} | ^ {2} / \delta) + \frac {3 \delta V _ {\max } ^ {2}}{n}
$$
Finally, apply a union bound over all  $t = 0 \dots K$ , we conclude the proof.
Now we are ready to show the main theorem. Given a dataset  $D$ . After generating virtual samples  $D_{k}$  we get a new combined dataset  $\hat{D}_k = D\cup D_k$  at each iteration  $k$  with virtual sample generating distribution  $\beta_{k}(x)$ . We first have
$$
\begin{array}{l} v ^ {\pi_ {f _ {k}}} - v ^ {*} = \frac {1}{1 - \gamma} \mathbb {E} _ {(s, x) \sim d _ {\eta_ {0}} ^ {\pi_ {f _ {k}}} (s, x)} \left[ Q ^ {*} (s, x, \pi_ {f _ {k}}) - V ^ {*} (s, x) \right] \\ \leq \frac {1}{1 - \gamma} \mathbb {E} _ {(s, x) \sim d _ {\eta_ {0}} ^ {\pi_ {f _ {k}}} (s, x)} \left[ Q ^ {*} (s, x, \pi_ {f _ {k}}) - f _ {k} (s, x, \pi_ {f _ {k}}) + f _ {k} (s, x, \pi^ {*}) - V ^ {*} (s, x) \right] \\ \leq \frac {1}{1 - \gamma} \left(\| Q ^ {*} - f _ {k} \| _ {1, d _ {\eta_ {0}} ^ {\pi f _ {k}} (s, x) \times \pi_ {f _ {k}}} + \| Q ^ {*} - f _ {k} \| _ {1, d _ {\eta_ {0}} ^ {\pi f _ {k}} (s, x) \times \pi^ {*}}\right) \\ \leq \frac {1}{1 - \gamma} \left(\| Q ^ {*} - f _ {k} \| _ {2, d _ {\eta_ {0}} ^ {\pi f _ {k}} (s, x) \times \pi_ {f _ {k}}} + \| Q ^ {*} - f _ {k} \| _ {2, d _ {\eta_ {0}} ^ {\pi f _ {k}} (s, x) \times \pi^ {*}}\right), \tag {20} \\ \end{array}
$$
where the first equality follows from the performance difference lemma (Kakade and Langford, 2002), the first inequality holds because  $\pi_{f_k} \in \arg \min_a f_k(s, x, a)$  and the last inequality is true by using the fact that for any vector  $a = (a_1, \ldots, a_n)$  and a valid distribution  $d = (d_1, \ldots, d_n)$ ,  $\sum_i d_i = 1$
$$
\begin{array}{l} \left\| a \right\| _ {1, d} = \sum_ {i} \left| a _ {i} \right| d _ {i} = \sum_ {i} \left| a _ {i} \right| \sqrt {d _ {i}} \sqrt {d _ {i}} \quad (\text {C a u c h y - S c h w a r z i n e q u a l i t y}) \\ \leq \sqrt {\sum_ {i} d _ {i}} \times \sqrt {\sum_ {i} | a _ {i} | ^ {2} d _ {i}} = \| a \| _ {2, d}. \\ \end{array}
$$
According to Assumption A.4, we know that  $\sum_{s\in \mathcal{S}}\sum_{x\in \mathcal{X},x\notin \mathcal{B}}d_{\eta_0}^{\pi_{f_k}}(s,x)\leq \sigma_2$  , which implies that  $\begin{array}{r}\sum_{(s,x,a)\in \mathcal{S}\times \mathcal{X}\times \mathcal{A},x\notin \mathcal{B}}\{d_{\eta_0}^{\pi_{f_k}}(s,x)\times \pi^*\} (s,x,a)\leq \sigma_2 \end{array}$  and  $\begin{array}{r}\sum_{(s,x,a)\in \mathcal{S}\times \mathcal{X}\times \mathcal{A},x\notin \mathcal{B}}\{d_{\eta_0}^{\pi_{f_k}}(s,x)\times \end{array}$ $\pi_{f_k}\} (s,x,a)\leq \sigma_2$  . Define  $\xi = \{d_{\eta_0}^{\pi_{f_k}}(s,x)\times \pi_{f_k}\} (s,x,a)$  , then we have
$$
\begin{array}{l} \| f _ {k} - Q ^ {*} \| _ {2, \xi} = \left(\sum_ {(s, x, a) \in \mathcal {S} \times \mathcal {X} \times \mathcal {A}} | f _ {k} (s, x, a) - Q ^ {*} (s, x, a) | ^ {2} \xi (s, x, a)\right) ^ {1 / 2} \\ \leq \left(\sum_ {(s, x, a) \in \mathcal {S} \times \mathcal {X} \times \mathcal {A}, x \in \mathcal {B}} | f _ {k} (s, x, a) - Q ^ {*} (s, x, a) | ^ {2} \xi (s, x, a)\right) ^ {1 / 2} \\ + \left(\sum_ {(s, x, a) \in \mathcal {S} \times \mathcal {X} \times \mathcal {A}, x \notin \mathcal {B}} | f _ {k} (s, x, a) - Q ^ {*} (s, x, a) | ^ {2} \xi (s, x, a)\right) ^ {1 / 2} \\ \leq \left\| f _ {k} - Q ^ {*} \right\| _ {2, \xi} + \sqrt {\sigma_ {2}} V _ {\max } \\ \leq \sqrt {\frac {(m + 1) C}{m \sigma_ {1}}} \| f _ {k} - \mathcal {T} f _ {k - 1} \| _ {2, \bar {\mu} _ {\beta_ {k}}} + \gamma \| f _ {k - 1} - Q ^ {*} \| _ {2, P (\xi) \times \pi_ {f _ {k - 1}, Q ^ {*}}} + \sqrt {\sigma_ {2}} V _ {\max }, \tag {21} \\ \end{array}
$$
where the first inequality holds because  $\sqrt{a + b} \leq \sqrt{a} + \sqrt{b}$  ( $a \geq 0, b \geq 0$ ) and the last inequality comes from Lemma A.4.
By using LemmaA.7 we have with at least probability  $1 - \delta$
$$
\left\| f _ {k} - \mathcal {T} f _ {k - 1} \right\| _ {\underline {{2}}, \bar {\mu} _ {\beta_ {k}}} ^ {2} \leq \left\| f _ {k} - \mathcal {T} f _ {k - 1} \right\| _ {2, \bar {\mu} _ {\beta_ {k}}} ^ {2} \leq \epsilon_ {1} \tag {22}
$$
where  $\epsilon_{1} = 5\left(\frac{1}{n} +\frac{1}{m}\right)V_{\max}^{2}\log (nK|\mathcal{F}|^{2} / \delta) + \frac{3\delta V_{\max}^{2}}{n}$ . Therefore, we obtain
$$
\left\| f _ {k} - Q ^ {*} \right\| _ {2, \xi} \leq \gamma \left\| f _ {k - 1} - Q ^ {*} \right\| _ {2, P (\xi) \times \pi_ {f _ {k - 1}, Q ^ {*}}} + \sqrt {\frac {(m + 1) C \epsilon_ {1}}{m \sigma_ {1}}} + \sqrt {\sigma_ {2}} V _ {\max } \tag {23}
$$
Now define  $\xi' = P(\xi) \times \pi_{f_{k-1},Q^*}$ . Then based on Assumption A.4, it is easy to obtain  $\sum_{(s,x,a) \in S \times \mathcal{X} \times \mathcal{A}, x \notin \mathcal{B}} \xi'(s,x,a) \leq c_0\sigma_2$ . Note that the distribution  $\xi'' = P(\xi') \times \pi_{f_{k-2},Q^*}$  still satisfies  $\sum_{(s,x,a) \in S \times \mathcal{X} \times \mathcal{A}, x \notin \mathcal{B}} \xi''(s,x,a) \leq c_0\sigma_2$  according to Assumption A.4, because  $\pi_{f_{k-1}}, \pi_{f_{k-2}}$  and  $Q^*$  are all assumed to be reasonable policies.
Repeating the expansion above for  $k$  times, we have
$$
\| f _ {k} - Q ^ {*} \| _ {2, \xi} \leq \frac {1 - \gamma^ {k}}{1 - \gamma} \left(\sqrt {\frac {(m + 1) C \epsilon_ {1}}{m \sigma_ {1}}} + \sqrt {c _ {0} \sigma_ {2}} V _ {\mathrm {m a x}}\right) + \gamma^ {k} V _ {\mathrm {m a x}}.
$$
All the above analyses are still applied to the case when  $\xi = \{d_{\eta_0}^{\pi f_k}(s,x)\times \pi^*\}$ . Therefore, it is straightforward to obtain
$$
v ^ {*} - v ^ {\pi_ {f _ {k}}} \leq \frac {2}{(1 - \gamma) ^ {2}} \left(\sqrt {\frac {(m + 1) C \epsilon_ {1}}{m \sigma_ {1}}} + \sqrt {c _ {0} \sigma_ {2}} V _ {\max } + \gamma^ {k} (1 - \gamma) V _ {\max }\right). \tag {24}
$$
Substituting  $\epsilon_{1}$  completes the proof.
A.2 Extension of Theorem 1
We repeat the assumption for extending our main results to the case when  $|\mathcal{X}|$  can be infinite such that  $f(s, x, a)$  may not be bounded by  $V_{\max}$ .
Repeat of Assumption 1: For the typical set  $\mathcal{B}$  of pseudo-stochastic states (defined in Assumption A.4), for any  $s, a \in S \times \mathcal{A}, f \in \mathcal{F}$ , if  $x \in \mathcal{B}$ , then  $f(s, x, a) \leq V_{\max}$  otherwise if  $x \notin \mathcal{B}$ , we have  $|f(s, x, a) - Q^{*}(s, x, a)| \leq V_{\max}$ . Furthermore, for any given  $f \in \mathcal{F}, (s, x), x \in \mathcal{B}$ , we have  $|V_{f}(s', x') - V_{f}(s'', x'')| \leq V_{\max}$ , where  $x' = g(s, x, \pi_{f}, s')$ ,  $x'' = g(s, x, \pi_{f}, s'')$ .
There are two places we need to pay attention to: (1) : a bound on  $\| f_0 - Q^*\|_{2,P(\xi)\times \pi_{f_0,Q^*}}$ , (2) : a bound on the variance of  $Y$  as shown in Eq. (15). For the first case, it automatically holds due to assumption 1. For the second term, we first have
$$
\begin{array}{l} \operatorname {V a r} \left(Y \left(f; f ^ {\prime}\right)\right) \leq \mathbb {E} \left[ Y \left(f; f ^ {\prime}\right) ^ {2} \right] \\ = \mathbb {E} \left[ \left(\left(f (s, x, a) - r - \gamma V _ {f ^ {\prime}} \left(s ^ {\prime}, x ^ {\prime}\right)\right) ^ {2} - \left(\mathcal {T} f ^ {\prime} (s, x, a) - r - \gamma V _ {f} \left(s ^ {\prime}, x ^ {\prime}\right)\right) ^ {2}\right) ^ {2} \right] \\ = \mathbb {E} \left[ \left(f (s, x, a) - \mathcal {T} f ^ {\prime} (s, x, a)\right) ^ {2} \left(f (s, x, a) + \mathcal {T} f ^ {\prime} (s, x, a) - 2 r - 2 \gamma V _ {f ^ {\prime}} \left(s ^ {\prime}, x ^ {\prime}\right)\right) ^ {2} \right] \\ = \mathbb {E} \left[ \left(f (s, x, a) - \mathcal {T} f ^ {\prime} (s, x, a)\right) ^ {2} \left(f (s, x, a) - \mathcal {T} f ^ {\prime} (s, x, a) + 2 \mathcal {T} f ^ {\prime} (s, x, a) - 2 r - 2 \gamma V _ {f ^ {\prime}} \left(s ^ {\prime}, x ^ {\prime}\right)\right) ^ {2} \right]. \\ \end{array}
$$
We also know that
$$
\begin{array}{l} \left(f (s, x, a) - \mathcal {T} f ^ {\prime} (s, x, a) + 2 \mathcal {T} f ^ {\prime} (s, x, a) - 2 r - 2 \gamma V _ {f ^ {\prime}} \left(s ^ {\prime}, x ^ {\prime}\right)\right) ^ {2} \\ \leq 2 (f (s, x, a) - \mathcal {T} f ^ {\prime} (s, x, a)) ^ {2} + 2 (2 \mathcal {T} f ^ {\prime} (s, x, a) - 2 r - 2 \gamma V _ {f ^ {\prime}} \left(s ^ {\prime}, x ^ {\prime}\right)) ^ {2} \\ = 2 \left(f (s, x, a) + Q ^ {*} (s, x, a) - Q ^ {*} (s, x, a) - \mathcal {T} f ^ {\prime} (s, x, a)\right) ^ {2} + 8 \gamma^ {2} \left(\mathbb {E} \left[ V _ {f ^ {\prime}} ^ {\prime} (\hat {s}, \hat {x}) \mid s, x, a \right] - V _ {f ^ {\prime}} \left(s ^ {\prime}, x ^ {\prime}\right)\right) ^ {2} \\ \leq 1 6 V _ {\max } ^ {2}. \\ \end{array}
$$
Therefore, we have  $\operatorname{Var}(Y(f; f')) \leq 16 V_{\max}^2 \mathbb{E}[Y(f; f')]$ .
Then we can obtain a similar result of the same order, which only differs for some constant  $\tilde{c}$  such that
$$
v ^ {*} - v ^ {\pi_ {f _ {k}}} \leq \frac {2 \tilde {c}}{(1 - \gamma) ^ {2}} \left(\sqrt {\frac {(m + 1) C \epsilon_ {1}}{m \sigma_ {1}}} + \sqrt {c _ {0} \sigma_ {2}} V _ {\max } + \gamma^ {k} (1 - \gamma) V _ {\max }\right). \tag {25}
$$
B Additional Simulations
B.1 Combining PSG with Policy Gradient-type algorithms
In this section, we investigate the possibilities of using ASG in policy gradient-type algorithms. In particular, we use ASG in the phase of policy evaluation. An algorithm (SAC-ASG) that incorporates ASG into SAC is presented in Alg. 3. We also compare our algorithm SAC-ASG with state-of-art Dyna-type model-based approaches, i.e., MBPO (Janner et al., 2019) on the phases criss-cross network environment (Fig. 2b). The simulation results are shown in Fig.4. We can observe that
Algorithm 3: SAC-ASG  
1 Input: Critic Networks:  $Q_{\theta_1}, Q_{\theta_2}$ , Target Critic Networks:  $Q_{\theta_1'}$ ,  $Q_{\theta_2'}$ ;  
2 Actor-Network:  $A_{\phi}$ , Empty sample reply buffer:  $\mathcal{D}$ , Learning rate:  $\lambda$ ;  
3 for each iteration do  
4 for each environment interaction do  
5 Take action  $a_t \sim A_{\phi}(a_t | s_t, x_t)$ , observe next state  $(s_{t+1}, x_{t+1})$ , and reward  $r_t$ ;  
6 Store the transition into replay buffer:  $\mathcal{D} \gets \mathcal{D} \cup \{s_t, x_t, a_t, r_t, s_{t+1}, x_{t+1}\}$ ;  
7 for each training step do  
8 Sample mini-batch  $d$  of  $n$  transitions from replay buffer  $\mathcal{D}$ ;  
9 for Each virtual training loop do  
10 Obtain virtual dataset  $d' := \mathbf{ASG}(\mathbf{d}, \mathbf{m})$ ;  
11 Combine training dataset  $d \cup d' := \{s, x, a, r, s', x'\}$ ;  
12  $\tilde{a} \gets A_{\phi}(s', x')$ ,  $y \gets r + \gamma (\min_{i=1,2} Q_{\theta_i'}(s', x', \tilde{a}) - \alpha \log(A_{\phi}(\tilde{a} | s', x'))$ ;  
13  $J_Q(\theta_i) = (nm + n)^{-1} \sum (y - Q_{\theta_i}(s, x, a))^2$  for  $i \in \{1,2\}$ ;  
14  $\theta_i \gets \theta_i - \lambda \nabla_{\theta_i} J_Q(\theta_i)$  for  $i \in \{1,2\}$  // Update Critic networks  
15  $J_{\pi}(\phi) = (nm + n)^{-1} \sum (\alpha \log A_{\phi}(a|s) - \min_{i=1,2} Q_{\theta_i}(s, a))$ ;  
16  $\phi \gets \phi - \lambda \nabla_{\phi} J_{\pi}(\phi)$  // Update Actor network  
17  $\theta_i' \gets \tau \theta_i + (1 - \tau) \theta_i'$  // Update target network weights  
Output: Actor Network  $A_{\phi}$ ;
the performance of our approach is significantly better than the baselines'. We also would like to emphasize that the training time of our approach is much less than that of MBPO (4 hours v.s. 3 days).
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
(b)	{0.6, 0.6}	{7, 3.5, 7}	5
(c)	{0.6, 0.6}	{2.5, 4.5, 2.5}	5
Table 3: Detailed Environment Parameters
B.3 Two-phase Wireless Network
The two-phase wireless network is a modified version of the downlink network in Section 2.1, which cannot be modeled as an MDP with exogenous inputs. In this example, each packet has two phases. When a packet from a mobile is scheduled, if it is in phase 1, it leaves with probability 0.2 or moves from phase 1 to phase 2; and if it is in phase 2, it leaves the system. We considered Poisson arrivals
with rates  $\Lambda = \{2,4,3\}$ , and the channel rates  $O = \{12,12,12\}$ . As the result shown in Figure. 5 ASG outperforms Max-weight in this environment as well.
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 a scaled version: divided by 20 in the wireless networks). 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 use a simple two layer neural network with hidden size 128 with learning rate 0.0003. The batch size used in all the simulations is 256.
C.1 Experimental settings
For all the simulations, We used a single NVIDIA GeForce RTX 2080 Super with AMD Ryzen 7 3700 8-Core Processor.