Bayesian Exploration Networks

We thank the reviewer for their comments and detailed review and thank them for pointing out typos and points of clarification. We respond to detailed feedback regarding the validity of our mathematical results over two separate comments:

The reviewer questions the validity of the final line and states: `The MDP parameter $\phi$ doesn't impact future decisions made by the contextual policy, which depend solely on history.'

We're not sure what the reviewer means by this. If they are claiming our formulation is myopic, this can be proved as not the case by constructing the equivalent formulation using the myopic transition distributions. In the myopic formulation, the current posterior at timestep $t$ is used to marginalise across all future timesteps $t'>t$:

\int  p^\pi_2(\tau_2\vert \phi)  p(\phi) d\phi =\int p_0(s_0)\pi(a_0\vert  s_0)\pi(a_1\vert h_1)   p(s_2,r_1\vert s_1,a_1,\phi) p(s_1,r_0\vert s_0,a_0,\phi) p(\phi) d\phi,$$

Applying the same proof but conditioning on $h_t,a_t$ recovers our result from Lemma 1. We are more than happy to include a full step by step derivation if the reviewer still feels it is necessary.

MLE Estimator:

We would like to emphasise that Theorem 2 is an example demonstrating the need to model aleatoric uncertainty. The issues that the reviewer takes can be resolved by assuming that the (empirical) projection exists and is unique:

$

\phi^\star_\textrm{MLE}(h_t)\in\arg\inf_{\phi\in\Phi}\left(\sum_{i=0}^{t-1} (b_i-B(h_i,a_i,\phi)^2\right).

$

This is an assumption that is commonly (and most typically implicitly [4,5,6]) made in works that analyse fitted/projection operator methods [1,2,3]. Again, we emphasise that the purpose of the theorem is to show why modelling aleatoric uncertainty is essential, as even under somewhat contrived assumptions, the best we can hope for is a policy that has no exploratory properties. We thank the reviewer for bringing this to our attention.

Sufficiency of $b_t$:

We would like to clarify what we mean here: our goal is to find an optimal policy and we propose that either $b_t$ or $r_{t}$ and $s_{t+1}$ is a sufficient statistic for solving the underlying RL problem (that is finding $\pi^\star$), not that $b_t$ is a sufficient statistic for $r_{t}$ and $s_{t+1}$. This must hold or else model-free methods in general, which estimate $b_t$, would not be sufficient for learning optimal policies. We recognise that this is confusingly worded and thank the reviewer for bringing this to our attention. We will amend this and emphasise this key difference in an updated version.

Related Work:

We are aware of the review of meta RL cited by the reviewer, however we are not sure which methods they are referring to. We take care to emphasise that naively applying ANY POMDP solver to solve a BAMDP when taking a model-based approach, regardless whether the solver itself uses a model-free or a model-based approach, is not feasible for reasons of tractability. The only tractable methods that attempt this are based on VariBAD, which, as we explain, forfeit being Bayesian over part of their model. We took the time to carefully explain that model-free Bayesian approaches model uncertainty in the Bellman operator (see Definition 1). None of the methods presented in the review do this. We emphasise that just because a POMDP solver uses model-free methods to solve the BAMDP does not mean that it is a model-free approach to Bayesian RL. Model-free Bayesian RL approaches characterise uncertainty in the variable $b_t$. We are happy to add the review to our list of citations.

[1] Antos et al.Learning near-optimal policies with Bellman-residual minimization based fitted policy iteration and a single sample path. Machine Learning, 07.

[2] Antos et al. Fitted Q-iteration in continuous action-space MDPs. NeurIPS 07

[3] A. Antos et al. Value-iteration based fitted policy iteration: learning with a single trajectory. IEEE 07

[4] Riedmiller, Neural Fitted Q Iteration - First Experiences with a Data Efficient Neural Reinforcement Learning Method, ECML 05

[5] Sutton et al. Fast Gradient-Descent Methods for Temporal-Difference Learning with Linear Function Approximation. ICML 09

[6] Maei et al. Convergent Temporal-Difference Learning with Arbitrary Smooth Function Approximation. NeurIPS 09

We think the confusion may be coming from the fact Lemma 1 holds only because we are restricting ourselves to contextual policies. For simplicity, consider the space of contextual optimal policies, as this is what Theorem 1 analyses. As we write in our paper, a contextual optimal policy is defined as:

$

\pi^\star(\cdot,\phi)\in \arg \max_{\pi\in\Pi_\Phi} J^\pi(\phi).  

$

Given an MDP indexed by $\phi$, we can always find a contextual optimal policy by taking the optimal action $a\in\arg\max_{a'} Q^\star(s,a, \phi)$, where:

$

Q^\star(s,a,\phi)=E_{\tau_\infty \sim P_\infty^\star (s,a,\phi)} \left[\sum_{i=0}^\infty \gamma^i r_i \right].

$

and:

$

p_\infty^\star (\tau_\infty\vert s,a,\phi)=\prod_{i=0}^\infty p(r_i\vert s_i,a_i,\phi)p(s_{i+1}\vert s_i,a_i,\phi) \pi^\star(a_{i+1}\vert s_{i+1},\phi).

$

Finding $Q^\star(s,a,\phi)$ can be achieved by solving the contextual optimal Bellman equation $B[Q^\star](s,a, \phi)=Q^\star(s,a, \phi)$ where:

$

B[Q^\star](s,a,
\phi)=E_{s',r\sim P_{R,S}(s,a,\phi)}[r+\gamma\sup_{a'} Q^\star(s',a',\phi)].

$

We can thus (in principle) learn a set of contextual optimal policies by learning $Q^\star(s,a, \phi)$ and $\pi^\star(\cdot,\phi)$ for all $\phi\in \Phi$ by solving the afformentioned optimal contextual Bellman equation. Existing model-free BRL methods then characterise the uncertainty in $B[Q^\star](s,a, \phi)$ by first sampling $s',r\sim P_{R,S}(s,a,\phi^\star)$ from the environment and then applying the transformation $b=r+\gamma\sup_{a'} Q^\star(s',a',\phi)$. A posterior over $\phi$ can be inferred $p(\phi\vert h_t)$, which is used to marginalise across all optimal Bellman operators $B[Q^\star](h_t, a_t)=E_{\phi \sim {P_\Phi}( h_t) } [ B [ Q^\star ] ( s_t,a_t, \phi) ]$

In practice, existing methods partially solve the corresponding contextual Bellman equation online $B[Q^\star(h_t,a_t)=Q^\star(h_t,a_t)$ after inferring each $p(\phi\vert h_t)$ using variational inference and employing posterior sampling to yield a tractable algorithm. Regardless, the goal of these approaches is still to find:

$

B[Q^\star](h_t,

a_t)= E_{\phi\sim P_\Phi( h_t) }[B[Q^\star](s_t,a_t, \phi)]=E_{\phi\sim P_\Phi( h_t)}[Q^\star(s_t,a_t,\phi)].

$

Now:

$

E_{\phi\sim P_\Phi( h_t) }[Q^\star(s_t,a_t,\phi)]=	E_{\phi\sim P_\Phi( h_t) }\bigg[E_{\tau_{t:\infty} \sim P_{t:\infty}^\star ( s_t,a_t,\phi)} \bigg[\sum_{i=t}^\infty \gamma^i r_i \bigg]\bigg]
=	E_{\tau_\infty \sim P_{t:\infty}^\star (h_t,a_t)} \bigg[\sum_{i=0}^\infty \gamma^i r_i \bigg],

$

where:

$

p_{t:\infty}^\star (\tau_{t:\infty}\vert h_t,a_t)= \int \prod_{i=0}^\infty p(r_i\vert s_i,a_i,\phi)p(s_{i+1}\vert s_i,a_i,\phi) \pi^\star(a_{i+1}\vert s_{i+1},\phi) dP_\Phi(h_t).

$

and so the proof of Lemma 1 holds, but only because we are limiting ourselves to the space of contextual optimal policies.

We thank the reviewer for their feedback, especially regarding the structure of the paper. Due to the novelty of our approach, we have had often contractionary feedback about how to best present our work. We do really like the suggestion to shorten the background material and splitting up the theoretical and implementation details and will re-structure the paper according to this proposal. With regards to experiments, not all of our implementations are one-shot tasks. As discussed in Section 6 we evaluate BEN in the episodic setting for the search and rescue task too. We focused on the one-shot setting for our ablations as this more closely mimics what a rescue robot would have to achieve in real life (including being given training knowledge in the form of simulated related situations) and it highlights the important failures of existing model-free approaches to BRL.

We really thank the reviewer for their feedback. We are glad that they enjoyed reading the paper and appreciated their thoughtful criticisms. In response to their question about when we would genuinely seek to avoid existing model-free BRL methods, we believe that BEN offers a Pareto improvement on these methods and many can be recovered by using a linear flow. We believe that there might be slight benefits from a practitioner's perspective in using existing approaches: implementing BEN requires expertise in normalising flows whereas existing methods are slightly simpler. For toy domains where exploration is difficult but the environment's dynamics are simple enough to avoid the pathological examples outlined in our work, for example in domains like Mountain Car, existing approaches will suffice.