Title: Large Language Models can Implement Policy Iteration
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

In-Context Learning

Foundation Models
Abstract: In this work, we demonstrate a method for implementing policy iteration using a large language model. While the application of foundation models to RL has received considerable attention, most approaches rely on either (1) the curation of expert demonstrations (either through manual design or task-specific pretraining) or (2) adaptation to the task of interest using gradient methods (either 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. All 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 rely heavily), we demonstrate our method using Codex (M. Chen et al. 2021b), a language model with no prior knowledge of the domains on which we evaluate it.

Large Language Models Can Implement Policy Iteration

Abstract

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.

1 Introduction

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).

st

LLM

a

1

s

1

A

(1)

=

st

=

For each action in

{A

(1)

, ...,

A

(

n

)

}

:

D

s

2

r

1

LLM

D

a

2

LLM

D

s

3

r

2

...

Q

(

st ,

A

(1)) =

∑

γ u

r u

...

...

...

...

...

LLM

s

1

a

1

A

(

n

)

=

st

=

D

s

2

r

1

LLM

D

a

2

LLM

D

s

3

r

2

...

Q

(

st ,

A

(

n

)) =

∑

γ u

r u

a

at

Environment

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 taskschain, 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.

2 Related Work

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.

3 Method

How can standard policy iteration make use of in-context learning? Policy iteration is either modelbased -using a world-model to plan future trajectories in the environment-or model-free -inferring value-estimates without explicit planning. Both methods can be realized with in-context learning.

We choose model-based learning because planned trajectories make the underlying logic of valueestimates 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 worldmodel 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 ∈A Q ( π t ) ( o t , a ) , where A = {A (1) , · · · , A ( n ) } denotes the set of n actions

Algorithm 1 Training Loop

available, π t denotes the policy of the agent at time step t , and Q ( π ) denotes the Q-estimate for policy π . Taking the greedy ( arg max ) actions with respect to Q π 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 D of transitions experienced by the agent. To compute Q ( π 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) , a (2) , r (2) , · · · , o ( T ) , a ( T ) , r ( T ) , o ( T +1) by predicting, at each simulation time step u : reward r ( u ) ∼ LLM ( D r , o ( u ) , a ( u ) ) ; termination b ( u ) ∼ LLM ( D b , o ( u ) , a ( u ) ) ; observation o ( u +1) ∼ LLM ( D o , o ( u ) , a ( u ) ) ; action a (1) ∼ LLM ( D a , o ( u ) ) . Termination b ( u ) decides whether the simulated trajectory ends at step u .

The prompts D r , 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 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. 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. 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 ) = true, and to a non-terminal time-step if b ( u ) = false. For D o , the prompt includes ( o k , a k o k +1 ) tuples with a k = a ( u ) and b k = 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 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.

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 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.

Algorithms

-ICPl-NoArgMax-TabularQ-Matching Model

Chain

Distractor Chain

Maze

MiniInvaders

Point Mass

0.8

gret

0.6

regr

0.4

0.2

0

20

40

60800

40

801200

40

80

120

0

80160240

0

80

160

0

80

160

step

step

step

step

step

step

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 == 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.

4 Experiments

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 steptaking 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 , . . . , 7 } . The agent must therefore learn to ignore the distractor integer and base its inferrences 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 × 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.

Mini Catch. The agent operates a paddle to catch a falling ball. The ball falls from a height of 5 units, descending one unit per time step. The paddle can stay in place (not move), or move left or right along the bottom of the 4-unit wide plane. The agent receives a reward of 1 for catching the ball and 0 for other time-steps. The episode ends when the ball's height reaches the paddle regardless of whether or not the paddle catches the ball. We chose this environment specifically to challenge the action-inference/rollout-policy component of our algorithm. Specifically, note that the success condition in Mini Catch allows the paddle to meander before moving under the ball-as long as it gets there on the final time-step. Successful trajectories that include movement away from the ball thus make a good rollout policies more challenging to learn (i.e., elicit from the LLM via prompts). Again, we represent both the paddle and the ball as namedtuples C(x, y) and we represent actions as methods of the paddle object: paddle.stay() , paddle.left() , and paddle.right() . For the hint, we call out the location of the paddle's x -position, the ball's x -position, the relation between these positions (which is larger than which, or whether they are equal) and the ball's y -position. Table 2 in the appendix provides an example. We also include the text ball.descend() to account for the change in the ball's position between states.

Mini Invaders. The agent operates a ship that shoots down aliens which descend from the top of the screen. At the beginning of an episode, two aliens spawn at a random location in two of four columns. The episode terminates when an alien reaches the ground (resulting in 0 reward) or when the ship shoots down both aliens (the agent receives 1 reward per alien). The agent can move left , right , or shoot . This domain highlights ICPI's capacity to learn incrementally, rather than discovering an optimal policy through random exploration and then imitating that policy, which is how our 'No arg max ' baseline learns (see Comparison of ICPI with baseline algorithms). ICPI initially learns to shoot down one alien, and then builds on this good but suboptimal policy to discover the better policy of shooting down both aliens. In contrast, random exploration takes much longer to discover the optimal policy and the 'No arg max ' baseline has only experienced one or two successful trajectories by the end of training.

