SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Intuitively,  $\mu_{\mathrm{min}}(\pi)$  is the occupancy of the bottleneck state that is hardest to reach by  $\pi$ , and  $t_{\mathrm{mix}}$  captures how fast the distribution generated by  $\pi$  decorrelates with any initial state distribution. In this setting, Li et al. [63] showed that with  $\tilde{O}\left(\frac{1}{\mu_{\mathrm{min}}(\pi_b)(1 - \gamma)^5\epsilon^2} + \frac{t_{\mathrm{mix}}(\pi_b)}{\mu_{\mathrm{min}}(\pi_b)(1 - \gamma)}\right)$  steps<sup>9</sup>, one can expect the Q-value to be  $\epsilon$ -close to the optimal Q-value, i.e.,  $\max_{(s,a) \in S \times A} |Q_t(s,a) - Q^\star(s,a)| \leq \epsilon$  with high probability. The second term of the complexity concerns the rate at which the state-action marginal converges to a steady state, and the first term characterizes the difficulty of learning once the empirical distribution has converged. The crucial point of this analysis is that the sample complexity is inversely proportional to the visitation probability of the hardest state to reach by  $\pi_b$ , so a better exploration policy that aims to explore the state space more uniformly (such as UCB) would improve the sample complexity, which in turn leads to better generalization (i.e., lower Q-value error for all state-action pairs).
The argument above shows that good exploration could ensure that the Q-values are accurate for any initial state distribution (which guarantees generalization in a uniform convergence sense), but it does not immediately explain why the Q-values are accurate around the training initial state distribution even though the training initial states further away have inaccurate Q-values, i.e., convergence at these states further away is slower. It turns out that this is related to resetting. It is perhaps worth noting that the analysis so far deals with the continuous RL (i.e., not episodic and does not have an end state). To account for the effect of reset distribution, we may assume that after the episode ends (e.g., reaching the goal in Figure 2), the agent is reset to a state drawn from the initial state distribution or that the agent may be reset to a state drawn from the initial state distribution at every timestep. We use this the continuous MDP as a proxy for the original MDP<sup>10</sup>. In this setting,  $d_{\pi}(s)$  and consequently  $d_{\pi}(s,a)$  will naturally have a higher density around the initial state distribution, so the Q-values at states further away from the initial state would converge slower, leading to these states having poor return even though the policy has high return around the states in the training initial state
<sup>8</sup>Notice that the stationary state-action distribution is independent of  $\mu$  since it is the stationary distribution of a Markov chain; however, it could still be close to a particular  $\mu_{\mathrm{reset}}$  if the environment resets to  $\mu_{\mathrm{reset}}$  and  $\pi_b$  does not explore efficiently (i.e., the Markov chain has a high probability of visiting the initial state distribution).
$9\tilde{O}$  hides poly-logarithmic factors.
<sup>10</sup>If the original episodic MDP is ergodic excluding the end state, then the continuous MDP will also be ergodic; however, this modified continuous MDP, in general, does not have the same optimal policy as the original episodic MDP, but it may be close for many common MDPs.
distribution. In other words, we can roughly say that the more a state-action pair is visited, the more likely that its Q-value has a low error for many MDPs.
To further formalize this intuition, we will leverage the concept of sub-MDP [51], where we prune  $(s,a)$  pairs from  $\mathcal{M}$  based on the behavioral policy. Kearns and Singh [51] prunes  $\mathcal{M}$  based on whether  $d_{\pi}(s,a)$  is larger than  $\alpha$ , a hyperparameter that is larger than  $\mu_{\mathrm{min}}(\pi_b)$ , and argues that one can obtain better sample complexity on the sub-MDP. This can be seen as throwing away state-action pairs that are not likely visited by the behavioral policy in the analysis. Concretely, the sub-MDP  $\mathcal{M}_{\pi}(\alpha)$  is an MDP where the state space is:
$$
\mathcal {G} (\alpha) = \{(s, a) \in \mathcal {S} \times \mathcal {A} \mid d _ {\pi} (s, a) > \alpha \}. \tag {8}
$$
If we further assume that all the state-action pairs visited by the optimal policy,  $\pi^{\star}$ , starting from  $\mu$  are contained within  $\mathcal{G}(\alpha)$ , then the optimal policy of the  $\mathcal{M}_{\pi_b}(\alpha)$  recovers that of the  $\mathcal{M}$  starting from  $\mu$ . This assumption generally holds for many reasonable MDPs (e.g., those with deterministic transition or bounded reward) and exploration policies  $\pi_{b}$  — in the grid world MDP (Figure 2a), the optimal policy starting from the top left only moves right and never visits any other states, and these states are well covered by  $\varepsilon$ -greedy exploration; for the training initial state, one can disregard all states that are not on the first row and still retain the optimal policy. By definition  $\alpha > \mu_{\min}(\pi_b)$  on  $\mathcal{M}_{\pi_b}(\alpha)$  and thus Q-values would converge to the optimal value on  $\mathcal{M}_{\pi_b}(\alpha)$  faster than the full MDP,  $\mathcal{M}$ . If  $\alpha$  is sufficiently larger than  $\mu_{\min}(\pi_b)$ , the sample complexity on  $\mathcal{M}_{\pi_b}(\alpha)$  could be much better than on  $\mathcal{M}$ , explaining the policy's superior performance near the initial state when the  $Q(s,a)$  further away from the initial states have not converged. Note that  $\mathcal{M}_{\pi_b}(\alpha)$  is only a construct for making the theoretical argument. The Q-values of the "pruned" state-action pairs are still updated as usual algorithmically. Even if the assumption is not true, the optimal policy on  $\mathcal{M}_{\pi_b}(\alpha)$  may still be good enough. If the optimal policy on  $\mathcal{M}_{\pi_b}(\alpha)$  or  $\alpha$  has to be small for the assumption to approximately hold, then one would need a better behavioral policy to solve the training environment efficiently, much less having a good generalization performance on the test environment.
Finally, Li et al. [63] showed that with high probability  $\min_{s,a} \frac{2}{3} N(s,a) / t \leq \mu_{\mathrm{min}}(\pi_b)$  which means that having a behavior policy that encourages visiting states with low visitation count,  $N(s,a)$ , should result in increasing  $\min_{s,a} \frac{2}{3} N(s,a) / t$  and thus increasing  $\mu_{\mathrm{min}}(\pi_b)$ . This indicates that UCB-based exploration should have a better sample complexity than  $\epsilon$ -greedy which does not explicitly target the visitation count and can be in general pretty bad at covering large state space. Indeed, for episodic RL, the state-of-the-art algorithms generally rely on UCB-based exploration [25] to ensure good sample complexity, which is more similar to what we do in the tabular tasks. Nonetheless, we focus on the continuous case because it highlights the importance of exploration via  $\mu_{\mathrm{min}}$ .
B Case Study: Simplified Bigfish
Figure 7: Simplified version of bigfish.
Figure 6: Example frames of bigfish. The goal for the agent (green fish) is to avoid fish bigger than itself and eat smaller fish. The spawning of enemy fish is different across different MDPs and the background can change.
To further illustrate the intuition that exploration on the training MDPs can help zero-shot generalization to other MDPs, we will use a slightly simplified example inspired by the game bigfish (Figure 6) from Procgen. As illustrated in Figure 7, in bigfish, the agent (green circle) needs to eat smaller fish (small red circle) and avoid bigger fish (large red circle). The images with solid borders are the observed frames and the images with dashed borders are the transitions.
If the agent collides with the small fish, it eats the small fish and gains  $+1$  reward; if the agent collides with the big fish, it dies and the episode ends. The blue rectangle (top) represents the training environment and the red rectangle (bottom) represents the test environment. In the training MDP, the agent always starts in the state shown in the  $T = 0$  of the greedy trajectory (top row). Using random exploration strategy (e.g.,  $\epsilon$ -greedy), the agent should be able to quickly identify that going to the top will be able to achieve the  $+1$  reward. However, there is an alternative trajectory (middle) where the agent can go down and right to eat the small fish on the right (orange border). This trajectory is longer (therefore harder to achieve via uniformly random exploration than the first one) and has a lower discounted return.
From the perspective of solving the training MDP, the trajectory is suboptimal. If this trajectory is sufficiently hard to sample, the agent with a naive exploration will most likely keep repeating the greedy trajectory. On the other hand, once the agent has sampled the greedy trajectory sufficiently, uncertainty-aware exploration (e.g., UCB or count-based exploration) will more likely sample these rarer trajectories since they have been visited less and thus have higher uncertainty. This has no impact on the performance in the training MDP because the agent has learned the optimal policy regardless of which exploration strategy is used.[11]
However, on the test MDP, the initial state is shown in  $T = 0$  of the test trajectory (bottom row). There is no guarantee that the agent knows how to behave correctly in this state because the agent has not seen it during training. One could hope that the neural network has learned a good representation so the agent knows what to do, but this does not necessarily have to be the case — the objective only cares about solving the training MDP and does not explicitly encourage learning representations that help generalization. Suppose that the agent keeps moving up due to randomness in the initialization of the weight, it will run into the big fish and die (bottom row,  $T = 2$ ). Notice that even though the two environments are distinct, the test trajectory (bottom row) and the explorative trajectory (middle row) share the same state at  $T = 1$  (pink arrow). Due to this structure, agents that have learned about
Using function approximators with insufficient capacity may have some effects on the optimal policy.
the explorative trajectory during training time will be able to recognize that the better action from  $T = 1$  is to move to the right which will avoid death and ultimately result in  $+1$  reward.
Once again, we emphasize that this is a simplified example for illustrative purposes. In reality, the sequences of states do not need to be exactly the same and the neural network can learn to map similar states to similar representations. In principle, the neural network can also learn to avoid the big fish but at  $T = 0$  of the training MDP, the behavior "moving up" is indistinguishable from the behavior "avoiding the big fish". Good exploration during training can lead the agent to run into the big fish from all directions which is much closer to "avoid the big fish", the intuitive generalizable behavior. Clearly, not all CMDPs will have this structure which allows us to improve generalization via exploration, but it should not hurt. Like most problems in deep RL, it is difficult to prove this hypothesis analytically but empirically it is indeed the case that EDE as better or the same performance as QR-DQN in 11 out of 16 Progen environments and perform approximately the same in the remaining environments (Figure 13).
This perspective allows us to tackle the problem with exploration that is in principle orthogonal to representation learning; however, since exploration naturally collects more diverse data, it may also be beneficial for representation learning as a side effect.
C Tabular Experiments
C.1 Generalization to Different Initial States
(a)  $\varepsilon$ -greedy
(b)UCB
(c) Policy Gradient
Figure 8: Mean test return for each method with different exploration hyperparameters.
Both the training MDP,  $\mu_{\mathrm{train}}$ , and the test MDP,  $\mu_{\mathrm{test}}$ , share the same state space  $\mathcal{S}$ ,  $[5] \times [5]$ , and action space  $\mathcal{A}$ , [4], corresponding to  $\{\mathrm{left}, \mathrm{right}, \mathrm{up}, \mathrm{down}\}$ . From each state, the transition,  $P(s' \mid s, a)$ , to the next state corresponding to an action is  $100\%$ .  $\gamma$  is 0.9. The two MDPs differ by only their initial state distribution:  $\rho_{\mathrm{train}}(s) = \mathbb{1}_{(0,0)}(s)$  and  $\rho_{\mathrm{test}}(s) = \mathbb{1}_{(4,0)}(s)$ .
For both policy gradient and Q-learning, we use a base learning rate of  $\alpha_0 = 0.05$  that decays as training progresses. Specifically, at time step  $t$ ,  $\alpha_{t} = \frac{1}{\sqrt{t}}\alpha_{0}$ .
For Q-learning, the Q-function is parameterized as  $\vartheta \in \mathbb{R}^{|S| \times |A|}$  where each entry is a state-action value. The update is:
$$
Q ^ {(t + 1)} \left(s _ {t}, a _ {t}\right) = Q ^ {(t)} \left(s _ {t}, a _ {t}\right) + \alpha_ {t} \left(r _ {t} + \gamma \max  _ {a} Q ^ {(t)} \left(s _ {t + 1}, a\right) - Q ^ {(t)} \left(s _ {t}, a _ {t}\right)\right), \tag {9}
$$
where  $Q(s,a) = \vartheta_{s,a}$ . For Q-learning with UCB, the agent keeps  $N(s,a)$  that keeps track of how many times the agent has taken action  $a$  from  $s$  over the course of training and explores according to:
$$
\pi_ {\mathrm {u c b}} (a \mid s) = \mathbb {1} _ {a ^ {\star}} (a) \text {w h e r e} a ^ {\star} = \underset {a} {\arg \max } Q (s, a) + c \sqrt {\frac {\log t}{N (s , a)}}. \tag {10}
$$
For Q-learning with  $\varepsilon$ -greedy, the exploration policy is:
$$
\pi_ {\text {e g r e e d y}} (a \mid s) = (1 - \varepsilon) \mathbb {1} _ {\arg \max  _ {a ^ {\prime}} Q (s, a ^ {\prime})} (a) + \varepsilon \left(\sum_ {a ^ {\prime} \in \mathcal {A}} \frac {1}{| \mathcal {A} |} \mathbb {1} _ {a ^ {\prime}} (a)\right). \tag {11}
$$
Ties in arg max are broken randomly which acts as the source of randomness for UCB.
C.2 Policy gradient
For policy gradient [114, 106], the policy is parameterized as  $\vartheta \in \mathbb{R}^{|S| \times |\mathcal{A}|}$  where each entry is the logit for the distribution of taking action  $a$  from  $s$ , i.e.,  $\pi(a \mid s) = \frac{\exp(\vartheta_{s,a})}{\sum_{a' \in \mathcal{A}} \exp(\vartheta_{s,a'})}$ . The overall gradient is:
$$
\nabla_ {\vartheta} \mathbb {E} _ {\pi_ {\vartheta}} [ J (\tau) ] = \mathbb {E} _ {\pi_ {\vartheta}} \left[ \sum_ {t = 0} ^ {T} \nabla_ {\vartheta} \log \pi_ {\vartheta} \left(a _ {t} \mid s _ {t}\right) \sum_ {k = t} ^ {T} \gamma^ {k - t} r _ {k} \right], \tag {12}
$$
$$
\nabla_ {\vartheta} \mathcal {H} _ {\pi_ {\vartheta}} = \nabla_ {\vartheta} \mathbb {E} _ {\pi_ {\vartheta}} \left[ \sum_ {t = 0} ^ {T} \sum_ {a \in \mathcal {A}} \pi_ {\vartheta} (a \mid s _ {t}) \log \pi_ {\vartheta} (a \mid s _ {t}) \right]. \tag {13}
$$
The update is:
$$
\vartheta^ {(t + 1)} = \vartheta^ {(t)} + \alpha_ {t} \left(\nabla_ {\vartheta} \mathbb {E} _ {\pi_ {\vartheta}} [ J (\tau) ] + \beta \nabla_ {\vartheta} \mathcal {H} _ {\pi_ {\vartheta}}\right) _ {| \vartheta = \vartheta^ {(t)}}, \tag {14}
$$
where the expectation is estimated with a single trajectory.
For each method, we repeat the experiment for the following 10 hyperparameter values:
$\varepsilon$  :  $\{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9\}$
-  $c$ :  $\{5, 10, 15, 20, 25, 30, 35, 40, 45, 50\}$
-  $\beta$ :  $\{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9\}$
Each experiment involves training the algorithm for 100 trials. Each trial consists of 1000 episodes and each episode ends when the agent reaches a terminal state or at 250 steps. At test time, all approaches are deterministic: Q-learning follows the action with highest Q-value and policy gradient follows the action with the highest probability. The mean test returns for each experiment are shown in Figure 8. We observe that both UCB and policy gradient are relatively robust to the choice of hyperparameters, whereas  $\varepsilon$ -greedy's generalization performance varies greatly with  $\varepsilon$ . Furthermore, even with extremely large  $\varepsilon$ , the generalization performance is still far worse than UCB and policy gradient. The best-performing values are respectively  $\varepsilon^{\star} = 0.9$ ,  $c^{\star} = 45$  and  $\beta^{\star} = 0.1$  and the train/test performance of this configuration are shown in Figure 9.PG's performance lies between Greedy and UCB. On average, it converges to a better solution with lower variance than Greedy but it is not as good as UCB.
Figure 9: Suboptimality of the best hyperparameters of Policy Gradient
C.3 Full trajectory update and Sarsa(λ)
Exploration is evidently not the only approach one can take to improve generalization in RL. In fact, we see that policy gradient with rather naive softmax exploration can also outperform  $\varepsilon$ -greedy exploration (Appendix C.2). We hypothesize that this improvement is due to the fact that these algorithms use whole-trajectory updates (i.e., Monte Carlo methods) rather than one-step updates like Q-learning. To study this, we also consider Sarsa( $\lambda$ ) [105, Ch 12.7] which is an on-policy RL method that learns a Q-value with whole-trajectory updates. The results are shown in Figure 10a.
While both on-policy methods use fairly naive exploration methods (softmax exploration for PG and  $\varepsilon$ -greedy for Sarsa(λ), we observe that both are able to outperform Greedy. This observation suggests that their improved performance could be attributed to something else. We believe this is due to the fact that when we use whole-trajectory updates (or n-step returns), the learning signal is propagated to states with low density as soon as a good trajectory is sampled. On the other hand, when we use off-policy approaches, the signal is only propagated when the state action pair is sampled multiple times by the behavioral policy, which may be exponential in the horizon. Nonetheless, there is still a gap between these two methods and UCB, suggesting that the generalization benefits of whole-trajectory updates may be distinct from that of exploration. Hence, better exploration can further improve generalization.
C.3.1 Details of Sarsa(λ)
Sarsa  $(\lambda)$  is an on-policy value-based method that learns Q-value with whole-trajectory update via eligibility traces. Since the algorithmic details are more involved than the other methods we consider, we refer the reader to Chapter 12.7 of Sutton and Barto [105] for the details. The Sarsa  $(\lambda)$  with  $\varepsilon$ -greedy experiments inherit all the hyperparameters of Q-learning  $+\varepsilon$ -greedy experiments, but have one more hyperparameter  $\lambda \in (0,1]$  that interpolates the method between Monte Carlo method and one-step TD update. We use a fixed  $\lambda = 0.9$  for all experiments. We sweep  $\varepsilon$  across 10 values (100 trials each):
$$
\cdot \varepsilon : \{0. 0, 0. 1, 0. 2, 0. 3, 0. 4, 0. 5, 0. 6, 0. 7, 0. 8, 0. 9 \}
$$
It can be seen from Figure 10a that even with whole-episode updates, the average return is still higher for larger exploration coefficients, which means exploration can still help. Furthermore, larger exploration coefficients also tend to have smaller variances in the test performance (Figure 10b). With the best hyperparameter configuration (Figure 10c),  $\mathrm{Sarsa}(\lambda)$  outperforms both Greedy and PG, but is still worse than UCB and has a much larger variance.
(a) Average of return.
(b) Standard deviation of return.
(c) Best hyperparameter.  $\varepsilon = 0.9$  and  $\lambda = 0.9$
Figure 10: Test performance of Sarsa(λ) across different  $\varepsilon$ . The results of Sarsa(λ). Sarsa(λ) is a value-based method, but it is able to achieve competitive performance because it uses whole-trajectory similar to policy gradient.
C.4 Generalization to Different Dynamics
In previous section, we presented a simple scenario that exploration helps generalization to different initial state distribution. Here, we demonstrate another scenario where exploration helps generalization to new dynamics. We use a  $7 \times 7$  grid as the state space (Figure 11a). The agent receives a small negative reward for every location except for the green square  $(4,2)$ , where the agent receives a reward of 2 and the episode ends. The difference between this environment and the previous one is that the agent always starts at  $(0,3)$  but at test time, a gust of wind can blow the agent off course. Specifically, with a probability of  $40\%$ , the agent will be blown to  $+2$  unit in the y-direction. This environment is an example of the windy world from Sutton and Barto [105]. The results are shown in Figure 11 and 12, where we see similar results as the previous section, but at test time the optimal expected return is generally worse and has higher variance for all methods since there is stochasticity (suboptimality is computed with respect to the return of an oracle that consistently moves towards
(b)  $\varepsilon$ -greedy
(d) Policy Gradient
Figure 11: (a) During both training and test time, the agent starts in the blue square, but at test time, with  $40\%$  probability, wind can blow the agent down by two units. In both cases, the goal is to get to the green square. The other plots show the mean and standard deviation of the train and test suboptimality (difference between optimal return and achieved return) over 100 runs for (b) Q-learning with  $\varepsilon$ -greedy exploration, (c) Q-learning with UCB exploration, and (d) policy gradient with entropy bonus.
(a)  $\varepsilon$ -greedy
Figure 12: Mean test return for each method with different exploration hyperparameters when the transition dynamic changes at test time.
(c) Policy Gradient
the goal). This setup is related to robust reinforcement learning [76], but exploration makes no assumptions about the type or magnitude of perturbation (thus making more complicated perturbation, e.g., different random initializations of the same game, possible) and does not require access to a simulator.
D Additional Background
D.1 Quantile Regression DQN
Quantile regression DQN (QR-DQN) is a variant of distributional RL that tries to learn the quantiles of the state-action value distribution with quantile regression rather than produce the actual distribution of Q-values. Compared to the original C51 [10], QR-DQN has many favorable theoretical properties. Here, we provide the necessary background for understanding QR-DQN and refer the readers to Dabney et al. [23] for a more in-depth discussion of QR-DQN.
Given a random variable  $Y$  with associated CDF  $F_{Y}(y) = \mathbb{P}(Y \leq y)$ , the inverse CDF of  $Y$  is
$$
F _ {Y} ^ {- 1} (\tau) = \inf  \left\{y \in \mathbb {R} \mid \tau \leq F _ {Y} (y) \right\}.
$$
To approximate  $Y$ , we may discretize the CDF with  $N$  qunatile values:
$$
\left\{F _ {Y} ^ {- 1} \left(\tau_ {1}\right), F _ {Y} ^ {- 1} \left(\tau_ {2}\right), \dots , F _ {Y} ^ {- 1} \left(\tau_ {N}\right) \right\} \quad \text {w h e r e} \quad \tau_ {i} = \frac {i}{N}.
$$
Using the same notation as main text, we use  $\pmb{\theta}:\mathcal{S}\times \mathcal{A}\to \mathbb{R}^N$  to denote a parametric function that maps a state-action pair,  $(s,a)$ , to the estimated quantiles  $\{\pmb{\theta}_j(\pmb{s},\pmb{a})\}_{j = 1}^N$ . The approximated CDF of
$Z(s, a)$  is thus:
$$
\widehat {Z} (\boldsymbol {s}, \boldsymbol {a}) \stackrel {{d}} {{=}} \frac {1}{N} \sum_ {j = 1} ^ {N} \delta \left(\boldsymbol {\theta} _ {j} (\boldsymbol {s}, \boldsymbol {a})\right),
$$
where  $\delta (\pmb{\theta}_j(\pmb{s},\pmb{a}))$  is the Dirac delta distribution centered at  $\pmb{\theta}_j(\pmb{s},\pmb{a})$ . To learn the quantiles, we use quantile regression which allows for unbiased estimate of the true quantile values. Concretely, for a target distribution  $Z$  and quantile  $\tau$ , the value of  $F_Z^{-1}(\tau)$  is the minimizer of the quantile regression loss:
$$
F _ {Z} ^ {- 1} (\tau) = \underset {\theta} {\arg \min } \mathcal {L} _ {\mathrm {Q R}} ^ {\tau} (\theta) = \mathbb {E} _ {z \sim Z} \left[ \rho_ {\tau} (z - \theta) \right] \quad \text {w h e r e} \quad \rho_ {\tau} (u) = u \left(\tau - \mathbb {1} (u <   0)\right).
$$