# Reason for Future, Act for Now: A Principled Framework for Autonomous LLM
Agents with Provable Sample Efficiency

Zhihan Liu Hao Hu11footnotemark: 1 Shenao Zhang11footnotemark: 1
22footnotemark: 2
Hongyi Guo22footnotemark: 2 Shuqi Ke Boyi Liu22footnotemark: 2 Zhaoran
Wang22footnotemark: 2 Equal contribution.Northwestern University.
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>University<EMAIL_ADDRESS>Chinese University of Hong Kong<EMAIL_ADDRESS>
###### Abstract

Large language models (LLMs) demonstrate impressive reasoning abilities, but
translating reasoning into actions in the real world remains challenging. In
particular, it remains unclear how to complete a given task provably within a
minimum number of interactions with the external environment, e.g., through an
internal mechanism of reasoning. To this end, we propose a principled
framework with provable regret guarantees to orchestrate reasoning and acting,
which we call “reason for future, act for now” (RAFA). Specifically, we design
a prompt template for reasoning that learns from the memory buffer and plans a
future trajectory over a long horizon (“reason for future”). At each step, the
LLM agent takes the initial action of the planned trajectory (“act for now”),
stores the collected feedback in the memory buffer, and reinvokes the
reasoning routine to replan the future trajectory from the new state.

The key idea is to cast reasoning in LLMs as learning and planning in Bayesian
adaptive Markov decision processes (MDPs). Correspondingly, we prompt LLMs to
form an updated posterior of the unknown environment from the memory buffer
(learning) and generate an optimal trajectory for multiple future steps that
maximizes a value function (planning). The learning and planning subroutines
are performed in an ”in-context” manner to emulate the actor-critic update for
MDPs. Our theoretical analysis proves that the novel combination of long-term
reasoning and short-term acting achieves a $\sqrt{T}$ regret. In particular,
the regret bound highlights an intriguing interplay between the prior
knowledge obtained through pretraining and the uncertainty reduction achieved
by reasoning and acting. Our empirical validation shows that it outperforms
various existing frameworks and achieves nearly perfect scores on a few
benchmarks. By incorporating “classical” MDP techniques, RAFA introduces the
first autonomous LLM agent with provable regret guarantees. Notably, LLMs do
not function as actors, critics, or learned world models, but rather as an
internal mechanism that improves them iteratively.

###### Contents

  1. 1 Introduction
    1. 1.1 Literature
  2. 2 Bridging LLM and RL
  3. 3 Algorithm
  4. 4 Theory
  5. 5 Experiment
    1. 5.1 Game of 24
    2. 5.2 ALFWorld
    3. 5.3 BlocksWorld
    4. 5.4 Tic-Tac-Toe
  6. A More Algorithms
    1. A.1 Proof of Proposition A.1
  7. B LLM Samples Posterior
  8. C Proof of Theorem 4.4
  9. D Linear Special Case
    1. D.1 Proof of Proposition D.3
  10. E More Experiments
    1. E.1 Game of 24
    2. E.2 ALFWorld
    3. E.3 BlocksWorld
    4. E.4 Tic-Tac-Toe
  11. F Prompts
    1. F.1 Game of 24
    2. F.2 ALFWorld
    3. F.3 Blocksworld
    4. F.4 Tic-Tac-Toe

##  1 Introduction

Large language models (LLMs) exhibit remarkable reasoning abilities, which
open a new avenue for agents to interact with the real world autonomously.
However, turning reasoning into actions remains challenging. Specifically,
although LLMs are equipped with the prior knowledge obtained through
pretraining, it is stateless in nature and ungrounded in the real world, which
makes the resulting action suboptimal. To bridge the reasoning-acting gap, we
aim to design an internal mechanism of reasoning on top of LLMs, which
optimizes actions iteratively by incorporating feedbacks from an external
environment. In particular, we focus on the sample efficiency of autonomous
LLM agents in interactive decision-making tasks, which plays a key role in
their practical adoption, especially when interactions are costly and risky.
In other words, our primary goal is to enable agents to complete a given task
in a guaranteed manner through reasoning within a minimum number of
interactions with the external environment.

Reinforcement learning (RL) is a well-studied paradigm for improving actions
by collecting feedbacks. However, to tailor existing RL techniques for
autonomous LLM agents, we lack a rigorous mapping between RL and LLMs, which
leads to various conceptual discrepancies. For example, RL operates in a
numerical system, where rewards and transitions are defined by scalars and
probabilities. In comparison, the inputs and outputs of LLMs are described by
tokens in a linguistic system. As another example, LLMs are trained on a
general-purpose corpus and remain fixed throughout the interactive process. In
contrast, RL trains actors and critics on the collected feedback iteratively.
Thus, it appears inappropriate to treat LLMs as actors or critics under the RL
framework, although all of them are parameterized by deep neural networks.
Moreover, it remains unclear what reasoning with LLMs means under the RL
framework, e.g., what are the inputs and outputs of a reasoning routine and
how reasoning should be coordinated with acting. Such conceptual discrepancies
prevent us from establishing a principled framework beyond borrowing the
“trial and error” concept from RL straightforwardly and make it difficult to
achieve provable sample efficiency guarantees. For instance, it is known in RL
that an improper design of agents may induce an exponential dependency on
horizons in the sample complexity. Without the RL-LLM correspondence, it is
hard to avoid the same flaw in autonomous LLM agents.

To address such conceptual discrepancies, we formalize reasoning and acting
with LLMs under a Bayesian adaptive Markov decision process (MDP) framework,
where the latent variable of interest is the unknown environment. The starting
point is to cast the full history of states (of the external environment),
actions, rewards, and their linguistic summaries in the memory buffer as the
information state of Bayesian adaptive MDPs. Throughout the interactive
process, the information state accumulates a growing collection of feedbacks
from the external environment, which is mapped to an optimized action at each
step by an internal mechanism of reasoning. As detailed below, we construct
the reasoning routine through two key subroutines, namely learning and
planning, which are instantiated by LLMs with specially designed prompts.

  * •

The learning subroutine forms an updated posterior of the unknown environment
from the memory buffer. Depending on whether we emulate the model-based or
model-free approach of RL, the learning subroutine infers the transition and
reward models (model) or/and the value function (critic).

  * •

The planning subroutine generates an optimal policy (actor) or trajectory for
multiple future steps, which maximizes the value function (up to a certain
error). Depending on the specific configuration of the state and action spaces
(continuous versus discrete) and the transition and reward models (stochastic
versus deterministic), the planning subroutine emulates the value iteration
algorithm, the random shooting algorithm, or the Monte-Carlo tree-search
algorithm.

Although LLMs remain fixed throughout the interactive process, they are
prompted to utilize the growing collection of feedbacks from the external
environment as contexts. Through the learning subroutine, the collected
feedback reduces the posterior uncertainty in models or values, which allows
the planning subroutine to obtain an improved policy at each step. In other
words, we emulate the actor-model or actor-critic update for Bayesian adaptive
MDPs in an in-context manner, where LLMs function as an internal mechanism
that improves models, values, and policies iteratively. Specifically, existing
RL methods use deep neural networks to parameterize models, values, and
policies, which map states (of the external environment) and actions to
scalars and probabilities. In comparison, we use LLMs to represent the
learning and planning algorithms in RL, which are composed to map data in the
memory buffer to actions. Here, data and actions are allowed to be tokens in a
linguistic system.

Figure 1: Illustration of the RAFA (“reason for future, act for now”)
framework.

By establishing the LLM-RL correspondence, we propose a principled framework
for orchestrating reasoning and acting, namely “reason for future, act for
now” (RAFA), which achieves provable sample efficiency guarantees in
autonomous LLM agents for the first time. At each step, the LLM agent invokes
the reasoning routine, which learns from the memory buffer and plans a future
trajectory over a long horizon (“reason for future”), takes the initial action
of the planned trajectory (“act for now”), and stores the collected feedback
in the memory buffer. Upon the state transition of the external environment,
the LLM agent reinvokes the reasoning routine to replan another future
trajectory from the new state. In other words, the LLM agent considers the
long-term consequence of actions from the current state before taking the
immediate action, which reduces the necessary number of interactions with the
external environment by optimizing actions. By combining long-term reasoning
with short-term acting, RAFA enjoys desirable theoretical and empirical
properties, which are detailed below.

  * •

Our theoretical analysis proves that RAFA achieves a $\sqrt{T}$ regret. In
particular, the regret bound highlights an intriguing interplay between the
prior knowledge obtained through pretraining and the uncertainty reduction
achieved by reasoning and acting.

  * •

Our empirical validation shows that RAFA outperforms various existing
frameworks in interactive decision-making tasks, including ALFWorld,
BlocksWorld, Game of 24, and a new benchmark based on TicTacToe. In a few
benchmarks, it achieves nearly perfect scores.

In summary, we contribute in the following aspects.

  * •

We establish a rigorous mapping between LLMs and RL, which enables us to
construct the reasoning routine through the learning and planning subroutines.

  * •

We propose a principled framework for autonomous LLM agents, namely RAFA,
which combines long-term reasoning with short-term acting.

  * •

We prove that RAFA achieves a $\sqrt{T}$ regret, which makes RAFA the first
autonomous LLM agent with provable sample efficiency guarantees.

  * •

We show that RAFA outperforms various existing frameworks, which validates the
theoretical analysis.

###  1.1 Literature

#### Large Language Model (LLM) and In-Context Learning (ICL).

LLMs (Radford et al., 2019; Brown et al., 2020; Hoffmann et al., 2022;
Chowdhery et al., 2022; OpenAI, 2023; Touvron et al., 2023) display notable
reasoning abilities. A pivotal aspect of reasoning is the ICL ability (Liang
et al., 2022; Razeghi et al., 2022; Shin et al., 2022; Olsson et al., 2022;
Akyürek et al., 2022; Kirsch et al., 2022; Garg et al., 2022; Von Oswald et
al., 2023; Li et al., 2023; Abernethy et al., 2023), which allows LLMs to
solve a broad range of tasks with only a few in-context examples instead of
finetuning parameters on a specific dataset. We focus on harnessing the ICL
ability of LLMs to optimize actions in the real world, which is crucial to
autonomous LLM agents. In particular, we build on a recent line of work (Xie
et al., 2021; Zhang et al., 2022, 2023b; Wang et al., 2023a; Wies et al.,
2023; Jiang, 2023; Lee et al., 2023) that attributes the ICL ability to
implicit Bayesian inference, i.e., an implicit mechanism that enables LLMs to
infer a latent concept from those in-context examples, which is verified both
theoretically and empirically. In RAFA, the latent concept is the transition
and reward models (model) of the unknown environment or/and the value function
(critic), which is inferred from the memory buffer in the learning subroutine.

#### Reinforcement Learning (RL) under a Bayesian Framework.

We build on a recent line of work on the infinite-horizon (Abbasi-Yadkori and
Szepesvári, 2015; Dong et al., 2019; Wei et al., 2020; Zhou et al., 2021a, b;
Chen et al., 2022) and Bayesian (Strens, 2000; Osband et al., 2013; Russo and
Van Roy, 2014a, b, 2016) settings of RL. The infinite-horizon setting allows
RAFA to interact with the external environment continuously without resetting
to an initial state, while the Bayesian setting allows RAFA to invoke the ICL
ability of LLMs to form an updated posterior of the unknown environment. RL
operates in a numerical system, where rewards and transitions are defined by
scalars and probabilities, and trains actors and critics on the collected
feedback iteratively. We focus on emulating the actor-model or actor-critic
update in RL through an internal mechanism of reasoning on top of LLMs, which
allows data and actions to be tokens in a linguistic system while bypassing
the explicit update of parameters. In particular, the learning and planning
subroutines of RAFA emulate the posterior update and various planning
algorithms in RL. Moreover, RAFA orchestrates reasoning (learning and
planning) and acting following the principled approach in RL, i.e.,
(re)planning a future trajectory over a long horizon (“reason for future”) at
the new state and taking the initial action of the planned trajectory (“act
for now”). As a result, RAFA inherits provable sample efficiency guarantees
from RL.

#### Reasoning with LLM.

We build on a recent line of work that develops various prompting schemes to
improve the reasoning performance of LLMs. “Chain of thoughts” (Wei et al.,
2022) decomposes a challenging problem into several reasoning stages and
guides LLMs to solve them one by one. As generalizations, “tree of thoughts”
(Yao et al., 2023a), “graph of thoughts” (Yao et al., 2023b), “algorithm of
thoughts” (Sel et al., 2023), and “cumulative reasoning” (Zhang et al., 2023a)
provide different graph-search schemes to guide LLMs. See also Wang et al.
(2022); Creswell et al. (2022); Creswell and Shanahan (2022). Also, “reasoning
via planning” (Hao et al., 2023) emulates the Monte-Carlo tree-search (MCTS)
algorithm to reduce the search complexity. For embodied LLM agents, Huang et
al. (2022a) propose to decompose a complex task into multiple executable
steps. Most of them focus on general reasoning tasks, e.g., solving a
mathematical or logic puzzle, where LLMs generate a detailed trace
(trajectory) of arguments through an internal mechanism to reach a final
answer. Here, LLMs play the same role as the planning subroutine in RAFA. In
contrast, we focus on interactive decision-making tasks, where autonomous LLM
agents collect feedbacks from the external environment to optimize actions
iteratively. In particular, we aim to complete a given task within a minimum
number of interactions with the external environment. To this end, it is
essential to operate three interleaved modules, namely learning, planning, and
acting, in a closed loop. While it is feasible to incorporate existing graph-
search or MCTS schemes as the planning subroutine for generating trajectories,
our core contribution is a principled framework that executes a selected
subset of the planned trajectory to collect feedbacks (“act for now”) and
replans an improved trajectory from the new state by learning from feedbacks
(“reason for future”). From an RL perspective, existing graph-search or MCTS
schemes are analogous to an open-loop method, e.g., motion planning or
trajectory optimization (Betts, 1998), which does not involve interactions
with the external environment. To integrate them into a closed-loop approach,
e.g., model predictive control (Rawlings, 2000), one has to specify how to act
given the planned trajectory and when to reinvoke the reasoning (learning and
planning) routine, which is the key technique of RAFA. Another recent line of
work tackles more complex tasks by allowing LLMs to access various additional
modules, e.g., tools, programs, and other learning algorithms (Ahn et al.,
2022; Shen et al., 2023; Lu et al., 2023; Liu et al., 2023; Cai et al., 2023),
or by finetuning LLMs on the collected feedback (Zelikman et al., 2022; Li et
al., 2022; Paul et al., 2023). Integrating them with RAFA is left as a future
direction of research.

#### Acting (and Reasoning) with LLM.

We build on a recent line of work that develops various closed-loop frameworks
for interacting with the external environment. “Inner monologue” (Huang et
al., 2022b) and “ReAct” (Yao et al., 2022) combine reasoning and acting to
refine each other for the first time. In comparison, RAFA provides a specific
schedule for orchestrating reasoning and acting (as discussed above). As
generalizations, “Reflexion” (Shinn et al., 2023) enables autonomous LLM
agents to revise the current action of a pregenerated trajectory by learning
from feedbacks, especially when they make mistakes. See also Kim et al.
(2023). However, making a local revision to the pregenerated trajectory is
myopic because it fails to consider the long-term consequence of actions.
Consequently, the obtained policy may get trapped by a local optimum. From an
RL perspective, “Reflexion” (Shinn et al., 2023) is an oversimplified version
of RAFA, where the planning subroutine revises the current action to maximize
the reward function (“reason for now”) instead of planning multiple future
steps to maximize the value function (“reason for future”), which measures the
expected cumulative future reward. To remedy this issue, “AdaPlanner” (Sun et
al., 2023) regenerates the whole trajectory at each step, which yields a
global improvement. See also Wang et al. (2023b). However, the reasoning
routine of “AdaPlanner” requires a handcrafted set of programs to reject
suboptimal candidate trajectories. Without the domain knowledge of a specific
task, the regenerated trajectory is not necessarily optimal, i.e., maximizing
the value function (up to a certain error). In contrast, the reasoning routine
of RAFA is designed following the principled approach in RL. In particular,
the learning subroutine infers the transition and reward models (model) or/and
the value function (critic), while the planning subroutine emulates the value
iteration algorithm, the random shooting algorithm, or the MCTS algorithm,
none of which use any domain knowledge. As a result, RAFA achieves provable
sample efficiency guarantees for the first time and outperforms those existing
frameworks empirically.

##  2 Bridging LLM and RL

#### Interaction Protocol.

We use Markov decision processes (MDPs) to model how autonomous LLM agents
interact with the external environment. We consider an infinite-horizon MDP
$M=(\mathcal{S},\mathcal{A},P,r,\rho,\gamma)$, where $\mathcal{S}$ is the
state space, $\mathcal{A}$ is the action space,
$P:\mathcal{S}\times\mathcal{A}\mapsto\Delta(\mathcal{S})$ is the transition
kernel, $r:\mathcal{S}\times\mathcal{A}\mapsto[0,1]$ is the reward function,
$\rho$ is the initial distribution of states, and $\gamma\in(0,1)$ is the
discount factor. Here, $P$ gives the probability distribution of the next
state given the current state and action, while $r$ is assumed to be
deterministic without loss of generality. For notational simplicity, we
parameterize $P$ and $r$ by a shared parameter $\theta^{\star}\in\Theta$ and
denote them as $P_{\theta^{\star}}$ and $r_{\theta^{\star}}$. At the $t$-th
step, the LLM agent receives a state $s_{t}\in{\mathcal{S}}$, takes an action
$a_{t}\in\mathcal{A}$ following the current policy
$\pi_{t}:{\mathcal{S}}\mapsto\mathcal{A}$, and receives a reward
$r_{t}=r_{\theta^{\star}}(s_{t},a_{t})$. Subsequently, the external
environment transits to the next state $s_{t+1}\sim
P_{\theta^{\star}}(\cdot|s_{t},a_{t})$, while the LLM agent computes the
updated policy $\pi_{t+1}$ through an internal mechanism of reasoning (as
discussed below). Note that ${\mathcal{S}}$ and $\mathcal{A}$ are represented
by tokens in a linguistic system. Here, $\pi\in\Pi$ is assumed to be
deterministic without loss of generality, where $\Pi$ is the feasible set of
policies.

#### Value Function.

For a policy $\pi$ and a parameter $\theta$ of the transition and reward
models, we define the state-value and action-value functions