We represent the ship by its namedtuple coordinate ( C(x, y) ) and the aliens as a list of these namedtuples. When an alien is shot down, we substitute None for the tuple, as in aliens == [C(x, y), None] . We add the text for a in aliens: a.descend() in order to account for the change in the alien's position between states.

Point-Mass. A point-mass spawns at a random position on a continuous line between -6 and +6 with a velocity of 0. The agent can either accelerate the point-mass (increase velocity by 1) or decelerate it (decrease the velocity by 1). The point-mass position changes by the amount of its velocity each timestep. The episode terminates with a reward of 1 once the point-mass is between -2 and +2 and its velocity is 0 once again. The episode also terminates after 8 time-steps. This domain tests the algorithm's ability to handle continuous states.

States are represented as assert pos == p and vel == v , substituting floats rounded to two decimals for p and v . The actions are accel(pos, vel) and decel(pos, vel) . The hint indicates whether the success conditions are met, namely the relationship of pos to -2 and +2 and the whether or not vel == 0 . The hint includes identification of the aliens' and the ship's x -positions as well as a comparison between them.

4.2 Experiment Methodology and Results

Methodology and evaluation. For the results, we record the agent's regret over the course of training relative to an optimal policy computed with a discount factor of 0.8. For all experiments

Ablations

·ICPl-NoHints·c=16-NoConstraints-NoBalance

Chain

Distractor Chain

Maze

MiniInvaders

Point Mass

0.8

gret

0.6

regr

0.4

0.2

0

20

40

60800

40

801200

40

80

120

0

80160240

0

80

160

0

80

160

step

step

step

step

step

step

c = 8 (the number of most recent successful trajectories to include in the prompt). We did not have time for hyperparameter search and chose this number based on intuition. However the c = 16 baseline demonstrates results when this hyperparameter is doubled. All results use 4 seeds.

For both versions of Codex, we used the OpenAI Beta under the API Terms of Use. For GPT-J (B. Wang et al. 2021) , InCoder (Fried et al. 2022) and OPT-30B (Zhang et al. 2022), we used the open-source implementations from Huggingface Transformers (Wolf et al. 2020), each running on one Nvidia A40 GPU. All language models use a sampling temperature of 0.1. Code for our implementation is available at https://github.com/ethanabrooks/icpi .

Comparison of ICPI with baseline algorithms. We compare ICPI with three baselines (Fig. 2).

The 'No arg max ' baseline learns a good policy through random exploration and then imitates this policy. This baseline assumes access to a 'success threshold' for each domain - an undiscounted cumulative return greater than which a trajectory is considered successful. The action selection mechanism emulates ICPI's rollout policy: prompting the LLM with a set of trajectories and eliciting an action as output. For this baseline, we only include trajectories in the prompt whose cumulative return exceeds the success threshold. Thus the policy improves as the number of successful trajectories in the prompt increases over time. Note that at the start of learning, the agent will have experienced too few successful trajectories to effectively populate the policy prompt. In order to facilitate exploration, we act randomly until the agent experiences 3 successes.

'Tabular Q' is a standard tabular Q-learning algorithm, which uses a learning rate of 1 . 0 and optimistically initializes the Q-values to 1 . 0 .

'Matching Model' is a baseline which uses the trajectory history instead of an LLM to perform modelling. This baseline searches the trajectory buffer for the most recent instance of the current state, and in the case of transition/reward/termination prediction, the current action. If a match is found, the model outputs the historical value (e.g. the reward associated with the state-action pair found in the buffer). If no match is found, the modelling rollout is terminated. Recall that ICPI breaks ties randomly during action selection so this will often lead to random action selection.

As our results demonstrate, only ICPI learns good policies on all domains. We attribute this advantage to ICPI's ability to generalize from its context to unseen states and state/action pairs (unlike 'Tabular Q' and 'Matching Model'). Unlike 'No arg max ' ICPI is able to learn progressively, improving the policy before experiencing good trajectories.

