SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Since we are using variational inference a natural question to ask is how well our variational solution approximates  $\mathbb{P}_{\Gamma^{\star}}$ . Here we provide a bound quantifying the dissimilarity between  $\mathbb{P}_{\Gamma^{\star}}$  and the solution of the VAPOR optimization problem, according to a weighted KL-divergence, where for any  $\tau \in \mathbb{R}_+^{L,S,A}$ ,  $\lambda, \lambda' \in \Lambda(P)$  we define  $\mathrm{KL}_{\tau}(\lambda || \lambda') \coloneqq \sum_{l,s,a} \tau_l(s,a) \lambda_l(s,a) \log (\lambda_l(s,a) / \lambda_l'(s,a))$ .
Lemma 5. Let  $\lambda^{*}$  solve the VAPOR optimization problem (4) and  $\tau^{*} \in \operatorname{argmin}_{\tau} \mathcal{V}_{\phi}(\lambda^{*}, \tau)$ , then
$$
\mathrm {K L} _ {\tau^ {*}} \left(\mathbb {P} _ {\Gamma^ {*}} | | \lambda^ {*}\right) \leq \mathcal {V} _ {\phi} \left(\lambda^ {*}\right) - \mathbb {E} _ {s \sim \rho} \mathbb {E} _ {\phi} V _ {1} ^ {\star} (s).
$$
In other words, the (weighted) KL-divergence between  $\mathbb{P}_{\Gamma^{\star}}$  and  $\lambda^{*}$  is upper bounded by how loose our upper bound in (4) is. In practice however, we care less about a bound on the divergence metric than we do about performance in the control problem. Our next result shows that the variational policy also satisfies a strong sub-linear Bayesian regret bound, i.e., the resulting policy performs well.
4.3 Bayesian Regret Analysis
Learning problem. We consider that the agent interacts with the MDP  $\mathcal{M}$  over a (possibly unknown) number of  $N$  episodes. We denote by  $\mathcal{F}_t$  the sigma-algebra generated by all the history (i.e., sequences of states, actions and rewards) before episode  $t$ , with  $\mathcal{F}_1 = \emptyset$ . We let  $\mathbb{E}^t[\cdot] = \mathbb{E}[\cdot \mid \mathcal{F}_t]$ ,  $\mathbb{P}^t(\cdot) = \mathbb{P}(\cdot \mid \mathcal{F}_t)$ . We denote by  $n_l^t(s, a)$  the visitation count to  $(s, a)$  at step  $l$  before episode  $t$ , and  $(\cdot \lor 1) := \max(\cdot, 1)$ . In the Bayesian approach where  $\mathcal{M}$  is sampled from a known prior  $\phi$ , we want to minimize the Bayes regret over  $T := NL$  timesteps of an algorithm alg producing policy  $\pi^t$  at each episode  $t$ , which is defined as its expected regret under that prior distribution
$$
\mathcal {R} _ {\mathcal {M}} (\mathrm {a l g}, T) := \sum_ {t = 1} ^ {N} \mathbb {E} _ {s \sim \rho} \left(V _ {1} ^ {\star} (s) - V _ {1} ^ {\pi^ {t}} (s)\right), \qquad \mathcal {B R} _ {\phi} (\mathrm {a l g}, T) := \mathbb {E} _ {\mathcal {M} \sim \phi} \mathcal {R} _ {\mathcal {M}} (\mathrm {a l g}, T).
$$
The VAPOR learning algorithm proceeds as follows: at the beginning of each episode, it solves the VAPOR optimization problem and executes the induced policy. We now show that it enjoys a sub-linear Bayesian regret bound under the following standard assumption [61, 47, 62].
Assumption 1. The mean rewards are bounded in  $[0,1]$  almost surely with independent priors and the reward noise is additive  $\nu$ -sub-Gaussian for a constant  $\nu > 0$ .
Assumption 1 implies that the mean rewards are sub-Gaussian under the posterior (see [70, App. D.2]), where we can upper bound the sub-Gaussian parameter  $\sigma_l^t(s,a) \leq \sqrt{(\nu^2 + 1) / (n_l^t(s,a) \vee 1)}$  at the beginning of each episode  $t$ .
Theorem 1. For known  $P$  and under Assumption 1, it holds that
$$
\mathcal {B R} _ {\phi} (\text {V A P O R}, T) \leq \sqrt {2 (\nu^ {2} + 1) T S A \log (S A) (1 + \log T / L)} = \tilde {O} (\sqrt {S A T}).
$$
In the above  $\widetilde{O}$  suppresses log factors. The same regret bound holds for the (intractable) algorithm that uses (1) as the policy each episode.
Corollary 1. Denote by  $\mathrm{alg}_{\Gamma^{\star}}$  the algorithm that produces policies based on  $\mathbb{P}_{\Gamma^{\star}}^{t}$  for each episode  $t$  using (1). Then, under the same conditions as Theorem 1, we have  $\mathcal{BR}_{\phi}(\mathrm{alg}_{\Gamma^{\star}},T)\leq \widetilde{O} (\sqrt{SAT})$ .
4.4 Interpretation of VAPOR
Inspecting  $\mathcal{V}_{\phi}$  (3) reveals that the VAPOR optimization problem is equivalent to solving a two-player zero-sum game between a 'policy' player  $\lambda \in \Lambda(P)$  and a 'temperature' player  $\tau \in \mathbb{R}_{+}^{L,S,A}$ . The latter finds the tightest  $\tau$  that best balances two exploration mechanisms on a per state-action basis: an optimism term  $\sigma^2/\tau$  that augments the expected reward  $\mathbb{E}_{\phi}r$ , and a  $\tau$ -weighted entropy regularization term. The policy player maximizes the entropy-regularized optimistic reward.
The fact that entropy regularization falls out naturally from VAPOR is interesting because the standard 'RL as inference' framework conveniently reveals policy entropy regularization (Section 2.1), which has been widely studied theoretically [45, 20] and empirically for deep RL exploration [42, 26]. The specificity here is that it is with respect to the occupancy measure instead of the policy, and it is adaptively weighted for each state-action pair. This enables us to obtain an 'RL as inference' framework with entropy regularization that explores provably efficiently. Regularizing with the (weighted) entropy of the occupancy measure is harder to implement in practice, therefore in Section 7 we propose a principled, albeit looser, upper bound of the VAPOR optimization problem that regularizes the optimistic reward with only the (weighted) entropy of the policy, making the resulting objective amenable to online optimization with policy gradients.
Algorithm 1 VAPOR learning algorithm
For episode  $t = 1,2,\ldots$  do
1. Compute expected rewards  $\mathbb{E}^t r$ , transitions  $\mathbb{E}^t P$ , uncertainty measure  $\hat{\sigma}^t$
2. Solve VAPOR optimization problem  $\lambda^t\in \mathrm{argmax}_{\lambda \in \Lambda (\mathbb{E}^t P)}\mathcal{V}_{\hat{\phi}^t}(\lambda)$  from Equation (3)
3. Execute policy  $\pi_l^t (s,a)\propto \lambda_l^t (s,a)$  , for  $l = 1,\dots ,L$
5 Extension to Unknown Transition Dynamics
So far we have focused on the special case of known transition dynamics  $P$ . In this section, we derive a generic reduction from the case of unknown  $P$  to known  $P$ , which may be of independent interest. We prove that any algorithm that enjoys a Bayesian regret bound in the known-  $P$  special case can be easily converted into one that enjoys a regret bound for the more challenging unknown-  $P$  case. Our analysis relies on the mean rewards being sub-Gaussian under the posterior (which holds e.g., under Assumption 1) and on the following standard assumption [61, 47, 62].
Assumption 2. The transition functions are independent Dirichlet under  $\phi$ , with parameter  $\alpha_{l}(s,a)\in \mathbb{R}_{+}^{S_{l + 1}}$  for each  $(s,a)\in S_l\times \mathcal{A}$  with  $\sum_{s^{\prime}\in S_{l + 1}}\alpha_{l}(s,a,s^{\prime})\geq 1$ .
At a high level, our reduction transfers the uncertainty on the transitions to additional uncertainty on the rewards in the form of carefully defined zero-mean Gaussian noise. It extends the reward perturbation idea of the RLSVI algorithm [62], by leveraging a property of Gaussian-Dirichlet optimism [60] and deriving a new property of Gaussian sub-Gaussian optimism (Appendix G) which allows to bypass the assumption of binary rewards in  $\{0,1\}$  from [62, Asm. 3]. Our reduction implies the following 'dominance' property of the expected  $V^{\star}$  under the transformed and original beliefs.
Lemma 6. Define the mapping  $\mathcal{T}:\phi \mapsto \widehat{\phi}$  that transforms the beliefs  $\phi$  into a distribution  $\widehat{\phi}$  on the same space, with transition dynamics equal to  $\mathbb{E}_{\phi}P$  and rewards distributed as  $\mathcal{N}(\mathbb{E}_{\phi}r,\widehat{\sigma}^{2})$  with
$$
\hat {\sigma} _ {l} ^ {2} (s, a) := 3. 6 ^ {2} \sigma_ {l} ^ {2} (s, a) + \frac {(L - l) ^ {2}}{\sum_ {s ^ {\prime}} \alpha_ {l} (s , a , s ^ {\prime})}.
$$
$$
\mathbb {E} _ {s \sim \rho} \mathbb {E} _ {\phi} V _ {1} ^ {\star} (s) \leq \mathbb {E} _ {s \sim \rho} \mathbb {E} _ {\hat {\phi}} V _ {1} ^ {\star} (s).
$$
Original beliefs φ	Transformed beliefs φ
Mean reward r	σ-sub-Gaussian	N(Εφr, θ2)
Transitions P	Dirichlet(α)	ΕφP
Now, let us return to the episodic interaction case and denote by  $\phi^t \coloneqq \phi(\cdot \mid \mathcal{F}_t)$  the posterior at the beginning of episode  $t$ . Under Assumptions 1 and 2, we can upper bound the uncertainty  $\hat{\sigma}^t$  of the transformed posteriors  $\widehat{\phi}^t$  as  $(\widehat{\sigma}_l^t)^2(s,a) \leq (3.6^2(\nu^2 + 1) + (L - l)^2) / (n_l^t(s,a) \vee 1)$ . This brings us to a key result necessary for a general regret bound.
Lemma 7. Let  $\mathrm{alg}$  be any procedure that maps posterior beliefs to policies, and denote by  $\mathcal{B}\mathcal{R}_{\mathcal{T},\phi}(\mathrm{alg},T)$  the Bayesian regret of  $\mathrm{alg}$  where at each episode  $t = 1,\dots ,T$  the policy and regret are computed by replacing  $\phi^t$  with  $\mathcal{T}(\phi^t)$  (Lemma 6), then under Assumptions 1 and 2,
$$
\mathcal {B R} _ {\phi} (\mathrm {a l g}, T) \leq \mathcal {B R} _ {\mathcal {T}, \phi} (\mathrm {a l g}, T).
$$
This tells us that if we have an algorithm with a Bayesian regret bound under known transitions  $P$ , then we can convert it directly into an algorithm that achieves a regret bound for the case of unknown  $P$ , simply by increasing the amount of uncertainty in the posterior of the unknown rewards and replacing the unknown  $P$  with its mean when executing the algorithm. We now instantiate this generic reduction to VAPOR in Algorithm 1. Combining Lemma 7 and Theorem 1 yields the following regret bound.
Theorem 2. Under Assumptions 1 and 2, it holds that  $\mathcal{BR}_{\phi}(\mathrm{VAPOR},T)\leq \widetilde{O} (L\sqrt{SAT})$
The Bayes regret bound in Theorem 2 is within a factor of  $\sqrt{L}$  of the known information theoretic lower bound [32, 13], up to constant and log terms. It matches the best known bound for K-learning [47] and Thompson sampling [61], under the same set of standard assumptions (see Appendix C). There exists a complementary line of work deriving minimax-optimal regret bounds in the frequentist setting, both in the model-based [3, 41] and model-free [86, 38] cases. We refer to e.g., [21] for discussion on the advantages and disadvantages of the Bayesian and frequentist approaches. Without
claiming superiority of either, it is worth highlighting that compared to frequentist algorithms that tend to be deterministic and non-stationary (i.e., explicitly dependent on the episode number), VAPOR is naturally stochastic and stationary (i.e., independent of the episode number).
6 Connections
Thompson Sampling. Thompson sampling (TS) or posterior sampling for RL (PSRL) [78, 75, 59] first samples an environment (i.e., rewards and transition dynamics) from the posterior and then computes the optimal policy for this sample to be executed during the episode. Our principled Bayesian inference approach to RL uncovers an alternative view of TS which
TS/PSRL	VAPOR
PΓ*approx.	Sample MDP M and compute π*M ⇌ λTS ~ PΓ* ⇒ E[λTS] = PΓ*	Variationally approximate λ* ≈ PΓ* and compute π*
Exploration mechanism	Stochastic optimism	Explicit optimism + Entropy regularization
deepens our understanding of this popular algorithm. The next lemma shows that TS implicitly approximates  $\mathbb{P}_{\Gamma^{\star}}$  by sampling, thus tightly connecting it to VAPOR's explicit approximation of  $\mathbb{P}_{\Gamma^{\star}}$ .
Lemma 8. Let  $\lambda^{\mathrm{TS}}$  be the occupancy measure of the TS policy, it holds that  $\mathbb{E}[\lambda^{\mathrm{TS}}] = \mathbb{P}_{\Gamma}$ .
Plugging this observation into Lemma 4 and retracing the proof of Theorem 2 immediately yields regret bounds for TS matching that of VAPOR. The explicit variational approximation of  $\mathbb{P}_{\Gamma^{\star}}$  by VAPOR has several advantages over the implicit sampling-based approximation of  $\mathbb{P}_{\Gamma^{\star}}$  by TS:
- Having direct (approximate) access to this probability can be desirable in some practical applications, for example to ensure safety constraints or to allocate budgets.
- TS suffers linear regret in the multi-agent and constrained cases [52], while we expect suitable modifications of VAPOR to be able to handle these cases, as suggested in the multi-armed bandit setting [51].
- The VAPOR optimization problem can naturally extend to parameterizing the occupancy measure to be in a certain family (e.g., linear in some basis).
- The objective we use in VAPOR is differentiable, opening the door to differentiable computing architectures such as deep neural network. In contrast, performing TS requires to maintain an explicit model over MDP parameters [59], which becomes prohibitively expensive as the MDP becomes large. To alleviate this computational challenge of TS, some works have focused on cheaply generating approximate posterior samples [57, 55, 62, 54], but it is not yet clear if this approach is better than the variational approximation we propose.
K-learning. K-learning [47] endows the agent with a risk-seeking exponential utility function and enjoys the same regret bound as Theorem 2. Its update rule is similar to soft Q-learning [19, 25, 44, 22], which emerges from inference under Approximation 1 [37]. Although soft Q-learning ignores epistemic uncertainty and thus suffers linear regret [53], K-learning learns a scalar temperature parameter which balances both an optimistic reward and a policy entropy regularization term. In contrast, VAPOR naturally produces a separate risk-seeking parameter for each state-action. This enables a fine-grained control of the exploration-exploitation trade-off depending on the region of the state-action space, which can make a large difference in practice. In the multi-armed bandit case it was shown that K-learning and VAPOR coincide when VAPOR is constrained to have all temperature parameters (denoted  $\tau$ ) equal [51]. It can be shown (see Appendix F.12) that the same observation holds in the more general MDP case, up to additional entropic terms added to the K-learning 'soft' value functions, which only contribute a small constant to the regret analysis. This sheds a new light on K-learning as a variational approximation of our probabilistic inference framework with the additional constraint that all 'temperature' parameters are equal.
Maximum Entropy Exploration. A recent line of work called maximum entropy exploration suggests that in the absence of reward the agent should cover the state-action space as uniformly as possible by solving a maximum-entropy problem over occupancy measures  $\max_{\lambda \in \Lambda(P)} \mathcal{H}(\lambda)$  [29, 9, 85, 79], where  $\mathcal{H}$  denotes the entropy summed over steps  $l$ . VAPOR can thus be interpreted as regularizing a reward-driven objective with a maximum entropy exploration term, that is weighted with vanishing temperatures. In other words, the principle of (weighted) maximum-entropy exploration can be derived by considering a variational approximation to Bayesian state-action optimality.
7 A Policy-Gradient Approximation
Up to this point, we instantiated our new Bayesian 'RL as inference' approach with a tabular, model-based algorithm that exactly solves the variational optimization problem and has a guaranteed regret bound (Algorithm 1). In this section, we derive a principled, albeit looser, upper-bound approximation of the variational optimization problem, which can be solved using policy-gradient techniques by representing the policy using a deep neural network. This will yield a scalable, model-free algorithm that we call VAPOR-lite.
Denoting by  $\mu^{\pi}$  the stationary state distribution of a policy  $\pi$ , note the relation  $\lambda_l^\pi (s,a) = \pi_l(s,a)\mu_l^\pi (s)$ . The regularization term of VAPOR can thus be decomposed in a (weighted) policy entropy regularization term and a (weighted) stationary state distribution entropy term, the latter of which is challenging to optimize in high-dimensional spaces [29, 36, 43]. In light of this, we introduce the following VAPOR-lite alternative optimization problem
$$
\max  _ {\pi \in \Pi} \sum_ {l, s} \mu_ {l} ^ {\pi} (s) \left(\sum_ {a} \pi_ {l} (s, a) \left(r _ {l} (s, a) + \hat {\sigma} _ {l} (s, a) + \mathcal {H} _ {\hat {\sigma} _ {l} (s, \cdot)} (\pi_ {l} (s, \cdot))\right)\right), \tag {5}
$$
where we define the weighted policy entropy  $\mathcal{H}_{\hat{\sigma}_l(s,\cdot)}(\pi_l(s,\cdot)) := -\sum_a\hat{\sigma}_l(s,a)\pi_l(s,a)\log \pi_l(s,a)$ . Akin to policy-gradient methods, VAPOR-lite now optimizes over the policies  $\pi$  which can be parametrized by a neural network. It depends on an uncertainty signal  $\hat{\sigma}$  which can also be parametrized by a neural network (e.g., an ensemble of reward predictors, as we consider in our experiments in Section 8). Compared to VAPOR, VAPOR-lite accounts for a weaker notion of weighted entropy regularization (i.e., in the policy space instead of the occupancy-measure space), which allows to solve the objective using policy-gradient techniques. For simplicity it also sets the temperatures  $\tau$  to the uncertainties  $\hat{\sigma}$  instead of minimizing for the optimal ones. Importantly, VAPOR-lite remains a principled approach. We indeed prove the following relevant properties of VAPOR-lite in Appendix H: (i) although the problem is no longer concave in  $\pi$ , solving it in the space of occupancy measures  $\lambda$  remains a computationally tractable, concave optimization problem, (ii) VAPOR-lite upper bounds the VAPOR objective (up to a multiplicative factor in the uncertainty measure and a negligible additive bias), and (iii) it yields the same  $\widetilde{O}(L\sqrt{SAT})$  regret bound as VAPOR for a careful schedule of uncertainty measures  $\hat{\sigma}$  over episodes.
Essentially, VAPOR-lite endows a policy-gradient agent with (i) an uncertainty reward bonus and (ii) an uncertainty-weighted policy entropy regularization. The latter exploration mechanism is novel and has an appealing interpretation: Unlike standard policy entropy regularization, it adaptively accounts for epistemic uncertainty, by eliminating actions from the entropy term where the agent has low uncertainty for each given state. In the limit of  $\hat{\sigma} \to 0$  (i.e., no uncertainty), the regularization of VAPOR-lite vanishes and we recover the original policy-gradient objective.
8 Numerical Experiments
GridWorld. We first study empirically how well VAPOR approximates  $\mathbb{P}_{\Gamma^{\star}}$ . Since  $\mathbb{E}[\lambda^{\mathrm{TS}}] = \mathbb{P}_{\Gamma^{\star}}$ , we estimate the latter by averaging over 1000 samples of the (random) TS occupancy measure (we denote it by  $\overline{\lambda}_{(1000)}^{\mathrm{TS}}$ ). We design simple  $10 \times 10$  GridWorld MDPs with four cardinal actions, known dynamics and randomly generated reward. Figure 2 suggests that VAPOR and the TS average output similar approximations of  $\mathbb{P}_{\Gamma^{\star}}$ , thus showing the accuracy of our variational approximation in this domain. Unlike TS which requires estimating it with many samples, VAPOR approximates  $\mathbb{P}_{\Gamma^{\star}}$  by only solving a single convex optimization problem.
DeepSea. We now study the learning performance of the VAPOR Algorithm 1 in the hard-exploration environment of DeepSea [62]. In this  $L \times L$  grid MDP, the agent starts at the top-left cell and must reach the lower-right cell. It can move left or right, always descending to the row below. Going left yields no reward, while going right incurs a small cost of  $0.01 / L$ . The bottom-right cell yields a reward of 1, so that the optimal policy is to always go right. Local dithering (e.g., Soft Q-learning,  $\epsilon$ -greedy) takes time exponential in  $L$ , so the agent must perform deep exploration to reach the goal. Figure 3 shows the time required to 'solve' the problem as a function of the depth  $L$ , averaged over 10 seeds. The 'time to solve' is defined as the first episode where the rewarding state has been found at least in  $10\%$  of the episodes so far [47]. We compare VAPOR to TS/PSRL [59], K-learning [47] and a variant of RLSVI [62] that runs PSRL under the transformed posteriors of Lemma 6. (Several
Figure 2: For 3 seeds of randomly generated reward means and noises in a GridWorld, visualization of the timestep-averaged stationary state distribution (i.e.,  $L^{-1} \sum_{l,a} \lambda_l(s,a)$ ) for  $\lambda^{\mathrm{VAPOR}}$  (left) and  $\overline{\lambda}_{(1000)}^{\mathrm{TS}}$  (right).
Figure 3: Learning time on DeepSea
Figure 4: Median human normalized score across 57 Atari games (with standard errors across 5 seeds).
optimistic methods like UCBVI [3] or Optimistic Q-learning [32] were additionally compared by [47] but they performed worse.) We see that VAPOR achieves the lowest learning time as  $L$  increases, thus displaying its ability to perform deep exploration.
Atari. Finally, we investigate the performance of VAPOR-lite on the Atari benchmark [5]. We consider a replay-based actor-critic agent with V-trace off-policy corrections [15] and an actor-learner decomposition [30] (see Appendix I for the full experimental details). We compare this agent without and with added fixed policy entropy regularization as commonly used [83, 42], to the same agent with the VAPOR-lite objective. We stress that unlike simply adding entropy regularization, which leads to local dithering and suffers linear regret, VAPOR-lite relies on a principled approach of casting RL as Bayesian inference. Figure 4 shows the performance advantage of augmenting a policy-gradient agent with VAPOR-lite on Atari. It reaches the peak performance of the replay actor-critic agent (resp. with entropy regularization) in about  $2 \times$  (resp.  $1.5 \times$ ) fewer environment frames, for essentially the same computational cost. This advantage comes from the fact that VAPOR-lite leads to deep exploration, which results in finding higher rewarding states and in better cumulative performance. In our ablation in Appendix I, we illustrate the isolated benefits of tuning entropy regularization on a per state-action basis with VAPOR-lite compared to tuning a single scalar for entropy regularization [49]. Our promising empirical results suggest that further gains could come from deriving a more accurate practical approximation of VAPOR.
RL is a statistical inference problem wrapped in a control problem. In this paper, we demonstrated that a single quantity can be used to handle both aspects of the problem — the posterior probability of each state-action pair being visited under the optimal policy,  $\mathbb{P}_{\Gamma^{\star}}$ . We can perform Bayesian inference to compute this, and then use it to generate a control policy that performs efficient exploration. This is a coherent and principled approach to 'RL as inference', rather than the heuristic approaches in prior work which could not compute valid posteriors. Unfortunately, computing  $\mathbb{P}_{\Gamma^{\star}}$  is intractable in practice so we derived a variational approximation that we showed also explores efficiently. We concluded with some numerical experiments showing improved performance in the challenging 'DeepSea' unit test, as well as on the Atari suite. We discuss some limitations of our work and directions for future investigation in Appendix D.