$\displaystyle
V^{\pi}_{\theta}(s)=\mathbb{E}\Bigl{[}\sum_{t=0}^{\infty}\gamma^{t}r_{\theta}(s_{t},a_{t})\Big{|}\,s_{0}=s\Bigl{]},\quad
Q_{\theta}^{\pi}(s,a)=\mathbb{E}\Bigl{[}\sum_{t=0}^{\infty}\gamma^{t}r_{\theta}(s_{t},a_{t})\Big{|}\,s_{0}=s,a_{0}=a\Bigr{]},$
(2.1)

where $\mathbb{E}$ is taken with respect to $a_{t}=\pi(s_{t})$ and
$s_{t+1}\sim P_{\theta}(\cdot|s_{t},a_{t})$ for all $t\geq 0$. In other words,
$V^{\pi}_{\theta}$ (and $Q_{\theta}^{\pi}$) gives the expected cumulative
future reward from the current state $s$ (and action $a$).

#### Planning Oracle.

To define the optimal policy $\pi_{\theta}^{\star}$ with respect to a given
parameter $\theta$, we define the Bellman optimality equation

$\displaystyle
Q^{\star}_{\theta}\left(s,a\right)=r_{\theta}\left(s,a\right)+\gamma\left(P_{\theta}V_{\theta}^{\star}\right)\left(s,a\right),\quad
V_{\theta}^{\star}(s)=\max_{a\in\mathcal{A}}Q_{\theta}^{\star}(s,a),$ (2.2)

where $Q^{\star}_{\theta}$ and $V_{\theta}^{\star}$ are the fixed-point
solution. Here, we define
$(P_{\theta}V_{\theta}^{\star})(s,a)=\mathbb{E}[V_{\theta}^{\star}(s^{\prime})]$,
where $\mathbb{E}$ is taken with respect to $s^{\prime}\sim
P_{\theta}(\cdot|s,a)$. Let $\pi_{\theta}^{\star}(s)=\mathop{\text{\rm
arg\,max}}_{a\in\mathcal{A}}Q_{\theta}^{\star}(s,a)$. We define
$\texttt{PL}^{\star}:\Theta\mapsto\Pi$ as the planning oracle that maps
$\theta$ to $\pi_{\theta}^{\star}$. See Sutton and Barto (2018) for the
existence and uniqueness guarantees for $Q^{\star}_{\theta}$,
$V_{\theta}^{\star}$, and $\pi_{\theta}^{\star}$.

#### Sample Efficiency.

Let $\theta^{\star}$ be the underlying parameter that generates states and
rewards. As the performance metric, we define the Bayesian regret

$\displaystyle\mathfrak{R}(T)=\mathbb{E}\Bigl{[}\sum_{t=0}^{T-1}V_{\theta^{\star}}^{\pi^{\star}}(s_{t})-V_{\theta^{\star}}^{\pi_{t}}(s_{t})\Bigr{]},\quad\text{where~{}~{}}\pi^{\star}=\texttt{PL}^{\star}(\theta^{\star}).$
(2.3)

Here, $\mathbb{E}$ is taken with respect to the prior distribution $p_{0}$ of
$\theta^{\star}$, the stochastic outcome of $s_{t}$, and the iterative update
of $\pi_{t}$, which involves states, actions, and rewards until the $t$-th
step, i.e., the full history
$\mathcal{D}_{t}=\\{(s_{i},a_{i},s_{i+1},r_{i})\\}_{i=0}^{t-1}$. We aim to
design a sample-efficient agent that satisfies $\mathfrak{R}(T)=o(T)$, i.e.,
the Bayesian regret is sublinear in the total number of interactions with the
external environment.

#### What Reasoning Means and Role of LLM.

We formalize reasoning and acting with LLMs under a Bayesian adaptive MDP
framework (Ghavamzadeh et al., 2015), where the underlying parameter
$\theta^{\star}$ is the latent variable of interest and the full history
$\mathcal{D}_{t}$ (and its linguistic summary) is the information state. In
particular, we aim to design an internal mechanism on top of LLMs that maps
$\mathcal{D}_{t}$ to an optimized action $a_{t}$ or the corresponding policy
$\pi_{t}$ (reasoning), which is executed in the external environment (acting).
To this end, we construct the reasoning routine through two key subroutines,
which emulate the learning and planning algorithms in RL. Specifically, the
learning subroutine maps $\mathcal{D}_{t}$ to the posterior distribution
$p_{t}$ of $\theta^{\star}$, while the planning subroutine maps $p_{t}$ or a
sampled parameter $\theta\sim p_{t}$ to $\pi_{t}$. In other words, the
learning subroutine forms an updated posterior of the unknown environment from
the memory buffer, while the planning subroutine approximates the planning
oracle $\texttt{PL}^{\star}$. As shown in Section 3, we invoke the ICL ability
of LLMs to achieve the formal goal (implicitly), while we design a prompt
template for LLMs to achieve the latter goal (explicitly). Following the
principled approach in RL, we develop a specific schedule for orchestrating
reasoning (learning and planning) and acting, which is proven as sample-
efficient in Section 4.

##  3 Algorithm

#### Architecture of RAFA.

By leveraging the LLM-RL correspondence in Section 2, we provide a principled
framework for orchestrating reasoning and acting, namely “reason for future,
act for now” (RAFA), in Algorithms 1 and 2. In Section 4, we present the RL
counterpart of RAFA in Algorithm 3 to illustrate the design rationale and
establish the theoretical foundation. At the $t$-th step of Algorithm 1, the
LLM agent invokes the reasoning routine, which learns from the memory buffer
and plans a future trajectory over a long horizon (“reason for future” in Line
6), takes the initial action of the planned trajectory (“act for now” in Line
7), and stores the collected feedback (state, action, and reward) in the
memory buffer (Line 8). Upon the state transition of the external environment,
the LLM agent reinvokes the reasoning routine to replan another future
trajectory from the new state (Line 6 following Line 9). To ensure the
learning and planning stability, we impose the switching condition (Line 10)
to decide whether to incorporate the newest chunk of history, i.e., the set
difference $\mathcal{D}_{t}-\mathcal{D}_{t_{k}}$, into the information state,
which is used in the reasoning routine as contexts. In other words, the
reasoning routine uses the same history $\mathcal{D}_{t_{k}}$ for all
$t_{k}\leq t<t_{k+1}$ until the $(k+1)$-th switch at the $(t_{k+1}-1)$-th
step, which guarantees that the posterior distribution and the optimized
action or the corresponding policy are updated in a conservative manner. We
specify the switching condition in Sections 4 and 5.

#### “Reason for Future” (Line 6 in Algorithm 1 and Lines 3-11 in Algorithm
2).

As detailed below, the reasoning routine composes the learning and planning
subroutines to map the full history $\mathcal{D}_{t_{k}}$ (until the
$t_{k}$-th step) to an optimized action $a_{t}$.

  * •

The learning subroutine (Lines 3-4 in Algorithm 2) maps $\mathcal{D}_{t_{k}}$
to a transition kernel (Model) and a value function (Critic), which are used
in the planning subroutine. Intuitively, we prompt LLMs to form an updated
posterior of the unknown environment from the memory buffer. Here, the updated
posterior is instantiated by Model and Critic, which estimate their ground-
truth counterparts in association with the data-generating parameter. From an
RL perspective (Sections 2 and 4), the learning subroutine maps
$\mathcal{D}_{t_{k}}$ to the posterior distribution $p_{t}$ of the underlying
parameter $\theta^{\star}$, which generates states and rewards, and returns
the transition kernel $P_{\theta}$ and the value function
$V^{\pi_{t}}_{\theta}$, where $\theta\sim p_{t}$ is the sampled parameter and
$\pi_{t}$ is the current policy. On the other hand, the ICL ability of LLMs
allows us to bypass the posterior update of $p_{t}$, sampling $\theta$ from
$p_{t}$, and the explicit parameterization of $P_{\theta}$ and
$V^{\pi_{t}}_{\theta}$ in RL. Instead, we represent $P_{\theta}$ and
$V^{\pi_{t}}_{\theta}$ using two LLM instances with specially designed
prompts, which instruct them to use $\mathcal{D}_{t_{k}}$ as contexts to
generate the next state and evaluate a given trajectory or the corresponding
policy. As $\mathcal{D}_{t_{k}}$ accumulates a growing collection of feedbacks
from the external environment, it reduces the posterior uncertainty about the
unknown environment, which yields more accurate versions of Model and Critic.
Consequently, the planning subroutine is able to use them to assess the long-
term outcome of actions with a higher accuracy. Depending on whether we
emulate the model-based or model-free approach of RL, we may choose to emulate
Model or Critic individually. For illustration, we consider a deterministic
setting of transitions and rewards with discrete state and action spaces,
where we emulate both of them in a tree-search example.

  * •

The planning subroutine (Lines 5-11 in Algorithm 2) maps Model and Critic to a
future trajectory
$(s_{0}^{\dagger},a_{0}^{\dagger},\ldots,s_{U}^{\dagger},a_{U}^{\dagger})$,
where $s_{0}^{\dagger}$ is the current state $s_{t}$ and $a_{0}^{\dagger}$ is
executed in the external environment as the current action $a_{t}$ during the
acting phase. Intuitively, we prompt LLMs to generate an optimal policy
(actor) for multiple future steps, which maximizes the value function
(Critic). From an RL perspective (Sections 2 and 4), the planning subroutine
approximates the planning oracle $\texttt{PL}^{\star}$, which maps a given
parameter $\theta$ to the optimal policy $\pi_{\theta}^{\star}$ or the
corresponding action $a_{t}=\pi_{\theta}^{\star}(s_{t})$. As two LLM instances
from the learning subroutine, Model and Critic instantiate the transition
kernel $P_{\theta}$ and the value function $V^{\pi_{t}}_{\theta}$ in
association with the sampled parameter $\theta\sim p_{t}$ (as discussed
above). Hence, we are able to simulate a given number of trajectories with
Model, evaluate them with Critic, and obtain an improved policy, which is
achieved by specially designed prompts instead of a numerical algorithm. By
maximizing the expected cumulative future reward (instead of the immediate
reward), the planning subroutine returns an optimized action that improves the
long-term outcome. In Section 4, we identify two error sources that affect the
planning subroutine, namely the posterior uncertainty, which is inherited from
Model and Critic due to the finite size of $\mathcal{D}_{t_{k}}$, and the
planning suboptimality, which is induced by the limited capacity for
computation, e.g., the bounded width and depth of tree-search (Lines 6-9 in
Algorithm 2). Depending on the specific configuration of the state and action
spaces (continuous versus discrete) and the transition and reward models
(stochastic versus deterministic), we may choose to emulate the value
iteration algorithm, the random shooting algorithm, or the Monte-Carlo tree-
search algorithm. All of them allow RAFA to achieve provable sample efficiency
guarantees as long as they satisfy a specific requirement of optimality
(Definition 4.2). For illustration, we emulate the tree-search algorithm and
defer its stochastic variant to Appendix A.

Note that the reasoning routine does not interact with the external
environment throughout the learning and planning subroutines.

Algorithm 1 Reason for future, act for now (RAFA): The LLM version.

1: input: An LLM learner-planner LLM-LR-PL, which aims at generating an
optimal trajectory given an initial state and returns the initial action
(e.g., Algorithm 2), and a switching condition If-Switch.

2: initialization: Sample the initial state $s_{0}\sim\rho$, set $t=0$, and
initialize the memory buffer $\mathcal{D}_{0}=\varnothing$.

3: for $k=0,1,\ldots,$ do

4: Set $t_{k}\leftarrow{t}$.

5: repeat

6: Learn and plan given memory $\mathcal{D}_{t_{k}}$ to get action
$a_{t}\leftarrow\texttt{LLM-LR-PL}(\mathcal{D}_{t_{k}},s_{t})$.(“reason for
future”)

7: Execute action $a_{t}$ to receive reward $r_{t}$ and state $s_{t+1}$ from
environment. (“act for now”)

8: Update memory
$\mathcal{D}_{t+1}\leftarrow\mathcal{D}_{t}\cup\\{(s_{t},a_{t},s_{t+1},r_{t})\\}$.

9: Set $t\leftarrow t+1$.

10: until $\texttt{If-Switch}(\mathcal{D}_{t})$ is True. (the switching
condition is satisfied)

11: end for

Algorithm 2 The LLM learner-planner (LLM-LR-PL): A tree-search example (for
the deterministic case).

1: input: The memory buffer $\mathcal{D}$, the initial state $s$, the search
breadth $B$, and the search depth $U$.

2: initialization: Initialize the state array
$\mathcal{S}_{0}\leftarrow\\{s\\}$ and the action array
$\mathcal{A}_{0}\leftarrow\varnothing$. ———————————————— (the learning
subroutine) ————————————————

3: Set Model as an LLM instance prompted to use $\mathcal{D}$ as contexts to
generate the next state.

4: Set Critic as an LLM instance prompted to use $\mathcal{D}$ as contexts to
estimate the value function. ———————————————— (the planning subroutine)
————————————————

5: Set Elite as an LLM instance prompted to use $\mathcal{D}$ as contexts to
generate multiple candidate actions.

6: for $u=0,\ldots,U$ do

7: For each current state in $\mathcal{S}_{u}$, invoke Elite to generate $B$
candidate actions and store them in $\mathcal{A}_{u}$.

8: For each candidate action in $\mathcal{A}_{u}$, invoke Model to generate
the next state and store it in $\mathcal{S}_{u+1}$.

9: end for

10: For all resulting rollouts in
$\mathcal{S}_{0}\times\mathcal{A}_{0}\times\cdots\times\mathcal{S}_{U}\times\mathcal{A}_{U}$,
invoke Critic to evaluate the expected cumulative future reward and select the
best one
$(s_{0}^{\dagger},a_{0}^{\dagger},\ldots,s_{U}^{\dagger},a_{U}^{\dagger})$,
where $s_{0}^{\dagger}=s$.

11: output: The initial action $a_{0}^{\dagger}$ of the selected rollout.

#### “Act for Now” (Lines 7-10 in Algorithm 1).

At the current state $s_{t}$, the LLM agent executes the optimized action
$a_{t}$ in the external environment, which is obtained from the reasoning
routine. Specifically, we take the initial action $a_{0}^{\dagger}$ of the
planned trajectory
$(s_{0}^{\dagger},a_{0}^{\dagger},\ldots,s_{U}^{\dagger},a_{U}^{\dagger})$,
where $s_{0}^{\dagger}=s_{t}$ and $a_{0}^{\dagger}=a_{t}$, and discard the
remaining subset. At the next state $s_{t+1}$, the LLM agent replans another
future trajectory
$(s_{0}^{\dagger},a_{0}^{\dagger},\ldots,s_{U}^{\dagger},a_{U}^{\dagger})$
with $s_{0}^{\dagger}=s_{t+1}$ and $a_{0}^{\dagger}=a_{t+1}$. In other words,
the acting phase follows a short-term subset of the long-term plan, which is
regenerated at every new state. The LLM agent stores the collected feedback
$(s_{t},a_{t},r_{t},s_{t+1})$ in the memory buffer $\mathcal{D}_{t}$ and
queries a switching condition If-Switch to decide when to update the
information state $\mathcal{D}_{t_{k}}\subseteq\mathcal{D}_{t}$, which is used
in the reasoning routine as contexts for learning and planning. Intuitively,
we incorporate the newest chunk of history
$\mathcal{D}_{t}-\mathcal{D}_{t_{k}}$ to improve the current policy only in
the case that it carries significant novel information, e.g., when the LLM
agent loses for the first time following a winning streak. In Section 4, we
provide a principled implementation of the switching condition, which measures
the posterior uncertainty given $\mathcal{D}_{t}$ with entropy and compares it
against that given $\mathcal{D}_{t_{k}}$. From an RL perspective, the lazy
update ensures the learning and planning stability and plays a pivotal role in
the regret analysis. In Section 5, we develop several practical variants that
achieve superior empirical performance.

##  4 Theory

We establish provable sample efficiency guarantees for RAFA (Algorithms 1 and
2) through its RL counterpart (Algorithm 3). In Line 6 of Algorithm 3, the
reasoning routine forms an updated posterior of the unknown environment
(learning) and generates an optimized action from an improved policy
(planning), which mirrors that of RAFA. Here, we emulate the model-based
approach of RL and cast RAFA as a Thompson sampling (TS) method. The following
assumption and definition formalize the learning and planning subroutines of
RAFA (Lines 3-4 and 5-11 in Algorithm 2).

#### Learning.

Let $\texttt{LLM}_{\mathcal{D},g}$ be an LLM instance with $\mathcal{D}$ as
contexts and $g$ as instructions to perform a specific task. Specifically,
$g^{\dagger}$ prompts LLMs to predict the next state $s^{\prime}$ and the
received reward $r$ from the current state $s$ and the current action $a$,
i.e., $\texttt{LLM}_{\mathcal{D},g^{\dagger}}:S\times A\rightarrow
S\times[0,1]$, where the generated state is stochastic. We denote the Markov
kernel in association with $\texttt{LLM}_{\mathcal{D},g^{\dagger}}$ as
$P_{\texttt{LLM}_{\mathcal{D},g^{\dagger}}}(s^{\prime},r|s,a)$. Also, we
denote the posterior distribution of the transition and reward models as
$\mathbb{P}_{\text{model}}(P_{\theta},r_{\theta}|\mathcal{D})$.

######  Assumption 4.1 (LLM Samples Posterior).

The Markov kernel $P_{\texttt{LLM}_{\mathcal{D},g^{\dagger}}}$ follows the
posterior distribution $\mathbb{P}_{\text{model}}(\cdot|\mathcal{D})$.