Ablation of ICPI components. With these experiments, we ablate those components of the algorithm which are not, in principle, essential to learning (Fig. 3). 'No Hints' ablates the hints described in the Prompt-Format paragraph. 'No Balance' removes the balancing of different kinds of time-steps described in the Computing Q-values paragraph (for example, D b is allowed to contain an unequal number of terminal and non-terminal time-steps). The 'No Constraints' baseline removes the constraints on these time-steps described in the same paragraph. For example, D r is allowed to contain a mixture of terminal and non-terminal time-steps (regardless of the model's termination prediction). Finally, ' c = 16 ' prompts the rollout policy with the last 16 trajectories (instead of the last 8, as in ICPI). We find that while some ablations match ICPI's performance in several domains, none match its performance on all six.

LanguageModels

code-davinci-002=code-cushman-001·OPT-30B-InCoder-GPT-J

Chain

DistractorChain

Maze

MiniInvaders

PointMass

0.8

gret

0.6

0.4

0.2

0

20

40

60800

40

80120

0

40

80

120

0

80

160

240

0

80

160

0

80

160

step

step

step

step

step

step

Comparison of Different Language Models. While our lab lacks the resources to do a full study of scaling properties, we did compare several language models of varying size (Fig. 4). See Table 1 for details about these models. Both code-davinci-002 and code-cushman-001 are variations of the Codex language model. The exact number of parameters in these models is proprietary according to OpenAI, but M. Chen et al. 2021a describes Codex as fine-tuned from GPT-3 (Brown et al. 2020a), which contains 185 billion parameters. As for the distinction between the variations, the OpenAI website describes code-cushman-001 as 'almost as capable as Davinci Codex, but slightly faster.'

We found that the rate of learning and final performance of the smaller models fell significantly short of Codex on all but the simplest domein, chain . Examining the trajectories generated by agents trained using these models, we noted that in several cases, they seemed to struggle to apprehend the underlying 'logic' of successful trajectories, which hampered the ability of the rollout policy to produce good actions. Since these smaller models were not trained on identical data, we are unable to isolate the role of size in these results. However, the failure of all of these smaller models to learn suggests that size has some role to play in performance. We conjecture that larger models developed in the future may demonstrate comparable improvements in performance over our Codex model.

Limitations ICPI can theoretically work on any control task with discrete actions, due to the guarantees associated with policy iteration. However, since our implementation uses Codex, the domains in our paper were limited by the ability to encode states as text and to fit these encodings in the model's context window. Moreover, Codex demonstrated a limited ability to predict transitions and actions in more complex domains. As sequence models mature, we anticipate that more domains will become tractable for ICPI. We also note that reliance on the proprietary OpenAPI API limits exact reproduction of these results.

Societal Impacts An extensive literature (Tamkin et al. 2021; Abid et al. 2021; Liang et al. 2021; Pan et al. 2023) has explored the possible positive and negative impacts of LLMs. Some of this work has explored mitigation strategies. In extending LLMs to RL, our work inherits these benefits and challenges. We highlight two concerns: the use of LLMs to spread misinformation and the detrimental carbon cost of training and using these models.

5 Conclusion

Our main contribution is a method for implementing policy iteration algorithm using Large Language Models and the mechanism of in-context learning. The algorithm uses a foundation models as both a world model and policy to compute Q-values via rollouts. Although we presented the method here as text-based, it is general enough to be applied to any foundation models that works through prompting, including multi-modal models like Reed et al. 2022 and Seo et al. 2022. In experiments we showed that the algorithm works in six illustrative domains imposing different challenges for ICPI, confirming the benefit of the LLM-rollout-based policy improvement. While the empirical results are preliminary, we believe the approach provides an important new way to use LLMs that will increase in effectiveness as the models themselves become more powerful.

References

Abubakar Abid et al. 'Persistent anti-muslim bias in large language models'. In: Proceedings of the 2021 AAAI/ACM Conference on AI, Ethics, and Society . 2021, pp. 298-306.

Michael Ahn et al. Do As I Can, Not As I Say: Grounding Language in Robotic Affordances . arXiv:2204.01691 [cs]. Apr. 2022. DOI: 10.48550/arXiv.2204.01691 . URL: http: //arxiv.org/abs/2204.01691 (visited on 08/09/2022).

Prithviraj Ammanabrolu et al. 'Learning Knowledge Graph-based World Models of Textual Environments'. In: Advances in Neural Information Processing Systems . Vol. 34. Curran Associates, Inc., 2021, pp. 3720-3731. URL: https://proceedings.neurips.cc/ paper/2021/hash/1e747ddbea997a1b933aaf58a7953c3c-Abstract.html (visited on 05/16/2022).

Bowen Baker et al. Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos . arXiv:2206.11795 [cs]. June 2022. DOI: 10.48550/arXiv.2206.11795 . URL: http://arxiv.org/abs/2206.11795 (visited on 08/09/2022).

Jason Baumgartner et al. The Pushshift Reddit Dataset . Number: arXiv:2001.08435 arXiv:2001.08435 [cs]. Jan. 2020. DOI: 10.48550/arXiv.2001.08435 . URL: http: //arxiv.org/abs/2001.08435 (visited on 08/02/2022).

Tom Brown et al. 'Language Models are Few-Shot Learners'. In: Advances in Neural Information Processing Systems . Vol. 33. Curran Associates, Inc., 2020, pp. 1877-1901. URL: https: / / papers . nips . cc / paper / 2020 / hash / 1457c0d6bfcb4967418bfb8ac142f64a Abstract.html (visited on 05/14/2022).

Tom Brown et al. 'Language models are few-shot learners'. In: Advances in neural information processing systems 33 (2020), pp. 1877-1901.

Stephanie C. Y. Chan et al. Data Distributional Properties Drive Emergent In-Context Learning in Transformers . arXiv:2205.05055 [cs]. May 2022. DOI: 10.48550/arXiv.2205.05055 . URL: http://arxiv.org/abs/2205.05055 (visited on 08/11/2022).

Lili Chen et al. 'Decision Transformer: Reinforcement Learning via Sequence Modeling'. In: arXiv:2106.01345 [cs] (June 2021). URL: http://arxiv.org/abs/2106.01345 (visited on 05/18/2022).

Mark Chen et al. Evaluating Large Language Models Trained on Code . Tech. rep. arXiv:2107.03374. arXiv, July 2021. URL: http://arxiv.org/abs/2107.03374 (visited on 05/15/2022).

Mark Chen et al. 'Evaluating large language models trained on code'. In: arXiv preprint arXiv:2107.03374 (2021).

Yanda Chen et al. On the Relation between Sensitivity and Accuracy in In-context Learning . arXiv:2209.07661 [cs]. Sept. 2022. URL: http://arxiv.org/abs/2209.07661 (visited on 09/28/2022).

Daniel Fried et al. 'InCoder: A Generative Model for Code Infilling and Synthesis'. In: arXiv:2204.05999 [cs] (Apr. 2022). URL: http://arxiv.org/abs/2204.05999 (visited on 05/15/2022).

Leo Gao et al. 'The Pile: An 800GB Dataset of Diverse Text for Language Modeling'. In: arXiv:2101.00027 [cs] (Dec. 2020). URL: http://arxiv.org/abs/2101.00027 (visited on 05/18/2022).

Divyansh Garg et al. 'LISA: Learning Interpretable Skill Abstractions from Language'. In: arXiv:2203.00054 [cs] (Feb. 2022). URL: http://arxiv.org/abs/2203.00054 (visited on 05/14/2022).

Shivam Garg et al. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes . arXiv:2208.01066 [cs]. Aug. 2022. DOI: 10.48550/arXiv.2208.01066 . URL: http://arxiv.org/abs/2208.01066 (visited on 09/28/2022).

Felix Hill et al. 'Human Instruction-Following with Deep Reinforcement Learning via TransferLearning from Text'. In: arXiv:2005.09382 [cs] (May 2020). URL: http://arxiv.org/ abs/2005.09382 (visited on 05/15/2022).

Wenlong Huang et al. Inner Monologue: Embodied Reasoning through Planning with Language Models . arXiv:2207.05608 [cs]. July 2022. DOI: 10.48550/arXiv.2207.05608 . URL: http://arxiv.org/abs/2207.05608 (visited on 08/09/2022).

Wenlong Huang et al. 'Language Models as Zero-Shot Planners: Extracting Actionable Knowledge for Embodied Agents'. In: arXiv:2201.07207 [cs] (Mar. 2022). URL: http: //arxiv.org/abs/2201.07207 (visited on 05/14/2022).

Michael Janner et al. Offline Reinforcement Learning as One Big Sequence Modeling Problem . arXiv:2106.02039 [cs]. Nov. 2021. DOI: 10.48550/arXiv.2106.02039 . URL: http: //arxiv.org/abs/2106.02039 (visited on 08/09/2022).

Rabeeh Karimi Mahabadi et al. 'Prompt-free and Efficient Few-shot Learning with Language Models'. In: Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) . Dublin, Ireland: Association for Computational Linguistics, May 2022, pp. 3638-3652. DOI: 10.18653/v1/2022.acl-long.254 . URL: https://aclanthology.org/2022.acl-long.254 (visited on 08/09/2022).

