Title: Probabilistic Inference in Reinforcement Learning Done Right
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement learning

Bayesian inference

Exploration
Abstract: A popular perspective in Reinforcement learning (RL) casts the problem as probabilistic inference on a graphical model of the Markov decision process (MDP). The core object of study is the probability of each state-action pair being visited under the optimal policy. Previous approaches to approximate this quantity can be arbitrarily poor, leading to algorithms that do not implement genuine statistical inference and consequently do not perform well in challenging problems. In this work, we undertake a rigorous Bayesian treatment of the posterior probability of state-action optimality and clarify how it flows through the MDP. We first reveal that this quantity can indeed be used to generate a policy that explores efficiently, as measured by regret. Unfortunately, computing it is intractable, so we derive a new variational Bayesian approximation yielding a tractable convex optimization problem and establish that the resulting policy also explores efficiently. We call our approach VAPOR and show that it has strong connections to Thompson sampling, K-learning, and maximum entropy exploration. We conclude with some experiments demonstrating the performance advantage of a deep RL version of VAPOR.

Probabilistic Inference in Reinforcement Learning Done Right

Abstract

A popular perspective in Reinforcement learning (RL) casts the problem as probabilistic inference on a graphical model of the Markov decision process (MDP). The core object of study is the probability of each state-action pair being visited under the optimal policy. Previous approaches to approximate this quantity can be arbitrarily poor, leading to algorithms that do not implement genuine statistical inference and consequently do not perform well in challenging problems. In this work, we undertake a rigorous Bayesian treatment of the posterior probability of state-action optimality and clarify how it flows through the MDP. We first reveal that this quantity can indeed be used to generate a policy that explores efficiently, as measured by regret. Unfortunately, computing it is intractable, so we derive a new variational Bayesian approximation yielding a tractable convex optimization problem and establish that the resulting policy also explores efficiently. We call our approach VAPOR and show that it has strong connections to Thompson sampling, Klearning, and maximum entropy exploration. We conclude with some experiments demonstrating the performance advantage of a deep RL version of VAPOR .

1 Introduction

Reinforcement learning (RL) is the problem of learning to control an unknown system by taking actions to maximize its cumulative reward through time [76]. As the agent navigates the environment, it receives noisy observations which it can use to update its (posterior) beliefs about the environment [21]. Unlike supervised learning, where the performance of an algorithm does not influence the data it will later observe, in RL the policy of the agent affects the data it will collect, which in turn affects the policy, and so on. As a result, an agent must sometimes take actions that lead to states where it has epistemic uncertainty about the value of those states, and sometimes take actions that lead to more certain payoff. The tension between these two modes is the exploration-exploitation trade-off [2, 35]. In this light, RL is a statistical inference problem wrapped in a control problem , and the two problems must be tackled simultaneously for good data efficiency.

It is natural to apply Bayesian probabilistic inference to the uncertain parameters in RL [21] and, since the goal of the agent is to find the optimal policy, a relevant object of study is the posterior probability of optimality for each state-action pair. In fact, the popular 'RL as inference' approach, most clearly summarized in the tutorial and review of [37], embeds the control problem into a graphical model by introducing optimality binary random variables that indicate whether a state-action is optimal. In the 'RL as inference' setup these optimality variables are assumed to be observed and take the value one with probability depending on the local reward. This surrogate potential seems like a peculiar and arbitrary choice, yet it has been widely accepted due to its mathematical convenience: for deterministic dynamics, the resulting probabilistic inference is equivalent to maximum entropy RL [80, 87]. However, this approximation leads to a critical shortcoming: The surrogate potential functions do not take into account the Bayesian (epistemic) uncertainty, and consequently do not perform genuine statistical inference on the unknowns in the MDP. This leads to 'posteriors' that are

in no way related to the true posteriors in the environment, and acting with those false posteriors leads to poor decision making and agents that require exponential time to solve even small problems [53].

Our contributions can be centered around a new quantity that we uncover as key for inference and control , denoted by P Γ ⋆ , which represents the posterior probability of each state-action pair being visited under the optimal policy .

We reveal that P Γ ⋆ formalizes from a Bayesian perspective what it means for a 'state-action pair to be optimal', an event at the core of the 'RL as inference' framework which had never been properly analyzed.

We establish that knowledge of P Γ ⋆ is sufficient to derive a policy that explores efficiently, as measured by regret (Section 3).

Since computing P Γ ⋆ is intractable, we propose a variational optimization problem that tractably approximates it (Section 4).

We first solve this optimization problem exactly, resulting in a new tabular model-based algorithm with a guaranteed regret bound (Section 5).

We then solve a variant of this optimization problem using policy-gradient techniques, resulting in a new scalable model-free algorithm (Section 7).

We show that both Thompson sampling [78, 75, 59, 71] and K-learning [47] can be directly linked to P Γ ⋆ , thus shedding a new light on these algorithms and tightly connecting them to our variational approach (Section 6).

Our approach has the unique algorithmic feature of adaptively tuning optimism and entropy regularization for each state-action pair, which is empirically beneficial as our experiments on 'DeepSea' and Atari show (Section 8).

2 Preliminaries

We model the RL environment as a finite state-action, time-inhomogeneous MDP given by the tuple M := {S , A , L, P, R, ρ } , where L is the horizon length, S = S 1 ∪ . . . ∪ S L is the state space with cardinality S = ∑ L l =1 S l where S l = |S l | , A is the action space with A possible actions, P l : S l ×A→ ∆( S l +1 ) denotes the transition dynamics at step l , R l : S l ×A→ ∆( R ) is the reward function at step l and ρ ∈ ∆( S 1 ) is the initial state distribution. Concretely, the initial state s 1 ∈ S 1 is sampled from ρ , then for steps l = 1 , . . . , L the agent is in state s l ∈ S l , selects action a l ∈ A , receives a reward sampled from R l ( s l , a l ) with mean r l ( s l , a l ) ∈ R and transitions to the next state s l +1 ∈ S l +1 with probability P l ( s l +1 | s l , a l ) . An agent following a policy π l ∈ ∆( A ) S l at state s at step l selects action a with probability π l ( s, a ) . The value functions are defined as

An optimal policy π ⋆ satisfies V l π ⋆ ( s ) = V ⋆ l ( s ) for each state s and step l . Let { π ⋆ l ( s ) = a } be the event that action a is optimal for state s at step l , with ties broken arbitrarily so that only one action is optimal at each state. Let R L,S,A + := { R S l × A + } L l =1 be the set of functions S l ×A→ R + for l ∈ [ L ] .

Occupancy measures. We denote the set of occupancy measures (a.k.a. stationary state-action distributions) with respect to the transition dynamics P as

Note the correspondence between Λ( P ) and the set of policies Π = { ∆( A ) S l } L l =1 [66]. Any policy π ∈ Π induces λ π ∈ Λ( P ) such that λ π l ( s, a ) denotes the probability of reaching ( s, a ) at step l under π . Conversely, any λ ∈ Λ( P ) induces π λ ∈ Π given by π λ l ( s, a ) := λ l ( s, a ) / ( ∑ a ′ λ l ( s, a ′ )) so long as ∑ a ′ λ l ( s, a ′ ) > 0 ; otherwise π λ l ( s, · ) can be any distribution, e.g. , uniform.

Bayesian RL. We consider the Bayesian view of the RL problem, where the agent has some beliefs about the MDP represented by a distribution ϕ on the space of all MDPs. We use the shorthand notation E ϕ and P ϕ to denote the expectation and probability under ϕ . In the Bayesian view, the mean rewards r and transition dynamics P are random variables, and thus so are Q ⋆ , V ⋆ , π ⋆ .

2.1 Previous Approach to 'RL as Inference'

The popular line of research casting 'RL as inference' is most clearly summarized in the tutorial and review of Levine, 2018 [37]. It embeds the control problem into a graphical model, mirroring the dual relationship between optimal control and inference [81, 34, 68]. As shown in Figure 5a, this approach defines an additional optimality binary random variable denoted by Γ ⋆ l , which indicates which stateaction at timestep l is 'optimal', i.e. , which state-action is visited under the (unknown) optimal policy. In these works the probability of Γ ⋆ l is then modeled as proportional to the exponentiated reward.

Approximation 1 (Inference over exponentiated reward [37]) . Denoting by Γ ⋆ l ( s, a ) the event that state-action pair ( s, a ) is 'optimal' at timestep l , set P (Γ ⋆ l ( s, a )) ∝ exp( r l ( s, a )) .

Under Approximation 1, the optimality variables become observed variables in the graphical model, over which we can apply standard tools for inference. This approach has the advantage of being mathematically and computationally convenient. In particular, under deterministic dynamics it directly leads to RL algorithms with 'soft' Bellman updates and added entropy regularization, thus recovering the maximum-entropy RL framework [80, 87] and giving a natural, albeit heuristic exploration strategy. Many popular algorithms lie within this class [65, 25, 1] and have demonstrated strong empirical performance in domains where efficient exploration is not a bottleneck.

Despite its popularity and convenience, the potential function of Approximation 1 ignores epistemic uncertainty [53]. Due to this shortcoming, the inference procedure produces invalid and arbitrarily poor posteriors. It can be shown that the resulting algorithms fail even in basic bandit-like problems ( e.g. , [53, Problem 1]), as they underestimate the probability of optimal actions and take actions which have no probability of optimality under the true posterior. This implies that the resulting agents ( e.g. , Soft Q-learning [25]) can perform poorly in simple domains that require deep exploration [57]. The fact that the 'RL as inference' framework does not perform valid inference on the optimality variables leads us to consider how a proper Bayesian inference approach might proceed. This brings us to the first contribution of this manuscript.

3 A Principled Bayesian Inference Approach to RL

In this section, we provide a principled Bayesian treatment of statistical inference over 'optimality' variables in an MDP. Although it is the main quantity of interest of the popular 'RL as inference' framework, a rigorous definition of whether a state-action pair is optimal has remained elusive. We provide a recursive one below.

Definition 1 (State-action optimality) . For any step l ∈ [ L ] , state-action ( s, a ) ∈ S l ×A , we define

In words, Γ ⋆ l ( s ) is the random event that state s is reached at step l after executing optimal actions from steps 1 to l -1 , and the event Γ ⋆ l ( s, a ) further requires that the action a is optimal at state s in step l . Equipped with this definition, 'optimality' in an MDP flows both in a backward way (via action optimality π ⋆ , which can be computed using dynamic programming) and in a forward way (via state-action optimality Γ ⋆ ). We illustrate this bidirectional property in the simplified graphical model of Figure 5b (which isolates a single trajectory). In general, state-action optimality encapsulates all possible trajectories leading to a given state being optimal (and there may be exponentially many of them). Thanks to the MDP structure and Bayes' rule, we show in Lemma 1 that the posterior probability of state-action optimality lies in the space of occupancy measures. To simplify exposition, we first assume that the transition dynamics P are known (only rewards are unknown), and we extend our analysis to unknown P in Section 5.

The significance of Lemma 1 is that we can readily sample from P Γ ⋆ with the induced policy

s

1

s

2

s

L

-

1

s

L

+1

w.p.

0

.

5

-

1

w.p.

0

.

5

-

ϵ

s

′

2

s

′

3

s

′

L

-

ϵ

0

. . .

-

ϵ

0

0

Policy

Agent starts at

s

1

π

1

(

s

1

, a

)

for

a

∈ {↓

,

→}

Agent reaches

s

l

,

l >

1

π

l

(

s

l

, a

)

for

a

∈ {↓

,

→}

Expected # episodes

to reach

s

L

Bayes-optimal

(0

,

1)

(0

,

1)

1

Approximation 1

(0

.

5

,

0

.

5)

(0

.

5

,

0

.

5)

Ω(2

L

)

Marginal

P

ϕ

(

π

⋆

)

(0

.

5

,

0

.

5)

(0

.

5

,

0

.

5)

Ω(2

L

)

Conditional

P

ϕ

(

π

⋆

|

Γ

⋆

)

(0

.

5

,

0

.

5)

(0

,

1)

2

Thompson Sampling

(0

,

1) w

.

p

.

0

.

5

(1

,

0) w

.

p

.

0

.