Assumption 4.1 states that LLMs perform implicit Bayesian inference, which is
verified both theoretically and empirically as the underlying mechanism of ICL
(Xie et al., 2021; Zhang et al., 2022, 2023b; Wang et al., 2023a; Wies et al.,
2023; Jiang, 2023; Lee et al., 2023). In particular, Xie et al. (2021); Wang
et al. (2023a) validate it in a general setting for generating texts, while
Lee et al. (2023) prove it in the imitation setting of RL to develop a new
framework for pretrained decision transformers. We consider a related setting
for predicting states and rewards that are described by texts. Here, the
pretraining dataset is a general-purpose corpus covering a wide variety of
$\mathcal{D}$ and $g$, whereas $(P_{\theta},r_{\theta})$ or $\theta$ is the
latent concept of interest. In comparison, Lee et al. (2023) consider the
imitation setting for predicting the optimal action without an explicit
planner, where the pretraining dataset contains the numerical trajectory
labeled by experts. In Appendix B, we prove that Assumption 4.1 holds for a
specific parameterization of $(P_{\theta},r_{\theta})$ under three regularity
conditions, namely (a) LLMs are trained to replicate the pretraining
distribution, which is assumed in Saunshi et al. (2020); Wei et al. (2021);
Xie et al. (2021) to simplify the statistical analysis, (b) the pretraining
dataset is generated through a Bayesian mechanism with a latent concept, which
is a simplified version of the latent variable model in Xie et al. (2021) and
resembles that in Wang et al. (2023a), and (c) LLMs are able to parameterize
an implicit Bayesian inference mechanism, which is proved in Zhang et al.
(2022, 2023b) for the attention architecture. Note that, if Assumption 4.1
holds approximately, the regret analysis can be relaxed to accommodate the
additional error in the posterior distribution.

#### Planning.

Assumption 4.1 allows us to bridge RAFA and TS. In the learning subroutine of
RAFA, we emulate $P_{\theta}$ with Model (Line 3 in Algorithm 2) and
$V^{\pi}_{\theta}$ with Critic (Line 4 in Algorithm 2), which is determined by
$P_{\theta}$, $r_{\theta}$, and $\pi$. At the $t$-th step, $\theta$ is sampled
from $p_{t}$, i.e., the updated posterior given the full history
$\mathcal{D}_{t_{k}}$ (until the $t_{k}$-th step). To formalize the planning
subroutine of RAFA, we define the planning suboptimality. Recall that $\Theta$
is the parameter space, $\Pi$ is the policy space, and $\texttt{PL}^{\star}$
is the planning oracle, which is defined in Section 2.

######  Definition 4.2 ($\epsilon$-Optimality of Planner).

A planning algorithm $\texttt{PL}^{\epsilon}:\Theta\mapsto\Pi$ is an
$\epsilon$-optimal planner if $\max_{s\in
S}[V^{\texttt{PL}^{\star}(\theta)}_{\theta}(s)-V^{\textrm{{PL}}^{\epsilon}(\theta)}_{\theta}(s)]\leq\epsilon$
for all $\theta\in\Theta$.

As a special case of Definition 4.2, we present the value iteration algorithm
in Appendix A, where we use a truncated horizon $U$, i.e., a finite length of
the lookahead window. Here, $\epsilon$ decreases as $U$ increases. In other
words, enlarging the lookahead window reduces the planning suboptimality
(“reason for future”). Alternatively, we may choose to emulate the tree-search
algorithm, the random shooting algorithm, or the Monte-Carlo tree-search
algorithm. In the tree-search example (Lines 5-11 in Algorithm 2), $\epsilon$
decreases as the search breadth $B$ and depth $U$ increase. Note that, as long
as we emulate an $\epsilon$-optimal planner, we are able to establish provable
sample efficiency guarantees.

#### Switching.

We consider a principled implementation of the switching condition (Line 10 in
Algorithms 1 and 3). Let $\mathfrak{H}(p)$ be the differential entropy of $p$.
We define the posterior entropy given $\mathcal{D}_{t}$ as

$\displaystyle H_{t}=\mathfrak{H}(p_{t})=-\int_{\Theta}p_{t}(\theta)\cdot\log
p_{t}(\theta){\mathrm{d}}\theta.$ (4.1)

As long as $H_{t_{k}}-H_{t}>\log 2$, i.e., the memory buffer accumulates one
extra bit of information, we incorporate $\mathcal{D}_{t}-\mathcal{D}_{t_{k}}$
into the information state and use it to improve the current policy. The
switching condition ensures that $\pi_{t}$ is switched for a logarithmic
number of times, which is a key step in establishing the sublinear regret.
Intuitively, the lazy update of policies ensures the learning and planning
stability. On the other hand, calculating the posterior entropy is challenging
in practice. In Section 5, we develop several practical variants that achieve
superior empirical performance.

Algorithm 3 Reason for future, act for now (RAFA): The RL counterpart.

1: input: An $\epsilon$-optimal planner $\texttt{PL}^{\epsilon}$, which
returns an $\epsilon$-optimal policy that maximizes the value function up to
an $\epsilon$ accuracy (Definition 4.2), and the prior distribution $p_{0}$ of
the model parameter $\theta$.

2: initialization: Sample the initial state $s_{0}\sim\rho$, set $t=0$, and
initialize the memory buffer $\mathcal{D}_{0}=\varnothing$.

3: for $k=0,1,\ldots,$ do

4: Set $t_{k}\leftarrow{t}$.

5: repeat

6: Update posterior $p_{t_{k}}$ given memory $\mathcal{D}_{t_{k}}$, sample
parameter $\theta_{t}\sim p_{t_{k}}$,obtain policy
$\pi_{t}\leftarrow\texttt{PL}^{\epsilon}(\theta_{t})$, and sample action
$a_{t}\leftarrow\pi_{t}(s_{t})$. (“reason for future”)

7: Execute action $a_{t}$ to receive reward $r_{t}$ and state $s_{t+1}$ from
environment. (“act for now”)

8: Update memory
$\mathcal{D}_{t+1}\leftarrow\mathcal{D}_{t}\cup\\{(s_{t},a_{t},s_{t+1},r_{t})\\}$.

9: Set $t\leftarrow t+1.$

10: until $H_{t_{k}}-H_{t}>\log 2$. (the switching condition is satisfied,
where $H_{t_{k}}$ and $H_{t}$ are defined in (4.1))

11: end for

#### Regret.

We define the information ratio to characterize the tail behavior of the
posterior distribution (Abbasi-Yadkori and Szepesvári, 2015; Osband et al.,
2013; Russo and Van Roy, 2014a, b, 2016; Lu and Van Roy, 2019). Let
$\delta\in(0,1)$ be the confidence level,
$\mathcal{D}_{T}=\\{(s_{t},a_{t},s_{t+1},r_{t})\\}_{t=0}^{T-1}$ be an
arbitrary dataset collected in the underlying MDP, and
$\\{V_{t}\\}_{t=0}^{T-1}$ be a value function sequence adapted to
$\\{\sigma(\mathcal{D}_{t})\\}_{t=0}^{T-1}$, where $\sigma(\mathcal{D}_{t})$
is the sigma-algebra of $\mathcal{D}_{t}\subseteq\mathcal{D}_{T}$. We define
the information gain as $I(\theta;\xi_{t+1}|\mathcal{D}_{t})=H_{t}-H_{t+1}$.
Here, $\xi_{t+1}$ denotes $(s_{t},a_{t},s_{t+1},r_{t})$ and $H_{t}$ is defined
in (4.1), where $p_{t}$ is the posterior distribution given $\mathcal{D}_{t}$.

######  Definition 4.3 (Information Ratio).

The information ratio $\Gamma_{t^{\dagger}}(\delta)$ is the smallest number
for which, if we have $H_{t^{\dagger}}-H_{t}\leq\log 2$, then it holds for all
$t\in\\{t^{\dagger},\ldots,T-1\\}$ with probability at least $1-\delta$ that

$\displaystyle\bigl{|}(r_{\theta^{\star}}-r_{\theta_{t^{\dagger}}})(s_{t},a_{t})+\bigl{(}(P_{\theta^{\star}}-P_{\theta_{t^{\dagger}}})V_{t}\bigr{)}(s_{t},a_{t})\bigr{|}\leq\Gamma_{t^{\dagger}}(\delta)\cdot\sqrt{I(\theta;\xi_{t+1}|\mathcal{D}_{t})},$
(4.2)

where $\theta^{\star}$ is the data-generating parameter and
$\theta_{t^{\dagger}}\sim p_{t^{\dagger}}$ is a sampled parameter.

Definition 4.3 quantifies the estimation error of the sampled parameter
$\theta_{t^{\dagger}}$ in terms of approximating the data-generating parameter
$\theta^{\star}$. To achieve this, we use the information gain
$I(\theta;\xi_{t+1}|\mathcal{D}_{t})$ as a benchmarking quantity. Intuitively,
the information ratio $\Gamma_{t^{\dagger}}(\delta)$ characterizes how
exploration reduces uncertainty. As long as $\Gamma_{t^{\dagger}}(\delta)$ is
finite, collecting a better dataset $\mathcal{D}_{t^{\dagger}}$ reduces the
estimation error of $\theta_{t^{\dagger}}$ by decreasing
$I(\theta;\xi_{t+1}|\mathcal{D}_{t})$. To see this, we consider the limiting
case that $\mathcal{D}_{t^{\dagger}}\subseteq\mathcal{D}_{t}$ is already
sufficiently informative. As the collected feedback $\xi_{t+1}$ provides
little new information, we know that $I(\theta;\xi_{t+1}|\mathcal{D}_{t})$ is
small, which implies that the estimation error of $\theta_{t^{\dagger}}$ is
small. Note that $\Gamma_{t^{\dagger}}(\delta)$ is an intrinsic property of
the underlying MDP, which is independent of the agent design. In Appendix D,
we study a general class of MDPs, namely linear kernel MDPs (Ayoub et al.,
2020; Cai et al., 2020; Zhou et al., 2021b), and prove that the information
ratio is $\mathcal{O}(d\cdot\log(dT/\delta))$, where $d$ is the feature
dimension.

We characterize the Bayesian regret of Algorithm 1 by connecting it to
Algorithm 3. Recall that the Bayesian regret is defined in (2.3) and
$\gamma\in(0,1)$ is the discount factor.

######  Theorem 4.4 (Bayesian Regret).

Under Assumption 4.1, the Bayesian regret of RAFA satisfies

$\displaystyle\mathfrak{R}(T)=\mathcal{O}\biggl{(}\frac{\gamma\cdot\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}[\sqrt{H_{0}-H_{T}}]}{1-\gamma}\cdot\sqrt{T}+\frac{\gamma\delta}{(1-\gamma)^{2}}\cdot
T+\epsilon\cdot
T+\frac{\gamma\cdot\mathbb{E}[H_{0}-H_{T}]}{(1-\gamma)^{2}}\biggr{)}.$ (4.3)

###### Proof.

See Appendix C for a detailed proof. ∎

Theorem 4.4 establishes the $\sqrt{T}$ regret of RAFA (Algorithms 1 and 3) for
a proper choice of the confidence level $\delta$ and the planning
suboptimality $\epsilon$, e.g., $\delta=\mathcal{O}(1/\sqrt{T})$ and
$\epsilon=\mathcal{O}(1/\sqrt{T})$. Here, the first term on the right-hand
side of (4.3) is the leading term and involves several multiplicative factors,
namely the effective horizon $1/(1-\gamma)$, the information ratio
$\Gamma_{t^{\dagger}}(\delta)$, and the information gain ${H_{0}-H_{T}}$
throughout the $T$ steps, which are common in the RL literature (Abbasi-
Yadkori and Szepesvári, 2015; Osband et al., 2013; Russo and Van Roy, 2014a,
b, 2016; Lu and Van Roy, 2019). In particular, $H_{0}$ highlights the prior
knowledge obtained through pretraining, as $H_{0}$ quantifies the prior
uncertainty of LLMs before incorporating any collected feedback. Hence,
$H_{0}-H_{T}$ highlights the uncertainty reduction achieved by reasoning and
acting, as $H_{T}$ quantifies the posterior uncertainty of LLMs after
incorporating the collected feedback. In Appendix D, we prove that
$H_{0}-H_{T}=\mathcal{O}(d\cdot\log T)$ for linear kernel MDPs, which implies
$\mathfrak{R}(T)=\widetilde{O}(\sqrt{T})$. Here $\widetilde{O}$ hides the
logarithmic factor.

Figure 2: RAFA for Game of 24. Actions are proposed (dotted) and selected
(green). Hallucinations that the same number can be reused are mitigated
through interactions.

##  5 Experiment

We evaluate RAFA in several text-based benchmarks, e.g., Game of 24, ALFWorld,
BlocksWorld, and Tic-Tac-Toe. The detailed setups, results, and ablations are
provided in Appendix E, while the detailed prompts are found in Appendix F.

###  5.1 Game of 24

Game of 24 (Yao et al., 2023a) is a mathematical puzzle to obtain 24 from four
natural numbers through basic arithmetic operations. The state is the
(possibly unfinished) current formula and the action is the next formula (or
the modified part).

#### Setup.

We emulate the tree-search algorithm to plan ($B\in\\{1,2\\}$). At the $t$-th
step, RAFA learns from the memory buffer and switches to a new policy upon
receiving an unexpected reward, which is the switching condition. After the
$t$-th step, RAFA digests the collected feedback and generates a linguistic
summary, which is saved into the memory buffer to avoid similar previous
mistakes.

|  RAFA $(B=1)$ |  RAFA $(B=2)$ |  ToT $(B=1)$ |  ToT $(B=2)$ | Reflexion
---|---|---|---|---|---
gpt-4 | 89% | 93% | 73% | 81% | 21%
gpt-3.5 | 29% | 46% | 10% | 17% | 16%

tableGame of 24 results.

#### Result.

RAFA attains SOTA performances as shown in Table 5.1. RAFA achieves superior
sample efficiency by mitigating hallucinations and avoid careless trials
(Figures 2 and 3).

(a)

(b)

Figure 3: Sample efficiency on Game of 24.

###  5.2 ALFWorld

Figure 4: An illustration of RAFA in the ALFWorld environment.

ALFWorld (Shridhar et al., 2020) is an interactive environment for embodied
agent simulations, which encompasses $134$ household tasks in six overall
categories (Table 1). We use gpt-3 (text-davinci-003).

Figure 5: Sample efficiency on ALFWorld.

#### Setup.

We emulate the tree-search algorithm to plan ($B=2$). RAFA invokes Critic to
evaluate the completed portion of the desired goal and switches to a new
policy after $20$ consecutive failures.

#### Result.

RAFA outperforms various existing frameworks (Figure 5). The better
performance of AdaPlanner at the initial episode is attributed to a
handcrafted set of programs for rejecting suboptimal candidate trajectories,
which is challenging to construct without the domain knowledge of a specific
task. One such example is the PickTwo category.

|  Pick |  Clean |  Heat |  Cool |  Examine |  PickTwo |  Total
---|---|---|---|---|---|---|---
BUTLER |  46.00 |  39.00 |  74.00 |  100.00 |  22.00 |  24.00 |  37.00
ReAct |  66.67 |  41.94 |  91.03 |  80.95 |  55.56 |  35.29 |  61.94
AdaPlanner |  100.00 |  96.77 |  95.65 |  100.00 |  100.00 |  47.06 |  91.79
Reflexion |  100.00 |  90.32 |  82.61 |  90.48 |  100.00 |  94.12 |  92.54
RAFA |  100.00 |  96.77 |  100.00 |  100.00 |  100.00 |  100.00 |  99.25
Table 1: ALFWorld results (success rates %).

###  5.3 BlocksWorld

BlocksWorld (Hao et al., 2023) contains tasks to arrange blocks in specific
configurations.

#### Setup.

We use the Vicuna (Zheng et al., 2023) model and emulate the MCTS algorithm to
plan.

Figure 6: RAFA for BlocksWorld.

#### Result.

RAFA achieves superior success rates across multiple Vicuna versions (Figure
7). Comparisons with CoT and RAP demonstrate how the learning subroutine
improves the planning optimality.

(a)  (b)

(c)  (d)

Figure 7: Sample efficiency on BlocksWorld ($4$ and $6$ are the minimum
numbers of steps for solving a specific task). CoT is prompted by four in-
context examples.

###  5.4 Tic-Tac-Toe

Figure 8: Sample efficiency on Tic-Tac-Toe ($0$ means tie).

Tic-Tac-Toe (Beck, 2008) is a competitive game where the X and O sides take
turns to place marks. RAFA invokes Model to simulate the transition and
opponent dynamics (Figure 9).

#### Setup.

We use gpt-4 and emulate the tree-search algorithm to plan ($B\in\\{3,4\\}$).
RAFA switches to a new policy when (a) the predicted state differs from the
observed one, (2) the predicted action of opponents differs from the observed
one, or (3) Critic gives the wrong prediction of the game status. Here, X has
an asymmetric advantage (winning surely if played properly).

#### Result.

RAFA (playing O) matches and beats gpt-4 for $T=5$ and $T=7$ (Table 5.4),
although O is destined to lose. The ablation study ($B=3$ versus $B=4$)
illustrates how the planning suboptimality affects the sample efficiency
(Figure 8).

Figure 9: RAFA (playing O) for Tic-Tac-Toe.

X O |  gpt-4 |  RAFA $(T\\!=\\!1)$ |  RAFA $(T\\!=\\!5)$ |  RAFA $(T\\!=\\!7)$
---|---|---|---|---
gpt-4 |  90%, 0%, 10% |  90%, 0%, 10% |  50%, 0%, 50% |  0%, 0%, 100%

tableTic-Tac-Toe Results. We set $B=4$ and report the winning rate of X, the
tie rate, and the winning rate of O.