Michael Laskin et al. 'In-context reinforcement learning with algorithm distillation'. In: arXiv preprint arXiv:2210.14215 (2022).

Kuang-Huei Lee et al. Multi-Game Decision Transformers . arXiv:2205.15241 [cs]. May 2022. DOI: 10.48550/arXiv.2205.15241 . URL: http://arxiv.org/abs/2205.15241 (visited on 08/09/2022).

Shuang Li et al. 'Pre-Trained Language Models for Interactive Decision-Making'. In: arXiv:2202.01771 [cs] (Feb. 2022). URL: http://arxiv.org/abs/2202.01771 (visited on 05/14/2022).

Xiang Lisa Li et al. 'Prefix-Tuning: Optimizing Continuous Prompts for Generation'. en. In: (Jan. 2021). DOI: 10.48550/arXiv.2101.00190 . URL: https://arxiv.org/abs/2101. 00190v1 (visited on 05/15/2022).

Paul Pu Liang et al. 'Towards understanding and mitigating social biases in language models'. In: International Conference on Machine Learning . PMLR. 2021, pp. 6565-6576.

Aman Madaan et al. 'Self-refine: Iterative refinement with self-feedback'. In: arXiv preprint arXiv:2303.17651 (2023).

Arjun Majumdar et al. 'Improving Vision-and-Language Navigation with Image-Text Pairs from the Web'. In: arXiv:2004.14973 [cs] (May 2020). URL: http://arxiv.org/abs/ 2004.14973 (visited on 05/14/2022).

Maxwell Nye et al. Show Your Work: Scratchpads for Intermediate Computation with Language Models . arXiv:2112.00114 [cs]. Nov. 2021. DOI: 10.48550/arXiv.2112.00114 . URL: http://arxiv.org/abs/2112.00114 (visited on 08/11/2022).

Yikang Pan et al. 'On the Risk of Misinformation Pollution with Large Language Models'. In: arXiv preprint arXiv:2305.13661 (2023).

Xiangyu Peng et al. 'Inherently Explainable Reinforcement Learning in Natural Language'. In: arXiv:2112.08907 [cs] (Dec. 2021). URL: http://arxiv.org/abs/2112.08907 (visited on 05/14/2022).

