SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Large Language Models Can Implement Policy Iteration
Ethan Brooks<sup>1</sup>, Logan Walls<sup>2</sup>, Richard L. Lewis<sup>2</sup>, Satinder Singh<sup>1</sup>
<sup>1</sup>Computer Science and Engineering, University of Michigan
$^{2}$ Department of Psychology, University of Michigan
{ethanbro, logwalls, rickl, baveja}@umich.edu
In this work, we demonstrate a method for implementing policy iteration using a large language model. While the application of foundation models to Reinforcement Learning (RL) has received considerable attention, most approaches rely on either (1) the curation of expert demonstrations (through manual design or task-specific pretraining) or (2) adaptation to the task of interest using gradient methods (fine-tuning or training of adapter layers). Both of these techniques have drawbacks. Collecting demonstrations is labor-intensive, and algorithms that rely on them do not outperform the experts from which the demonstrations were derived. Gradient techniques are inherently slow, sacrificing the "few-shot" quality that makes in-context learning attractive to begin with. Our method demonstrates that a large language model can be used to implement policy iteration using the machinery of in-context learning, enabling it to learn to perform RL tasks without expert demonstrations or gradients. Our approach iteratively updates the contents of the prompt from which it derives its policy through trial-and-error interaction with an RL environment. In order to eliminate the role of in-weights learning (on which approaches like Decision Transformer (L. Chen et al. 2021) rely heavily), we demonstrate our method using Codex (M. Chen et al. 2021b), a language model that was not trained to perform the tasks which we are evaluating.
In many settings, models implemented using a transformer or recurrent architecture will improve their performance as information accumulates in their context or memory. We refer to this phenomenon as "in-context learning." Brown et al. 2020b demonstrated a technique for inducing form of learning by prompting a large language model with a small number of input/output exemplars. An interesting property of in-context learning in the case of large pre-trained models (or "foundation models") is that the models are not directly trained to optimize a meta-learning objective, but demonstrate an emergent capacity to generalize (or at least specialize) to diverse downstream task-distributions (Wei et al. 2022b).
This remarkable capability has been studied in a variety of settings, including Reinforcement Learning (RL). However, most work in this area either (1) assumes access to expert demonstrations — collected either from human experts (Huang et al. 2022b; Baker et al. 2022), or domain-specific pre-trained RL agents (L. Chen et al. 2021; Lee et al. 2022; Janner et al. 2021; Reed et al. 2022; Xu et al. 2022). — or (2) relies on gradient-based methods — e.g. fine-tuning of the foundation models parameters as a whole (Lee et al. 2022; Reed et al. 2022; Baker et al. 2022) or newly training an adapter layer or prefix vectors while keeping the original foundation models frozen (X. L. Li et al. 2021; A. K. Singh et al. 2022; Karimi Mahabadi et al. 2022).
Figure 1: For each possible action  $\mathcal{A}(1),\ldots ,\mathcal{A}(n)$ , the LLM generates a rollout by alternately predicting transitions and selecting actions. Q-value estimates are discounted sums of rewards. The action is chosen greedily with respect to Q-values. Both state/reward prediction and next action selection use trajectories from  $\mathcal{D}$  to create prompts for the LLM. Changes to the content of  $\mathcal{D}$  change the prompts that the LLM receives, allowing the model to improve its behavior over time.
Our work demonstrates an approach to in-context learning which relaxes these assumptions. Our method, In-Context Policy Iteration (ICPI), implements policy iteration using the prompt content, instead of the model parameters, as the locus of learning, thereby avoiding gradient methods. Furthermore, the use of policy iteration frees us from expert demonstrations because suboptimal prompts can be improved over the course of training.
We illustrate the method empirically on six small illustrative RL tasks—chain, distractor-chain, maze, mini-catch, mini-invaders, and point-mass—in which the method very quickly finds good policies. We also compare five pretrained Large Language Models (LLMs), including two different size models trained on natural language—OPT-30B and GPT-J—and three different sizes of a model trained on program code—two sizes of Codex as well as InCoder. On our six domains, we find that only the largest model (the code-davinci-001 variant of Codex) consistently demonstrates learning.
A common application of foundation models to RL involves tasks that have language input, for example natural language instructions/goals (D. Garg et al. 2022; Hill et al. 2020) or text-based games (Peng et al. 2021; I. Singh et al. 2021; Majumdar et al. 2020; Ammanabrolu et al. 2021). Another approach encodes RL trajectories into token sequences, and processes them with a foundation model, model representations as input to deep RL architectures (S. Li et al. 2022; Tarasov et al. 2022; Tam et al. 2022). Finally, a recent set of approaches (which we will focus on in this Related Work section) treat RL as a sequence modeling problem and use the foundation models itself to predict states or actions. In this related work section, we will focus a third set of recent approaches that treat reinforcement learning (RL) as a sequence modeling problem and utilize foundation models for state prediction, action selection, and task completion. We will organize our survey of these approaches based on how they elicit these RL-relevant outputs from the foundation models. In this respect the approaches fall under three broad categories: learning from demonstrations, specialization (via training or finetuning), and context manipulation (in-context learning).
2.1 Learning from demonstrations
Many recent sequence-based approaches to reinforcement learning use demonstrations that come either from human experts or pretrained RL agents. For example, Huang et al. 2022b use a frozen LLM as a planner for everyday household tasks by constructing a prefix from human-generated task instructions, and then using the LLM to generate instructions for new tasks. This work is extended by Huang et al. 2022a. Similarly, Ahn et al. 2022 use a value function that is trained on
human demonstrations to rank candidate actions produced by an LLM. Baker et al. 2022 use human demonstrations to train the foundation model itself: they use video recordings of human Minecraft players to train a foundation models that plays Minecraft. Works that rely on pretrained RL agents include Janner et al. 2021 who train a "Trajectory Transformer" to predict trajectory sequences in continuous control tasks by using trajectories generated by pretrained agents, and L. Chen et al. 2021, who use a dataset of offline trajectories to train a "Decision Transformer" that predicts actions from state-action-reward sequences in RL environments like Atari. Two approaches build on this method to improve generalization: Lee et al. 2022 use trajectories generated by a DQN agent to train a single Decision Transformer that can play many Atari games, and Xu et al. 2022 use a combination of human and artificial trajectories to train a Decision Transformer that achieves few-shot generalization on continuous control tasks. Reed et al. 2022 take task-generality a step farther and use datasets generated by pretrained agents to train a multi-modal agent that performs a wide array of RL (e.g. Atari, continuous control) and non-RL (e.g. image captioning, chat) tasks.
Some of the above works include non-expert demonstrations as well. L. Chen et al. 2021 include experiments with trajectories generated by random (as opposed to expert) policies. Lee et al. 2022 and Xu et al. 2022 also use datasets that include trajectories generated by partially trained agents in addition to fully trained agents. Like these works, our proposed method (ICPI) does not rely on expert demonstrations—but we note two key differences between our approach and existing approaches. Firstly, ICPI only consumes self-generated trajectories, so it does not require any demonstrations (like L. Chen et al. 2021 with random trajectories, but unlike Lee et al. 2022, Xu et al. 2022, and the other approaches reviewed above). Secondly, ICPI relies primarily on in-context learning rather than in-weights learning to achieve generalization (like Xu et al. 2022, but unlike L. Chen et al. 2021 & Lee et al. 2022). For discussion about in-weights vs. in-context learning see Chan et al. 2022.
2.2 Gradient-based training & finetuning on RL tasks
Many approaches that use foundation models for RL involve specifically training or fine-tuning on RL tasks. For example, Janner et al. 2021; L. Chen et al. 2021; Lee et al. 2022; Xu et al. 2022; Baker et al. 2022; Reed et al. 2022 all use models that are trained from scratch on tasks of interest, and A. K. Singh et al. 2022; Ahn et al. 2022; Huang et al. 2022a combine frozen foundation models with trainable components or adapters. In contrast, Huang et al. 2022b use frozen foundation models for planning, without training or fine-tuning on RL tasks. Like Huang et al. 2022b, ICPI does not update the parameters of the foundation model, but relies on the frozen model's in-context learning abilities. However, ICPI gradually builds and improves the prompts within the space defined by the given fixed text-format for observations, actions, and rewards (in contrast to Huang et al. 2022b, which uses the frozen model to select good prompts from a given fixed library of goal/plan descriptions).
2.3 In-Context learning
Several recent papers have specifically studied in-context learning. Laskin et al. 2022 demonstrates an approach to performing in-context reinforcement learning by training a model on complete RL learning histories, demonstrating that the model actually distills the improvement operator of the source algorithm. Chan et al. 2022 and S. Garg et al. 2022 provide analyses of the properties that drive in-context learning, the first in the context of image classification, the second in the context of regression onto a continuous function. These papers identify various properties, including "burstiness," model-size, and model-architecture, that in-context learning depends on. Y. Chen et al. 2022 studies the sensitivity of in-context learning to small perturbations of the context. They propose a novel method that uses sensitivity as a proxy for model certainty. Some recent work has explored iterative forms of in-context learning, similar to our own. For example, Shinn et al. 2023 and Madaan et al. 2023 use iterative self-refinement to improve the outputs of a large language model in a natural language context. These approaches rely on the ability of the model to examine and critique its own outputs, rather than using policy iteration as our method does.
How can standard policy iteration make use of in-context learning? Policy iteration is either model-based—using a world-model to plan future trajectories in the environment—or model-free—inflerring value-estimates without explicit planning. Both methods can be realized with in-context learning.
Algorithm 2 Computing Q-values
1: function  $Q(o_{t},a,\mathcal{D})$    
2:  $u\gets t$    
3:  $o^{(1)} = o_t$    
4:  $a^{(1)} = a$    
5: repeat ▷ All samples come from the experience buffer  $\mathcal{D}$    
6:  $\mathcal{D}_b\sim$  time-steps with action  $a^{(u)}$  balancing terminal and non-terminal   
7:  $b^{(u)}\sim$  LLM  $(\mathcal{D}_b,o^{(u)},a^{(u)})$    
8:  $\mathcal{D}_r\sim$  time-steps with action  $a^{(u)}$  and termination  $b^{(u)}$  balancing reward   
9:  $r^{(u)}\sim$  LLM  $(\mathcal{D}_r,o^{(u)},a^{(u)})$    
10:  $\mathcal{D}_o\sim$  time-steps with action  $a^{(u)}$  and termination  $b^{(u)}$  no balancing   
11:  $o^{(u + 1)}\sim$  LLM  $(\mathcal{D}_o,o^{(u)},a^{(u)})$    
12:  $\mathcal{D}_a\sim c$  recent trajectories   
13:  $a^{(u + 1)}\sim$  LLM  $(o^{(u + 1)},\mathcal{D}_a)$    
14:  $u\gets u + 1$    
15: until  $b^{(u)}$  is terminal   
16: return  $\sum_{k = 1}^{u}\gamma^{k - 1}r^{(k)}$    
17: end function
We choose model-based learning because planned trajectories make the underlying logic of value-estimates explicit to our foundation model backbone, providing a concrete instantiation of a trajectory that realizes the values. This ties into recent work (Wei et al. 2022a; Nye et al. 2021) demonstrating that "chains of thought" can significantly improve few-shot performance of foundation models.
Model-based RL requires two ingredients, a rollout-policy used to act during planning and a world-model used to predict future rewards, terminations, and states. Since our approach avoids any mutation of the foundation model's parameters (this would require gradients), we must instead induce the rollout-policy and the world-model using in-context learning, i.e. by selecting appropriate prompts. We induce the rollout-policy by prompting the foundation model with trajectories drawn from the current (or recent) behavior policy (distinct from the rollout-policy). Similarly, we induce the world-model by prompting the foundation models with transitions drawn from the agent's history of experience. Note that our approach assumes access to some translation between the state-space of the environment and the medium (language, images, etc.) of the foundation models. This explains how an algorithm might plan and estimate values using a foundation model. It also explains how the rollout-policy approximately tracks the behavior policy.
How does the policy improve? When acting in the environment (as opposed to planning), we choose the action that maximizes the estimated Q-value from the current state (see Training Loop pseudocode, line 6). At time step  $t$ , the agent observes the state of the environment (denoted  $s_t$ ) and executes action  $a_t = \arg \max_{a \in \mathcal{A}} Q^{(\pi_t)}(o_t, a)$ , where  $\mathcal{A} = \{\mathcal{A}(1), \dots, \mathcal{A}(n)\}$  denotes the set of  $n$  actions
Algorithm 1 Training Loop
1: function TRAIN(environment)  
2: initialize  $\mathcal{D}$   
3: while training do  
4:  $o_0 \gets$  Reset environment.  
5: while episode is not done do  
6:  $a_t \gets$  arg max  $a$ $Q(o_t, a, \mathcal{D})$   
7:  $o_{t+1}, r_t, b_t \gets$  Execute  $a_t$  in environment.  
8:  $t \gets t + 1$   
9: end while  
10:  $\mathcal{D} \gets \mathcal{D} \cup (o_0, a_0, r_0, b_0, o_1, \ldots, o_t, a_t, r_t, b_t, o_{t+1})$   
11: end while  
12: end function
available,  $\pi_t$  denotes the policy of the agent at time step  $t$ , and  $Q^{(\pi)}$  denotes the Q-estimate for policy  $\pi$ . Taking the greedy (arg max) actions with respect to  $Q^{\pi_t}$  implements a new and improved policy.
Computing Q-values This section provides details on the prompts that we use in our computation of Q-values (see Computing Q-values pseudocode & Figure 1). During training, we maintain a buffer  $\mathcal{D}$  of transitions experienced by the agent. To compute  $Q^{(\pi_t)}(o_t,a)$  at time step  $t$  in the real-world we rollout a simulated trajectory  $o^{(1)} = o_t$ ,  $a^{(1)} = a$ ,  $r^{(1)}$ ,  $o^{(2)}$ ,  $r^{(2)}$ ,  $\dots$ ,  $o^{(T)}$ ,  $a^{(T)}$ ,  $r^{(T)}$ ,  $o^{(T + 1)}$  by predicting, at each simulation time step  $u$ : reward  $r^{(u)} \sim \text{LLM}(\mathcal{D}_r, o^{(u)}, a^{(u)})$ ; termination  $b^{(u)} \sim \text{LLM}(\mathcal{D}_b, o^{(u)}, a^{(u)})$ ; observation  $o^{(u + 1)} \sim \text{LLM}(\mathcal{D}_o, o^{(u)}, a^{(u)})$ ; action  $a^{(1)} \sim \text{LLM}(\mathcal{D}_a, o^{(u)})$ . Termination  $b^{(u)}$  decides whether the simulated trajectory ends at step  $u$ .
The prompts  $\mathcal{D}_r$ ,  $\mathcal{D}_b$  contain data sampled from the replay buffer. For each prompt, we choose some subset of replay buffer transitions, shuffle them, convert them to text (examples are provided in table 4.1) and clip the prompt at the 4000-token Codex context limit. We use the same method for  $\mathcal{D}_a$ , except that we use random trajectory subsequences.
In order to maximize the relevance of the prompt contents to the current inference we select transitions using the following criteria.  $\mathcal{D}_b$  contains  $(o_k, a_k, b_k)$  tuples such that  $a_k$  equals  $a^{(u)}$ , the action for which the LLM must infer termination.  $\mathcal{D}_r$  contains  $(o_k, a_k, r_k)$  tuples, again constraining  $a_k = a^{(u)}$  but also constraining  $b_k = b^{(k)}$  — that the tuple corresponds to a terminal time-step if the LLM inferred  $b^{(u)} = \text{true}$ , and to a non-terminal time-step if  $b^{(u)} = \text{false}$ . For  $\mathcal{D}_o$ , the prompt includes  $(o_k, a_k o_{k+1})$  tuples with  $a_k = a^{(u)}$  and  $b_k = \text{false}$  (only non-terminal states need to be modelled).
We also maintain a balance of certain kinds of transitions in the prompt. For termination prediction, we balance terminal and non-terminal time-steps. Since non-terminal time-steps far outnumber terminal time-steps, this eliminates a situation wherein the randomly sampled prompt time-steps are entirely non-terminal, all but ensuring that the LLM will predict non-termination. Similarly, for reward prediction, we balance the number of time-steps corresponding to each reward value stored in  $\mathcal{D}$ . In order to balance two collections of unequal size, we take the smaller and duplicate randomly chosen members until the sizes are equal.
In contrast to the other predictions, we condition the rollout policy on trajectory subsequences, not individual time-steps. Prompting with sequences better enables the foundation model to apprehend the logic behind a policy. Trajectory subsequences consist of  $(o_k, a_k)$  pairs, randomly clipped from the  $c$  most recent trajectories. More recent trajectories will, in general demonstrate higher performance, since they come from policies that have benefited from more rounds of improvement.
Finally, the Q-value estimate is simply the discounted sum of rewards for the simulated episode. Given this description of Q-value estimation, we now return to the concept of policy improvement.
Policy-Improvement The arg max (line 6 of Algorithm 1) drives policy improvement in ICPI. Critically this is not simply a one-step improvement but a mechanism that builds improvement on top of improvement. This occurs through a cycle in which the arg max improves behavior. The improved behavior is stored in the buffer  $\mathcal{D}$ , and then used to condition the rollout policy. This improves the returns generated by the LLM during planning rollouts. These improved rollouts improve the Q-estimates for each action. Completing the cycle, this improves the actions chosen by the arg max. Because this process feeds into itself, it can drive improvement without bound until optimality is achieved.
Note that this process takes advantage of properties specific to in-context learning. In particular, it relies on the assumption that the rollout policy, when prompted with trajectories drawn from a mixture of policies, will approximate something like an average of these policies. Given this assumption, the rollout policy will improve with the improvement of the mixture of policies from which its prompt-trajectories are drawn. This results in a kind of rapid policy improvement that works without any use of gradients.
Figure 2: Comparison of ICPI with three baselines, "No arg max", "Tabular Q," and "Nearest Neighbor." The  $y$ -axis depicts regret (normalized between 0 and 1), computed relative to an optimal return with a discount-factor of 0.8. The  $x$ -axis depicts time-steps during training. Error bars are standard errors from four seeds.
Prompt-Format The LLM cannot take non-linguistic prompts, so our algorithm assumes access to a textual representation of the environment—of states, actions, terminations, and rewards—and some way to recover the original action, termination, and reward values from their textual representation (we do not attempt to recover states). Since our primary results use the Codex language model (see Table 1), we use Python code to represent these values (examples are available in Table 2 in the appendix).
In our experiments, we discovered that the LLM world-model was unable to reliably predict rewards, terminations, and next-states on some of the more difficult environments. We experimented with providing domain hints in the form of prompt formats that make explicit useful information — similar to Chain of Thought Prompting (Wei et al. 2022a). For example, for the chain domain, the hint includes an explicit comparison  $(==$  or  $!=$ ) of the current state with the goal state. Note that while hints are provided in the initial context, the LLM must infer the hint content in rollouts generated from this context.
We use a consistent idiom for rewards and terminations, namely assert reward  $== \mathbf{x}$  and assert done or assert not done. Some decisions had to be made when representing states and actions. In general, we strove to use simple, idiomatic, concise Python. On the more challenging environments, we did search over several options for the choice of hint. For examples, see Table 2. We anticipate that in the future, stronger foundation models will be increasingly robust to these decisions.
We have three main goals in our experiments: (1) Demonstrate that the agent algorithm can in fact quickly learn good policies, using pretrained LLMs, in a set of six simple illustrative domains of increasing challenge; (2) provide evidence through an ablation that the policy-improvement step—taking the arg max over Q-values computed through LLM rollouts—accelerates learning; and (3) investigate the impact of using different LLMs (see Table 1)—different sizes and trained on different data, in particular, trained on (mostly) natural language (GPT-3 and GPT-J) vs. program code (Codex and InCoder). We next describe the six domains and their associated prompt formats, and then describe the experimental methodology and results.
4.1 Domains and prompt format
Chain. In this environment, the agent occupies an 8-state chain. The agent has three actions: Left, right, and try goal. The try goal action always terminates the episode, conferring a reward of 1 on state 4 (the goal state) and 0 on all other states. Episodes also terminate after 8 time-steps. States are represented as numbers from 0 to 7, as in assert state  $= =$  n, with the appropriate integer substituted for n. The actions are represented as functions left(), right(), and try_goal(). For the hint, we simply indicate whether or not the current state matches the goal state, 4.
Distractor Chain. This environment is an 8-state chain, identical to the chain environment, except that the observation is a pair of integers, the first indicating the true state of the agent and the second acting as a distractor which transitions randomly within  $\{0,\dots ,7\}$ . The agent must therefore learn to ignore the distractor integer and base its inferences on the information contained in the first integer.
Aside from the addition of this distractor integer to the observation, all text representations and hints are identical to the chain environment.
Maze. The agent navigates a small  $3 \times 3$  gridworld with obstacles. The agent can move up, down, left, or right. The episode terminates with a reward of 1 once the agent navigates to the goal grid, or with a reward of 0 after 8 time-steps. This environment tests our algorithms capacity to handle 2-dimensional movement and obstacles, as well as a 4-action state-space. We represent the states as namedtuples — C(x, y), with integers substituted for x and y. Similar to chain, the hint indicates whether or not the state corresponds to the goal state.