## References

  * Abbasi-Yadkori et al. (2011) Abbasi-Yadkori, Y., Pál, D. and Szepesvári, C. (2011).  Improved algorithms for linear stochastic bandits.  In Advances in Neural Information Processing Systems.
  * Abbasi-Yadkori and Szepesvári (2015) Abbasi-Yadkori, Y. and Szepesvári, C. (2015).  Bayesian optimal control of smoothly parameterized systems.  In Uncertainty in Artificial Intelligence.
  * Abernethy et al. (2023) Abernethy, J., Agarwal, A., Marinov, T. V. and Warmuth, M. K. (2023).  A mechanism for sample-efficient in-context learning for sparse retrieval tasks.  arXiv preprint arXiv:2305.17040.
  * Ahn et al. (2022) Ahn, M., Brohan, A., Brown, N., Chebotar, Y., Cortes, O., David, B., Finn, C., Fu, C., Gopalakrishnan, K., Hausman, K. et al. (2022).  Do as I can, not as I say: Grounding language in robotic affordances.  arXiv preprint arXiv:2204.01691.
  * Akyürek et al. (2022) Akyürek, E., Schuurmans, D., Andreas, J., Ma, T. and Zhou, D. (2022).  What learning algorithm is in-context learning? Investigations with linear models.  arXiv preprint arXiv:2211.15661.
  * Ayoub et al. (2020) Ayoub, A., Jia, Z., Szepesvari, C., Wang, M. and Yang, L. (2020).  Model-based reinforcement learning with value-targeted regression.  In International Conference on Machine Learning.
  * Beck (2008) Beck, J. (2008).  Combinatorial games: Tic-Tac-Toe theory.
  * Betts (1998) Betts, J. T. (1998).  Survey of numerical methods for trajectory optimization.  Journal of Guidance, Control, and Dynamics.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A. et al. (2020).  Language models are few-shot learners.  Advances in neural information processing systems, 33 1877–1901.
  * Cai et al. (2020) Cai, Q., Yang, Z., Jin, C. and Wang, Z. (2020).  Provably efficient exploration in policy optimization.  In International Conference on Machine Learning.
  * Cai et al. (2023) Cai, T., Wang, X., Ma, T., Chen, X. and Zhou, D. (2023).  Large language models as tool makers.  arXiv preprint arXiv:2305.17126.
  * Chen et al. (2022) Chen, Y., He, J. and Gu, Q. (2022).  On the sample complexity of learning infinite-horizon discounted linear kernel MDPs.  In International Conference on Machine Learning.
  * Chowdhery et al. (2022) Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S. et al. (2022).  PaLM: Scaling language modeling with pathways.  arXiv preprint arXiv:2204.02311.
  * Creswell and Shanahan (2022) Creswell, A. and Shanahan, M. (2022).  Faithful reasoning using large language models.  arXiv preprint arXiv:2208.14271.
  * Creswell et al. (2022) Creswell, A., Shanahan, M. and Higgins, I. (2022).  Selection-inference: Exploiting large language models for interpretable logical reasoning.  arXiv preprint arXiv:2205.09712.
  * Dong et al. (2019) Dong, K., Wang, Y., Chen, X. and Wang, L. (2019).  Q-learning with UCB exploration is sample efficient for infinite-horizon MDP.  arXiv preprint arXiv:1901.09311.
  * Gao et al. (2021) Gao, M., Xie, T., Du, S. S. and Yang, L. F. (2021).  A provably efficient algorithm for linear markov decision process with low switching cost.  arXiv preprint arXiv:2101.00494.
  * Garg et al. (2022) Garg, S., Tsipras, D., Liang, P. S. and Valiant, G. (2022).  What can transformers learn in-context? A case study of simple function classes.  In Advances in Neural Information Processing Systems.
  * Ghavamzadeh et al. (2015) Ghavamzadeh, M., Mannor, S., Pineau, J., Tamar, A. et al. (2015).  Bayesian reinforcement learning: A survey.  Foundations and Trends® in Machine Learning, 8 359–483.
  * Ghosh (2021) Ghosh, M. (2021).  Exponential tail bounds for chisquared random variables.  Journal of Statistical Theory and Practice, 15 35\.
  * Hao et al. (2023) Hao, S., Gu, Y., Ma, H., Hong, J. J., Wang, Z., Wang, D. Z. and Hu, Z. (2023).  Reasoning with language model is planning with world model.  arXiv preprint arXiv:2305.14992.
  * Hoffmann et al. (2022) Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A. et al. (2022).  Training compute-optimal large language models.  arXiv preprint arXiv:2203.15556.
  * Huang et al. (2022a) Huang, W., Abbeel, P., Pathak, D. and Mordatch, I. (2022a).  Language models as zero-shot planners: Extracting actionable knowledge for embodied agents.  In International Conference on Machine Learning.
  * Huang et al. (2022b) Huang, W., Xia, F., Xiao, T., Chan, H., Liang, J., Florence, P., Zeng, A., Tompson, J., Mordatch, I., Chebotar, Y. et al. (2022b).  Inner monologue: Embodied reasoning through planning with language models.  arXiv preprint arXiv:2207.05608.
  * Jiang (2023) Jiang, H. (2023).  A latent space theory for emergent abilities in large language models.  arXiv preprint arXiv:2304.09960.
  * Jin et al. (2021) Jin, Y., Yang, Z. and Wang, Z. (2021).  Is pessimism provably efficient for offline RL?  In International Conference on Machine Learning.
  * Kim et al. (2023) Kim, G., Baldi, P. and McAleer, S. (2023).  Language models can solve computer tasks.  arXiv preprint arXiv:2303.17491.
  * Kirsch et al. (2022) Kirsch, L., Harrison, J., Sohl-Dickstein, J. and Metz, L. (2022).  General-purpose in-context learning by meta-learning transformers.  arXiv preprint arXiv:2212.04458.
  * Lee et al. (2023) Lee, J. N., Xie, A., Pacchiano, A., Chandak, Y., Finn, C., Nachum, O. and Brunskill, E. (2023).  Supervised pretraining can learn in-context reinforcement learning.  arXiv preprint arXiv:2306.14892.
  * Li et al. (2022) Li, B. Z., Nye, M. and Andreas, J. (2022).  Language modeling with latent situations.  arXiv preprint arXiv:2212.10012.
  * Li et al. (2023) Li, Y., Ildiz, M. E., Papailiopoulos, D. and Oymak, S. (2023).  Transformers as algorithms: Generalization and implicit model selection in in-context learning.  arXiv preprint arXiv:2301.07067.
  * Liang et al. (2022) Liang, P., Bommasani, R., Lee, T., Tsipras, D., Soylu, D., Yasunaga, M., Zhang, Y., Narayanan, D., Wu, Y., Kumar, A. et al. (2022).  Holistic evaluation of language models.  arXiv preprint arXiv:2211.09110.
  * Liu et al. (2023) Liu, B., Jiang, Y., Zhang, X., Liu, Q., Zhang, S., Biswas, J. and Stone, P. (2023).  LLM+P: Empowering large language models with optimal planning proficiency.  arXiv preprint arXiv:2304.11477.
  * Lu et al. (2023) Lu, P., Peng, B., Cheng, H., Galley, M., Chang, K.-W., Wu, Y. N., Zhu, S.-C. and Gao, J. (2023).  Chameleon: Plug-and-play compositional reasoning with large language models.  arXiv preprint arXiv:2304.09842.
  * Lu and Van Roy (2019) Lu, X. and Van Roy, B. (2019).  Information-theoretic confidence bounds for reinforcement learning.  Advances in Neural Information Processing Systems.
  * Muandet et al. (2017) Muandet, K., Fukumizu, K., Sriperumbudur, B., Schölkopf, B. et al. (2017).  Kernel mean embedding of distributions: A review and beyond.  Foundations and Trends® in Machine Learning, 10 1–141.
  * Olsson et al. (2022) Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A. et al. (2022).  In-context learning and induction heads.  arXiv preprint arXiv:2209.11895.
  * OpenAI (2023) OpenAI (2023).  GPT-4 technical report.
  * Osband et al. (2013) Osband, I., Russo, D. and Van Roy, B. (2013).  (More) efficient reinforcement learning via posterior sampling.  In Advances in Neural Information Processing Systems.
  * Paul et al. (2023) Paul, D., Ismayilzada, M., Peyrard, M., Borges, B., Bosselut, A., West, R. and Faltings, B. (2023).  REFINER: Reasoning feedback on intermediate representations.  arXiv preprint arXiv:2304.01904.
  * Radford et al. (2019) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D. and Sutskever, I. (2019).  Language models are unsupervised multitask learners.
  * Rawlings (2000) Rawlings, J. B. (2000).  Tutorial overview of model predictive control.  IEEE Control Systems Magazine.
  * Razeghi et al. (2022) Razeghi, Y., Logan IV, R. L., Gardner, M. and Singh, S. (2022).  Impact of pretraining term frequencies on few-shot reasoning.  arXiv preprint arXiv:2202.07206.
  * Russo and Van Roy (2014a) Russo, D. and Van Roy, B. (2014a).  Learning to optimize via information-directed sampling.  In Advances in Neural Information Processing Systems.
  * Russo and Van Roy (2014b) Russo, D. and Van Roy, B. (2014b).  Learning to optimize via posterior sampling.  Mathematics of Operations Research.
  * Russo and Van Roy (2016) Russo, D. and Van Roy, B. (2016).  An information-theoretic analysis of Thompson sampling.  Journal of Machine Learning Research.
  * Saunshi et al. (2020) Saunshi, N., Malladi, S. and Arora, S. (2020).  A mathematical exploration of why language models help solve downstream tasks.  arXiv preprint arXiv:2010.03648.
  * Sel et al. (2023) Sel, B., Al-Tawaha, A., Khattar, V., Wang, L., Jia, R. and Jin, M. (2023).  Algorithm of thoughts: Enhancing exploration of ideas in large language models.  arXiv preprint arXiv:2308.10379.
  * Shen et al. (2023) Shen, Y., Song, K., Tan, X., Li, D., Lu, W. and Zhuang, Y. (2023).  HuggingGPT: Solving AI tasks with ChatGPT and its friends in HuggingFace.  arXiv preprint arXiv:2303.17580.
  * Shin et al. (2022) Shin, S., Lee, S.-W., Ahn, H., Kim, S., Kim, H., Kim, B., Cho, K., Lee, G., Park, W., Ha, J.-W. et al. (2022).  On the effect of pretraining corpora on in-context learning by a large-scale language model.  arXiv preprint arXiv:2204.13509.
  * Shinn et al. (2023) Shinn, N., Cassano, F., Labash, B., Gopinath, A., Narasimhan, K. and Yao, S. (2023).  Reflexion: Language agents with verbal reinforcement learning.  arXiv preprint arXiv:2303.11366.
  * Shridhar et al. (2020) Shridhar, M., Yuan, X., Côté, M.-A., Bisk, Y., Trischler, A. and Hausknecht, M. (2020).  Alfworld: Aligning text and embodied environments for interactive learning.  arXiv preprint arXiv:2010.03768.
  * Song et al. (2013) Song, L., Fukumizu, K. and Gretton, A. (2013).  Kernel embeddings of conditional distributions: A unified kernel framework for nonparametric inference in graphical models.  IEEE Signal Processing Magazine, 30 98–111.
  * Strens (2000) Strens, M. (2000).  A Bayesian framework for reinforcement learning.  In International Conference on Machine Learning.
  * Sun et al. (2023) Sun, H., Zhuang, Y., Kong, L., Dai, B. and Zhang, C. (2023).  AdaPlanner: Adaptive planning from feedback with language models.  arXiv preprint arXiv:2305.16653.
  * Sutton and Barto (2018) Sutton, R. S. and Barto, A. G. (2018).  Reinforcement learning: An introduction.
  * Touvron et al. (2023) Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F. et al. (2023).  LLaMa: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971.
  * Von Oswald et al. (2023) Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A. and Vladymyrov, M. (2023).  Transformers learn in-context by gradient descent.  In International Conference on Machine Learning.
  * Wang et al. (2022) Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A. and Zhou, D. (2022).  Self-consistency improves chain of thought reasoning in language models.  arXiv preprint arXiv:2203.11171.
  * Wang et al. (2023a) Wang, X., Zhu, W. and Wang, W. Y. (2023a).  Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning.  arXiv preprint arXiv:2301.11916.
  * Wang et al. (2023b) Wang, Z., Cai, S., Liu, A., Ma, X. and Liang, Y. (2023b).  Describe, explain, plan and select: Interactive planning with large language models enables open-world multi-task agents.  arXiv preprint arXiv:2302.01560.
  * Wei et al. (2021) Wei, C., Xie, S. M. and Ma, T. (2021).  Why do pretrained language models help in downstream tasks? An analysis of head and prompt tuning.  arXiv preprint arXiv:2106.09226.
  * Wei et al. (2020) Wei, C.-Y., Jahromi, M. J., Luo, H., Sharma, H. and Jain, R. (2020).  Model-free reinforcement learning in infinite-horizon average-reward Markov decision processes.  In International Conference on Machine Learning.
  * Wei et al. (2022) Wei, J., Wang, X., Schuurmans, D., Bosma, M., Xia, F., Chi, E., Le, Q. V. and Zhou, D. (2022).  Chain-of-thought prompting elicits reasoning in large language models.  In Advances in Neural Information Processing Systems.
  * Wies et al. (2023) Wies, N., Levine, Y. and Shashua, A. (2023).  The learnability of in-context learning.  arXiv preprint arXiv:2303.07895.
  * Xie et al. (2021) Xie, S. M., Raghunathan, A., Liang, P. and Ma, T. (2021).  An explanation of in-context learning as implicit Bayesian inference.  arXiv preprint arXiv:2111.02080.
  * Yang and Wang (2019) Yang, L. and Wang, M. (2019).  Sample-optimal parametric $q$-learning using linearly additive features.  In International Conference on Machine Learning.
  * Yang and Wang (2020) Yang, L. and Wang, M. (2020).  Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound.  In International Conference on Machine Learning.
  * Yao et al. (2023a) Yao, S., Yu, D., Zhao, J., Shafran, I., Griffiths, T. L., Cao, Y. and Narasimhan, K. (2023a).  Tree of thoughts: Deliberate problem solving with large language models.  arXiv preprint arXiv:2305.10601.
  * Yao et al. (2022) Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K. and Cao, Y. (2022).  ReAct: Synergizing reasoning and acting in language models.  arXiv preprint arXiv:2210.03629.
  * Yao et al. (2023b) Yao, Y., Li, Z. and Zhao, H. (2023b).  Beyond chain-of-thought, effective graph-of-thought reasoning in large language models.  arXiv preprint arXiv:2305.16582.
  * Zelikman et al. (2022) Zelikman, E., Wu, Y., Mu, J. and Goodman, N. (2022).  STaR: Bootstrapping reasoning with reasoning.  In Advances in Neural Information Processing Systems.
  * Zhang et al. (2022) Zhang, Y., Liu, B., Cai, Q., Wang, L. and Wang, Z. (2022).  An analysis of attention via the lens of exchangeability and latent variable models.  arXiv preprint arXiv:2212.14852.
  * Zhang et al. (2023a) Zhang, Y., Yang, J., Yuan, Y. and Yao, A. C.-C. (2023a).  Cumulative reasoning with large language models.  arXiv preprint arXiv:2308.04371.
  * Zhang et al. (2023b) Zhang, Y., Zhang, F., Yang, Z. and Wang, Z. (2023b).  What and how does in-context learning learn? Bayesian model averaging, parameterization, and generalization.  arXiv preprint arXiv:2305.19420.
  * Zheng et al. (2023) Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. et al. (2023).  Judging LLM-as-a-judge with MT-bench and chatbot arena.  arXiv preprint arXiv:2306.05685.
  * Zhou et al. (2021a) Zhou, D., Gu, Q. and Szepesvári, C. (2021a).  Nearly minimax optimal reinforcement learning for linear mixture Markov decision processes.  In Conference on Learning Theory.
  * Zhou et al. (2021b) Zhou, D., He, J. and Gu, Q. (2021b).  Provably efficient reinforcement learning for discounted MDPs with feature mapping.  In International Conference on Machine Learning.

##  Appendix A More Algorithms

Depending on the specific configuration of the state and action spaces
(continuous versus discrete) and the transition and reward models (stochastic
versus deterministic), we may choose to emulate the tree-search algorithm, the
value iteration algorithm, the random shooting algorithm, or the MCTS
algorithm. All of them allow RAFA to achieve provable sample efficiency
guarantees as long as they satisfy a specific requirement of optimality
(Definition 4.2). For illustration, we emulate the beam-search algorithm (an
advanced version of the tree-search algorithm) in Algorithm 4 and the MCTS
algorithm in Algorithm 5. For the theoretical discussion, we present the value
iteration algorithm in Algorithm 6.

Algorithm 4 The LLM learner-planner (LLM-LR-PL): A beam-search example (for
the deterministic case).

1: input: The memory buffer $\mathcal{D}$, the initial state $s$, the proposal
width $L$, the search breadth $B$, and the search depth $U$.

2: initialization: Initialize the state array
$\mathcal{S}_{0}\leftarrow\\{s\\}$ and the action array
$\mathcal{A}_{0}\leftarrow\varnothing$. ———————————————— (the learning
subroutine) ———————————–

3: Set Model as an LLM instance prompted to use $\mathcal{D}$ as contexts to
generate the next state.

4: Set Critic as an LLM instance prompted to use $\mathcal{D}$ as contexts to
estimate the value function. ———————————————— (the planning subroutine)
———————————–

5: Set Elite as an LLM instance prompted to use $\mathcal{D}$ as contexts to
generate multiple candidate actions.

6: for $u=0,\ldots,U$ do

7: For each current state $s_{u}$ in $\mathcal{S}_{u}$, invoke Elite to
generate $L$ candidate actions.

8: For each candidate action $a_{u}^{(\ell)}$, invoke Model to generate the
next state $s_{u+1}^{(\ell)}$ and the received reward $r_{u}^{(\ell)}$.

9: For each resulting tuple
$(s_{u},a_{u}^{(\ell)},s_{u+1}^{(\ell)},r_{u}^{(\ell)})$, invoke Critic to
evaluate the expected cumulative future reward
$\widehat{Q}(s_{u},a^{(\ell)}_{u})\leftarrow
r_{u}^{(\ell)}+\gamma\widehat{V}(s^{(\ell)}_{u+1})$, where $\widehat{V}$ is
given by Critic.

10: Select $B$ best tuples $(s_{u},a_{u}^{(\ell)},s_{u+1}^{(\ell)})$ with the
highest value $\widehat{Q}(s_{u},a^{(\ell)}_{u})$ and write them to
$\mathcal{S}_{u}\times\mathcal{A}_{u}\times\mathcal{S}_{u+1}$.

11: end for

12: For $B$ preserved rollouts in
$\mathcal{S}_{0}\times\mathcal{A}_{0}\times\cdots\times\mathcal{S}_{U}\times\mathcal{A}_{U}\times\mathcal{S}_{U+1}$,
invoke Critic to evaluate the expected cumulative future reward
$\sum_{u=0}^{U}\gamma^{u}r_{u}^{(b)}+\gamma^{U+1}\widehat{V}(s_{U+1}^{(b)})$
and select the best one
$(s_{0}^{\dagger},a_{0}^{\dagger},\ldots,s_{U}^{\dagger},a_{U}^{\dagger},s_{U+1}^{\dagger})$,
where $\widehat{V}$ is given by Critic and $s_{0}^{\dagger}=s$.

13: output: The initial action $a_{0}^{\dagger}$ of the selected rollout.

Algorithm 5 LLM learner-planner (LLM-PL) for RAFA: A Monte-Carlo tree-search
example (for the stochastic case).