5

(0

,

1)

2

VAPOR

(0

,

1)

(0

,

1)

1

which is the probability under the posterior that action a is optimal at state s conditioned on all actions taken before timestep l being optimal. It turns out that this policy explores efficiently, as measured by regret (see Corollary 1). Concretely, we have shown that the two facets of RL (inference and control) can be distilled into a single quantity: P Γ ⋆ . We use Bayesian inference to compute P Γ ⋆ and then use that to derive a good control policy. This is a principled and consistent Bayesian 'RL as inference' approach. The remaining challenge thus lies in computing P Γ ⋆ . Unfortunately, doing so involves computing several complicated integrals with respect to the posterior and is intractable in most cases. In Section 4, we introduce the second contribution of this manuscript, which is a way to approximate P Γ ⋆ using a computationally tractable variational approach.

3.1 Instructive Example

To gain intuition on the difference between action optimality ( i.e. , P ϕ ( π ⋆ ) ) and state-action optimality ( i.e. , P Γ ⋆ ), we consider the simple decision problem of Figure 1. Each episode starts in the leftmost state and the agent can traverse a chain of states in order to reach the rightmost state, where the reward is always either +1 or -1 with equal probability. Once the agent reaches the rightmost state it will resolve its uncertainty. At each state the agent has the option to move → , paying a cost of ϵ where ϵL ≪ 1 , or move ↓ with no cost. We see that there are only 2 consistent trajectories: either take action ↓ at state s 1 (optimal for M -), or repeat actions → to reach s L (optimal for M + ). We detail how a selection of algorithms perform on this problem in Table 1.

It is worth examining in detail what has gone wrong for the policy that samples according to P ϕ ( π ⋆ ) . Even if the agent has traveled to the right a few times it will exit the chain with probability 0 . 5 at each new state, which leads to poor exploration. Taking action → a few times and then exiting the chain has zero probability of being optimal under the posterior, so this sequence of actions is inconsistent with the posterior. Conditioning on the event Γ ⋆ ( s ) when at state s forces the policy to be consistent with the set of the beliefs that would get the agent to state s in the first place, which, as we prove later, yields deep exploration. In our chain example, the only belief that takes action → is the one in which the end state has reward +1 , so the only consistent policy conditioned on having taken → once is to continue all the way to the end state. The same issue arises in Thompson sampling, where in order to achieve consistency the agent samples a policy from its posterior and keeps it fixed for the entire episode since resampling at every timestep leads to inefficient exploration [71]. As we shall show in Section 6, Thompson sampling is exactly equivalent to sampling from P Γ ⋆ . In conclusion:

4 A Variational Bayesian Approach

Having access to P Γ ⋆ is sufficient to enable deep exploration, but computing this probability is intractable in general. In this section, we derive a variational , i.e. , optimization-based, approach to approximate P Γ ⋆ . In most variational Bayesian approaches [18], the optimization objective is to find a surrogate probability measure that minimizes some dissimilarity metric ( e.g. , a KL-divergence) from the intractable probability measure of interest. Yet unlike standard variational inference techniques, we cannot minimize the KL-divergence to this distribution as we do not have access to samples. However, we do have a downstream control objective. That is to say, we are not simply interested in approximating the distribution, but also in our approximation doing well in the control problem. This insight provides an alternative objective that takes into account the rewards and uncertainties that the agent is likely to encounter. The optimal value function V ⋆ is a random variable under the beliefs ϕ . Our next lemma relates the expected value of V ⋆ under ϕ to our state-action optimality event Γ ⋆ .

4.1 Information Theoretic Upper Bound

Lemma 2 depends on P Γ ⋆ , but the conditional expectation is not easy to deal with. To handle this, we upper bound each E ϕ [ r l ( s, a ) | Γ ⋆ l ( s, a )] in Lemma 2 using tools from information theory. To simplify the exposition, we consider the following standard sub-Gaussian assumption [70, 61, 47] and we defer to Appendix F.3 the treatment of the general case. We say that X : Ω → R is υ -sub-Gaussian for υ > 0 if E exp( c ( X -E X )) ≤ exp( c 2 υ 2 / 2) , for all c ∈ R .

Lemma 3. If r l ( s, a ) is σ l ( s, a ) -sub-Gaussian under ϕ for l ∈ [ L ] and ( s, a ) ∈ S l ×A , then

4.2 Optimization Problem

Finally, we combine Lemmas 2 and 3 to reveal a concave function in P Γ ⋆ that upper bounds the value function objective E s ∼ ρ E ϕ V ⋆ 1 ( s ) . For any τ ∈ R L,S,A + and occupancy measure λ ∈ Λ( P ) , we define the τ -weighted entropy of λ (summed over steps l ) as H τ ( λ ) := -∑ l,s,a τ l ( s, a ) λ l ( s, a ) log λ l ( s, a ) . In the following definition we take division and the square-root to be applied elementwise, and ◦ denotes elementwise multiplication. With this we can define the following optimistic value functions

Lemma 4. For known P and σ -sub-Gaussian r , we have

where the VAPOR optimization problem is concave.

The above optimization problem is our variational objective, the solution of which yields an occupancy measure that approximates P Γ ⋆ . We call the approach ' VAPOR ', for variational approximation of the posterior probability of optimality in RL. We discuss various properties of the optimization problem in Appendix E (its unconstrained dual problem and message passing interpretation).

Since we are using variational inference a natural question to ask is how well our variational solution approximates P Γ ⋆ . Here we provide a bound quantifying the dissimilarity between P Γ ⋆ and the solution of the VAPOR optimization problem, according to a weighted KL-divergence, where for any τ ∈ R L,S,A + , λ, λ ′ ∈ Λ( P ) we define KL τ ( λ || λ ′ ) := ∑ l,s,a τ l ( s, a ) λ l ( s, a ) log( λ l ( s, a ) /λ ′ l ( s, a )) .

Lemma 5. Let λ ∗ solve the VAPOR optimization problem (4) and τ ∗ ∈ argmin τ V ϕ ( λ ∗ , τ ) , then

In other words, the (weighted) KL-divergence between P Γ ⋆ and λ ∗ 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 M over a (possibly unknown) number of N episodes. We denote by F t the sigma-algebra generated by all the history ( i.e. , sequences of states, actions and rewards) before episode t , with F 1 = ∅ . We let E t [ · ] = E [ · | F t ] , P t ( · ) = P ( · | F t ) . We denote by n t l ( s, a ) the visitation count to ( s, a ) at step l before episode t , and ( · ∨ 1) := max( · , 1) . In the Bayesian approach where M is sampled from a known prior ϕ , we want to minimize the Bayes regret over T := NL timesteps of an algorithm alg producing policy π t at each episode t , which is defined as its expected regret under that prior distribution

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 ν -sub-Gaussian for a constant ν > 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 σ t l ( s, a ) ≤ √ ( ν 2 +1) / ( n t l ( s, a ) ∨ 1) at the beginning of each episode t .

Theorem 1. For known P and under Assumption 1, it holds that

In the above ˜ 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 alg Γ ⋆ the algorithm that produces policies based on P t Γ ⋆ for each episode t using (1) . Then, under the same conditions as Theorem 1, we have BR ϕ (alg Γ ⋆ , T ) ≤ ˜ O ( √ SAT ) .

4.4 Interpretation of VAPOR

Inspecting V ϕ (3) reveals that the VAPOR optimization problem is equivalent to solving a two-player zero-sum game between a 'policy' player λ ∈ Λ( P ) and a 'temperature' player τ ∈ R L,S,A + . The latter finds the tightest τ that best balances two exploration mechanisms on a per state-action basis: an optimism term σ 2 /τ that augments the expected reward E ϕ r , and a τ -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 , . . . do

Compute expected rewards E t r , transitions E t P , uncertainty measure p σ t

Solve VAPOR optimization problem λ t ∈ argmax λ ∈ Λ( E t P ) V p ϕ t ( λ ) from Equation (3)

Execute policy π t l ( s, a ) ∝ λ t l ( s, a ) , for l = 1 , . . . , 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 knownP special case can be easily converted into one that enjoys a regret bound for the more challenging unknownP 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 ϕ , with parameter α l ( s, a ) ∈ R S l +1 + for each ( s, a ) ∈ S l ×A with ∑ s ′ ∈S l +1 α l ( s, a, s ′ ) ≥ 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 ⋆ under the transformed and original beliefs.

Lemma 6. Define the mapping T : ϕ ↦→ p ϕ that transforms the beliefs ϕ into a distribution p ϕ on the same space, with transition dynamics equal to E ϕ P and rewards distributed as N ( E ϕ r, p σ 2 ) with

Now, let us return to the episodic interaction case and denote by ϕ t := ϕ ( · | F t ) the posterior at the beginning of episode t . Under Assumptions 1 and 2, we can upper bound the uncertainty p σ t of the transformed posteriors p ϕ t as ( p σ t l ) 2 ( s, a ) ≤ (3 . 6 2 ( ν 2 +1) + ( L -l ) 2 ) / ( n t l ( s, a ) ∨ 1) . This brings us to a key result necessary for a general regret bound.

Lemma 7. Let alg be any procedure that maps posterior beliefs to policies, and denote by BR T ,ϕ ( alg , T ) the Bayesian regret of alg where at each episode t = 1 , . . . , T , the policy and regret are computed by replacing ϕ t with T ( ϕ t ) (Lemma 6), then under Assumptions 1 and 2,

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 BR ϕ ( VAPOR , T ) ≤ ˜ O ( L √ SAT ) .

√

The Bayes regret bound in Theorem 2 is within a factor of 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

deepens our understanding of this popular algorithm. The next lemma shows that TS implicitly approximates P Γ ⋆ by sampling, thus tightly connecting it to VAPOR 's explicit approximation of P Γ ⋆ .