Scott Reed et al. A Generalist Agent . arXiv:2205.06175 [cs]. May 2022. DOI: 10.48550/ arXiv.2205.06175 . URL: http://arxiv.org/abs/2205.06175 (visited on 08/09/2022).

John Schulman et al. 'Proximal policy optimization algorithms'. In: arXiv preprint arXiv:1707.06347 (2017).

Younggyo Seo et al. HARP: Autoregressive Latent Video Prediction with High-Fidelity Image Generator . arXiv:2209.07143 [cs]. Sept. 2022. URL: http://arxiv.org/abs/2209.07143 (visited on 09/28/2022).

Noah Shinn et al. Reflexion: Language Agents with Verbal Reinforcement Learning . 2023. arXiv: 2303.11366 [cs.AI] .

Aaditya K. Singh et al. Know your audience: specializing grounded language models with the game of Dixit . arXiv:2206.08349 [cs]. June 2022. DOI: 10.48550/arXiv.2206.08349 . URL: http://arxiv.org/abs/2206.08349 (visited on 08/09/2022).

Ishika Singh et al. 'Pre-trained Language Models as Prior Knowledge for Playing Text-based Games'. In: arXiv:2107.08408 [cs] (Dec. 2021). URL: http://arxiv.org/abs/2107. 08408 (visited on 05/14/2022).

Allison C. Tam et al. 'Semantic Exploration from Language Abstractions and Pretrained Representations'. In: arXiv:2204.05080 [cs] (Apr. 2022). URL: http://arxiv.org/abs/ 2204.05080 (visited on 05/14/2022).

Alex Tamkin et al. 'Understanding the capabilities, limitations, and societal impact of large language models'. In: arXiv preprint arXiv:2102.02503 (2021).

Denis Tarasov et al. 'Prompts and Pre-Trained Language Models for Offline Reinforcement Learning'. en. In: Mar. 2022. URL: https://openreview.net/forum?id=Spf4TE6NkWq (visited on 05/14/2022).

Ben Wang et al. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model . May 2021. URL: https://github.com/kingoflolz/mesh-transformer-jax .

Jason Wei et al. 'Chain of Thought Prompting Elicits Reasoning in Large Language Models'. In: arXiv:2201.11903 [cs] (Apr. 2022). URL: http://arxiv.org/abs/2201.11903 (visited on 05/15/2022).

Jason Wei et al. Emergent Abilities of Large Language Models . arXiv:2206.07682 [cs]. June 2022. DOI: 10.48550/arXiv.2206.07682 . URL: http://arxiv.org/abs/2206.07682 (visited on 08/11/2022).

Thomas Wolf et al. 'Transformers: State-of-the-Art Natural Language Processing'. In: Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations . Online: Association for Computational Linguistics, Oct. 2020, pp. 38-45. DOI: 10.18653/v1/2020.emnlp-demos.6 . URL: https://aclanthology.org/2020.emnlpdemos.6 (visited on 01/22/2022).

Mengdi Xu et al. Prompting Decision Transformer for Few-Shot Policy Generalization . arXiv:2206.13499 [cs]. June 2022. DOI: 10.48550/arXiv.2206.13499 . URL: http: //arxiv.org/abs/2206.13499 (visited on 08/09/2022).

Susan Zhang et al. OPT: Open Pre-trained Transformer Language Models . _eprint: 2205.01068. 2022.

Appendix

A Example prompts for each domain

B Tabular Q-Learning Hyperparameter Sweeps

Tabular Q learning rate

-0.01-0.05-0.1-0.25-0.5-1-2

Chain

DistractorChain

Maze

MiniCatch

MiniInvaders

0.8

regret

0.6

0.4

MW

0.2

MW

0

200

400

0

6,000

0

6,000

0

6,000

0

6,000

step

step

step

step

step

This figure presents learning curves using the standard Tabular Q-Learning algorithm. Given an appropriately chosen learning rate, this algorithm does converge, but this can take several thousand timesteps on all but the simplest domains.

Tabular Q exploration bonus

-00.01-0.1=1

Chain

DistractorChain

Maze

MiniInvaders

0.8

regret

0.6

0.4

W

0.2

0

200

400

0

6,000

0

6,000

0

6,000

0

6,000

step

step

step

step

step

This figure presents learning curves for Tabular Q-Learning with a count-based exploration bonus. Actions are chosen using the following formulation: a ← arg max a ′ Q ( s, a ′ ) + k/ (1 + N ( s, a ′ )) where k is the exploration bonus hyperparameter, assuming values { 0 , 0 . 01 , 0 . 1 , 1 } in the graph, and N ( s, a ′ ) is the visitation count for the state-action ( s, a ′ ) .

C Proximal Policy Optimization Hyperparameter Sweeps

Chain

Distractor-chain

Maze

Mean return per episode

1.0

Mean return per episode

1.0

Mean return per episode

1.0

0.8

0.8-

0.8

0.6

0.6

0.6-

0.4

0.4-

0.4

0.2

0.2-

0.2-

0.0

00

0.0-

Mini-catch

Point-mass

Space-invaders

Mean return per episode

1.0

Mean return per episode

1.0

Mean return per episode

2.0-

0.8

80

1.5

0.6-

0.6-

1.0

0.4-

0.5

0.2

00