1: input: The memory buffer $\mathcal{D}$, the initial state $s$, the proposal
width $L$, $L^{\prime}$, and the expansion budget $E$.

2: initialization: Initialize the root node $n\leftarrow s$ and the child
function $c(\cdot)\leftarrow\varnothing$. ———————————————— (the learning
subroutine) ———————————–

3: Set Model as an LLM instance prompted to use $\mathcal{D}$ as contexts to
generate the next state.

4: Set Critic as an LLM instance prompted to use $\mathcal{D}$ as contexts to
estimate the value function. ———————————————— (the planning subroutine)
———————————–

5: Set Elite as an LLM instance prompted to use $\mathcal{D}$ as contexts to
generate multiple candidate actions.

6: for $e=0,\ldots,E$ do

7: Set $s_{e}\leftarrow n$.

8: while $s_{e}$ is not a leaf node, i.e., $c(s_{e})\neq\varnothing$, do

9: Invoke Critic to evaluate the expected cumulative future reward and select
the child node $a_{e}$ in $c(s_{e})$ with the highest value
$\widehat{Q}(s_{e},a_{e})$.

10: Set $s_{e}$ as a child node in $c(a_{e})$.

11: end while

12: For the current state $s_{e}$, invoke Elite to generate $L$ candidate
actions.

13: Write each candidate action $a^{(\ell)}_{e}$ to $c(s_{e})$, i.e.,
$c(s_{e})\leftarrow\\{a^{(\ell)}_{e}\\}_{\ell=1}^{L}$.

14: For each candidate action $a^{(\ell)}_{e}$, invoke Model to sample
$L^{\prime}$ next states.

15: Write each next state $s^{(\ell,\ell^{\prime})}_{e}$ to
$c(a^{(\ell)}_{e})$, i.e.,
$c(a^{(\ell)}_{e})\leftarrow\\{s^{(\ell,\ell^{\prime})}_{e}\\}_{\ell^{\prime}=1}^{L^{\prime}}$.

16: For each generated state $s^{(\ell,\ell^{\prime})}_{e}$, invoke Critic to
evaluate the expected cumulative future reward and update the estimated value
$\widehat{V}$ for all ancestor nodes.(Optional)

17: end for

18: Set $s^{\dagger}_{0}\leftarrow n$ and $i\leftarrow 0$.

19: while $s^{\dagger}_{i}$ is not a leaf node, i.e.,
$c(s^{\dagger}_{i})\neq\varnothing$, do

20: Invoke Critic to evaluate the expected cumulative future reward and select
the child node $a^{\dagger}_{i+1}$ in $c(s^{\dagger}_{i})$ with the highest
value $\widehat{Q}(s^{\dagger}_{i},a^{\dagger}_{i})$.

21: Set $s^{\dagger}_{i+1}$ as a child node in $c(a^{\dagger}_{i})$ and
$i\leftarrow i+1$.

22: end while

23: output: The initial action $a_{0}^{\dagger}$ of the selected rollout
$(s_{0}^{\dagger},a_{0}^{\dagger},\ldots,s_{i}^{\dagger},a_{i}^{\dagger})$.

We also present the value iteration algorithm (Algorithm 6) with a truncated
horizon $U$, i.e., a finite length of the lookahead window as the
$\epsilon$-optimal planner in Algorithm 3. The following proposition ensures
that Algorithm 6 satisfies Definition 4.2.

######  Proposition A.1.

Algorithm 6 is an $\epsilon$-optimal planner as long as we set
$U\geq\lceil\log((2L/(\epsilon(1-\gamma))))/\log(1/\gamma)\rceil+1$ and any
value function is upper bounded by $L\geq 0$.

###### Proof.

See Appendix A.1 for a detailed proof. ∎

Algorithm 6 $\epsilon$-Optimal planner: The value iteration algorithm with a
truncated horizon.

1: input: The model parameter $\theta$ and the truncated horizon $U$.

2: initialization: Set the value function $V_{\theta}^{(U)}(\cdot)\leftarrow
0$.

3: for $u=U-1,\ldots,1$ do

4: Set the value function
$V_{\theta}^{(u)}(\cdot)\leftarrow\max_{a\in\mathcal{A}}Q_{\theta}^{(u)}(\cdot,a)$,
where $Q^{(u)}_{\theta}(\cdot,\cdot)\leftarrow
r_{\theta}(\cdot,\cdot)+\gamma(P_{\theta}V_{\theta}^{(u+1)})(\cdot,\cdot)$.

5: end for

6: output: The greedy policy
$\pi(\cdot)=\arg\max_{a\in\mathcal{A}}Q^{(1)}_{\theta}(\cdot,a)$.

###  A.1 Proof of Proposition A.1

###### Proof of Proposition A.1.

We prove that Algorithm 6 satisfies Definition 4.2, where $U$ is dependent on
$\epsilon$. Let $\pi_{\theta}(\theta)$ be the output policy of Algorithm 6 for
the given $\theta$. For notational simplicity, we denote
$\max_{s\in{\mathcal{S}}}$ and $\max_{a\in\mathcal{A}}$ as $\max_{s}$ and
$\max_{a}$. Then we have

$\displaystyle\max_{s}\bigl{[}V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s)-V_{\theta}^{\pi_{\theta}}(s)\bigr{]}$
$\displaystyle\leq\underbrace{\max_{s}\bigl{[}V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s)-\max_{a}Q^{(1)}_{\theta}(s,a)\bigr{]}}_{\text{term
(A) }}$
$\displaystyle\qquad+\underbrace{\max_{s}\bigl{[}\max_{a}Q^{(1)}_{\theta}(s,a)-V_{\theta}^{\pi_{\theta}}(s)\bigr{]}}_{\text{term
(B)}},$ (A.1)

where $Q^{(1)}_{\theta}$ is defined in Algorithm 6.

#### Analysis of Term (A).

For any $1\leq u<U$ and $s\in\mathcal{S}$, we have

$\displaystyle
V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s)-\max_{a}Q^{(u)}_{\theta}(s,a)$
$\displaystyle=\max_{a}Q_{\theta}^{\texttt{PL}^{\star}(\theta)}(s,a)-\max_{a}Q^{(u)}_{\theta}(s,a)$
$\displaystyle\leq\max_{a}\bigl{|}Q_{\theta}^{\texttt{PL}^{\star}(\theta)}(s)-Q^{(u)}_{\theta}(s,a)\bigr{|}$
$\displaystyle\leq\gamma\cdot\max_{a}\bigl{|}\mathbb{E}_{s^{\prime}\sim
P_{\theta}(\cdot\,|\,s,a)}\bigl{[}V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s^{\prime})-\max_{a^{\prime}}Q_{\theta}^{(u+1)}(s^{\prime},a^{\prime})\bigr{]}\bigr{|}$
$\displaystyle\leq\gamma\cdot\max_{s^{\prime}}\bigl{|}V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s^{\prime})-\max_{a^{\prime}}Q_{\theta}^{(u+1)}(s^{\prime},a^{\prime})\bigr{|},$
(A.2)

where the first equality uses the definition of $\texttt{PL}^{\star}$ in
Section 2, the first inequality uses the contraction property of the $\max$
operator, and the second equality uses the Bellman equation in (C.7) and
Algorithm 6. By induction, we have

term (A)
$\displaystyle\leq\gamma^{U-1}\max_{s}\cdot\bigl{|}V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s)-\max_{a}Q_{\theta}^{(U)}(s,a)\bigr{|}$
$\displaystyle\leq\gamma^{U-1}L,$ (A.3)

where the last equality uses $Q_{\theta}^{(U)}=0$ and the fact that any value
function is upper bounded by $L$.

#### Analysis of Term (B).

Recall that $\pi_{\theta}(s)=\arg\max_{a}Q^{(1)}_{\theta}(s,a)$ in Algorithm
6. We have

term (B)
$\displaystyle=\max_{s}\bigl{[}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-Q_{\theta}^{\pi_{\theta}}(s,\pi_{\theta}(s))\bigr{]}.$
(A.4)

Let

$\displaystyle\epsilon^{\dagger}=\max_{s}\bigl{[}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-r_{\theta}(s,\pi_{\theta}(s))-\gamma(P_{\theta}V^{(1)}_{\theta})(s,\pi_{\theta}(s))\bigr{]}.$
(A.5)

We have

$\displaystyle
Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-Q_{\theta}^{\pi_{\theta}}(s,\pi_{\theta}(s))$
$\displaystyle\leq\epsilon^{\dagger}+r_{\theta}(s,\pi_{\theta}(s))+\gamma(P_{\theta}V^{(1)}_{\theta})(s,\pi_{\theta}(s))$
$\displaystyle\qquad-
r_{\theta}(s,\pi_{\theta}(s))-\gamma(P_{\theta}V^{\pi_{\theta}}_{\theta})(s,\pi_{\theta}(s))$
$\displaystyle=\epsilon^{\dagger}+\gamma\cdot\mathbb{E}_{s^{\prime}\sim
P_{\theta}(\cdot\,|\,s,\pi_{\theta}(s))}\bigl{[}Q_{\theta}^{(1)}(s^{\prime},\pi_{\theta}(s^{\prime}))-Q_{\theta}^{\pi_{\theta}}(s^{\prime},\pi_{\theta}(s^{\prime}))\bigr{]}$
$\displaystyle\leq\epsilon^{\dagger}+\gamma\cdot\max_{s}\bigl{|}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-Q_{\theta}^{\pi_{\theta}}(s,\pi_{\theta}(s))\bigr{|}.$
(A.6)

Here, the first inequality is based on the definition of $\epsilon^{\dagger}$
and the Bellman equation in (C.7) in the proof of Theorem 4.4, where the first
equality uses fact that
$Q_{\theta}^{(1)}(s,\pi_{\theta}(s))=\max_{a}Q^{(1)}_{\theta}(s,a)=V_{\theta}^{(1)}(s)$
for all $s\in\mathcal{S}$. Taking the $\sup_{s\in{\mathcal{S}}}$ in the left-
hand side of (A.6), we obtain

$\displaystyle\max_{s}\bigl{|}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-Q_{\theta}^{\pi_{\theta}}(s,\pi_{\theta}(s))\bigr{|}\leq\epsilon^{\dagger}+\gamma\cdot\max_{s}\bigl{|}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-Q_{\theta}^{\pi_{\theta}}(s,\pi_{\theta}(s))\bigr{|},$

which implies

term (B)
$\displaystyle\leq\max_{s}\bigl{|}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-Q_{\theta}^{\pi_{\theta}}(s,\pi_{\theta}(s))\bigr{|}$
$\displaystyle\leq\epsilon^{\dagger}/(1-\gamma).$ (A.7)

Meanwhile, the convergence analysis of the value iteration algorithm in Sutton
and Barto (2018) gives

$\displaystyle\max_{s,a}\bigl{|}Q_{\theta}^{(1)}(s,a)-Q_{\theta}^{(2)}(s,a)\bigr{|}\leq\gamma^{U-2}\max_{s,a}\bigl{|}Q_{\theta}^{(U-1)}(s,a)-Q_{\theta}^{(U)}(s,a)\bigr{|},$

which implies

$\displaystyle\max_{s,a}\bigl{|}Q_{\theta}^{(1)}(s,a)-Q_{\theta}^{(2)}(s,a)\bigr{|}\leq\gamma^{U-1}L.$

We have

$\displaystyle\epsilon^{\dagger}$
$\displaystyle=\max_{s}\bigl{[}Q_{\theta}^{(1)}(s,\pi_{\theta}(s))-r_{\theta}(s,\pi_{\theta}(s))-\gamma(P_{\theta}V^{(2)}_{\theta})(s,\pi_{\theta}(s))$
$\displaystyle\qquad\qquad+\gamma\mathbb{E}_{s^{\prime}\sim
P_{\theta}(\cdot\,|\,s,\pi_{\theta}(s))}\bigl{[}V_{\theta}^{(1)}(s^{\prime})-V_{\theta}^{(2)}(s^{\prime})\bigr{]}\bigr{]}$
$\displaystyle=\gamma\cdot\mathbb{E}_{s^{\prime}\sim
P_{\theta}(\cdot\,|\,s,\pi_{\theta}(s))}\bigl{[}V_{\theta}^{(1)}(s^{\prime})-V_{\theta}^{(2)}(s^{\prime})\bigr{]}$
$\displaystyle=\gamma\cdot\mathbb{E}_{s^{\prime}\sim
P_{\theta}(\cdot\,|\,s,\pi_{\theta}(s))}\bigl{[}\max_{a}Q_{\theta}^{(1)}(s^{\prime},a)-\max_{a}Q_{\theta}^{(2)}(s^{\prime},a)\bigr{]}$
$\displaystyle\leq\gamma\cdot\mathbb{E}_{s^{\prime}\sim
P_{\theta}(\cdot\,|\,s,\pi_{\theta}(s))}\bigl{[}\max_{a}\bigl{|}Q_{\theta}^{(1)}(s^{\prime},a)-Q_{\theta}^{(2)}(s^{\prime},a)\bigr{|}\bigr{]}$
$\displaystyle\leq\gamma^{U-1}L,$ (A.8)

where the first and third equalities are based on Algorithm 6, the second last
inequality uses the contraction property of the $\max$ operator.

Plugging (A.3), (A.7), and (A.8) into (A.1), we obtain

$\displaystyle\max_{s}\bigl{[}V_{\theta}^{\texttt{PL}^{\star}(\theta)}(s)-V_{\theta}^{\pi_{\theta}}(s)\bigr{]}\leq{\gamma^{U-1}L}+\gamma^{U-1}L/(1-\gamma)\leq\epsilon,$

for $U\geq\lceil\log((2L/(\epsilon(1-\gamma))))/\log(1/\gamma)\rceil+1$. Thus,
we prove Proposition A.1. ∎

##  Appendix B LLM Samples Posterior

###### Verification of Assumption 4.1.

We prove that Assumption 4.1 holds for a specific parameterization of
$(P_{\theta},r_{\theta})$ under three regularity conditions, namely (a) LLMs
are trained to replicate the pretraining distribution, which is assumed in
Saunshi et al. (2020); Wei et al. (2021); Xie et al. (2021) to simplify the
statistical analysis, (b) the pretraining dataset is generated through a
Bayesian mechanism with a latent concept, which is a simplified version of the
latent variable model in Xie et al. (2021) and resembles that in Wang et al.
(2023a), and (c) LLMs are able to parameterize an implicit Bayesian inference
mechanism, which is proved in Zhang et al. (2022, 2023b) for the attention
architecture. Note that, if Assumption 4.1 holds approximately, the regret
analysis can be relaxed to accommodate the additional error in the posterior
distribution.

######  Assumption B.1 (Model Class).

We assume that $(P_{\theta},r_{\theta})$ is parameterized under the following
conditions.

  * Factorization:

For any
$(s^{\prime},s,a,g)\in\mathcal{S}\times\mathcal{S}\times\mathcal{A}\times\mathcal{G}$,
it holds that

$\displaystyle
P_{\theta}(s^{\prime}\,|\,s,a;g)=\phi_{g}(s^{\prime},s,a)^{\top}f_{g}(\theta),$

where
$\phi_{g}:\mathcal{S}\times\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}^{d_{g}}$
is the feature mapping, $f_{g}:\Theta\mapsto\mathbb{R}^{d_{g}}$ is the
parameter transformation, and $d_{g}$ is the feature dimension.

  * No Redundant Feature:

For any $g\in\mathcal{G}$, there exists a set
$\\{(s_{(i)}^{\prime},s_{(i)},a_{(i)})\\}_{i=1}^{d_{g}}\subseteq\mathcal{S}\times\mathcal{S}\times\mathcal{A}$
such that the set
$\\{\phi_{g}(s_{(i)}^{\prime},s_{(i)},a_{(i)})\\}_{i=1}^{d_{g}}$ forms a basis
of $\mathbb{R}^{d_{g}}$.

  * Identification:

For two probability distributions $q_{1}$ and $q_{2}$ over $\Theta$, if it
holds for all $g\in\mathcal{G}$ that

$\displaystyle\mathbb{E}_{\theta\sim
q_{1}}[f_{g}(\theta)]=\mathbb{E}_{\theta\sim q_{2}}[f_{g}(\theta)],$

then we have $q_{1}=q_{2}$.

The factorization condition in Assumption B.1 is common in the RL literature
(Yang and Wang, 2020; Cai et al., 2020; Zhou et al., 2021b). As long as
$\phi_{g}$ has sufficient expressive power, $(P_{\theta},r_{\theta})$ is able
to represent a broad range of models.

The redundancy condition states that $\phi_{g}$ does not have any redundant
dimension. In other words, we are able to construct a basis of
$\mathbb{R}^{d_{g}}$ by choosing $d_{g}$ different tuples $(s^{\prime},s,a)$.
We provide the following example. Let $\mathcal{G}$ be
$\\{1,\ldots,|\mathcal{G}|\\}$ and $d_{g}=d$ for all $g\in\mathcal{G}$. We
assume that there exists a surjective mapping
$k:\mathcal{S}\times\mathcal{S}\times\mathcal{A}\mapsto\mathbb{B}^{|\mathcal{G}|d}$,
where $\mathbb{B}^{|\mathcal{G}|d}$ is the unit ball under the Euclidean norm
in $\mathbb{R}^{|\mathcal{G}|d}$. Let
$p_{g}:\mathbb{R}^{|\mathcal{G}|d}\mapsto\mathbb{R}^{d}$ be a projection
operator, which outputs the $((g-1)d+1)$-th to $gd$-th entries of the input
vector. We define $\phi_{g}$ as $p_{g}\circ k$. Since $k$ is surjective, we
are able to construct
$\\{(s^{\prime}_{(i)},s_{(i)}a_{(i)})\\}_{i=1}^{|\mathcal{G}|d}$ as the tuple
set in association with the canonical basis of $\mathbb{R}^{|\mathcal{G}|d}$.
As a result,
$\\{\phi_{g}(s_{(i)}^{\prime},s_{(i)},a_{(i)})\\}_{i=(g-1)d+1}^{gd}$ forms the
canonical basis of $\mathbb{R}^{d_{g}}$.

The identification condition ensures that the parameter distribution has a
one-to-one correspondence to its mean embedding (Song et al., 2013; Muandet et
al., 2017). For example, the Gaussian distribution is determined by its first-
order and second-order moments.

The following assumption specifies the data-generating process, which is a
simplified version of the latent variable model in Xie et al. (2021) and
resembles that in Wang et al. (2023a).