Lemma 8. Let λ TS be the occupancy measure of the TS policy, it holds that E [ λ TS ] = P Γ ⋆ .

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 P Γ ⋆ by VAPOR has several advantages over the implicit sampling-based approximation of P Γ ⋆ 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 Qlearning 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 τ ) 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 λ ∈ Λ( P ) H ( λ ) [29, 9, 85, 79], where 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, modelbased 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 µ π the stationary state distribution of a policy π , note the relation λ π l ( s, a ) = π l ( s, a ) µ π l ( 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

where we define the weighted policy entropy H p σ l ( s, · ) ( π l ( s, · )) := -∑ a p σ l ( s, a ) π l ( s, a ) log π l ( s, a ) . Akin to policy-gradient methods, VAPOR-lite now optimizes over the policies π which can be parametrized by a neural network. It depends on an uncertainty signal p σ 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 τ to the uncertainties p σ 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 π , solving it in the space of occupancy measures λ 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 ˜ O ( L √ SAT ) regret bound as VAPOR for a careful schedule of uncertainty measures p σ 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 p σ → 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 P Γ ⋆ . Since E [ λ TS ] = P Γ ⋆ , we estimate the latter by averaging over 1000 samples of the (random) TS occupancy measure (we denote it by λ TS (1000) ). We design simple 10 × 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 P Γ ⋆ , thus showing the accuracy of our variational approximation in this domain. Unlike TS which requires estimating it with many samples, VAPOR approximates P Γ ⋆ 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 × 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, ϵ -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

Randomrewardseed1

Randomrewardseed2

Randomrewardseed3

1e0

1e-1

1e-2

1e-3

1e-4

VAPOR

TSaverage

VAPOR

TSaverage

VAPOR

TS average

5000

2L

failure

1000

%06>

500

100

RLSVI

K-learning

TS/PSRL

10

VAPOR

10

20

30

40

50

DeepSeadepthL

Median human-normalized score

1.5

0.5

Replay Actor-Critic(RAC)

RAC+ entropy regularization

RAC+VAPOR-lite

0

0e+00

1e+08

2e+08

3e+08

4e+08

5e+08

Frames

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 × (resp. 1 . 5 × ) 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 .

9 Conclusion

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, P Γ ⋆ . 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 P Γ ⋆ 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.

References

A. Abdolmaleki, J. T. Springenberg, Y. Tassa, R. Munos, N. Heess, and M. Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint arXiv:1806.06920 , 2018.

P. Auer. Using confidence bounds for exploitation-exploration trade-offs. Journal of Machine Learning Research , 3(Nov):397-422, 2002.

M. G. Azar, I. Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning , pages 263-272. PMLR, 2017.

M. Bellemare, S. Srinivasan, G. Ostrovski, T. Schaul, D. Saxton, and R. Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems , 29, 2016.

M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research , 47:253-279, 2013.

S. P. Boyd and L. Vandenberghe. Convex optimization . Cambridge university press, 2004.

Y. Burda, H. Edwards, A. Storkey, and O. Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.12894 , 2018.

C. B˘ alc˘ au. Weighted entropic perturbation for linear programming. Analele Universit˘ a¸ tii Bucure¸ sti. Matematic˘ a-Informatic˘ a , 55, 01 2006.

W. C. Cheung. Exploration-exploitation trade-off in reinforcement learning on online markov decision processes with global concave rewards. arXiv preprint arXiv:1905.06466 , 2019.

T. M. Cover and J. A. Thomas. Information theory and statistics. Elements of information theory , 1(1):279-335, 1991.

S. Diamond and S. Boyd. CVXPY: A Python-embedded modeling language for convex optimization. Journal of Machine Learning Research , 17(83):1-5, 2016.

A. Domahidi, E. Chu, and S. Boyd. ECOS: An SOCP solver for embedded systems. In European Control Conference (ECC) , pages 3071-3076, 2013.

O. D. Domingues, P. Ménard, E. Kaufmann, and M. Valko. Episodic reinforcement learning in finite mdps: Minimax lower bounds revisited. In Algorithmic Learning Theory , pages 578-598. PMLR, 2021.

V. Dwaracherla, Z. Wen, I. Osband, X. Lu, S. M. Asghari, and B. Van Roy. Ensembles for uncertainty estimation: Benefits of prior functions and bootstrapping. arXiv preprint arXiv:2206.03633 , 2022.

L. Espeholt, H. Soyer, R. Munos, K. Simonyan, V. Mnih, T. Ward, Y. Doron, V. Firoiu, T. Harley, I. Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning , pages 1407-1416. PMLR, 2018.

S.-C. Fang and H.-S. J. Tsao. Linear programming with entropic perturbation. Zeitschrift für Operations Research , 37(2):171-186, 1993.

M. Fellows, A. Mahajan, T. G. Rudner, and S. Whiteson. Virel: A variational inference framework for reinforcement learning. Advances in neural information processing systems , 32, 2019.

C. W. Fox and S. J. Roberts. A tutorial on variational bayesian inference. Artificial intelligence review , 38:85-95, 2012.

R. Fox, A. Pakman, and N. Tishby. Taming the noise in reinforcement learning via soft updates. arXiv preprint arXiv:1512.08562 , 2015.

M. Geist, B. Scherrer, and O. Pietquin. A theory of regularized markov decision processes. In International Conference on Machine Learning , pages 2160-2169. PMLR, 2019.

M. Ghavamzadeh, S. Mannor, J. Pineau, A. Tamar, et al. Bayesian reinforcement learning: A survey. Foundations and Trends® in Machine Learning , 8(5-6):359-483, 2015.

J. Grau-Moya, F. Leibfried, and P. Vrancx. Soft Q-learning with mutual-information regularization. In International conference on learning representations , 2019.

R. M. Gray. Entropy and information theory . Springer Science & Business Media, 2011.

D. Grytskyy, J. Ramírez-Ruiz, and R. Moreno-Bote. A general markov decision process formalism for action-state entropy-regularized reward maximization. arXiv preprint arXiv:2302.01098 , 2023.

T. Haarnoja, H. Tang, P. Abbeel, and S. Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning , pages 1352-1361. PMLR, 2017.

T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018.

S. Han and Y. Sung. Diversity actor-critic: Sample-aware entropy regularization for sampleefficient exploration. In International Conference on Machine Learning , pages 4018-4029. PMLR, 2021.

B. Hao and T. Lattimore. Regret bounds for information-directed reinforcement learning. arXiv preprint arXiv:2206.04640 , 2022.

E. Hazan, S. Kakade, K. Singh, and A. Van Soest. Provably efficient maximum entropy exploration. In International Conference on Machine Learning , pages 2681-2691. PMLR, 2019.

M. Hessel, M. Kroiss, A. Clark, I. Kemaev, J. Quan, T. Keck, F. Viola, and H. van Hasselt. Podracer architectures for scalable reinforcement learning. arXiv preprint arXiv:2104.06272 , 2021.

M. Jaggi. Revisiting frank-wolfe: Projection-free sparse convex optimization. In International conference on machine learning , pages 427-435. PMLR, 2013.

C. Jin, Z. Allen-Zhu, S. Bubeck, and M. I. Jordan. Is Q-learning provably efficient? Advances in neural information processing systems , 31, 2018.

H. Juergensen and D. E. Matthews. Entropy and higher moments of information. Journal of Universal Computer Science , 16(5):749-794, 2010.

H. J. Kappen, V. Gómez, and M. Opper. Optimal control as a graphical model inference problem. Machine learning , 87:159-182, 2012.

M. Kearns and S. Singh. Near-optimal reinforcement learning in polynomial time. Machine learning , 49:209-232, 2002.

L. Lee, B. Eysenbach, E. Parisotto, E. Xing, S. Levine, and R. Salakhutdinov. Efficient exploration via state marginal matching. arXiv preprint arXiv:1906.05274 , 2019.

S. Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909 , 2018.

G. Li, L. Shi, Y. Chen, Y. Gu, and Y. Chi. Breaking the sample complexity barrier to regretoptimal model-free reinforcement learning. Advances in Neural Information Processing Systems , 34:17762-17776, 2021.

X. Lu, B. Van Roy, V. Dwaracherla, M. Ibrahimi, I. Osband, Z. Wen, et al. Reinforcement learning, bit by bit. Foundations and Trends® in Machine Learning , 16(6):733-865, 2023.

C. E. Luis, A. G. Bottero, J. Vinogradska, F. Berkenkamp, and J. Peters. Model-based uncertainty in value functions. arXiv preprint arXiv:2302.12526 , 2023.

P. Ménard, O. D. Domingues, X. Shang, and M. Valko. UCB momentum Q-learning: Correcting the bias without forgetting. In International Conference on Machine Learning , pages 7609-7618. PMLR, 2021.

V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning , pages 1928-1937. PMLR, 2016.

M. Mutti, L. Pratissoli, and M. Restelli. Task-agnostic exploration via policy gradient of a non-parametric state entropy estimate. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 9028-9036, 2021.

O. Nachum, M. Norouzi, K. Xu, and D. Schuurmans. Bridging the gap between value and policy based reinforcement learning. In Advances in Neural Information Processing Systems , pages 2772-2782, 2017.

G. Neu, A. Jonsson, and V. Gómez. A unified view of entropy-regularized markov decision processes. arXiv preprint arXiv:1705.07798 , 2017.

B. O'Donoghue. Operator splitting for a homogeneous embedding of the linear complementarity problem. SIAM Journal on Optimization , 31(3):1999-2023, 2021.

B. O'Donoghue. Variational Bayesian reinforcement learning with regret bounds. Advances in Neural Information Processing Systems , 34:28208-28221, 2021.

B. O'Donoghue. On the connection between Bregman divergence and value in regularized markov decision processes. arXiv preprint arXiv:2210.12160 , 2022.

B. O'Donoghue. Efficient exploration via epistemic-risk-seeking policy optimization. arXiv preprint arXiv:2302.09339 , 2023.

B. O'Donoghue, E. Chu, N. Parikh, and S. Boyd. Conic optimization via operator splitting and homogeneous self-dual embedding. Journal of Optimization Theory and Applications , 169(3):1042-1068, June 2016.

B. O'Donoghue and T. Lattimore. Variational Bayesian optimistic sampling. Advances in Neural Information Processing Systems , 34:12507-12519, 2021.

B. O'Donoghue, T. Lattimore, and I. Osband. Matrix games with bandit feedback. In Uncertainty in Artificial Intelligence , pages 279-289. PMLR, 2021.

B. O'Donoghue, I. Osband, and C. Ionescu. Making sense of reinforcement learning and probabilistic inference. In International Conference on Learning Representations , 2020.

B. O'Donoghue, I. Osband, R. Munos, and V. Mnih. The uncertainty Bellman equation and exploration. In International Conference on Machine Learning , pages 3836-3845, 2018.

I. Osband. Deep Exploration via Randomized Value Functions . PhD thesis, Stanford University, 2016.

I. Osband, J. Aslanides, and A. Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems , 31, 2018.

I. Osband, C. Blundell, A. Pritzel, and B. Van Roy. Deep exploration via bootstrapped DQN. In Advances In Neural Information Processing Systems , pages 4026-4034, 2016.

I. Osband, Y. Doron, M. Hessel, J. Aslanides, E. Sezener, A. Saraiva, K. McKinney, T. Lattimore, C. Szepesvari, S. Singh, et al. Behaviour suite for reinforcement learning. arXiv preprint arXiv:1908.03568 , 2019.

I. Osband, D. Russo, and B. Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems , 26, 2013.

I. Osband and B. Van Roy. Gaussian-dirichlet posterior dominance in sequential learning. arXiv preprint arXiv:1702.04126 , 2017.

I. Osband and B. Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In International conference on machine learning , pages 2701-2710. PMLR, 2017.

I. Osband, B. Van Roy, D. J. Russo, and Z. Wen. Deep exploration via randomized value functions. Journal of Machine Learning Research , 20:1-62, 2019.

G. Ostrovski, M. G. Bellemare, A. Oord, and R. Munos. Count-based exploration with neural density models. In International conference on machine learning , pages 2721-2730. PMLR, 2017.

D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell. Curiosity-driven exploration by selfsupervised prediction. In International conference on machine learning , pages 2778-2787. PMLR, 2017.

J. Peters, K. Mulling, and Y. Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 24, pages 1607-1612, 2010.

M. L. Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014.

J. Qian, R. Fruit, M. Pirotta, and A. Lazaric. Concentration inequalities for multinoulli random variables. arXiv preprint arXiv:2001.11595 , 2020.

K. Rawlik, M. Toussaint, and S. Vijayakumar. On stochastic optimal control and reinforcement learning by approximate inference. Proceedings of Robotics: Science and Systems VIII , 2012.

D. Russo and B. Van Roy. Learning to optimize via information-directed sampling. Advances in Neural Information Processing Systems , 27, 2014.

D. Russo and B. Van Roy. An information-theoretic analysis of thompson sampling. The Journal of Machine Learning Research , 17(1):2442-2471, 2016.

D. J. Russo, B. Van Roy, A. Kazerouni, I. Osband, Z. Wen, et al. A tutorial on thompson sampling. Foundations and Trends® in Machine Learning , 11(1):1-96, 2018.

T. Schaul, J. Quan, I. Antonoglou, and D. Silver. Prioritized experience replay. arXiv preprint arXiv:1511.05952 , 2015.

M. J. Schervish. Theory of statistics . Springer Science & Business Media, 2012.

S. A. Serrano. Algorithms for unsymmetric cone optimization and an implementation for problems with the exponential cone . PhD thesis, Stanford University, 2015.

M. Strens. A bayesian framework for reinforcement learning. In ICML , volume 2000, pages 943-950, 2000.

R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018.

J. Tarbouriech and A. Lazaric. Active exploration in markov decision processes. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 974-982. PMLR, 2019.

W. R. Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika , 25(3-4):285-294, 1933.

D. Tiapkin, D. Belomestny, D. Calandriello, E. Moulines, R. Munos, A. Naumov, P. Perrault, Y. Tang, M. Valko, and P. Menard. Fast rates for maximum entropy exploration. arXiv preprint arXiv:2303.08059 , 2023.

E. Todorov. Linearly-solvable markov decision problems. Advances in neural information processing systems , 19, 2006.

E. Todorov. Efficient computation of optimal actions. Proceedings of the national academy of sciences , 106(28):11478-11483, 2009.

M. J. Wainwright, M. I. Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning , 1(1-2):1-305, 2008.

R. J. Williams and J. Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science , 3(3):241-268, 1991.

H. Yu, H. Zhang, and W. Xu. Do you need the entropy reward (in practice)? arXiv preprint arXiv:2201.12434 , 2022.

T. Zahavy, B. O'Donoghue, G. Desjardins, and S. Singh. Reward is enough for convex mdps. Advances in Neural Information Processing Systems , 34:25746-25759, 2021.

Z. Zhang, Y. Zhou, and X. Ji. Almost optimal model-free reinforcement learning via referenceadvantage decomposition. Advances in Neural Information Processing Systems , 33:1519815207, 2020.

B. D. Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy . Carnegie Mellon University, 2010.

Appendix

Table of Contents

a

1

a

2

a

3

a

4

s

1

s

2

s

3

s

4

r

1

r

2

r

3

r

4

Γ

⋆

1

Γ

⋆

2

Γ

⋆

3

Γ

⋆

4

(a) graphical model of the standard

'RL as inference' framework [37]

a

1

a

2

a

3

a

4

s

1

s

2

s

3

s

4

r

1

r

2

r

3

r

4

Q

⋆

1

Q

⋆

2

Q

⋆

3

Q

⋆

4

π

⋆

1

π

⋆

2

π

⋆

3

π

⋆

4

Γ

⋆

1

Γ

⋆

2

Γ

⋆

3

Γ

⋆

4

(b) graphical model for principled

Bayesian 'RL as inference'

A Graphical Models

We illustrate in Figure 5 the graphical models representing the standard 1 'RL as inference' framework [37] and our principled, Bayesian 'RL as inference' framework. For ease of presentation, we consider here a simplified representation where we isolate a single trajectory ( s 1 , a 1 , . . . , s 4 , a 4 ) and consider that Γ ⋆ does not encompass the randomness in the transition dynamics (which is e.g. , the case for deterministic dynamics).

VAPOR-lite

T(s,α)=T

K-learning,

VAPOR

ERSAC

variational

P

IDS

T*

V-IDS

sampling

TS

RVFs, RLSVI

(PSRL)

Figure 6: VAPOR is an explicit, variational approximation of P Γ ⋆ , while Thompson sampling (TS) implicitly approximates P Γ ⋆ by sampling, similarly to the TS approximations Randomized value functions (RVFs), RLSVI [62]. A special case of VAPOR with equal temperatures approximately recovers K-learning and ERSAC [49]. We conjecture that there also exist connections between Γ ⋆ and information-theoretic approaches such as information-directed sampling (IDS), variance-IDS (V-IDS) [69, 28, 39].

B Algorithmic Connections to Γ ⋆

Figure 6 illustrates how numerous exploration algorithms can be linked to Γ ⋆ . This sheds a new light on existing algorithms, tightly connects them to our variational approach, and unifies these approaches within our principled Bayesian framework.

C Assumptions

In this section, we review and discuss the assumptions made throughout the paper. We stress that these assumptions are not required computationally for our variational Bayesian approach, but only for the analysis to obtain VAPOR 's regret bound (Theorem 2). These are:

(A) Layered, time-inhomogoneous MDP,

(B) The reward noise is additive sub-Gaussian and the mean rewards are bounded almost surely with independent priors,

(C) The prior over transition functions is independent Dirichlet.

(A) This assumption is non-restrictive in the sense that any finite-horizon MDP with cycles can be transformed into an equivalent MDP without cycles by adding a timestep variable l ∈ [ L ] to the state space (picking up an additional factor of √ L in the regret bound). It implies the transition function and the value function at the next state are conditionally independent, which is a technical property that simplifies the regret analysis.

(B) The sub-Gaussian assumption is standard in the literature and arises commonly in practice. For instance, it holds in any environment with bounded rewards (e.g., DeepSea, Atari). In Appendix F.3 (see Lemma 11) we extend Lemma 3 to the general case, and we refer to Appendix F.5 for the resulting VAPOR optimization problem (this would give a generic regret bound for the VAPOR learning algorithm which would depend on how fast the posteriors concentrate).

(C) The Dirichlet prior over the transition functions is the canonical prior for transitions in Bayesian RL, see [21]. This is because visiting a state-action pair simply increments by 1 the appropriate entry of the Dirichlet vector - note the ease of the update because the Dirichlet is the conjugate prior of the categorical distribution (which models transition probabilities of discrete-state-action-MDPs).

Note that we make the same assumptions as existing state-of-the-art Bayesian regret analyses, including PSRL [61] 3 , K-learning [47], RLSVI [62]. To the best of our knowledge, it is an open question how to obtain an algorithm with a ˜ O ( L √ SAT ) Bayes regret bound without these assumptions.

D Limitations

We now discuss some limitations of our work, which constitute relevant directions for future investigation.

The VAPOR convex optimization problem grows with L, S, A . Although large-scale exponential cone solvers exist that can handle this type of problem, it does not immediately lend itself to an online approximation.

The regret guarantee of VAPOR relies on Assumptions 1 and 2 which, although standard (Appendix C), can be strong and it would be relevant to (partly) relax them.

VAPOR-lite makes some upper-bound approximations of our variational optimization problem to use policy-gradient techniques, and there may be a tighter way to approximate VAPOR in the case where the policy is parameterized using a deep neural network that is updated using online stochastic gradients.

E Properties of the VAPOR Optimization Problem

The VAPOR optimization problem is an exponential cone program that can be solved efficiently using modern optimization methods [50, 12, 74, 46]. In Appendix E.1, we show that its dual is a convex optimization problem that is unconstrained in a 'value' variable and a 'temperature' variable. In Appendix E.2, we relate VAPOR to message passing , a popular technique to perform inference in graphical models [82], by showing that VAPOR can be cast as passing some forward and backward messages back and forth.

E.1 Dual Problem of the VAPOR Optimization Problem

We derive the dual problem of the VAPOR optimization problem (4) and show that it is a convex optimization problem in a 'value' variable V ∈ R L,S := { R S l } L l =1 and a 'temperature' variable τ ∈ R L,S,A + . Even though the 'primal' admits flow constraints in λ , this dual is unconstrained in V (the only constraint is the non-negativity of τ ). In the case of unknown P (Section 5), the dual is the same as in Lemma 9 except that P is replaced by E ϕ P and the uncertainty σ is replaced by p σ .

Lemma 9. The λ ∗ ∈ Λ( P ) solving the VAPOR optimization problem (4) is unique and satisfies

where we define the 'advantage' function

and where V ∗ , τ ∗ minimize the (convex) dual function

The induced policy π ∗ is uniquely defined as

if the denominator is positive, otherwise π ∗ l ( s, · ) can be any distribution.

Proof. The proof follows from applying the method of Lagrange multipliers to the VAPOR optimization problem using the definition of V in (3). Derivations of dual problems of related optimization problems in the space of occupancy measures can be found in e.g. , [65, 24]. Note that for a fixed temperature function τ ∈ R L,S,A + , VAPOR is an instance of linear programming with weighted entropic perturbation, which is a class of optimization problems shown to admit an unconstrained convex dual [8, 16].

Algorithm 2 Frank-Wolfe algorithm for the VAPOR optimization problem (Lemma 10)

Require: Accuracy level ε > 0

Set the smoothing parameter δ = ε/ ( σ max LSA )

Set the number of iterations K = ( σ max Lε -1 ) 5 ( SA ) 4

Initialize any λ (0) ∈ Λ

for iterations k = 1 , 2 , . . . , K do

Backward message pass: Compute the optimal policy π ( k -1) for the reward ∇V δ ( λ ( k -1) ) Forward message pass: Compute d ( k -1) the stationary state distribution induced by π ( k -1) , i.e. ,

Combine: Compute for step size γ k = 2 k +1 the occupancy measure

end for return λ ( K )

E.2 Message Passing View of the VAPOR Optimization Problem

Message passing is a popular technique to perform inference in graphical models [82]. In fact, as shown in [37, Section 2.3], the standard 'RL as inference' framework (Section 2.1) can recover the optimal policy using standard sum-product message passing, relying on a single backward message pass. While VAPOR does not have such a simplified derivation, the next lemma shows that it can be cast as passing some forward and backward messages back and forth. In this view, a backward message amounts to computing the optimal policy for a given reward (via backward induction), and a forward message amounts to computing the associated stationary state distribution (via the flow equations). This view stems from iteratively solving VAPOR using the Frank-Wolfe (a.k.a. conditional gradient) algorithm, where an update is equivalent to solving an MDP [29, 77, 85].

Lemma 10. For any accuracy ε > 0 , there exist forward-backward messages such that after poly ( S, A, L, σ max , ε -1 ) passes they output a λ MP ∈ Λ( P ) with max λ ∈ Λ( P ) V ϕ ( λ ) -V ϕ ( λ MP ) ≤ ε .

Proof. We consider in Algorithm 2 the Frank-Wolfe algorithm applied to maximizing the smoothed proxy V δ over λ ∈ Λ (throughout the proof we write Λ = Λ( P ) and omit the dependence on the beliefs ϕ for ease of notation). Each iteration represents one forward-backward message pass: the backward message computes the optimal policy for the reward given by the gradient of the objective V δ evaluated at the current iterate (via backward induction), while the forward message computes the associated stationary state distribution (via the flow equations). Note that the chosen number of iterations K is polynomial in the inverse of the accuracy level ε -1 , the maximum uncertainty σ max := max l,s,a σ l ( s, a ) and the MDP parameters L, S, A . To prove the ε -convergence in the original objective V , we decompose

① and ③ can be bounded using the fact that for any λ ∈ Λ , | V δ ( λ ) -V ( λ ) |≤ σ max LSAδ . ② can be bounded from the Frank-Wolfe convergence rate for smooth objective functions [31], since the proxy V δ is C δ -smooth by construction, with C δ = O ( σ max δ -4 ) . As a result, ② ≤ M δ / ( K +2) , where M δ is the curvature constant of V δ over Λ , with M δ ≤ C δ diam 2 (Λ) , where diam 2 (Λ) ≤ 2 L . Therefore, choosing δ = ε/ ( σ max LSA ) and K = ( σ max Lε -1 ) 5 ( SA ) 4 gives

F Proofs

F.1 Proof of Lemma 1

Lemma 1. For known P , it holds that P Γ ⋆ := { P ϕ (Γ ⋆ l ) } L l =1 ∈ Λ( P ) .

Proof. We omit the dependence on the beliefs ϕ for ease of notation. At the initial step,

and from Bayes' rule and the fact that P is assumed known, we have the recursive equation

Summing (7) over a ′ ∈ A implies the non-negativity and flow conservation properties of P Γ ⋆ . Moreover, when ∑ a ′ ∈A P (Γ ⋆ l ( s, a ′ )) > 0 , it holds that

which corresponds to (1).

F.2 Proof of Lemma 2

Proof. Denote by Z ⋆ the random variable of the total reward accumulated along a single rollout of the optimal policy π ⋆ , and let Z ⋆ l ( s ) be the total reward accumulated by π ⋆ from state s at layer l and Z ⋆ l ( s, a ) be the total reward accumulated by π ⋆ from state s at layer l after taking action a . We can write the total reward accumulated as

Now for arbitrary l ∈ { 1 , . . . , L -1 } ,

and unrolling we obtain

Denoting by F r the sigma-algebra generated by r , it holds from the tower property of conditional expectation that

Therefore it holds that

Now note that

Putting it all together,

F.3 Proof of Lemma 3

We prove in Lemma 11 a generalization of Lemma 3 that does not rely on the assumption of σ -sub-Gaussian mean rewards. It builds on the Donsker-Varadhan representation [23, 70, 69, 28] (Proposition 1), see also [51, Theorem 1]. We denote by KL( · || · ) the Kullback-Leibler divergence, by Ψ X : R → R the cumulant generating function of random variable X -E X and by f ⋆ : R → R the convex conjugate of function f : R → R , i.e. ,

Lemma 11. Let X : Ω → R be a random variable on (Ω , F , P ) satisfying X ∈ L 1 such that the interior of the domain of Ψ X is non-empty, and let B ∈ F be an event with P ( B ) > 0 . Then, 4

Proof. Applying Proposition 1 to λ ( X -E X ) for arbitrary λ ∈ R + with the choices Q := P ( X ∈ · ) , P := P ( X ∈ · | B ) gives

by definition of Ψ X the cumulant generating function of X -E X and by definition of the convex conjugate of Ψ X which satisfies Ψ ⋆ X ( y ) = sup λ ≥ 0 { λy -Ψ X ( λ ) } . Since Ψ ⋆ X is strictly increasing, its inverse (Ψ ∗ X ) -1 is also strictly increasing, so

which proves the first inequality. We now derive the second inequality. Note the definition of the KL divergence

where dP dQ is the Radon-Nikodym derivative of P with respect to Q . We can then write

where (i) is from the formulation of Bayes' theorem stated in [73, Theorem 1.31] and (ii) applies Bayes' rule.

Proposition 1 (Variational form of the KL-Divergence given in Theorem 5.2.1 of [23]) . Fix two probability distributions P and Q such that P is absolutely continuous with respect to Q . Then,

where E P and E Q denote the expectation operator under P and Q respectively, and the supremum is taken over all real valued random variables X such that E P [ X ] is well defined and E Q [exp X ] < ∞ .

Proof of Lemma 3. Under the σ l ( s, a ) -sub-Gaussian assumption of r l ( s, a ) , we can bound the cumulant generating function of r l ( s, a ) as

We therefore have that for any y ≥ 0 ,

Plugging this bound into Lemma 11 yields Lemma 3.

F.4 Proof of Lemma 4

Lemma 4. For known P and σ -sub-Gaussian r , we have

where the VAPOR optimization problem is concave.

Proof. The first inequality in (4) combines Lemmas 2 and 3. The second inequality comes from the fact that P Γ ⋆ ∈ Λ( P ) from Lemma 1. The VAPOR optimization problem is concave since (i) the objective function V is concave in λ by concavity of the function x ∈ (0 , 1) → x √ -log x , and (ii) the constraints are a convex set due to the fact that the state-action polytope Λ( P ) is closed, bounded and convex [66, Theorem 8.9.4]. Finally, note that the maximum exists and is attained, since Λ( P ) is closed, bounded and nonempty due to the fact that P Γ ⋆ ∈ Λ( P ) .

F.5 VAPOR without the Sub-Gaussian Assumption

From (8), for any set of reward beliefs, the VAPOR optimization problem (Lemma 4) becomes

This objective remains a concave optimization problem since the function (Ψ ∗ r l ( s,a ) ) -1 is always concave (as the inverse of the strictly increasing convex function Ψ r l ( s,a ) ).

F.6 Proof of Lemma 5

Lemma 5. Let λ ∗ solve the VAPOR optimization problem (4) and τ ∗ ∈ argmin τ V ϕ ( λ ∗ , τ ) , then

Proof. The proof extends the relationship between Bregman divergence and value in policyregularized MDPs [48] to the setting of occupancy-measure-regularized MDPs. The Bregman divergence generated by Ω between two points λ, λ ′ ∈ Λ( P ) is defined as

For ease of notation, we omit the dependence on the beliefs ϕ throughout the proof. We write

where Ω : Λ( P ) → R is a continuously differentiable strictly convex regularizer defined as

By definition, λ ∗ must satisfy the first-order optimality conditions for the maximum of V , and since λ ∗ ∈ relint (Λ( P )) the following holds

Then,

We now notice that

For any positive τ ∈ R L,S,A + , we define

Using that the KL-divergence is the Bregman divergence generated by the negative-entropy function, (9) implies that

Moreover, from strong duality (which holds because Λ( P ) is convex),

Denote τ ∗ ∈ argmin τ V ∗ τ , then it holds that

which concludes the proof.

F.7 Proof of Theorem 1 and Corollary 1

Below we prove a more general statement that will imply Theorem 1 and Corollary 1.

t

Lemma 12. Let alg produce any sequence of F t -measurable policies π , t = 1 , . . . , N , whose induced occupancy measures λ π t satisfy

where we assume that the uncertainty measure σ t of the V ϕ t function (3) satisfies

for some c 0 > 0 and any t ∈ [ N ] , l ∈ [ L ] , ( s, a ) ∈ S l ×A , where we recall that n t l ( s, a ) denotes the visitation count to ( s, a ) at step l before episode t and ( · ∨ 1) := max( · , 1) . Then it holds that

BR

(

alg

, T

)

=

√

O

˜

2

(

c

0

c

2

TSA

√

0

log(

SAT

)

.

Proof. Note that E t λ π t corresponds to the occupancy measure of policy π t on the expected transition dynamics E t P , i.e. , E t λ π t ∈ Λ( E t P ) . We bound the Bayes regret of alg as

where (i) follows from the tower property of conditional expectation where the outer expectation is with respect to F 1 , F 2 , . . . , (ii) is from the dual formulation of the value function averaged with respect to the initial state distribution [66] and the fact that λ π t l ( s, a ) and r l ( s, a ) are conditionally independent given F t , (iii) is due to condition (10), (iv) is from the definition of V ϕ t , (v) and (vi) are by applying the Cauchy-Schwarz inequality, where H denotes the entropy function [10].

We bound the left term by using the fact that H ( E t λ π t l ) ≤ log SA . We bound the right term using condition (11) and by applying the pigeonhole principle ( e.g. , [47, Lemma 6]) which gives

SA

)(1 + log

T/L

)

ϕ

which follows from the tower property of conditional expectation and since the counts at time t are F t -measurable. Taking T = NL finally yields

Theorem 1. For known P and under Assumption 1, it holds that

Proof. 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 σ t l ( s, a ) ≤ √ ( ν 2 +1) / ( n t l ( s, a ) ∨ 1) at the beginning of each episode t . Thus the condition (11) holds for c 0 = √ ν 2 +1 . Moreover, Lemma 4 implies that the occupancy measures that solve the VAPOR optimization problem satisfy condition (10). The result thus follows from applying Lemma 12.

Corollary 1. Denote by alg Γ ⋆ the algorithm that produces policies based on P t Γ ⋆ for each episode t using (1) . Then, under the same conditions as Theorem 1, we have BR ϕ (alg Γ ⋆ , T ) ≤ ˜ O ( √ SAT ) .

Proof. The result comes from combining Lemmas 4 and 12, as done in the proof of Theorem 1.

F.8 Proof of Lemma 6

Definition 2 (Transformed beliefs) . Consider any beliefs ϕ such that the mean rewards are σ -subGaussian and the transition dynamics are α -Dirichlet. We define the transformed beliefs p ϕ according to which the transition dynamics are known and equal to E ϕ P and the rewards are distributed as E ϕ r + ω , where ω l ( s, a ) follows an independent zero-mean Gaussian distribution with variance p σ 2 defined for each step l ∈ [ L ] and state-action pair ( s, a ) ∈ S l ×A as

We prove in Appendix F.8.1 the following stronger result which implies Lemma 6.

Lemma 13. For any step l ∈ [ L ] and state s ∈ S l , it holds that

F.8.1 Proof of Lemma 13

Overview. The proof relies on studying properties of two stochastic Bellman operators, drawing inspiration from [62, Section 6.5]. In particular we will make use of the notion of stochastic optimism (Definition 3). To deal with the unknown transitions, we will use the property of Gaussian-Dirichlet optimism [60]. To deal with unknown rewards, we will derive a property of Gaussian-sub-Gaussian optimism (Lemma 19) which holds when the variance of the Gaussian distribution is multiplied by a small universal constant.

Definition 3. A random variable X is stochastically optimistic with respect to another random variable Y , written X ≥ SO Y , if E u ( X ) ≥ E u ( Y ) for all convex increasing functions u : R → R .

Likewise, we will say that X is stochastically optimistic with respect to Y under ϕ and write that X | ϕ ≥ SO Y | ϕ if E ϕ u ( X ) ≥ E ϕ u ( Y ) for all convex increasing functions u : R → R .

For any step l ∈ [ L ] , we introduce the Bellman operator B l : R S l +1 × A → R S l +1 × A which is defined for any Q l +1 ∈ R S l +1 × A and state-action pair ( s, a ) ∈ S l ×A as

We also define the 'transformed' Bellman operator p B l : R S l +1 × A → R S l +1 × A as

where we recall that ω l ( s, a ) ∼ N ( 0 , p σ l ( s, a ) 2 ) is distributed independently across the state-action pairs ( s, a ) ∈ S l ×A under ϕ . Note that B l and p B can be viewed as randomized Bellman operators due to the randomness in the MDP M and in the distribution ω , respectively. We rely on two key properties.

Lemma 14 (Property 1: monotonicity) . For any step l ∈ [ L ] , consider two random Q functions Q l +1 , Q ′ l +1 ∈ R S l +1 ×A such that conditioned on ϕ the entries of Q l +1 (respectively Q ′ l +1 ) are drawn independently across state-action pairs ( s, a ) ∈ S l +1 ×A and drawn independently of the noise terms ω l ( s, a ) . Then

implies

Proof. The proof follows exactly the steps of [62, Lemma 3]. It relies on the fact that conditioned on p ϕ , p B l Q l +1 ( s, a ) is a convex increasing function of ( Q l +1 ( s, a )) s,a convolved with the independent noise term ω l ( s, a ) . The result therefore follows from the fact that stochastic optimism is preserved under convex increasing operations [62, Lemma 2].

Lemma 15 (Property 2: stochastic optimism) . For any l ∈ [ L ] and ( s, a ) ∈ S l +1 ×A , it holds that

for any fixed Q l +1 ∈ R S l +1 ×A such that Span( Q l +1 ) ≤ L -l .

Proof. For any fixed Q l +1 , it holds that

where we define

Lemma 16 ([60]) . Let Y = ∑ n i =1 A i b i for fixed b ∈ R n and random variable A , where A is Dirichlet with parameter α ∈ R n , and let X ∼ N ( µ X , σ 2 X ) with µ X ≥ ∑ i α i b i ∑ i α i and σ 2 X ≥ ( ∑ i α i ) -1 Span( b ) 2 , where Span( b ) = max i b i -min j b j , then X ≥ SO Y .

In our case, in the notation of Lemma 16, A will represent the transition function probabilities, and b will represent V := max a Q ( · , a ) , i.e. , for a given ( s, a ) ∈ S l × A let X be a random variable distributed as N ( µ X , σ 2 X ) where

due to Assumption 2. Moreover, the lemma assumes that Span( V l +1 ) ≤ L -l , so we choose σ 2 X = ( L -l ) 2 / ( ∑ s ′ α l ( s, a, s ′ )) . As a result, it holds that

Moreover, from the assumption of σ l ( s, a ) -sub-Gaussian r l ( s, a ) and from the Gaussian SubGaussian stochastic optimism property of Lemma 19,

which means that

As a result, since the random variables r l ( s, a ) , P l ( · | s, a ) ⊤ max a ′ Q l +1 ( · , a ′ ) , N ( E ϕ r l ( s, a ) , 3 . 6 2 σ 2 l ( s, a ) ) and N ( E ϕ P l ( · | s, a ) ⊤ max a ′ Q l +1 ( · , a ′ ) , ( L -l ) 2 ∑ s ′ α l ( s,a,s ′ ) ) are mutually conditionally independent, we have

which concludes the proof.

Lemma 17. For any l ∈ [ L ] and ( s, a ) ∈ S l ×A ,

Proof. Note that under ϕ , Q ⋆ 1 = B 1 B 2 . . . B L 0 and under p ϕ , Q ⋆ 1 = p B 1 p B 2 . . . p B L 0 . By Lemma 15,

Proceeding by induction, suppose that for some l ≤ L ,

Combining this with Lemma 14 shows

where the final step uses Lemma 15 and the fact that Span( B l +1 . . . B L 0) ≤ L -l .

We are now ready to prove Lemma 13. We apply the property of stochastic optimism derived in Lemma 17 to the convex increasing function u s ( Q ) := max a ∈A Q ( s, a ) for every step l ∈ [ L ] and state s ∈ S l , which yields the desired inequality E ϕ V ⋆ l ( s ) ≤ E p ϕ V ⋆ l ( s ) .

F.9 Proof of Lemma 7

Lemma 7. Let alg be any procedure that maps posterior beliefs to policies, and denote by BR T ,ϕ ( alg , T ) the Bayesian regret of alg where at each episode t = 1 , . . . , T , the policy and regret are computed by replacing ϕ t with T ( ϕ t ) (Lemma 6), then under Assumptions 1 and 2,

Proof. We recall that we denote by ϕ the prior on the MDP, by ϕ t := ϕ ( · | F t ) the posterior beliefs at the beginning of each episode t and by T the mapping that transforms the posteriors according to Lemma 6. Using the tower property of conditional expectation (where the outer expectation is with respect to F 1 , F 2 , . . . ), we can define the Bayes regret over T timesteps of alg under the sequence of original and transformed posteriors respectively as

The result then immediately comes from combining Lemmas 13 and 18.

Lemma 18. For any fixed policy π , it holds that

Proof. We prove by induction on the step l that E ϕ Q π l ( s, a ) = E p ϕ Q π l ( s, a ) for any ( s, a ) ∈ S l ×A . This is true at step L +1 . Assuming that it is true at step l +1 , then

because E ϕ Q π l ( s ′ , a ′ ) = E p ϕ Q π l ( s ′ , a ′ ) for every ( s ′ , a ′ ) ∈ S l +1 ×A by induction hypothesis. In the above we used that E ϕ [ P l ( s ′ | s, a ) Q π l +1 ( s ′ , a ′ )] = E ϕ P l ( s ′ | s, a ) E ϕ Q π l +1 ( s ′ , a ′ ) , which holds due to the time-inhomogeneity of the MDP, since the future return from a fixed state-action pair cannot be influenced by the dynamics that gets the agent to that state-action. As a result,

F.10 Proof of Theorem 2

Theorem 2. Under Assumptions 1 and 2, it holds that BR ϕ ( VAPOR , T ) ≤ ˜ O ( L √ SAT ) .

Proof. We first explicitly bound the uncertainty measure p σ t given our assumptions. The prior over the transition function P l ( · | s, a ) is assumed Dirichlet, and let us denote the parameter of the Dirichlet distribution α 0 l ( s, a ) ∈ R S l +1 + for each ( s, a ) , where ∑ s ′ ∈S l +1 α 0 l ( s, a, s ′ ) ≥ 1 , i.e. , we start with a total pseudo-count of at least one for every state-action (as done in [47]). Since the likelihood for the transition function is a Categorical distribution, conjugacy of the categorical and Dirichlet distributions implies that the posterior over P l ( · | s, a ) at time t is Dirichlet with parameter α t l ( s, a ) , where α t l ( s, a, s ′ ) = α 0 l ( s, a, s ′ ) + n t l ( s, a, s ′ ) for each s ′ ∈ S l +1 , where n t l ( s, a, s ′ ) ∈ N is the number of times the agent has been in state s , taken action a , and transitioned to state s ′ at timestep l , and note that ∑ s ′ ∈S l +1 n t l ( s, a, s ′ ) = n t l ( s, a ) , the total visit count to ( s, a ) . Meanwhile, the reward noise is assumed additive ν -sub-Gaussian, so we can upper bound the uncertainty measure σ t ∈ R L,S,A + as σ t l ( s, a ) ≤ √ ( ν 2 +1) / ( n t l ( s, a ) ∨ 1) . As a result, we can bound

The result follows from

where the first inequality stems from Lemma 7 and the second inequality applies Lemma 12 whose condition (11) holds for c 0 = √ 3 . 6 2 ( ν 2 +1) + L 2 .

F.11 Proof of Lemma 8

Lemma 8. Let λ TS be the occupancy measure of the TS policy, it holds that E [ λ TS ] = P Γ ⋆ .

Proof. Given any environment M sampled from the posterior ϕ , the optimal policy π ⋆ M induces the occupancy measure λ π ⋆ M l ( s, a ) = P M (Γ ⋆ l ( s, a )) . Marginalizing over the environment M yields E [ λ TS l ( s, a )] = ∫ M λ π ⋆ M l ( s, a ) d P ( M ) = P ϕ (Γ ⋆ l ( s, a )) .

F.12 Connection between K-learning and VAPOR

We sketch a proof of the connection between K-learning and VAPOR here. It proceeds via the dual in Lemma 9, by setting τ l ( s, a ) = τ for all l, s, a and explicitly minimizing over the V ∈ R L,S variable. First, we define the following K -function as done in the original K-learning paper [47]

for each l, s, a . After setting the τ parameters to be equal, we obtain the following dual over K ∈ R L,S × A , V ∈ R L,S , τ ∈ R ,

where we have introduced the K -function as an inequality constraint that becomes tight at the optimum and preserves convexity. Explicitly minimizing over V 1 we obtain

for each s ∈ S 1 , and doing the same for V l ( s ) , l > 1 we have

for each s ∈ S l , where µ l ( s ) is the dual variable associated with the inequality constraint and corresponds to the stationary state distribution at s . Substituting in for V we can write the K-values as

Thus we have derived a variant of K-learning, with additional terms corresponding to the stationary state distribution. These additional terms becomes entropy in the objective function, and since the sum of these entropy terms can be bounded above by L log S , it only contributes a small additional term to the regret bound. Practically speaking this version of K-learning requires knowing the stationary state distribution µ ( s ) , so it is more challenging to implement than the original K-learning. This analysis is not to suggest this algorithm as a practical alternative, but instead to demonstrate the connection between K-learning and VAPOR .

G Gaussian Sub-Gaussian Stochastic Optimism

The following technical lemma, which may be of independent interest, establishes a relation of stochastic optimism (Definition 3) between a sub-Gaussian and a Gaussian whose variance is inflated by a small universal constant.

Lemma 19. Let Y be σ -sub-Gaussian. Let X ∼ N (0 , κ 2 σ 2 ) with κ ≥ 3 . 6 . Then X ≥ SO Y .

Proof. Let u : R → R be a convex (increasing) function. We aim to show that E [ u ( X )] ≥ E [ u ( Y )] . Let s 0 ∈ R be a subgradient of u at 0 (which exists since u is convex). Then by convexity of u , we have for all x ∈ R , u ( x ) ≥ u (0)+ s 0 x . Define the function f : R → R as f ( x ) := u ( x ) -u (0) -s 0 x . We have that f ( x ) ≥ 0 for all x ∈ R and f (0) = 0 . Since E X = E Y = 0 , it suffices to show that E [ f ( X )] ≥ E [ f ( Y )] . For any ϵ ≥ 0 , define the ϵ -sublevel set of the convex function f as K ϵ := { x ∈ R : f ( x ) ≤ ϵ } . Since K ϵ is convex and by property of f , there exists a ( ϵ ) ∈ R -∪{-∞} and b ( ϵ ) ∈ R + ∪ { + ∞} such that K ϵ = [ a ( ϵ ) , b ( ϵ )] . Then, for any random variable Z ,

By the σ -sub-Gaussian property of Y ,

where we define

Let d := 2 . 69 and c := κ d ≥ 1 . 338 . Let Z ∼ N (0 , c 2 σ 2 ) and F ( x ) := P ( Z ≥ x ) . Note that

where erf denotes the error function, i.e. , erf( y ) = 2 π -1 / 2 ∫ y 0 exp( -t 2 ) d t . From Lemma 20, by choice of constants c, d , it holds that G ( x ) ≤ dF ( x ) . Therefore,

where the last inequality applies Lemma 21 to the convex function f that satisfies f (0) = 0 , and the last equality is because κ = cd which means that X and dZ follow the same distribution.

Lemma 20. Denote by c 0 > 0 the (unique) solution of the equation

and let d 0 > 0 be defined as

Note that c 0 ≈ 1 . 33016 and d 0 ≈ 2 . 68271 .

Then for any c ≥ c 0 , d ≥ d 0 , the function ϕ : R → R defined as

is non-negative everywhere.

Proof. Define

It suffices to prove that ϕ 0 ( x ) ≥ 0 for all x ∈ R . By the choices of c 0 , d 0 , it holds that

Moreover, analyzing the variations of ϕ yields that there exists x 0 > 1 such that ϕ ′ 0 ( x ) ≥ 0 for x ∈ [1 , x 0 ] , and ϕ ′ 0 ( x ) ≤ 0 otherwise. Therefore ϕ 0 is non-decreasing on [1 , x 0 ] and non-increasing elsewhere. Since lim x → + ∞ ϕ 0 ( x ) = 0 and ϕ 0 (1) = 0 , we get that ϕ 0 ( x ) ≥ 0 for all x ∈ R .

Lemma 21. For any convex function g : R → R such that g (0) = 0 , for any s ≥ 1 and x ∈ R , it holds that sg ( x ) ≤ g ( sx ) .

Proof. By convexity of g , for any x ′ ∈ R , we have g ( s -1 x ′ ) ≤ s -1 g ( x ′ ) + (1 -s -1 ) g (0) = s -1 g ( x ′ ) , so by setting x ′ = sx it holds that sg ( x ) = sg ( s -1 x ′ ) ≤ g ( x ′ ) = g ( sx ) .

H VAPOR-lite Analysis

Denoting the uncertainty signal of the VAPOR optimization problem by p σ , we define the VAPOR-lite ( c ) optimization problem for any multiplicative scalar c > 0 as follows

Note that this corresponds to the same objective as (5) except that it is optimized over the occupancy measures instead of the policies.

Lemma 22. The VAPOR-lite ( c ) optimization problem is concave in λ .

Proof. The objective is concave due to the fact that the normalized entropy function h of u ∈ R n > 0 given by h ( u ) = ∑ n i =1 u i log( 1 ⊤ u/u i ) is concave (see [6, Example 3.19]) since the perspective operation preserves convexity.

Lemma 23. For any η > 0 and λ ∈ Λ( E ϕ P ) , it holds that

Proof. Any occupancy measure λ ∈ Λ( E ϕ P ) can be decomposed as λ l ( s, a ) = µ π l ( s ) π l ( s, a ) , where µ π denotes the stationary state distribution under π , i.e. , µ π l ( s ) := ∑ a ∈A λ l ( s, a ) , and π is the policy given by π l ( s, a ) := λ l ( s, a ) /µ π l ( s ) so long as µ π l ( s ) > 0 , otherwise π l ( s, · ) can be any distribution, e.g. , uniform. It holds that

where (i) uses that √ -log x ≤ 1 -log( x ) for any 0 < x ≤ 1 , and (ii) uses that -x log x ≤ x log 1 /η + η for any 0 ≤ x ≤ 1 and η > 0 . Hence, we have

Lemma 24. Consider the algorithm VAPOR-lite that at each episode t = 1 , . . . , N solves the optimization problem VAPOR-lite ( c = √ 2(1 + log SLt )) . Then under Assumptions 1 and 2, it holds that

Proof. Denote by π t the policy executed by the algorithm at episode t , and by E t P Γ ⋆ the posterior probability of state-action optimality in the expected model E t P . It holds that

where (i) combines Lemmas 4 and 6, (ii) applies Lemma 23 to E t P Γ ⋆ ∈ Λ( E t P ) with the choice η = 1 / ( SLt ) , and (iii) comes from the fact that E t λ π t ∈ argmax λ ∈ Λ( E t P ) U p ϕ t ( λ, 1 / ( SLt )) . Hence, retracing the steps of the regret derivation of the proof of Lemma 12, we can bound the Bayes regret of VAPOR-lite as

To bound Z 1 we apply the Cauchy-Schwarz inequality and the pigeonhole principle, as in the proof of Lemma 12, which gives that Z 1 = ˜ O ( p σ max √ SAT ) . To bound Z 2 , we introduce the second moment of the information [33] of a discrete random variable X supported on finite set X as

and note that H (2) ( X ) ≤ log 2 ( |X| ) for |X| ≥ 3 [33, Proposition 8]. We bound

where (i) uses that E t λ π t l ( s, a ) ≤ π t l ( s, a ) , (ii) applies the Cauchy-Schwarz inequality and (iii) uses the pigeonhole principle (as in the Z 1 bound) combined with the aforementioned logarithmic upper bound on H (2) . Finally, we bound the last term by using that ∑ N t =1 1 t ≤ 1 + log N . Putting everything together and using that p σ max = O ( L ) yields the desired bound BR ϕ ( VAPOR-lite , T ) ≤ ˜ O ( L √ SAT ) .

10000

MaxEnt

Weighted MaxEnt

State-action coverage time

8000

6000

4000

2000

0

0

50

100

150

200

GridWorld size

I Experimental Details

I.1 Reward-Free GridWorld

We empirically complement our discussion in Section 6 on the connection of VAPOR to maximum entropy exploration. We consider a simple 4-room GridWorld domain with no rewards, known deterministic transitions, four cardinal actions and varying state space size. We measure the ability of the agent to cover the state-action space as quickly as possible, i.e. , visit each state-action at least once. We consider the state-of-the-art algorithm of [79] for the original maximum entropy exploration objective max λ ∈ Λ( P ) H ( λ ) , and compare it to the same algorithm that optimizes for weighted entropy max λ ∈ Λ( P ) H σ ( λ ) , where we define the uncertainty as σ ( s, a ) = 1 { n ( s, a ) = 0 } , with n ( s, a ) the visitation count to state-action ( s, a ) . We see in Figure 7 that optimizing for this weighted entropy objective results in faster coverage time. This illustrates the exploration benefits of the state-action weighted entropy regularization of VAPOR ( e.g. , weighted by the uncertainty as done in VAPOR-lite ).

I.2 DeepSea (tabular)

We provide details on the experiment of Figure 3. For the transition function we use a prior Dirichlet (1 / √ S ) and for rewards a standard normal N (0 , 1) , as done by [53]. Similar to [62, 47, 40], we accelerate learning by imagining that each experienced transition ( s, a, s ′ , r ) is repeated 100 times. Effectively, this strategy forces the MDP posterior to shrink faster without favoring any algorithm, making them all converge in fewer episodes.

VAPOR implementation. The VAPOR optimization problem is an exponential cone program that can be solved efficiently using modern optimization methods. We point out that we do not need to solve the two-player zero-sum game between λ and τ since the minimization over τ conveniently admits a closed-form solution, see Equation (3). In our experiments, we use CVXPY [11], specifically the ECOS solver [12] (with 1e-8 absolute tolerance). In terms of runtime, this took a few seconds on the largest DeepSeas, which was sufficient for our purposes. There is a natural trade-off between a less accurate optimization solution (i.e., better computational complexity) and a more accurate policy (i.e., better sample complexity), which can be balanced with the choice of CVXPY solver, its desired accuracy, its maximum number of iterations, etc. The implementation of the VAPOR optimization problem in CVXPY is straightforward using the library's atomic functions. Just note that although the expression x √ -log x that appears in the V ϕ function (3) is not an atomic function in CVXPY, a simple trick is to use the equivalence

where -x log x (cvxpy.entr ( x ) ) and y 2 /x (cvxpy.quad_over_lin ( y, x ) ) are atomic functions.

VAPOR-lite replay actor-critic

100000

90%bad episodes

75000-

solved

V

50000

True

#episodes until

25000

25

50

75

100

deep sea problem size

I.3 DeepSea (neural network)

We consider the DeepSea domain where instead of using a tabular state representation, we feed a one-hot representation of the agent location into a neural network, using bsuite [58]. As discussed in prior works [62, 53, 49], agents that do not adequately perform deep exploration are far from being able to solve depths of up to 100 within 10 5 episodes. This includes vanilla actor-critic or Soft Q-learning (which can only solve depths up to around 14 , showcasing a learning time that suffers from an exponential dependence on depth), but also agents with some exploration mechanisms such as optimistic actor-critic or Bootstrapped DQN (depths up to around 50 ). In contrast, Figure 8 shows that VAPOR-lite is able to solve DeepSea instances out to size 100 within 4 × 10 4 episodes, without a clear performance degradation. We chose the exact same algorithmic configuration and hyperparameter choices as our experiments on Atari (Appendix I.4), except σ scale = 3 . 0 and Replay fraction = 0 . 995 . This experiment shows that VAPOR-lite is capable of deep exploration , as suggested by the regret analysis in Appendix H.

I.4 Atari

RL agent. Our setup involves actors generating experience and sending them back to a learner, which mixes online data and offline data from a replay buffer to to update the network weights [30]. Our underlying agent is an actor-critic algorithm with use of replay data, thus we call it 'Replay Actor-Critic (RAC)'. We also made use of V-trace clipped importance sampling to the off-policy trajectories [15]. Replay was prioritized by TD-error and when sampling the replay prioritization exponent was 1 . 0 [72]. We note that, as commonly done on Atari, we consider the discounted RL setting (rather than the finite-horizon setting used to derive our theoretical results). We refer to Table 2 for hyperparameter details.

Uncertainty measure in VAPOR-lite . For the uncertainty measure p σ , we use an ensemble of reward predictors. Specifically, we set it to the standard deviation of H = 10 randomly initialized reward prediction heads p r ( i ) with randomized prior functions [56], i.e. ,

with σ scale > 0 a scaling hyperparameter. We also added a small amount of noise to the rewards in the replay buffer for the targets for the reward prediction ensemble (specifically, 0 . 1 times an independent standard normal distribution) [14], to prevent the ensemble from potentially collapsing due to the use of replay data.

0.75

0.50

1

alien

0.4

0.6-

assault

60-

40

asterix

0.08

0.06-

0.04

asteroids

40 -

60

atlantis

0.25

0.2

2

20

0.02

20

0

0

0

0

0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

bank_heist

battle_zone

beam_rider

bowling

boxing

1.2

1

80

1.5

0.10

0.8

0.75

0.6 -

1

0.05

Mw

6

0.4

0.25 -

0.50 -

0.4-

0.2-

0.5

F0

0-

3

0

0

0

0

0.05

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

10-

breakout

2

centipede

chopper_command

5

crazy_climber

12.5

demon_attack

7.5

1.5

3

10-

30-

5 -

1

7.5

5

20-

2.5

0.5

1-

2.5

10

-0

0

0

0

0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

double_dunk

enduro

fishing_derby

1.2

freeway

frostbite

20

gopher

15 -

10-

5

2

1

2-

1

0.9

0.6

0.3

2-

15

10

5

0-

-0

0

0-

0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

gravitar

hero

ice_hockey

jamesbond

kangaroo

score

0.2

1 -

0.5

1

20

3

2

20

Human-normalized

0.1

0.5

F0

10-

1

10

Mmy

0-

0

0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

8

kung_fu_master

montezuma_revenge

0.8-

ms_pacman

name_this_game

0.04

pitfall

4

0.2

0.4

0.3

0.1-

0.6

0.2

0.4 -

3

2

1-

7.5 -

2.5 -

5

00

0.02-

-0

0

0

0

0.01

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0.9

0.3

0.6-

1.2

pong

0.10

0.05

private_eye

0.75 -

0.50 -

0.25-

W

qbert

1.2 -

0.8-

0.4

riverraid

20

15

10-

5

road_runner

F9

4 -

2

robotank

0

0-

0

F0

0

0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0.05

00

0.02

0.04

0.01 -

seaquest

0.4

0.4

0-

M

sking

0.1 -

F0

solaris

space_invaders

7.5

2.5

10-

5

star_gunner

0.75 -

0.50

0.25

1

surround

0

0.8

0.1

0

0

0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

3

tennis

5

time_pilot

1.5

15

up_n_down

60

venture

video_pinball

2

2.5

10-

0.6-

20

FL

F0

0.5 -

5

W

0.3

10 -

0-

0

0

-0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

wizard_of_wor

yars_revenge

zaxxon

4-

1.5

0.5

2

1-

2

2

3

1-

RAC+VAPOR-lite

Replay Actor-Critic(RAC)

o

0

-0

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

0e+00

2e+08

5e+08

Frames

Median human-normalized score

1.5

1 -

Entropyregularization

0.5

none

fixed，scalar

tuned,scalar

tuned,state-action

0-

0e+00

1e+08

2e+08

3e+08

4e+08

5e+08

Frames

Baselines. In Figure 4, we compare VAPOR-lite , whose policy loss may be condensely expressed as

to a standard actor-critic objective

as well as to an actor-critic with fixed scalar entropy regularization

where β > 0 is a hyperparameter ( e.g. , β = 0 . 01 in [42], which is also used in Figure 4 as it performed best). We later compare to additional baselines (Figure 10).

On entropy regularization for policy improvement/evaluation. We refer to [84] for in-depth discussion on the difference of using entropy for policy improvement and/or for policy evaluation in entropy-regularized RL. Theory prescribes both [87] and some algorithms such as Soft Actor-Critic [26] implement both, yet as remarked by [84], omitting entropy for policy evaluation ( i.e. , not having entropy as an intrinsic reward) tends to result in more stable and efficient learning. We also observed the same and thus only use entropy for policy improvement.

Related works. We mention here a couple of related works to our proposed algorithm. [22] argue that standard policy entropy regularization can be harmful when the RL problem contains actions that are rarely useful, and propose a method that uses mutual-information regularization to optimize a prior action distribution. [27] discuss the limitation of policy entropy regularization being 'sampleunaware', and propose to regularize with respect to the entropy of a weighted sum of the policy action distribution and the sample action distribution from the replay buffer.

Per-game results. Figure 9 reports the performance of the Replay Actor-Critic (RAC) baseline compared to RAC augmented with VAPOR-lite, on each Atari game. We see that VAPOR-lite yields improvements in many games, especially in hard-exploration games such as Montezuma's Revenge or Venture.

Ablation on entropy regularization. VAPOR-lite advocates for two ingredients: optimism and weighted entropy regularization . While the first has already been investigated in deep RL for various choices of uncertainty bonus [4], the second is novel so we focus our attention on it here. For a fair comparison, we augment all agents with the optimistic reward component, to isolate the effect of the adaptive state-action entropy regularization of VAPOR-lite . We consider the following policy losses for various choices of entropy regularization

In particular, we have further compared to an optimistic actor-critic objective with tuned, scalar entropy regularization, specifically the epistemic-risk-seeking objective of [49] which corresponds to solving a principled saddle-point problem between a policy player and a scalar temperature player τ parameterized by a neural network. We refer to Table 2 for hyperparameter details. We observe in Figure 10 the isolated benefits of VAPOR-lite 's tuned, state-action entropy regularization.

Future directions. Below we list some relevant directions for future empirical investigation:

Although our simple p σ uncertainty measure using an ensemble of reward predictors worked well, more sophisticated domain-specific uncertainty signals could be used [63, 64, 7].

A first 'looseness' of VAPOR-lite with respect to VAPOR is that it sets the temperatures τ to the uncertainties p σ (which yields the same ˜ O regret bound). An interesting extension for VAPOR-lite could be to optimize directly the temperatures by minimizing VAPOR 's saddle-point problem with a state-action dependent temperature player τ parameterized by an independent neural network.

A second 'looseness' of VAPOR-lite with respect to VAPOR is that it bypasses the challenging optimization of the entropy of the state visitation distribution, by only (weight-)regularizing with the entropy of the policy. Various attempts have turned to density models or non-parametric entropy estimation [29, 36, 43], and it could be relevant to incorporate such techniques to get closer to the VAPOR objective.

Our agent is relatively simple compared to modern state-of-the-art Atari agents since it is missing components like model-based rollouts, distributional heads, auxiliary tasks, etc. An interesting extension would be to incorporate the techniques discussed in this paper into the most effective policy-based agents.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a model-based RL algorithm for MDPs with unknown rewards and transition dynamics, which approximates the posterior probability of an action being optimal in a given state to ensure efficient exploration. The approach is mostly contrasted with model-free RL as inference, but there is also an experiment comparing with a Thompson sampling-based baseline.\n\nEDIT AFTER DISCUSSION\nI revised my score from 3 to 7.\n\nSTRENGTHS:\nThe presentation is clear and the proposed method is easy to understand.\n\nWEAKNESSES:\nI think the abstract is rather misleading and the authors are missing a large portion of existing literature on the topic. RL as inference, as presented by Levine [30], is a model-free algorithm. Any shortcomings of this approach notwithstanding, it is simple and computationally efficient, as model-free methods generally are. The abstract of this submission promises that the authors fix some problems with RL as inference, when in fact they just use a model to compute optimal actions instead. Sure, in some sense one could argue that this is better, but not exactly novel. The RL as inference framework is thus only tangentially relevant to the proposed method and the authors should compare with model-based approaches instead. There is some discussion of the relationship to PSRL in Section 6,  but way too little and too late.\n\nI don't do model-based RL myself, so I'm not really in a position to evaluate how the proposed method relates to existing approaches in that space. I also don't really know the right papers to use as baselines, but here are a few articles I found with a quick search in case that's helpful.\nAsmuth, J. and Littman, M. Learning is planning: Near Bayes-optimal reinforcement learning via Monte-Carlo tree search. In UAI, 2011.\nAsmuth, J., Li, L., Littman, M. L., Nouri, A., and Wingate, D. A Bayesian sampling approach to exploration in reinforcement learning. In UAI 2009, Proceedings of the Twenty-Fifth Conference on Uncertainty in Artificial Intelligence, 2009, pp. 19–26. AUAI Press, 2009.\nGuez, A., Silver, D., and Dayan, P. Efficient bayes-adaptive reinforcement learning using sample-based search. In NIPS, 2012.\nStrens, M. J. A. A Bayesian framework for reinforcement learning. In Proceedings of the Seventeenth International Conference on Machine Learning (ICML 2000), pp. 943–950, 2000.\n\nQUESTIONS:\nWhere does this paper stand with respect to existing literature on Bayesian approaches to model-based RL? Only once this question is answered exhaustively would it be possible to review this paper. I don't see how that's going to be possible within this cycle, though.\n\nLIMITATIONS:\nUnclear, as this paper is not properly positioned with respect to existing work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper views reinforcement learning as a Bayesian variational inference problem, and proposes VAPOR, an algorithm which computes an approximately optimal occupancy measure and its corresponding policy. Bayesian regret of VAPOR is analyzed to have a sub-linear bound. In numerical experiments, VAPOR shows good performance in simple GridWord and DeepSea environments. For more complex Atari environments, a further approximated VAPOR-lite is used and is compared with its entropy-regularized counterpart.\n\nSTRENGTHS:\n- While the idea of \"RL as inference\" has been considered, this paper proposes a new approach which aims to directly approximate the optimal occupancy measure in the Bayesian setting. This approach has the potential to better approximate Bayesian optimal reinforcement learning policy which might handle hard exploration situations.\n\n- The motivation of the algorithm is explained well. The main idea of VAPOR is based on inequality (5) in Lemma 4. This inequality provides an upper bound of the optimal Bayesian performance, and the objective of VAPOR is basically finding the policy achieving the upper bound.  Some analysis of the approximation is provided in Lemma 5.\n\n- The performance of VAPOR is analyzed in terms of Bayesian regret in Theorem 1 with known transitions, and then the analysis is extended in Theorem 2 under the independent Dirichlet assumption on the transition dynamics. Bayesian regret is shown to be sub-linear in both cases.\n\n- VAPOR shows good performance compared with Thompson sampling based methods in simple GridWord and DeepSea environments. \n\n- For more complex Atari environments, a further approximated VAPOR-lite is used and is compared with its entropy-regularized counterpart.\n\nWEAKNESSES:\n- Although Lemma 5 provides a bound on the KL-divergence between the VAPOR solution and the optimal occupancy measure, the bound is in terms of the gap of the VAPOR approximation of the optimal Bayesian performance in (5). So this bound is basically using a property of VAPOR itself to bound another property of VAPOR and none of them can be evaluated. Therefore, beside numerical experiments, we have no idea how loose the upper bound (5) can.\n\n- Since the considered dynamics are time-inhomogeneous, the space of occupancy measure VAPOR is solving has dimension growing in the time horizon. This means that VAPOR might not be able to handle problems with long horizon, and this contradicts the statement at the end of Section 5 that VAPOR produces stationary policies, while it doesn't due to the time-inhomogeneous nature of the occupancy measure.\n\n- In VAPOR, $\\sigma$ seems to be a hyper-parameter, but it's not clearly defined. In particular, $\\sigma$ seems to be a constant in most lemma and theorem statements, but $\\sigma$ looks like a multi-dimensional vector with entry $\\sigma_l(s, a)$ in (2) and (4). This confusion makes the description of the algorithm inconsistent and it might lead to some errors.\n\n- The issue of $\\sigma$ mentioned above might lead to a major error in the regret analysis and the VAPOR algorithm. In the proof of Lemma 12, equation (14) implies $c_l(s, a) = \\sigma_l(s, a) \\sqrt{n^t_l(s, a) \\vee 1}$. According to (14), $c_{\\max} = \\max_{l, s, a} \\sigma_l(s, a) \\sqrt{n^t_l(s, a) \\vee 1}$. Therefore, unless $\\sigma_l(s, a)$ decrease at least at the rate of $1/\\sqrt{n^t_l(s, a) \\vee 1}$, the value of $c_{\\max}$ will not be a constant.  But from the statement of Theorem 1, $\\sigma$ seems to be treated as constant which is inconsistent with the proof of Lemma 12. It seems like, the uncertainty measure $\\sigma$ needs to follow a decay procedure. Gradually decreased uncertainty is generally required for a policy to have sub-linear regret because non-decreasing randomness would lead to linear regret due to the constant noise in policy. \n\n- Lemma 7 is very counter-intuitive. If it's true in general, uncertainty in the transition dynamics can be easily handled by replacing it with its mean. This would establish a kind of certainty equivalence for MDP which seems too strong to be true. I try to follow the proof of Lemma 7 and found several issues. Those issues may be correctable, but even if Lemma 7 is correct, it is probably too strong due to Assumption 1. Although Assumption 1 has been used in the literature, this assumption is restrictive and is not applicable in most applications. For example, any MDP with an unknown but fixed (over times/steps) transition function violates the assumption because the transition functions are not independent over times/steps. \n\n- VAPOR-lite sounds promising with its application in complex environments like Atari games, but the paper only provides very limited information for VAPOR-lite and there is no details available for its implementation. The numerical results are also very limited, with no results on individual games and no comparison with popular existing algorithms.\n\nQUESTIONS:\n- As described in the weakness, is $\\sigma_l(s, a)$ a constant or it follows a decay process? This seems critical to the algorithm and its analysis. It would be great if the authors can clarify the role of $\\sigma_l(s, a)$ and correct the proofs one way or another.\n\n- There are several issues in the lemmas that lead to the proof of Lemma 7. \n  - In the proof of Lemma 13, summations are missing from Bellman operators. It seems like the authors may want to use some short-hand notations for those summations, what are the definitions for those notations? It's not possible to follow the proof without proper descriptions and definitions.\n  - In the proof of Lemma 14, the second sentence claims to be conditioned on $\\phi$, but $\\hat{\\mathcal B}_l$ is defined using the distribution of $\\hat \\phi$ instead of $\\phi$. Is there some typos for $\\hat \\phi$ and $\\phi$? Either way, the proof requires more steps transiting from $\\phi$ to $\\hat\\phi$ to be true.\n  - In the proof of Lemma 18, why can one brings the expectation $\\mathbb E_\\phi$ inside the Q-function? This seems to be a consequence of independent Dirichlet, but one needs to show that it leads to some kind of independence for the Q-function.\n\nLIMITATIONS:\nBeside possible errors in the analysis, the assumption makes the analysis in the paper limited and there is no discussion on this aspect.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a Bayesian treatement for the reinforcement learning as probabilistic inference framework for the discrete state/action space. As extension to the standard formulation, posterior probabilities of state-action optimalities are regarded and formally defined. For tractability, an approximation of the probabilities via variational inference is derived. Further, the case of unknown dynamics is considered and links to other approaches discussed. The authors evaluate their method on a grid world problem, deep sea, and (in an reduced form) to atari.\n\nSTRENGTHS:\nThe problem which is considered in the paper is interesting for the use in state-of-the-art reinforcement learning. While not being obvious at first sight why a Bayesian treatment is advantageous, the simple decision problem makes that very clear. \n\nBy providing a formalization of the Bayesian formulation and the variational approximation and details about the algorithm the contribution of the paper is very good.\n\nThe paper is well-written and easy to follow.\n\nWhile just briefly looking at the appendix, all theorems seem to be proven and assumptions listed.\n\nWEAKNESSES:\nThere is no code provided in the supplementary material or paper. In these days, I would expect this for an accepted paper at NeurIPS.\n\nFuture work and limitations are not sufficiently discussed in my opinion.\n\nFor atari evaluation, only a reduced version of the algorithm is considered (without weighted entropy)\n\nInformation on how to solve the VAPOR objective are only given in the appendix.\n\nThe paper seems very crowded, for example the table besides lemma 6 is a bit confusing. It is good, however, that it has a lot of content/contribution.\n\nMinor:\nThe quantity $\\lambda$ seems not to be well-defined for me as it is used a function but not defined this way. It took me a while to understand what it really is. This should be written in a more consistent way.\n\nline 110: Add \"pair\" after state-action.\n\nQUESTIONS:\nHow strong is the assumption for VAPOR that $r_l$ is a sub-Gaussian (in Lemma 3)?\n\nIs it possible to apply the method to continuous state/action spaces?\n\nWhat are the runtimes for the experiments?\n\nLIMITATIONS:\nThe formulation is limited to discrete/action spaces and does not seem to be scalable.\n\nNo runtimes are provided. It would have been very interesting to see how much time it takes to solve the two-player zero-sum geme in every iteration.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper undertakes a rigorous Bayesian treatment of the posterior probability of state-action optimality. It proposes a variational approach to approximate the state-action optimality. The proposed method involves a tractable convex optimization problem and is provably efficient. This paper also conducts experiments showing that the proposed method compares favorably with previous methods.\n\nSTRENGTHS:\nThis paper, for the first time, undertakes a rigorous Bayesian treatment of the posterior probability of state-action optimality. It proposes a novel method having deep connections to previous work. Related work has been adequately cited.\n\nThis paper seems to be technically sound, though I did not check the proof in Appendix. All claims are well supported by theoretical analysis and experimental results. It is also clearly written and well organized.\n\nI believe the result presented in this paper is significant. What I find most interesting is that the proposed method finds a balance between optimism and entropy regularization, and the resulting reward bonus is the product of surprise and uncertainty.\n\nWEAKNESSES:\nThis paper formally treats the posterior probability of state-action optimality and demonstrates its connection to Thompson sampling. Similar analysis has been given in [O’Donoghue et al.](https://openreview.net/pdf?id=S1xitgHtvS). This should be mentioned in the paper.\n\nThe proposed method is novel. However, its connection to previous work, especially K-learning, should be discussed in more depth. Specifically, K-learning could be understood as a variational approximation of the policy induced by the state-action optimality ([O’Donoghue et al.](https://openreview.net/pdf?id=S1xitgHtvS)). VAPOR proposed in this paper additionally optimizes the stationary state distribution, which, to my understanding, complicates the optimization by introducing the need for forward message passing. It is not discussed in this paper whether and why this is beneficial.\n\nIn Figure 2, this paper compares the stationary state distribution induced by VAPOR and TS. I believe it should additionally provide a comparison with K-learning as K-learning also approximates the policy induced by the state-action optimality.\n\nThis paper provides a performance comparison in environments like DeepSea and Atari. As Bayesian approaches, it should also compare the Bayesian regret by simulating in MDPs randomly generated from prior.\n\nQUESTIONS:\nIs there a theoretical benefit for optimizing both the policy and stationary state distribution?\nIs there an intuitive explanation for the balance between optimism and entropy regularization?\n\nLIMITATIONS:\nThis paper has not discussed the limitations of the proposed method and the theoretical analysis.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
The paper identifies "the posterior probability of state-action optimality" denoted $\mathbb{P}_{\Gamma^*}$ as a key object for inference and control and provides a variational optimization approach to estimate it. Clear presentation, insightful analysis, and experiments on GridWorld, DeepSea, and Atari are provided.

STRENGTHS:
Outstanding paper in all respects. A fresh and principled approach to RL as inference. High quality execution both on the theory and experimental side.

WEAKNESSES:
None

QUESTIONS:
None

LIMITATIONS:
There is no dedicated section for limitations, but limitations are sufficiently discussed in the conclusion.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06