0.0

0.0

0

40,000

80,000

40,000

80,000

0

40,000

80,000

Steps

Steps

Steps

This figure presents learning curves for Proximal Policy Optimization (PPO) (Schulman et al. 2017), a deep RL algorithm trained using gradient descent. These runs search the following hyperparameters:

As the graphs indicate, PPO takes several thousand steps to converge, even on simple domains.

D Random Prediction Ablation

This ablation studies the contribution that each prediction type makes to ICPI's ability to learn. To this end, for prediction type - observation, termination, action, and reward - we ran an ablation that substitutes a random, valid prediction for the prediction made by the LLM. For example, for actions, we first sample a random action from the action space of the environment and then apply the standard string formatting used in other parts of the algorithm. As figure 5 indicates, on Chain, the simplest of our environments, none of the ablations are able to learn.

E Error Analysis

In this section we present examples of errors made by ICPI on each domain. In each subsection, we will examine one rollout generated by the Codex LLM that causes the agent to choose suboptimal behavior in the actual environment.

Chain

Here, we analyze an example in which the rollouts incorrectly estimate equal values for all three actions. Each line in the code snippet represents actual output from the LLM, excluding comments which are added for clarity. The agent's state is 5 (recall that the goal state is 4). For the left() action, the model generates the following rollout:

Note that the model outputs the action right() when the agent is in goal state 4, and then outputs action try_goal() . As a result the rollout estimates value 0 for an action which has value γ under the optimal policy (one timestep away from a terminal state with reward 1).

Distracter Chain

In this example the agent starts in state (2, 3) . Recall that in this environment, the second digit in the tuple is chosen randomly and does not affect environment transitions or reward.

In the first of the three rollouts, the first action is left() . Since the goal is state ( 4, *), this action is already suboptimal. However, we see that the agent continues to behave sub-optimally throughout the rest of the rollout by continuously choosing action left() instead of the optimal action right() . This occurs because at this stage of training, the agent has mostly taken left() actions and as a result, the rollout policy prompt mostly contains left() actions.

Here the initial action is optimal: right() . However, we once again see that the rollout policy is suboptimal because it favors the left() action over the optimal right() action.

This causes the model to inaccurately produce a value estimate of 0 for the right() action.

Maze

Similar to the previous example, we study an example for maze early in training, when the rollout policy is still quite suboptimal. As a result, the rollout policy does not discover the reward, causing the rollouts to estimate equal value for each of the four actions.

The rollouts all start at state ( 1, 1).

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper examines the possibility of taking advantage of the ability of in-context learning to use large language models to implement policy iteration so that large language models can be used to solve reinforcement learning problems without querying the gradients. Experiments have been conducted on 5 artificial environments showing that the proposed method ICPI can indeed achieve decent performance across all considered environments.\n\nSTRENGTHS:\nGiven the success of in-context learning of large language models, applying in-context learning to policy iteration is a well-motivated idea. The proposed ICPI is sufficiently different from the original policy iteration and the other recent attempts to combine the success of large language models with reinforcement learning techniques. In the experiments, 5 artificial environments are proposed and ICPI achieves better results than baselines in all environments. Detailed ablations are also provided to justify the importance of each component of ICPI.\n\nWEAKNESSES:\nThe applications of ICPI are limited. Right now it can only be applied to handcrafted small environments like chain and mini catch. It is unclear if large language models can do complex planning in more practical tasks such as Atari and mujoco tasks. Additionally, the current formulation of ICPI can only handle discrete state and action space. Still, it would be interesting to see if large language models can be used to handle realistic continuous-action robotics tasks. It also seems that a curated hint is required for more complicated tasks like mini invaders and point mass.\n\nComparisons to other more recent methods are missing. To justify the advantage of ICPI, it would be interesting to see how ICPI compares to popular deep RL baselines (either state-based or text-based) such as PPO [1], SAC [2], and CQL [3]/IQL [4] since some recent works [5, 6] have shown that they can be used for large language models too. \n\nSome of the arguments are not supported. For example, in line 8-10, \"All gradient techniques are inherently slow, sacrificing the 'few-shot' quality that makes in-context learning attractive to begin with\" is not supported. ICPI also involves inferencing the entire trajectory to calculate the q function for each action, so it seems unclear whether it is faster than gradient-based techniques. A comparison of ICPI and gradient-based methods in terms of computation time is needed.\n\nThe writing can be improved. The related works and methodology sections can be more concise. There are some references typos and the paragraph from 160 - 164 is repeated.\n\n[1] Proximal Policy Optimization Algorithms, https://arxiv.org/abs/1707.06347\n\n[2] Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, https://arxiv.org/abs/1801.01290\n\n[3] Conservative Q-Learning for Offline Reinforcement Learning, https://arxiv.org/abs/2006.04779\n\n[4] Offline Reinforcement Learning with Implicit Q-Learning, https://arxiv.org/abs/2110.06169\n\n[5] Is Reinforcement Learning (Not) for Natural Language Processing: Benchmarks, Baselines, and Building Blocks for Natural Language Policy Optimization, https://arxiv.org/abs/2210.01241 \n\n[6] Offline RL for Natural Language Generation with Implicit Language Q Learning, https://arxiv.org/abs/2206.11871\n\nQUESTIONS:\nPlease see weakness.\n\nLIMITATIONS:\nThe limitation that current large language models lack the ability to plan in complex domains accurately is discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-22"}