######  Assumption B.2 (Data Generation).

We assume that each individual rollout
$\\{(s_{t},a_{t},r_{t})\\}_{t=0}^{\infty}$ in the pretraining dataset
$\mathcal{D}$ is generated through the following process.

  1. Instruct:

Choose a specific instruction $g\in\mathcal{G}$ and sample an underlying
parameter $\theta^{\star}$ from the prior distribution $p_{0}$. Store $g$ in
$\mathcal{D}$.

  2. Initialize:

Sample the initial state $s_{0}$ from the initial distribution $\rho$.

  3. Interact:

At the $t$-th step, take an action $a_{t}$ given the history
$\\{(s_{i},a_{i},s_{i+1},r_{i})\\}_{i=0}^{t}$ and receive the reward $r_{t}$
and the next state $s_{t+1}$, where
$r_{t}=r_{\theta^{\star}}^{g}(s_{t},a_{t})$ and $s_{t+1}\sim
P_{\theta^{\star}}^{g}(\cdot\,|\,s_{t},a_{t})$.

The following assumption states that LLMs parameterize an implicit Bayesian
inference mechanism, which is proved in Zhang et al. (2022, 2023b) for the
attention architecture.

######  Assumption B.3 (Parameterization of Implicit Bayesian Inference).

We assume that LLMs parameterize the following mechanism, which is summarized
by the conditional distribution
$P_{\text{LLM}}(s^{\prime}\,|\,s,a,\mathcal{D},g)$.

  * Infer:

Sample a latent concept $\theta\sim P_{\text{LLM}}(\cdot\,|\,\mathcal{D})$,
where $P_{\text{LLM}}$ is the conditional distribution in association with the
given LLM.

  * Generate:

Generate the next state $s^{\prime}$ and the received reward $r$ via
$s^{\prime}\sim P_{\theta}(\cdot\,|\,s,a;g)$ and $r=r_{\theta}(s,a;g)$ given a
specific instruction $g$ and the current state $s$ and action $a$.

The following assumption states that LLMs are trained to replicate the
pretraining distribution, which covers a wide variety of $\mathcal{D}$ and
$g$. The same assumption is employed by Saunshi et al. (2020); Wei et al.
(2021); Xie et al. (2021); Lee et al. (2023) to simplify the statistical
analysis. Let $\mathbb{P}_{\text{Data}}$ be the pretraining distribution.

######  Assumption B.4 (Pretraining Consistency).

We assume that

$\displaystyle
P_{\text{LLM}}(s^{\prime}\,|\,s,a,\mathcal{D},g)=\mathbb{P}_{\text{Data}}(s^{\prime}\,|\,s,a,\mathcal{D},g)$

for all $\mathcal{D}$ and
$(s^{\prime},s,a,g)\in\mathcal{S}\times\mathcal{S}\times\mathcal{A}\times\mathcal{G}$.

In the following, we prove that Assumption 4.1 holds. By the implicit Bayesian
inference mechanism in Assumption B.3, we have

$\displaystyle P_{\text{LLM}}(s^{\prime}\,|\,s,a,\mathcal{D},g)=\int
P_{\theta}(s^{\prime}\,|\,s,a;g)\cdot
P_{\text{LLM}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta,$ (B.1)

for all $\mathcal{D}$ and
$(s^{\prime},s,a,g)\in\mathcal{S}\times\mathcal{S}\times\mathcal{A}\times\mathcal{G}$.
By the data-generating mechanism in Assumption B.2, we have

$\displaystyle\mathbb{P}_{\text{Data}}(s^{\prime}\,|\,s,a,\mathcal{D},g)$
$\displaystyle=\int
P_{\theta^{\star}}(s^{\prime}\,|\,s,a;g)\cdot\mathbb{P}_{\text{Data}}(\theta^{\star}\,|\,s,a,\mathcal{D},g){\mathrm{d}}\theta^{\star}$
$\displaystyle=\int
P_{\theta^{\star}}(s^{\prime}\,|\,s,a;g)\cdot\mathbb{P}_{\text{Data}}(\theta^{\star}\,|\,\mathcal{D}){\mathrm{d}}\theta^{\star}$
(B.2)

for all $\mathcal{D}$ and
$(s^{\prime},s,a,g)\in\mathcal{S}\times\mathcal{S}\times\mathcal{A}\times\mathcal{G}$.
Let $\mathbb{P}_{\text{Data}}(\cdot\,|\,\mathcal{D})$ be the posterior
distribution of the underlying parameter $\theta^{\star}$. By the pretraining
consistency in Assumption B.4, we have

$\displaystyle\int P_{\theta}(s^{\prime}\,|\,s,a;g)\cdot
P_{\text{LLM}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta=\int
P_{\theta}(s^{\prime}\,|\,s,a;g)\cdot\mathbb{P}_{\text{Data}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta$
(B.3)

for all $\mathcal{D}$ and
$(s^{\prime},s,a,g)\in\mathcal{S}\times\mathcal{S}\times\mathcal{A}\times\mathcal{G}$.
By the factorization condition in Assumption B.1, we transform (B.3) into

$\displaystyle\phi_{g}(s^{\prime},s,a)^{\top}\int f_{g}(\theta)\cdot
P_{\text{LLM}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta\
=\phi_{g}(s^{\prime},s,a)^{\top}\int
f_{g}(\theta)\cdot\mathbb{P}_{\text{Data}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta\,$
(B.4)

for all $\mathcal{D}$ and
$(s^{\prime},s,a,g)\in\mathcal{S}\times\mathcal{S}\times\mathcal{A}\times\mathcal{G}$.
By the redundancy condition in Assumption B.1 and the property of a basis, we
have

$\displaystyle\int f_{g}(\theta)\cdot
P_{\text{LLM}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta$
$\displaystyle=\sum_{i=1}^{d_{g}}\Bigl{(}\phi_{g}(s_{(i)}^{\prime},s_{(i)},a_{(i)})^{\top}\int
f_{g}(\theta)\cdot
P_{\text{LLM}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta\Bigr{)}\cdot\phi_{g}(s_{(i)}^{\prime},s_{(i)},a_{(i)})$
$\displaystyle=\sum_{i=1}^{d_{g}}\Bigl{(}\phi_{g}(s_{(i)}^{\prime},s_{(i)},a_{(i)})^{\top}\int
f_{g}(\theta)\cdot\mathbb{P}_{\text{Data}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta\Bigr{)}\cdot\phi_{g}(s_{(i)}^{\prime},s_{(i)},a_{(i)})$
$\displaystyle=\int
f_{g}(\theta)\cdot\mathbb{P}_{\text{Data}}(\theta\,|\,\mathcal{D}){\mathrm{d}}\theta,$
(B.5)

which implies

$\displaystyle\mathbb{E}_{\theta\sim
P_{\text{LLM}}(\cdot\,|\,\mathcal{D})}\bigl{[}f_{g}(\theta)\bigr{]}=\mathbb{E}_{\theta\sim\mathbb{P}_{\text{Data}}(\cdot\,|\,\mathcal{D})}\bigl{[}f_{g}(\theta)\bigr{]}$

for all $g\in\mathcal{G}$ and $\mathcal{D}$. By the identifiability condition
in Assumption B.1, we have

$\displaystyle
P_{\text{LLM}}(\theta\,|\,\mathcal{D})=\mathbb{P}_{\text{Data}}(\theta\,|\,\mathcal{D})$

which shows that the implicit Bayesian inference mechanism in Assumption B.3
matches the data-generating process in Assumption B.2. Thus, we verify
Assumption 4.1. ∎

##  Appendix C Proof of Theorem 4.4

###### Proof of Theorem 4.4.

We specify the terminating condition for Algorithm 3 (and Algorithm 1). Let
$(K-1)$ be the total number of switches until $t$ reaches $(T-1)$. Let
$t_{K}=T$. At the $(T-1)$-th step, Algorithm 3 samples $\theta_{T-1}$ from
$p_{t_{K-1}}$ and executes $a_{T-1}=\pi_{T-1}(s_{T-1})$ , where we have
$\pi_{T-1}=\texttt{PL}^{\epsilon}(\theta_{T-1})$. Upon receiving $r_{T-1}$ and
$s_{T}$ from the external environment, Algorithm 3 updates
$\mathcal{D}_{T}=\\{(s_{t},a_{t},s_{t+1},r_{t})\\}_{t=0}^{T-1}$ and
terminates. Throughout the following proof, we denote the upper bound of the
value function as $L$, i.e.,
$\max_{\theta\in\Theta,\pi\in\Pi,s\in{\mathcal{S}}}|V_{\theta}^{\pi}(s)|\leq
L$, for a positive constant $L$.

Recall that the Bayesian regret $\mathfrak{R}(T)$ and the optimal planner
$\texttt{PL}^{\star}(\theta)$ are defined in Section 2. By the tower property
of the conditional expectation, we have

$\displaystyle\mathfrak{R}(T)$
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{\star}}^{\pi^{\star}}(s_{t})-V_{\theta^{\star}}^{\pi_{t}}(s_{t})\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}\Bigl{[}\mathbb{E}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{\star}}^{\pi^{\star}}(s_{t})-V_{\theta^{\star}}^{\pi_{t}}(s_{t})\Big{|}\,\mathcal{F}_{k}\Bigr{]}\Bigr{]}\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}\Bigl{[}\mathbb{E}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{\star}}^{\texttt{PL}^{\star}(\theta^{\star})}(s_{t})-V_{\theta^{\star}}^{\pi_{t}}(s_{t})\Big{|}\,\mathcal{F}_{k}\Bigr{]}\Bigr{]}\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}\Bigl{[}\mathbb{E}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V^{\texttt{PL}^{\star}(\theta_{t})}_{\theta_{t}}(s_{t})-V_{\theta^{\star}}^{\pi_{t}}(s_{t})\Big{|}\,\mathcal{F}_{k}\Bigr{]}\Bigr{]}\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}V^{\texttt{PL}^{\star}(\theta_{t})}_{\theta_{t}}(s_{t})-V_{\theta^{\star}}^{\pi_{t}}(s_{t})\Bigr{]},$
(C.1)

where we define

$\displaystyle\mathcal{F}_{k}=\sigma\Bigl{(}\bigl{\\{}(s_{t},a_{t},s_{t+1},r_{t})\bigr{\\}}_{t=0}^{t_{k}-1}\Bigr{)}$
(C.2)

and the fourth equality uses the fact that
$\theta^{\star}\,|\,\mathcal{F}_{k}$ and $\theta_{t}\,|\,\mathcal{F}_{k}$
follows the same distribution for all $t_{k}\leq t<t_{k+1}$. Note that the
parameter sequence $\\{\theta_{t}\\}_{t=t_{k}}^{t_{k+1}-1}$ is sampled from
$p_{t_{k}}$ independently and identically, while we have
$\pi_{t}=\texttt{PL}^{\epsilon}(\theta_{t})$, where $\texttt{PL}^{\epsilon}$
is an $\epsilon$-optimal planner in Definition 4.2. Let
$\theta^{k}=\theta_{t_{k}}$ and $\pi^{k}=\pi_{t_{k}}$. We are able to rewrite
(C) as

$\displaystyle\mathfrak{R}(T)=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}V^{\texttt{PL}^{\star}(\theta^{k})}_{\theta^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}.$
(C.3)

Meanwhile, Definition 4.2 gives

$\displaystyle
V^{\texttt{PL}^{\star}(\theta^{k})}_{\theta^{k}}(s_{t})-V^{\pi^{k}}_{\theta^{k}}(s_{t})\leq\epsilon.$
(C.4)

Here, we use $\pi^{k}=\texttt{PL}^{\epsilon}(\theta^{k})$.

To connect (C.3) and (C.4), we introduce the regret decomposition. For
notational simplicity, we denote by

$\displaystyle(B_{\theta}V)(s,a)=r_{\theta}(s,a)+\gamma(P_{\theta}V)(s,a)$
(C.5)

for all $(s,a)\in{\mathcal{S}}\times\mathcal{A}$ and value function $V$.

######  Lemma C.1 (Regret Decomposition).

For Algorithm 3, it holds that

$\displaystyle\frac{1-\gamma}{\gamma}\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}\Bigr{]}$
$\displaystyle\quad=\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}({B}_{\theta^{k}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})-({B}_{\theta^{\star}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})\Bigr{]}\Bigr{]}}_{\text{term
(A): information gain}}$
$\displaystyle\quad\qquad+\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t_{k+1}})-V^{\pi^{k}}_{\theta^{\star}}(s_{t_{k+1}})\bigr{)}-\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t_{k}})-V^{\pi^{k}}_{\theta^{\star}}(s_{t_{k}})\bigr{)}\Bigr{]}\Bigr{]}}_{\text{term
(B): value inconsistency}},$ (C.6)

where $\mathbb{E}_{\pi^{k}}$ is taken with respect to the state-action
sequence following $s_{t+1}\sim P_{\theta^{\star}}(\cdot\,|\,s_{t},a_{t})$ and
$a_{t}=\pi^{k}(s_{t})$ for all $t_{k}\leq t<t_{k+1}$, while $\mathbb{E}$ is
taken with respect to the prior distribution $p_{0}$ of $\theta^{\star}$, the
posterior distribution $p_{t_{k}}$ of $\theta^{k}$, and the iterative update
of $\pi^{k}$. Here, $\pi^{k}$ is determined by $\theta^{k}$ and $p_{t_{k}}$ is
conditioned on
$\mathcal{D}_{t_{k}}=\\{(s_{t},a_{t},s_{t+1},r_{t})\\}_{t=0}^{t_{k}-1}$, which
is generated by $\theta^{\star}$ and Algorithm 3.

###### Proof of Lemma C.1.

The Bellman equation (Sutton and Barto, 2018) connects $Q_{\theta}^{\pi}(s,a)$
and $V^{\pi}_{\theta}(s)$ by

$\displaystyle
Q^{\pi}_{\theta}\left(s,a\right)=r_{\theta}\left(s,a\right)+\gamma\left(P_{\theta}V_{\theta}^{\pi}\right)\left(s,a\right),\quad
V_{\theta}^{\pi}(s)=Q_{\theta}^{\pi}\bigl{(}s,\pi(s)\bigr{)}.$ (C.7)

By (C.5), we rewrite (C.7) as
$Q^{\pi}_{\theta}\left(s,a\right)=(B_{\theta}V_{\theta}^{\pi})(s,a)$. For the
left-hand side of (C.1), we have

$\displaystyle\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}\Bigr{]}$
$\displaystyle\quad=\gamma\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}({B}_{\theta^{k}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})-({B}_{\theta^{\star}}V_{\theta^{\star}}^{\pi^{k}})(s_{t},a_{t})\bigr{]}\bigr{]}$
$\displaystyle\quad=\gamma\cdot\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}({B}_{\theta^{k}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})-({B}_{\theta^{\star}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})\Bigr{]}\Bigr{]}}_{\text{term
(A)}}$
$\displaystyle\quad\qquad+\gamma\cdot\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}({B}_{\theta^{\star}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})-V_{\theta^{k}}^{\pi^{k}}(s_{t+1})\Bigr{]}\Bigr{]}}_{\text{term
(C1)}}$
$\displaystyle\quad\qquad+\gamma\cdot\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{\star}}^{\pi^{k}}(s_{t+1})-({B}_{\theta^{\star}}V^{\pi^{k}}_{\theta^{\star}})(s_{t},a_{t})\Bigr{]}\Bigr{]}}_{\text{term
(C2)}}$
$\displaystyle\quad\qquad+\gamma\cdot\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t+1})-V_{\theta^{\star}}^{\pi^{k}}(s_{t+1})\Bigr{]}\Bigr{]}}_{\text{term
(D)}},$ (C.8)

where the first equality uses $a_{t}=\pi^{k}(s_{t})$. Since we have

$\displaystyle(P_{\theta^{\star}}V)(s_{t},a_{t})=\mathbb{E}_{s_{t+1}\sim
P_{\theta^{\star}}(\cdot\,|\,s_{t},a_{t})}\bigl{[}V(s_{t+1})\bigr{]},$

terms (C1) and (C2) in (C.8) are zero.

Meanwhile, term (D) in (C.8) satisfies

term (D)
$\displaystyle={\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\bigr{)}\Bigr{]}\Bigr{]}}$
$\displaystyle\qquad+\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t_{k+1}})-V^{\pi^{k}}_{\theta^{\star}}(s_{t_{k+1}})\bigr{)}-\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t_{k}})-V_{\theta^{\star}}^{\pi^{k}}(s_{t_{k}})\bigr{)}\Bigr{]}\Bigr{]}}_{\text{term
(B)}},$ (C.9)

where term (B) is defined in Lemma C.1. Rearranging (C.8) and (C.9), we prove
Lemma C.1. ∎

We characterize terms (A) and (B) in (C.1). Let $\mathcal{E}$ be the high-
probability event in Definition 4.3. For the information gain in term (A), we
have

term (A)
$\displaystyle\leq\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\operatorname{\mathds{1}}_{\mathcal{E}}\cdot\sum_{t=t_{k}}^{t_{k+1}-1}({B}_{\theta^{k}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})-({B}_{\theta^{\star}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})\Bigr{]}\Bigr{]}+L\delta
T$
$\displaystyle\leq\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}\Gamma_{t_{k}}(\delta)\cdot\sqrt{I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})}\Bigr{]}\Bigr{]}+L\delta
T$
$\displaystyle\leq\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}\sqrt{I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})}\Bigr{]}\Bigr{]}+L\delta
T$
$\displaystyle\leq\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}\Bigl{[}\sqrt{T}\cdot\Bigl{(}{\sum_{t=0}^{T-1}I(\theta;\xi_{t+1}\Big{|}\,\mathcal{D}_{t})}\Bigr{)}^{1/2}\Bigr{]}+L\delta
T$
$\displaystyle=\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}\bigl{[}\sqrt{T(H_{0}-H_{T})}\bigr{]}+L\delta
T.$ (C.10)

The first inequality uses the fact that the value function is upper bounded by
$L$ and $\mathbb{P}(\mathcal{E})\geq 1-\delta$ in Definition 4.3.

The last inequality invokes the Cauchy-Schwarz inequality. The last equality
is derived from $I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})=H_{t}-H_{t+1}$.

We explain the second inequality as follows. For any $0\leq k<K$, we consider
$t^{\dagger}=t_{k}$ and $\theta^{k}=\theta_{t_{k}}$, which is sampled from
$p_{t_{k}}$ given $\mathcal{D}_{t_{k}}$, and the value function sequence
$\\{V_{t}\\}_{t=t_{k}}^{t_{k+1}-1}=\\{V_{\theta^{k}}^{\pi^{k}}\\}_{t=t_{k}}^{t_{k+1}-1}$,
which is adapted to $\\{\sigma(\mathcal{D}_{t})\\}_{t=t_{k}}^{t_{k+1}-1}$
since $\mathcal{D}_{t_{k}}\subseteq\mathcal{D}_{t}$ for all $t_{k}\leq
t<t_{k+1}$. Correspondingly, we have

$\displaystyle\bigl{|}(r_{\theta^{\star}}-r_{\theta_{t_{k}}})(s_{t},a_{t})+\bigl{(}(P_{\theta^{\star}}-P_{\theta_{t_{k}}})V_{t}\bigr{)}(s_{t},a_{t})\bigr{|}\leq\Gamma_{t_{k}}(\delta)\cdot\sqrt{I(\theta;\xi_{t+1}|\mathcal{D}_{t})}.$

Hence, we obtain the second inequality in (C.10).

As the switching condition in Algorithm 3 implies
$H_{t_{k}}-H_{t_{k+1}}\geq\log 2$, we have

$\displaystyle
H_{0}-H_{t_{K-1}}=\sum_{k=0}^{K-2}H_{t_{k}}-H_{t_{k+1}}\geq(K-1)\cdot\log 2,$

which implies

$\displaystyle K-1\leq{(H_{0}-H_{t_{K-1}})}/{\log 2}\leq(H_{0}-H_{T})/\log 2.$
(C.11)

Recall that the value function is upper bounded by $L$. For the value
inconsistency in term (B), we have

$\displaystyle\text{term (B)}\leq{(4L/{\log
2})\cdot\mathbb{E}[H_{0}-H_{T}]}+4L.$ (C.12)

Plugging (C.10) and (C.12) into Lemma C.1, we have

$\displaystyle\frac{1-\gamma}{\gamma}\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}\Bigr{]}$
$\displaystyle\quad\leq L\delta
T+{\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}\bigl{[}\sqrt{T(H_{0}-H_{T})}\bigr{]}}+(4L/{\log
2})\cdot\mathbb{E}[H_{0}-H_{T}]+4L.$ (C.13)

Combining (C.3), (C.4), and (C.13), we have

$\displaystyle\mathfrak{R}(T)$
$\displaystyle=\mathcal{O}\biggl{(}\frac{\gamma\cdot\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}[\sqrt{H_{0}-H_{T}}]}{1-\gamma}\cdot\sqrt{T}$
$\displaystyle\qquad\qquad+\frac{\gamma\delta L}{1-\gamma}\cdot
T+\epsilon\cdot T+\frac{\gamma
L\cdot\mathbb{E}[H_{0}-H_{T}]}{(1-\gamma)}\biggr{)}.$ (C.14)

Since the reward function is bounded in $[0,1]$, we take
$L=\sum_{t=0}^{\infty}\gamma^{t}=1/(1-\gamma)$ and obtain

$\displaystyle\mathfrak{R}(T)=\mathcal{O}\biggl{(}\frac{\gamma\cdot\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)\cdot\mathbb{E}[\sqrt{H_{0}-H_{T}}]}{1-\gamma}\cdot\sqrt{T}+\frac{\gamma\delta}{(1-\gamma)^{2}}\cdot
T+\epsilon\cdot
T+\frac{\gamma\cdot\mathbb{E}[H_{0}-H_{T}]}{(1-\gamma)^{2}}\biggr{)}.$

Thus, we prove Theorem 4.4. ∎

##  Appendix D Linear Special Case

In this section, we first propose a liner function approximation setting, and
then we propose the implementation of Algorithm 3 (RAFA) in this setting. We
also prove the Bayesian regret of Algorithm 3 in this setting. Formally, we
introduce the definition of linear kernel MDP, where the expected state-value
function is linearly parameterized with respect to an unknown $d$-dimensional
parameter $\theta^{\star}$ and we assume that both the prior of the parameter
$\theta^{\star}$ and the state-value function follow Gaussian distributions
(Lu and Van Roy, 2019; Jin et al., 2021; Ghavamzadeh et al., 2015).

######  Definition D.1 (Linear Kernel MDP (Ghavamzadeh et al., 2015; Yang and
Wang, 2020, 2019; Cai et al., 2020; Zhou et al., 2021b)).

Suppose the state-value function in a MDP $M$ follows a Gaussian distribution
whose mean is linearly parameterized with respect to a $d$-dimensional
parameter $\theta$. At the initial state $s_{0}\in\mathcal{S}$, we assume that

$\displaystyle
V(s_{0})\sim\mathcal{N}(\phi_{V}^{\top}\theta,1)\qquad\text{and}\qquad
V(s_{t+1})\sim\mathcal{N}(\psi_{V}(s_{t},a_{t})^{\top}\theta,1),$

for all $t\geq 0$, $(s_{t},a_{t})\in\mathcal{S}\times\mathcal{A}$, any value
function $V$, and $s_{t+1}$ generated from $s_{t+1}\sim
P_{\theta^{\star}}(\cdot\,|\,s_{t},a_{t})$. Here $\phi_{V}$ and
$\psi_{V}(s,a)$ are both $d$-dimensional known maps satisfying
$\max\\{\|\psi_{V}(s,a)\|_{2},\|\phi_{V}\|_{2}\\}\leq C$ for a positive
constant $C$, any value function $V$, and
$(s,a)\in{\mathcal{S}}\times\mathcal{A}$. The prior distribution of the
parameter $\theta^{\star}\in\mathbb{R}^{d}$ is $\mathcal{N}(0,\lambda I_{d})$
with a positive constant $\lambda$. For simplicity, we consider the reward to
be deterministic and known.

Let $(K-1)$ be the total number of switches until $t$ reaches $(T-1)$ in
Algorithm 3. By Definition D.1 and Algorithm 3, we take
$V_{t}=V_{\theta^{k}}^{\pi^{k}}$ for $t_{k}\leq t<t_{k+1}$ and then obtain the
closed form of the posterior distribution of the parameter as

$\displaystyle\theta\,|\,\mathcal{D}_{t}\sim\mathcal{N}(\widehat{\theta}_{t};\Sigma_{t}^{-1}),$

where

$\displaystyle\widehat{\theta}_{t}=\Bigl{(}\lambda
I_{d}+\sum_{i=0}^{t-1}\psi_{V_{i}}(s_{i},a_{i})\psi_{V_{i}}(s_{i},a_{i})^{\top}\Bigr{)}^{-1}\Bigl{(}\sum_{i=0}^{t-1}\psi_{V_{i}}(s_{i},a_{i})V_{i}(s_{i+1})\Bigr{)},$
(D.1)

and

$\displaystyle\Sigma_{t}=\lambda
I_{d}+\sum_{i=0}^{t-1}\psi_{V_{i}}\left(s_{i},a_{i}\right)\psi_{V_{i}}\left(s_{i},a_{i}\right)^{\top}.$
(D.2)

Hence the entropy of the posterior distribution of the parameter is

$\displaystyle H_{t}=1/2\cdot\log(\det(\Sigma_{t}))+d/2\cdot(1+\log(2\pi)).$
(D.3)

Then we specify the switching condition in Algorithm 3 in the linear kernel
MDP setting. As the switching condition in Line 10 of Algorithm 3 is
$H_{t_{k}}-H_{t}>\log 2$, we use the closed form of the posterior distribution
of the parameter in (D.3) to specify the switching condition as
$1/2\cdot\log(\det(\Sigma_{t_{k}}))-1/2\cdot\log(\det(\Sigma_{t}))>\log 2$,
that is, $\det(\Sigma_{t_{k}})>4\det(\Sigma_{t})$, which implies

$\displaystyle\det(\Sigma_{t_{k}})\leq 4\det(\Sigma_{t}),$ (D.4)

for all $t_{k}\leq t<t_{k+1}$ and $k<K$.

By the definition of $\Sigma_{t}$ in (D.2), we have
$\log\det(\Sigma_{0})=d\log\lambda$ and

$\displaystyle\log\det(\Sigma_{T})$ $\displaystyle=\log\det\Bigl{(}\lambda
I_{d}+\sum_{t=0}^{T-1}\psi_{V_{t}}(s_{t},a_{t})\psi^{\top}_{V_{t}}(s_{t},a_{t})\Bigr{)}$
$\displaystyle\leq d\log\Bigl{(}1/d\cdot\mathrm{Tr}\bigl{(}\lambda
I_{d}+\sum_{t=0}^{T-1}\psi_{V_{t}}(s_{t},a_{t})\psi^{\top}_{V_{t}}(s_{t},a_{t})\bigr{)}\Bigr{)}$
$\displaystyle=d\log\Bigl{(}1/d\cdot\bigl{(}\lambda
d+\sum_{t=0}^{T-1}\|\psi_{V_{t}}(s_{t},a_{t})\|_{2}^{2}\bigr{)}\Bigr{)}$
$\displaystyle\leq d\log\bigl{(}\lambda+TC^{2}/d\bigr{)},$ (D.5)

almost surely, where the first inequality is derived from the relationship
between the trace and the determinant of a square matrix, the second equality
uses $\mathrm{Tr}(a^{\top}b)=\mathrm{Tr}(b^{\top}a)$ for two arbitrary vectors
$a$ and $b$, and the last inequality uses the fact that
$\|\psi_{V}(s,a)\|_{2}$ is bounded by $C$ for any value function $V$, and
$(s,a)\in{\mathcal{S}}\times\mathcal{A}$. Thus, we have

$\displaystyle
H_{0}-H_{T}=\mathcal{O}\big{(}d\log(1+TC^{2}/(d\lambda)\big{)},$ (D.6)

almost surely. By (C.11) in the proof of Theorem 4.4, we obtain the almost
surely upper bound of the number of switches $K-1$ as

$\displaystyle K-1=\mathcal{O}\big{(}d\log(1+TC^{2}/(d\lambda))\big{)}.$ (D.7)

In the following proposition, we characterize the upper bound of the
information ratio in Definition 4.3 for the linear kernel MDP.

######  Proposition D.2.

The information ratio $\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta)$ for
the linear kernel MDP is $\mathcal{O}(d\log(dT/\delta))$.

###### Proof of Proposition D.2.

For any $t^{\dagger}$ and $t\geq t^{\dagger}$, if
$H_{t^{\dagger}}-H_{t}\leq\log 2$, then we have

$\displaystyle(1/2)\cdot\log\det{(\Sigma_{t^{\dagger}})}-(1/2)\cdot\log\det{(\Sigma_{t})}\leq\log
2,$

which implies

$\displaystyle\det{(\Sigma_{t^{\dagger}})}\leq 4\det{(\Sigma_{t})}.$ (D.8)

By Definition D.1, we obtain

$\displaystyle\Bigl{|}\bigl{(}(P_{\theta^{\star}}-P_{\theta_{t^{\dagger}}})V_{t}\bigr{)}(s_{t},a_{t})\Bigr{|}$
$\displaystyle=(\theta^{\star}-\theta_{t^{\dagger}})^{\top}\big{(}\psi_{V_{t}}(s_{t},a_{t})\big{)}$
$\displaystyle\leq\|\theta^{\star}-\theta_{t^{\dagger}}\|_{\Sigma_{t}}\cdot\|\psi_{V_{t}}(s_{t},a_{t})\|_{\Sigma_{t}^{-1}}$
$\displaystyle\leq
4\cdot\|\theta^{\star}-\theta_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}\cdot\|\psi_{V_{t}}(s_{t},a_{t})\|_{\Sigma_{t}^{-1}},$
(D.9)

where the first inequality invokes the Cauchy–Schwarz inequality and the
second inequality is derived from Lemma 12 in Abbasi-Yadkori et al. (2011) and
(D.8). Based on the fact that $\theta_{t^{\dagger}}$ and $\theta^{\star}$
share the same posterior distribution
$\mathcal{N}(\widehat{\theta}_{t^{\dagger}},\Sigma_{t^{\dagger}}^{-1})$ given
$\mathcal{D}_{t^{\dagger}}$, we know that
$\|\theta-\widehat{\theta}_{t^{\dagger}}\|^{2}_{\Sigma_{t^{\dagger}}}$ follows
$\mathcal{X}_{d}^{2}$, where
$\theta\in\\{\theta^{\star},\theta_{t^{\dagger}}\\}$ and $\mathcal{X}_{d}^{2}$
is the chi-square distribution with $d$ degrees of freedom. By the tail
behavior of the chi-square distribution (Ghosh, 2021),

we have

$\displaystyle\mathbb{P}\big{(}\|\theta^{\star}-\widehat{\theta}_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}>2d\log(2d/\delta)\big{)}\leq\delta/2$
(D.10)

and

$\displaystyle\mathbb{P}\big{(}\|\theta_{t^{\dagger}}-\widehat{\theta}_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}>2d\log(2d/\delta)\big{)}\leq\delta/2,$
(D.11)

for a sufficiently large $d$ and a sufficiently small $\delta$. With
probability at least $1-\delta$, we have

$\displaystyle\|\theta^{\star}-\theta_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}\leq\sqrt{2}\cdot\sqrt{\|\theta_{t^{\dagger}}-\widehat{\theta}_{t^{\dagger}}\|^{2}_{\Sigma_{t^{\dagger}}}+\|\theta^{\star}-\widehat{\theta}_{t^{\dagger}}\|^{2}_{\Sigma_{t^{\dagger}}}}=\mathcal{O}(\sqrt{d}\log(d/\delta)),$
(D.12)

where we use the triangle inequality of $\|\cdot\|_{\Sigma_{t^{\dagger}}}$ and
the inequality $(a+b)^{2}\leq 2(a^{2}+b^{2})$ for $a,b\geq 0$. By the
definition of information gain and the closed form of the entropy in (D.3), we
have

$\displaystyle{I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})}$
$\displaystyle=H_{t}-H_{t+1}=1/2\cdot\log\bigl{(}{\det(\Sigma_{t+1})}/{\det(\Sigma_{t})}\bigr{)}$
$\displaystyle=1/2\cdot\log(1+\psi_{V_{t}}(s_{t},a_{t})^{\top}\Sigma_{t}^{-1}\psi_{V_{t}}(s_{t},a_{t}))$
$\displaystyle\geq\log(1+d)/(2d)\cdot\psi_{V_{t}}(s_{t},a_{t})^{\top}\Sigma_{t}^{-1}\psi_{V_{t}}(s_{t},a_{t})$
$\displaystyle=\log(1+d)/(2d)\cdot\|\psi_{V_{t}}(s_{t},a_{t})\|_{\Sigma_{t}^{-1}}^{2},$
(D.13)

where the first equality uses the matrix determinant lemma. Here, the last
inequality uses that $\log(1+x)/x$ is an increasing function for $x\geq 0$ and

$\displaystyle 0$
$\displaystyle\leq\psi_{V_{t}}(s_{t},a_{t})^{\top}\Sigma_{t}^{-1}\psi_{V_{t}}(s_{t},a_{t})$
$\displaystyle\leq\psi_{V_{t}}(s_{t},a_{t})^{\top}\big{(}{\psi_{V_{t}}(s_{t},a_{t})\psi_{V_{t}}(s_{t},a_{t})^{\top}}\big{)}^{-1}\psi_{V_{t}}(s_{t},a_{t})$
$\displaystyle=\mathrm{Tr}\Big{(}\psi_{V_{t}}(s_{t},a_{t})\psi_{V_{t}}(s_{t},a_{t})^{\top}\big{(}{\psi_{V_{t}}(s_{t},a_{t})\psi_{V_{t}}(s_{t},a_{t})^{\top}}\big{)}^{-1}\Big{)}$
$\displaystyle=d,$ (D.14)

where the first inequality relies on the nonnegativity of a quadratic form,
the first equality uses $\mathrm{Tr}(AB)=\mathrm{Tr}(BA)$ for two arbitrary
matrices $A$ and $B$, and the second inequality is based on the definition of
$\Sigma_{t}^{-1}$ in (D.2). Then we upper bound the right side of (D.9) as

$\displaystyle
4\cdot\|\theta^{\star}-\theta_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}\cdot\|\psi_{V_{t}}(s_{t},a_{t})\|_{\Sigma_{t}^{-1}}$
$\displaystyle\leq{4\sqrt{2d}}/{\sqrt{\log(1+d)}}\cdot\|\theta^{\star}-\theta_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}\cdot\sqrt{I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})}$
$\displaystyle\leq
4{\sqrt{2(d+1)}}\cdot\|\theta^{\star}-\theta_{t^{\dagger}}\|_{\Sigma_{t^{\dagger}}}\cdot\sqrt{I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})}$
$\displaystyle=\mathcal{O}\bigl{(}\sqrt{d(d+1)}\cdot\log(d/\delta)\cdot\sqrt{I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})}\bigr{)},$
(D.15)

for all $t<T$ such that $H_{t^{\dagger}}-H_{t}\leq\log 2$. Here the first
inequality is the consequence of (D.12), the second inequality uses
$\log(1+x)\geq x/(1+x)$ for $x\geq 0$, and the last inequality is the result
of (D.13). By scaling $\delta^{\prime}=\delta/T$ and applying the union bound
over $0\leq t<T$, we derive the upper bound of
$\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta^{\prime})$ as

$\displaystyle\sup_{t^{\dagger}<T}\Gamma_{t^{\dagger}}(\delta^{\prime})$
$\displaystyle=\mathcal{O}\big{(}\sqrt{d(d+1)}\cdot\log(dT/\delta^{\prime})\big{)}$
$\displaystyle=\mathcal{O}\big{(}d\cdot\log(dT/\delta^{\prime})\big{)},$

with probability at least $1-\delta^{\prime}$ for a sufficiently large enough
$d$ and a sufficiently small $\delta$. ∎

Based on Proposition D.2, we can prove Theorem 4.4 in the linear kernel MDP
given the upper bound $L$ of the value function. However, since the state-
value function is Gaussian distributed, it is hard to directly derive the
almost surely upper bound of the value function. Instead of directly applying
the definition of information ratio, we can use another method to deal with
term (A) in the regret decomposition (Lemma C.1) by taking the conditional
expectation for both $\theta^{\star}$ and $\theta^{k}$ to obtain the desired
result, which does not rely on the almost surely upper bound of the value
function. We present the upper bound of the Bayesian regret of Algorithm 3 on
the linear kernel MDP as follows.