comments: SUMMARY:
This paper formalizes a framework where language model can use its in-context learning capability to predict actions and Q-values, and the training occurs by updating the prompt with the latest history. As a result, the LM can perform reinforcement learning without gradient. The paper explored a few language models and a few different toy tasks and found that Codex can learn successfully in this environment that it has no prior knowledge.

STRENGTHS:
- The paper formalizes a framework where LLM can perform RL using its in-context learning capability. Such a framework might allow future works to connect LLM research to existing RL concepts to further improvement performance.
- The empirical results look promising.

WEAKNESSES:
I'm not an expert in RL, so I am not confident in my judgement.

- While the formalization seems nice, it is not particularly novel; the empirical results/datasets presented also seem toy and low-hanging fruits.

QUESTIONS:
N/A

LIMITATIONS:
The author properly listed their limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-28

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors show how LMs can solve simple MDPs by implementing policy iteration in-context. They do that by constructing prompts that simulate policy rollouts using partial trajectories from an experience buffer as few-shot examples. They empirically validate their approach on six toy MDPs that test for specific failure modes.\n\nSTRENGTHS:\n1. I find the result from the paper title impressive. It’s interesting from at least four different angles: (i) understanding the capabilities of LMs as general-purpose simulators, (ii) sketching the design space of hybrid systems combining RL with LMs, (iii) exploring prompting strategies for representing and rolling-out MDPs in token space, (iv) understanding the trade-offs between in-context learning and standard in-weights learning. The paper is a very original and solid contribution to all four.\n2. The paper is well-written and easy to follow. \n3. The experiments are pretty extensive, covering different MDPs (carefully designed to test specific failure modes), well-motivated ablations and different model families and sizes.\n\nWEAKNESSES:\nThe weaknesses listed below are relatively minor and are greatly outweighed by strengths. The authors shouldn’t feel obliged to respond to them in detail.\n\n1. The authors attribute the relatively good performance of ICPI its “ability to generalize from its context to unseen states and state/action pairs (unlike “Tabular Q” and “Matching Model”)” (lines 296-298). How is that different from function approximation in more traditional RL? It could've been good to also compare ICPI with some baseline using function approximation; both a variant with a small MLP and with a linear head on top of an LLM could be interesting.\n2. The claim in the abstract that “algorithms that rely on [demonstrations] do not outperform the experts from which the demonstrations were derived” is too strong or too vague. It is possible in principle to generalise to rewards higher than those seen in the dataset and the DecisionTransformer paper demonstrated that in practice (in some cases; Fig. 4 in Chen et al.). \n3. The Related work section is very comprehensive but also somewhat repetitive. If the authors’ are short of space, I think it could be shortened without much harm.\n\nQUESTIONS:\n1. What roles does the choice of LM sampling temperature=0.1 play in the setup? Was temperature=1 causing too much hallucination when generating rollouts? Does the theoretical convergence of policy iteration still hold when Q values are estimated from a  biased (due to temperature) distribution of rollouts?\n2. What discount factor $\\gamma$ is used in experiments? Is it 0.8, the same as for the optimal policy relative to which regret is computed?\n\nLIMITATIONS:\nLimitations and societal impacts are discussed. One class of risks that could be discussed is those associated with combining powerful world models with optimisation pressure for agentic behaviour (e.g., Ngo et al., 2023).\n\nNgo et al., 2023. The alignment problem from a deep learning perspective", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper demonstrates a way to implement policy iteration by means of in-context learning with a large language model. Authors suggest ICPI, a reinforcement learning algorithm that performs policy evaluation and policy improvement solely by encoding tokens and inferencing the model, without updating model parameters or adapters. They proceed by evaluating ICPI on a collection of simple RL problems, where each problem was manually converted to an LLM-digestible text representation. Authors show that certain ICPI configurations can achieve comparable or smaller regret than traditional RL algorithms, emphasizing the first few hundred steps. Finally, the paper analyzes the learned behavior produced by ICPI in the tested environments.\n\nSTRENGTHS:\nThe main strength of the paper is that it provides a new way to think about policy iteration from an LLM perspective. While ICPI follows the general PI framework, the way it does so certainly \"feels different\" than traditional (e.g. tabular) PI. Applying similar ideas to other RL areas  may inspire interesting future research.\n\nWhile the evaluation of ICPI is rather minimalistic, these experiments clearly showcase how the algorithm performs and explain the reported results in due detail.\n\nFinally, the paper is generally well-written and easy to follow. The presentation quality is high, both for text and the accompanying visualizations.\n\nWEAKNESSES:\n### 1. Other methods have drawbacks... so what about them?\n\nIn the abstract / intro, authors review existing LLM applications in RL and broadly group them in two categories: those based on expert demosntrations, and those that use gradient-based updates. As a motivation, they state that each of these methods have drawbacks, and the drawback of gradient-based methods is that they are slow (e.g. L9).\n\nTo the best of my knowledge, there is nothing in the paper that evaluates if ICPI is \"faster\" than equivalently tuned gradient-based methods. \nFurthermore, I have a hunch that ICPI scales worse than linearly with the number of demonstrations, i.e. 10000th episode is processed slower than the first one, due to the nature of Transformer attention.\n\nMy concern is that, since authors use speed as one of the motivating factors, this claim should be addressed somewhere in the paper. If ICPI is in fact slower than gradient-based methods, this should also be stated as a limitation. Please note that being slower does not discredit this work in any way -- but not exploring speed after using it as motivation does.\n\n### 2. Comparison to other in-context reinforcement learning methods\n\nThere are at least several prior works that also fall in the category of \"reinforcement learning in-context without updating model parameters\".\nOne prior work that fits this definition is [1], where the LLM-based agent \"reflects\" on it's performance in natural language, which is shown to improve performance. Another, similar work [2] tasks the large language model with providing feedback on model performance in natural language, which also leads to improved behavior. It is likely that there are more prior & concurrent papers that are also relevant.\n\n[1] https://arxiv.org/abs/2303.11366\n[2] https://arxiv.org/abs/2303.17651\n\nOn a surface level, it seems that the above two papers take a less formalistic stance on policy improvement, relying on instruction-following LLMs to do most of their work. However, they have similar requirements and solve similar problems. Both papers use text-based RL environments of similar complexity to the ones used in this submission.\n\nSince both these works were available for several months before submission, I would argue that this paper would benefit from comparing against them: both in terms of theoretical properties and in empirical regret on the same problem set.\n\nQUESTIONS:\n> L185 - since our primary results use codex language model, we use python code to represent these values\n\nMinor question: was it an educated guess or does python outperform some alternative languages when used in this capacity?\n\n> six domains and their associated prompt formats\n\nMinor question: how robust is ICPI performance to the choice of prompts?\n\n\n> L21: improve … as information accumulated\n\nPossibly missing a verb, e.g. as information is accumulated / accumulates\n\n\n\n> Fig 4 caption: Error bars are standard errors from four seeds\n\nAssuming you use standard deviation, you may want to specify whether or not this is an adjusted estimate (i.e. whether or not you use Bessel’s correction). If you deliberately opt not to specify this, please feel free to ignore the question.\n\nLIMITATIONS:\nAuthors properly address most of the main limitations of this work, as well as it's societal limitations.\n\nOne potentially missing limitation is the reproducibility of the main experiments. As it stands, ICPI relies on proprietary language models that can change without warning. To the best of my understanding, the current OpenAI policy allows them to alter the model behavior without warning users in advance, e.g. to patch a potential misuse. This could inadvertently affect the performance of ICPI, making it difficult to reproduce the main results in several years. While there is no guaranteed way to fix this issue, I would argue that it is worth acknowledging it.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe problem of using large-pretrained foundation models (LLMs) like codex as backbones for learning RL policies without gradient updates or expert demonstrations is addressed by the authors. They propose an algorithm called In-Context Policy Iteration (ICPI) to perform model-based RL policy iterations using a pretrained LLM to predict state/rewards as well as the next-action. \n\nTo make the environment prompt-ready, the task is translated into a text-based description. The authors focus on code-pre-trained backbones and consider translating to Python code in their experiments. They demonstrate that ICPI can learn to minimize regret over time in simple tabular environments such as maze and point-mass.\n\nSTRENGTHS:\nSome strengths of the work are outlined below:\n1. The proposed algorithm utilizes In-Context learning to implement policy-iteration with LLMs. This parameter-free approach has the potential to enhance the task-specific behavior of an LLM agent by designing a prompt-ready text-based interface for any MDP.\n2. The authors present a clear explanation of the algorithm for ICPI, including the necessary design choices. They also provide empirical evidence of policy improvement over time in six MDPs that are both simple and tabular.\n\nWEAKNESSES:\n1. Based on the results, it appears that the algorithm works efficiently when combined with code-pre-training. However, relying solely on text-based MDP descriptions may not yield consistent outcomes. It is uncertain whether this is due to the limitations in reasoning and in-context learning capabilities of current LLMs, or the difficulty in creating effective intermediate representations of the environment.\n2. To perform Q-value computation, the algorithm needs to unroll action sequences for all actions. Unfortunately, this greatly restricts the algorithm's usefulness to finite action-spaces.\n\nQUESTIONS:\n1. The success of policy-iteration using replay-buffer enhanced prompting is likely influenced by the information content of the demonstration set $\\mathcal{D}$. It's worth noting that these trajectories may not be on-policy and are sourced from previous policy versions. It's important to investigate how the algorithm responds to the quality of the replay-buffer, including factors like the number of trajectories and whether failed examples are necessary.\n\nLIMITATIONS:\nThe authors have acknowledged certain limitations of their work and discussed them thoroughly. One such limitation is that the ICPI method requires exhaustive unrolling of the action sequence for all actions, making it suitable only for domains with finite action space. Furthermore, it appears that only code-pre-trained models are capable of reliably performing ICPI, indicating that translating the MDP to an appropriate code format is crucial. Additionally, the method often relies on additional hints to ground the reward function/terminal states, hence a more comprehensive discussion on potential ways to mitigate this would be beneficial.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}