######  Proposition D.3 (Linear Kernel MDP).

If we apply Algorithm 3 on the linear kernel MDP defined in Definition D.1,
then we can upper bound of the Bayesian regret as

$\displaystyle\mathfrak{R}(T)\leq\mathcal{O}\left(\frac{\gamma
d^{3/2}}{1-\gamma}\cdot\sqrt{T}+\epsilon\cdot T\right).$

Proposition D.3 demonstrates the $\sqrt{T}$ Bayesian regret if we have
$\epsilon=\mathcal{O}(1/\sqrt{T})$, which characterizes the sample efficiency
of Algorithm 3 (RAFA) in the linear kernel MDP setting. As a variant of
Theorem 4.4, the upper bound in Prop D.3 is irrelevant to the upper bound of
the value function. The $\sqrt{T}$ Bayesian regret of Algorithm 3 also matches
the result of other RL algorithms in linear kernel MDPs (Zhou et al., 2021b;
Cai et al., 2020; Gao et al., 2021) without an explicit parameter update.

###  D.1 Proof of Proposition D.3

###### Proof of Proposition D.3.

Similar to the arguments in the proof of Theorem 4.4, it suffices to bound
$\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}\Bigr{]}$.
By Lemma C.1, we have that

$\displaystyle\frac{1-\gamma}{\gamma}\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}\Bigr{]}$
$\displaystyle\quad=\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}({B}_{\theta^{k}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})-({B}_{\theta^{\star}}V^{\pi^{k}}_{\theta^{k}})(s_{t},a_{t})\Bigr{]}\Bigr{]}}_{\text{term
(A): information gain}}$
$\displaystyle\quad\qquad+\underbrace{\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t_{k+1}})-V^{\pi^{k}}_{\theta^{\star}}(s_{t_{k+1}})\bigr{)}-\bigl{(}V_{\theta^{k}}^{\pi^{k}}(s_{t_{k}})-V^{\pi^{k}}_{\theta^{\star}}(s_{t_{k}})\bigr{)}\Bigr{]}\Bigr{]}}_{\text{term
(B): value inconsistency}},$ (D.16)

where all the expectations are specified in Lemma C.1.

#### Analysis of Term (A).

By Definition D.1, we have

term (A)
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}(\psi_{V_{t}}(s_{t},a_{t}))^{\top}(\theta^{k}-\theta^{\star})\Bigr{]}\Bigr{]}$
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}\mathbb{E}\bigl{[}(\psi_{V_{t}}(s_{t},a_{t}))^{\top}(\theta^{k}-\theta^{\star})\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{]}\Bigr{]}$
$\displaystyle\leq\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}\Bigl{(}\mathbb{E}\bigl{[}\|\psi_{V_{t}}(s_{t},a_{t})\|^{2}_{\Sigma_{t}^{-1}}\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}^{1/2}\cdot\Bigl{(}\mathbb{E}\bigl{[}\|\theta^{k}-\theta^{\star}\|^{2}_{\Sigma_{t}}\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}^{1/2}\Bigr{]}$
$\displaystyle\leq
4\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}\Bigl{(}\mathbb{E}\bigl{[}\|\psi_{V_{t}}(s_{t},a_{t})\|^{2}_{\Sigma_{t}^{-1}}\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}^{1/2}\cdot\Bigl{(}\mathbb{E}\bigl{[}\|\theta^{k}-\theta^{\star}\|^{2}_{\Sigma_{t_{k}}}\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}^{1/2}\Bigr{]},$
(D.17)

where the first equality uses Definition D.1, the filtration $\mathcal{F}_{k}$
is defined in (C.2), and the second equality uses the tower property of the
conditional expectation.

Here the first inequality invokes the Holder’s inequality and the second
inequality is derived from Lemma 12 in Abbasi-Yadkori et al. (2011) and (D.4).
Based on the fact that $\theta^{k}$ and $\theta^{\star}$ share the same
posterior distribution
$\mathcal{N}(\widehat{\theta}_{t_{k}},\Sigma_{t_{k}}^{-1})$ given
$\mathcal{D}_{t_{k}}$, we have that
$\|\theta-\widehat{\theta}_{t_{k}}\|^{2}_{\Sigma_{t_{k}}}$ follows
$\mathcal{X}_{d}^{2}$, where $\theta\in\\{\theta^{\star},\theta^{k}\\}$ and
$\mathcal{X}_{d}^{2}$ is the chi-square distribution with $d$ degrees of
freedom, whose expectation is $d$. Thus, we have

$\displaystyle\mathbb{E}\bigl{[}\|\theta^{k}-\theta^{\star}\|^{2}_{\Sigma_{t_{k}}}\big{|}\,\mathcal{F}_{k}\bigr{]}$
$\displaystyle\leq
2\Bigl{(}\mathbb{E}\bigl{[}\|\theta^{k}-\widehat{\theta}_{t_{k}}\|^{2}_{\Sigma_{t_{k}}}\big{|}\,\mathcal{F}_{k}\bigr{]}+\mathbb{E}\bigl{[}\|\theta^{\star}-\widehat{\theta}_{t_{k}}\|^{2}_{\Sigma_{t_{k}}}\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}=2d,$
(D.18)

where the first inequality invokes the triangle inequality of
$\|\cdot\|_{\Sigma_{t_{k}}}$ and the fact that $(a+b)^{2}\leq 2(a^{2}+b^{2})$
for $a,b\geq 0$. Plugging (D.18) into (D.17), we have that

term (A) $\displaystyle\leq
4\sqrt{2d}\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}\Bigl{(}\mathbb{E}\bigl{[}\|\psi_{V_{t}}(s_{t},a_{t})\|^{2}_{\Sigma_{t}^{-1}}\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}^{1/2}\Bigr{]}$
$\displaystyle\leq
8\sqrt{2(d+1)d}\cdot\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\sum_{t=t_{k}}^{t_{k+1}-1}\Bigl{(}\mathbb{E}\bigl{[}I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})\big{|}\,\mathcal{F}_{k}\bigr{]}\Bigr{)}^{1/2}\Bigr{]}$
$\displaystyle\leq
8\sqrt{2(d+1)d}\cdot\sqrt{T}\cdot\mathbb{E}\Bigl{[}\Bigl{(}\sum_{t=0}^{T-1}I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})\Bigr{)}^{1/2}\Bigr{]}$
$\displaystyle=8\sqrt{2(d+1)d}\cdot\sqrt{T}\cdot\mathbb{E}\bigl{[}\sqrt{H_{0}-H_{T}}\bigr{]},$
(D.19)

where the second inequality follows the same argument in (D.13) and (D.15),
the last inequality invokes the Cauchy-Schwarz inequality, and the last
equality is derived from
$I(\theta;\xi_{t+1}\,|\,\mathcal{D}_{t})=H_{t}-H_{t+1}$. Recall that we prove
that the almost surely upper bound for $H_{0}-H_{T}$ is
$\mathcal{O}(d\cdot\log(1+TC^{2}/(d\lambda)))$ in (D.6). Plugging the almost
surely upper bound for $H_{0}-H_{T}$ into the right side of (D.20), we derive
the upper bound of term (A) in (D.1) as

$\displaystyle\text{term
(A)}=\mathcal{O}\bigl{(}d^{3/2}\cdot\log(TC^{2}/d)\cdot\sqrt{T}\bigr{)}.$
(D.20)

#### Analysis of Term (B).

For the simplicity of discussions, we define
$\psi_{V}(s_{-1},a_{-1})=\phi_{V}$ for any value function $V$, where
$\phi_{V}$ is defined in Definition D.1. By Definition D.1, we obtain

term (B)
$\displaystyle=\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\bigl{(}\psi_{V_{\theta^{k}}^{\pi^{k}}}(s_{t_{k+1}-1},a_{t_{k+1}-1})-\psi_{V_{\theta^{\star}}^{\pi^{k}}}(s_{t_{k+1}-1},a_{t_{k+1}-1})\bigr{)}^{\top}\theta^{\star}\Bigr{]}\Bigr{]}$
$\displaystyle\qquad+\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\bigl{(}\psi_{V_{\theta^{\star}}^{\pi^{k}}}(s_{t_{k}-1},a_{t_{k}-1})-\psi_{V_{\theta^{k}}^{\pi^{k}}}(s_{t_{k}-1},a_{t_{k}-1})\bigr{)}^{\top}\theta^{\star}\Bigr{]}\Bigr{]}$
$\displaystyle\leq\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\big{\|}\psi_{V_{\theta^{k}}^{\pi^{k}}}(s_{t_{k+1}-1},a_{t_{k+1}-1})-\psi_{V_{\theta^{\star}}^{\pi^{k}}}(s_{t_{k+1}-1},a_{t_{k+1}-1})\big{\|}_{2}\cdot\|\theta^{\star}\|_{2}\Bigr{]}\Bigr{]}$
$\displaystyle\qquad+\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\big{\|}\psi_{V_{\theta^{\star}}^{\pi^{k}}}(s_{t_{k}-1},a_{t_{k}-1})-\psi_{V_{\theta^{k}}^{\pi^{k}}}(s_{t_{k}-1},a_{t_{k}-1})\big{\|}_{2}\cdot\|\theta^{\star}\|_{2}\Bigr{]}\Bigr{]}$
$\displaystyle\leq 4C\cdot\mathbb{E}\bigl{[}K\|\theta^{\star}\|_{2}\bigr{]},$
(D.21)

where the first equality is based on the fact that $s_{t+1}\sim
P_{\theta^{\star}}(\cdot\,|\,s_{t},a_{t})$ and $a_{t}=\pi^{k}(s_{t})$ for all
$t_{k}\leq t<t_{k+1}$ and $k<K$. Here the first inequality invokes the Cauchy-
Schwarz inequality, and the last inequality uses the triangle inequality of
$\|\cdot\|_{2}$ and the fact that
$\max\\{\|\psi_{V}(s,a)\|_{2},\phi_{V}\\}\leq C$ for any value function $V$
and $(s,a)\in{\mathcal{S}}\times\mathcal{A}$. From (D.7), we plug the almost
surely upper bound of $K-1$ into the right-hand side of (D.21) to obtain

$\displaystyle 4C\cdot\mathbb{E}\bigl{[}K\|\theta^{\star}\|_{2}\bigr{]}$
$\displaystyle=\mathcal{O}\Bigl{(}\bigl{(}1+d\log
2\cdot\log(1+TC^{2}/(d\lambda))\bigr{)}\cdot\mathbb{E}\bigl{[}\|\theta^{\star}\|_{2}\bigr{]}\Bigr{)}$
$\displaystyle=\mathcal{O}\Bigl{(}d^{1/2}\cdot\bigl{(}1+d\log
2\cdot\log(1+TC^{2}/(d\lambda))\bigr{)}\Bigr{)},$

where the last equality is derived from the fact that the prior distribution
of $\theta^{\star}$ is $\mathcal{N}(0,\lambda I_{d})$ and
$\mathbb{E}[\|\theta^{\star}\|_{2}]\leq\sqrt{\mathbb{E}[\|\theta^{\star}\|_{2}^{2}]}\leq\lambda\sqrt{d}$
as $\|\theta^{\star}\|_{2}^{2}/\lambda^{2}$ follows the chi-square
distribution with $d$ degrees of freedom. If $T$ and $d$ are large enough, we
have

term (B) $\displaystyle=\mathcal{O}\Bigl{(}d^{1/2}\cdot\bigl{(}1+d\log
2\cdot\log(1+TC^{2}/(d\lambda))\bigr{)}\Bigr{)}$
$\displaystyle=\mathcal{O}\Bigl{(}d^{3/2}\cdot\log(TC^{2}/d)\Bigr{)}.$ (D.22)

Combining Lemma C.1, (D.20), and (D.22), we obtain

$\displaystyle\mathbb{E}\Bigl{[}\sum_{k=0}^{K-1}\mathbb{E}_{\pi^{k}}\Bigl{[}\sum_{t=t_{k}}^{t_{k+1}-1}V_{\theta^{k}}^{\pi^{k}}(s_{t})-V_{\theta^{\star}}^{\pi^{k}}(s_{t})\Bigr{]}\Bigr{]}=\mathcal{O}\Bigl{(}\frac{\gamma
d^{3/2}\cdot\log(TC^{2}/d)}{1-\gamma}\cdot\sqrt{T}\Bigr{)}.$

By Definition 4.2 and a similar argument in (C) in the proof of Theorem 4.4,
we obtain

$\displaystyle\mathfrak{R}(T)=\mathcal{O}\Bigl{(}\frac{\gamma
d^{3/2}\cdot\log(TC^{2}/d)}{1-\gamma}\cdot\sqrt{T}+\epsilon\cdot T\Bigr{)}.$

Then we finish the proof of Proposition D.3. ∎

##  Appendix E More Experiments

In what follows, we provide the detailed setups and additional results of our
experiments.

###  E.1 Game of 24

#### Task Setup.

Game of 24 (Yao et al., 2023a) is a mathematical puzzle where the player uses
basic arithmetic operations (i.e., addition, subtraction, multiplication,
division) with four given numbers to get 24. Figure 10 gives an illustrative
example for Game of 24.

[Illustrative example for Game of 24] • Numbers: [2, 5, 8, 11] • Arithmetic
Operations: [$+$, $-$, $\times$, $/$, $($, $)$] • Solution:
$\displaystyle(11-5)\times 8/4=24$

Figure 10: An illustrative example of the Game of 24. The player uses
combinations of basic arithmetic operations with four given numbers to get 24.

Figure 11:  Sample efficiency on Game of 24. RAFA agent achieves strong
performance due to an orchestration of reasoning and acting. The success rate
at a given step is the number of tasks that is solved within the given step.

Following Yao et al. (2023a), we use the same subset indexed 901-1,000 from a
total of 1,362 tasks collected from 4nums.com. The index is arranged from easy
to hard by human solving time so the subset is relatively challenging. The
agent receives a reward of 1 if the proposed formula is correct and the
proposed formula is accepted and concatenated into the state; if the final
result is exactly 24, the agent receives a reward of 10, and the episode
terminates. Otherwise, the agent receives a reward of 0, and the proposed
formula is not accepted. We limit the maximum trials for each task to 20 to
avoid meaningless retries. The task is successful if the agent receives a
return larger than 10 111For gpt-3.5-turbo, we report the success rate when
the agent receives a return no less than 3 (i.e., find all sub-steps to get 24
but not necessarily generate a whole correct formula). This is because ToT
with gpt-3.5-turbo is known to suffer from correctly get a whole formula due
to limited reasoning ability and non-perfect prompts. See
https://github.com/princeton-nlp/tree-of-thought-llm/issues/24 for more
details. (i.e., find a valid solution within 20 steps). We report the final
success rate and sample efficiency for each method on the subset of 100 tasks.
Notably, a task is considered successful if the RAFA agent returns one and
only one correct formula, which is more strictly evaluated than Tree of
Thoughts (ToT, Yao et al. (2023a)): we allow open-loop agents like ToT to
retry 20 times and consider them successful if they generate a valid solution
in any of the 20 trials. For CoT (Wei et al., 2022) and Reflexion (Shinn et
al., 2023) agents, we allow them to reflect on the environment’s feedback but
require them to generate a plan immediately without sophisticated reasoning.

#### RAFA Setup.

In the Game of 24, the RAFA agent uses ToT as the planner, regenerates a plan
when the agent receives a zero reward and continues acting according to the
previous plan when the agent receives a positive reward. We set the base ToT
planner with beam search width $b=1,2$ and use both gpt-3.5-turbo and gpt-4 to
test the RAFA’s boost-up over LLM agents with different reasoning abilities.
We set the temperature $t=0.2$ by default to favor rigorous reasoning and
$t=0.7$ for majority voting.

#### Reduced Hallucination Through Interaction.

A comprehensive review of various method proposals revealed significant
hallucination, especially with gpt-3.5-turbo. A common hallucination is that
the agent believes she can reuse the same number (e.g. using the number $2$
twice as illustrated in Figure 2). RAFA efficiently mitigates such
hallucination by actively interacting with the environment, displaying
exceptional hallucination resistance and improved performance.

#### Enhanced Efficiency Through Planning.

Evidenced in Figure 3, the RAFA agent substantially surpasses the Reflexion
baseline, reflecting heightened efficiency and minimized regret by negating
careless trials. For example, without carefully planning, agent may give
negative answers, e.g,“Impossible to obtain 24 with the given numbers, or
unchecked answers, e.g.,“Answer: 6 * 9 / (3 - 2) = 24”. This reduction of
careless trails is especially achieved when a strong backbone LLMs (e.g.,
gpt-4) is used, even with a basic planning method, such as BFS with $B=1$.

#### Ablation Study.

The RAFA agent’s performance is dissected by individually examining its
components: (1) Planning modules or model/elite LLM, (2) Reflection modules or
critic LLM, and (3) Different LLMs. Results, displayed in Table 5.1 and Figure
3, affirm the substantial contribution of each segment to the aggregate
performance. Compared to absent or rudimentary zero-shot planning, a basic
planner markedly enhances overall performance. However, augmenting planner
strength only offers marginal performance enhancements. Both critic LLM and
robust LLM usage emerge as pivotal for optimal performance.

###  E.2 ALFWorld

#### Task Setup.

The action space of ALFWorld consists of high-level actions such as “heat a
potato with a microwave”, which is executed in the underlying embodied
simulator through low-level action primitives. The egocentric visual
observations of the simulator are translated into natural language before
being provided to the agent. The state is the history of the observations. If
a task goal can be precisely achieved by the agent, it will be counted as a
success.

#### RAFA Setup.

In the ALFWorld environment, the RAFA planner is instantiated as Breadth First
Search (BFS). Specifically, $B$ and $U$ are both set to 2, and we use gpt-3
(text-davinci-003) for the Critic, Model, and Elite modules. Besides, since it
is challenging to prompt the LLM with the stored full trajectories in the
memory buffer due to the token limit, we make the following modifications: the
Model LLM instance uses only the partial trajectory executed so far in the
current episode, and the Elite LLM instance uses the same partial executed
trajectory with additional model-generated state-action pairs during the
planning subroutine. When switching is triggered after $20$ failed timesteps
(i.e., an episode), a summary from the failure trajectory is generated by
gpt-4 and added to the Critic prompt.

#### Reduced Hallucination Through Interaction.