Original: I N -C ONTEXT P OLICY I TERATION
Revised: I N -C ONTEXT P OLICY I TERATION

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: This work presents In-Context Policy Iteration, an algorithm for performing Reinforcement Learning (RL), in-context, using foundation models.
Revised: This work presents In-Context Policy Iteration, an algorithm for performing Reinforcement Learning (RL), in-context, using foundation models.

Original: 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).
Revised: 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).

Original: Both of these techniques have drawbacks.
Revised: Both of these techniques have drawbacks.

Original: Collecting demonstrations is labor-intensive, and algorithms that rely on them do not outperform the experts from which the demonstrations were derived.
Revised: Collecting demonstrations is labor-intensive, and algorithms that rely on them do not outperform the experts from which the demonstrations were derived.

Original: All gradient techniques are inherently slow, sacrificing the “few-shot” quality that made in-context learning attractive to begin with.
Revised: All gradient techniques are inherently slow, sacrificing the “few-shot” quality that made in-context learning attractive to begin with.

Original: In this work, we present an algorithm, ICPI, that learns to perform RL tasks without expert demonstrations or gradients.
Revised: In this work, we present an algorithm, ICPI, that learns to perform RL tasks without expert demonstrations or gradients.

Original: Instead we present a policy-iteration method in which the prompt content is the entire locus of learning.
Revised: Instead we present a policy-iteration method in which the prompt content is the entire locus of learning.

Original: ICPI iteratively updates the contents of the prompt from which it derives its policy through trial-and-error interaction with an RL environment.
Revised: ICPI iteratively updates the contents of the prompt from which it derives its policy through trial-and-error interaction with an RL environment.

Original: In order to eliminate the role of in-weights learning (on which approaches like Decision Transformer rely heavily), we demonstrate our algorithm using Codex Chen et al.
Revised: In order to eliminate the role of in-weights learning (on which approaches like Decision Transformer rely heavily), we demonstrate our algorithm using Codex Chen et al.

Original: (2021b), a language model with no prior knowledge of the domains on which we evaluate it.
Revised: (2021b), a language model with no prior knowledge of the domains on which we evaluate it.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In-context learning describes the ability of sequence-prediction models to generalize to novel downstream tasks when prompted with a small number of exemplars (Lu et al., 2021; Brown et al., 2020).
Revised: In-context learning describes the ability of sequence-prediction models to generalize to novel downstream tasks when prompted with a small number of exemplars (Lu et al., 2021; Brown et al., 2020).

Original: The introduction of the Transformer architecture (Vaswani et al., 2017) has significantly increased interest in this phenomenon, since this architecture demonstrates much stronger generalization capacity than its predecessors (Chan et al., 2022).
Revised: The introduction of the Transformer architecture (Vaswani et al., 2017) has significantly increased interest in this phenomenon, since this architecture demonstrates much stronger generalization capacity than its predecessors (Chan et al., 2022).

Original: Another 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
Revised: Another 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

Original: downstream task-distributions
Revised: downstream task-distributions

Original: (Brown et al., 2020; Wei et al., 2022a).
Revised: (Brown et al., 2020; Wei et al., 2022a).

Original: A litany of existing work has explored methods for applying this remarkable capability to downstream tasks (see Related Work), including Reinforcement Learning (RL).
Revised: A litany of existing work has explored methods for applying this remarkable capability to downstream tasks (see Related Work), including Reinforcement Learning (RL).

Original: Most work in this area either (1) assumes access to expert demonstrations — collected either from human experts (Huang et al., 2022a;Ahn et al., 2022;
Revised: Most work in this area either (1) assumes access to expert demonstrations — collected either from human experts (Huang et al., 2022a;

Original: Huang et al., 2022b; Baker et al., 2022), or domain-specific pre-trained RL agents (Chen et al., 2021a; Lee et al., 2022; Janner et al., 2021;
Revised: Baker et al., 2022), or domain-specific pre-trained RL agents (Chen et al., 2021a; Lee et al., 2022; Janner et al., 2021;

Original: Reed et al., 2022; Xu et al., 2022).
Revised: Reed et al., 2022; Xu et al., 2022).

Original: — or (2) relies on gradient-based methods — e.g. fine-tuning of the foundation models parameters as a whole (Lee et al., 2022;
Revised: — or (2) relies on gradient-based methods — e.g. fine-tuning of the foundation models parameters as a whole (Lee et al., 2022;

Original: Reed et al., 2022; Baker et al., 2022) or newly training an adapter layer or prefix vectors while keeping the original foundation models frozen (Li & Liang, 2021; Singh et al., 2022; Karimi Mahabadi et al., 2022).
Revised: Reed et al., 2022; Baker et al., 2022) or newly training an adapter layer or prefix vectors while keeping the original foundation models frozen (Li & Liang, 2021; Singh et al., 2022; Karimi Mahabadi et al., 2022).

Original: Our work presents an algorithm, In-Context Policy Iteration (ICPI) which relaxes these assumptions.
Revised: Our work presents an algorithm, In-Context Policy Iteration (ICPI) which relaxes these assumptions.

Original: ICPI is a form of policy iteration in which the prompt content is the locus of learning.
Revised: ICPI is a form of policy iteration in which the prompt content is the locus of learning.

Original: Because our method operates on the prompt itself rather than the model parameters, we are able to avoid gradient methods.
Revised: Because our method operates on the prompt itself rather than the model parameters, we are able to avoid gradient methods.

Original: Furthermore, the use of policy iteration frees us from expert demonstrations because suboptimal prompts can be improved over the course of training.
Revised: Furthermore, the use of policy iteration frees us from expert demonstrations because suboptimal prompts can be improved over the course of training.

Original: We illustrate the algorithm empirically on six small illustrative RL tasks— chain, distractor-chain, maze, mini-catch, mini-invaders , and point-mass —in which the algorithm very quickly finds good policies.
Revised: We illustrate the algorithm empirically on six small illustrative RL tasks— chain, distractor-chain, maze, mini-catch, mini-invaders , and point-mass —in which the algorithm very quickly finds good s t

Original: We also compare five pretrained Large Language Models (LLMs), including two different s t

Original: LLM as
Revised: LLM as

Original: A (1) =
Revised: A (1) =

Original: s t =
Revised: s t =

Original: For each action in {A (1) ,..., A ( n ) } :
Revised: For each action in {A (1) ,..., A ( n ) } :

Original: D sr
Revised: D sr

Original: LLM
Revised: LLM

Original: D a
Revised: D a

Original: LLM
Revised: LLM

Original: D sr...
Revised: D sr...

Original: Q ( s t , A (1)) =
Revised: Q ( s t , A (1)) =

Original: P γ u r u ...
Revised: P γ u r u ...

Original: LLM sa 1 A ( n ) = s t =
Revised: LLM sa 1 A ( n ) = s t =

Original: D sr
Revised: D sr

Original: LLM
Revised: LLM

Original: D a
Revised: D a

Original: LLM
Revised: LLM

Original: D sr...
Revised: D sr...

Original: Q ( s t , A ( n )) =
Revised: Q ( s t , A ( n )) =

Original: P γ u r u argmax a a t
Revised: P γ u r u argmax a a t

Original: Environment  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.
Revised: Environment 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.

Original: On our six domains, we find that only the largest model (the code-davinci-001 variant of Codex) consistently demonstrates learning.
Revised: On our six domains, we find that only the largest model (the code-davinci-001 variant of Codex) consistently demonstrates learning.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: A common application of foundation models to RL involves tasks that havesome kind of language input, for example natual language instructions/goals (Garg et al., 2022a; Hill et al., 2020) or textbased games (Peng et al., 2021; Singh et al., 2021; Majumdar et al., 2020; Ammanabrolu & Riedl, 2021).
Revised: A common application of foundation models to RL involves tasks that have language input, for example natural language instructions/goals (Garg et al., 2022a; Hill et al., 2020) or text-based games (Peng et al., 2021; Singh et al., 2021; Majumdar et al., 2020; Ammanabrolu & Riedl, 2021).

Original: Anothercommon approach encodes RL trajectories into token sequences, and processes the sequences with a foundation models—using the foundation model’s representations as input to deep 
Revised: Another approach encodes RL trajectories into token sequences, and processes them with a foundation model, model representations as input to deep RL architectures (Li et al., 2022;

Original: RL architectures (Li et al., 2022; Tarasov et al., 2022; Tam et al., 2022).
Revised: Tarasov et al., 2022; Tam et al., 2022).

Original: Finally, a recent set of approaches treat RL as a sequence modeling problem and use the foundation models itself to predict states or actions.
Revised: Finally, a recent set of approaches treat RL as a sequence modeling problem and use the foundation models itself to predict states or actions.

Original: This section will focus on this last category.
Revised: This section will focus on this last category.

Original: 2. L EARNING FROM DEMONSTRATIONS
Revised: 2. L EARNING FROM DEMONSTRATIONS

Original: Many recent sequence-based approaches to reinforcement learning use demonstrations that come either from human experts or pretrained RL agents.
Revised: Many recent sequence-based approaches to reinforcement learning use demonstrations that come either from human experts or pretrained RL agents.

Original: For example, Huang et al.
Revised: For example, Huang et al.

Original: (2022a) use a frozen
Revised: (2022a) use a frozen

Original: 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.
Revised: 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.

Original: This work is extended by Huang et al.
Revised: This work is extended by Huang et al.

Original: (2022b).
Revised: (2022b).

Original: Similarly, Ahn et al.
Revised: Similarly, Ahn et al.

Original: (2022) use a value function that is trained on human demonstrations to rank candidate actions produced by an LLM.
Revised: (2022) use a value function that is trained on human demonstrations to rank candidate actions produced by an LLM.

Original: Baker et al.
Revised: Baker et al.

Original: (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.
Revised: (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.

Original: Works that rely on pretrained RL agents include Janner et al.
Revised: Works that rely on pretrained RL agents include Janner et al.

Original: (2021) who train a “Trajectory Transformer” to predict trajectory sequences in continuous control tasks by using trajectories generated by pretrained agents, and Chen et al.
Revised: (2021) who train a “Trajectory Transformer” to predict trajectory sequences in continuous control tasks by using trajectories generated by pretrained agents, and Chen et al.

Original: (2021a), 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.
Revised: (2021a), 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.

Original: Two approaches build on this method to improve generalization: Lee et al.
Revised: Two approaches build on this method to improve generalization: Lee et al.

Original: (2022) use trajectories generated by a DQN agent to train a single
Revised: (2022) use trajectories generated by a DQN agent to train a single

Original: Decision Transformer that can play many Atari games, and Xu et al.
Revised: Decision Transformer that can play many Atari games, and Xu et al.

Original: (2022) use a combination of human and artificial trajectories to train a Decision Transformer that achieves few-shot generalization on continuous control tasks.
Revised: (2022) use a combination of human and artificial trajectories to train a Decision Transformer that achieves few-shot generalization on continuous control tasks.

Original: Reed et al.
Revised: Reed et al.

Original: (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.
Revised: (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.

Original: Atari, continuous control) and non-RL (e.g. image captioning, chat) tasks.
Revised: Atari, continuous control) and non-RL (e.g. image captioning, chat) tasks.

Original: Some of the above works include non-expert demonstrations as well.
Revised: Some of the above works include non-expert demonstrations as well.

Original: Chen et al.
Revised: Chen et al.

Original: (2021a) include experiments with trajectories generated by random (as opposed to expert) policies.
Revised: (2021a) include experiments with trajectories generated by random (as opposed to expert) policies.

Original: Lee et al.
Revised: Lee et al.

Original: (2022) and Xu et al.
Revised: (2022) and Xu et al.

Original: (2022) also use datasets that include trajectories generated by partially trained agents in addition to fully trained agents.
Revised: (2022) also use datasets that include trajectories generated by partially trained agents in addition to fully trained agents.

Original: 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.
Revised: 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.

Original: Firstly, ICPI only consumes self-generated trajectories, so it does not require any demonstrations (like
Revised: Firstly, ICPI only consumes self-generated trajectories, so it does not require any demonstrations (like

Original: Chen et al.
Revised: Chen et al.

Original: (2021a) with random trajectories, but unlike Lee et al.
Revised: (2021a) with random trajectories, but unlike Lee et al.

Original: (2022), Xu et al.
Revised: (2022), Xu et al.

Original: and the other approaches reviewed above).
Revised: and the other approaches reviewed above).

Original: Secondly, ICPI relies primarily on in-context learning rather than in-weights learning to achieve generalization (like Xu et al.
Revised: Secondly, ICPI relies primarily on in-context learning rather than in-weights learning to achieve generalization (like Xu et al.

Original: , but unlike Chen et al.
Revised: , but unlike Chen et al.

Original: (2021a)
Revised: (2021a)

Original: Lee et al.
Revised: Lee et al.

Original: For discussion about in-weights vs. in-context learning see Chan et al.
Revised: For discussion about in-weights vs. in-context learning see Chan et al.

Original: 2. G RADIENT - BASED T RAINING & F INETUNING ON RL T ASKS
Revised: 2. G RADIENT - BASED T RAINING & F INETUNING ON RL T ASKS

Original: Most approaches involve training or fine-tuning foundation models on RL tasks.
Revised: Most approaches involve training or fine-tuning foundation models on RL tasks.

Original: For example, Janner et al.
Revised: For example, Janner et al.

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2021a)
Revised: (2021a)

Original: ; Lee et al.
Revised: ; Lee et al.

Original: ; Xu et al. (2022); Baker et al.
Revised: ; Xu et al. (2022); Baker et al.

Original: ; Reed et al.
Revised: ; Reed et al.

Original: (2022) all use models that are trained from scratch on tasks of interest, and Singh et al.
Revised: (2022) all use models that are trained from scratch on tasks of interest, and Singh et al.

Original: ; Ahn et al.
Revised: ; Ahn et al.

Original: ; Huang et al.
Revised: ; Huang et al.

Original: (2022b) combine frozen foundation models with trainable components or adapters.
Revised: (2022b) combine frozen foundation models with trainable components or adapters.

Original: In contrast, Huang et al.
Revised: In contrast, Huang et al.

Original: (2022a) use frozen foundation models for planning, without training or fine-tuning on RL tasks.
Revised: (2022a) use frozen foundation models for planning, without training or fine-tuning on RL tasks.

Original: Like Huang et al.
Revised: Like Huang et al.

Original: (2022a), ICPI does not update the parameters of the foundation model, but relies on the frozen model’s in-context learning abilities.
Revised: (2022a), ICPI does not update the parameters of the foundation model, but relies on the frozen model’s in-context learning abilities.

Original: 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.
Revised: 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.

Original: (2022a), which uses the frozen model to select good prompts from a given fixed library of goal/plan descriptions).
Revised: (2022a), which uses the frozen model to select good prompts from a given fixed library of goal/plan descriptions).

Original: 2. I N -C ONTEXT L EARNING
Revised: 2. I N -C ONTEXT L EARNING

Original: Several recent papers have specifically studied in-context learning.
Revised: Several recent papers have specifically studied in-context learning.

Original: Min et al.
Revised: Min et al.

Original: (2022) demonstrates that LLMs can learn in-context, even when the labels in the prompt are randomized, problemetizing the conventional understanding of in-context learning and showing that label distribution is more important than label correctness.
Revised: (2022) demonstrates that LLMs can learn in-context, even when the labels in the prompt are randomized, problemetizing the conventional understanding of in-context learning and showing that label distribution is more important than label correctness.

Original: Chan et al.
Revised: Chan et al.

Original: and
Revised: and

Original: Garg et al.
Revised: Garg et al.

Original: (2022b) 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.
Revised: (2022b) 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.

Original: These papers identify various properties, including “burstiness,” model-size, and model-architecture, that in-context learning depends on.
Revised: These papers identify various properties, including “burstiness,” model-size, and model-architecture, that in-context learning depends on.

Original: Chen et al.
Revised: Chen et al.

Original: (2022) studies the sensitivity of in-context learning to small perturbations of the context.
Revised: (2022) studies the sensitivity of in-context learning to small perturbations of the context.

Original: They propose a novel method that uses sensitivity as a proxy for model certainty.
Revised: They propose a novel method that uses sensitivity as a proxy for model certainty.

Original:  A LGORITHM
Revised:  A LGORITHM

Original: How can standard policy iteration make use of in-context learning?
Revised: How can standard policy iteration make use of in-context learning?

Original: In general, policy iteration is either model-based —using a world-model to plan future trajectories in the environment—or modelfree —inferring value-estimates withoutreliance on explicit planning.
Revised: 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.

Original: Both methods can be realized with in-context learning.
Revised: Both methods can be realized with in-context learning.

Original: We choose model-based learning because planned trajectories make the Algorithm 1 Training Loop 1: function T RAIN (environment) 2: initialize D ▷ replay buffer containing prompt "parameters" 3: while training do 4: s 0 ← Reset environment.
Revised: We choose model-based learning because planned trajectories make the underlying logic of valueAlgorithm 1 Training Loop 1: function T RAIN (environment) 2: initialize D ▷ replay buffer containing full history of behavior 3: while training do 4: s 0 ← Reset environment.

Original: 5: while episode is not done do 6: a t ← arg max a Q ( s t , a, D )
Revised: 5: while episode is not done do 6: a t ← arg max a Q ( s t , a, D )

Original: ▷ policy improvement 7: s t +1 , r t , b t ← Execute a t in environment.
Revised: ▷ policy improvement 7: s t +1 , r t , b t ← Execute a t in environment.

Original: 8: t ← t + 1 9: end while 10:
Revised: 8: t ← t + 1 9: end while 10:

Original: D ← D ∪ ( s 0 , a 0 , r 0 , b 0 , s 1 , . . .
Revised: D ← D ∪ ( s 0 , a 0 , r 0 , b 0 , s 1 , . . .

Original: , s t , a t , r t , b t , s t +1 )
Revised: , s t , a t , r t , b t , s t +1 )

Original: ▷ add trajectory to buffer 11: end while 12: end function
Revised: ▷ add trajectory to buffer 11: end while 12: end function

Original: 1: function Q ( s t , a , D ) 2: u ← t 3: s 

Revised: 1: function Q ( s t , a , D ) 2: u ← t 3: s

Original: = s t 4: a 1 = a [-
-]
Revised: = s t 4: a 1 = a [-
-]5: repeat

Original: : repeat ▷ All samples come from the experience buffer D 6:

Original: D b ∼ time-steps with action a u ▷ balancing terminal and non-terminal 7: b u ∼ LLM ( D b , s u , a u ) 8:
Revised: D b ∼ time-steps with action a u ▷ balancing terminal and non-terminal 7: b u ∼ LLM ( D b , s u , a u ) 8:

Original: D r ∼ time-steps with action a u and termination b u ▷ balancing reward 9: r u ∼ LLM ( D r , s u , a u ) 10:
Revised: D r ∼ time-steps with action a u and termination b u ▷ balancing reward 9: r u ∼ LLM ( D r , s u , a u ) 10:

Original: D s ∼ time-steps with action a u and termination b
Revised: D s ∼ time-steps with action a u and termination b

Original: u ▷ no balancing 11: s u +1 ∼ LLM
Revised: u ▷ no balancing 11: s u +1 ∼ LLM

Original: ( D s , s u , a u ) 12:
Revised: ( D s , s u , a u ) 12:

Original: D a ∼ c recent trajectories 13: a u +1 ∼ LLM (cid:0) s u +1 , D a (cid:1) 14: u ← u + 1 15: until b u is terminal 16: return (cid:80) uk =1 γ k − 1 r k 17: end function underlying logic of value-estimates explicit to our foundation model backbone, providing a concrete instantiation of how the values might be realized.
Revised: D a ∼ c recent trajectories 13: a u +1 ∼ LLM (cid:0) s u +1 , D a (cid:1) 14: u ← u + 1 15: until b u is terminal 16: return (cid:80) uk =1 γ k − 1 r k 17: end function estimates explicit to our foundation model backbone, providing a concrete instantiation of a trajectory that realizes the values.

Original: This ties into  work such as Wei et al. (2022b) and Nye et al. (2021) which demonstrates thatexplicit “chains of thought” can significantly improve few-shot performance of foundation models.
Revised: This ties into recent work (Wei et al., 2022b; Nye et al., 2021) demonstrating that “chains of thought” can significantly improve few-shot performance of foundation models.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Similarly, we induce the world-model by prompting the foundation models with transitions drawn from the agent’s history of experience.
Revised: Similarly, we induce the world-model by prompting the foundation models with transitions drawn from the agent’s history of experience.

Original: 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.
Revised: 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.

Original: This explains how an algorithm might plan and estimate values using a foundation model.
Revised: This explains how an algorithm might plan and estimate values using a foundation model.

Original: It also explains how the rollout-policy approximately tracks the behavior policy.
Revised: It also explains how the rollout-policy approximately tracks the behavior policy.

Original: How does  policy improvement occur?
Revised: How does the policy improve?

Original: When acting in the environment (as opposed to planning), we use the standard Q-learning technique of choosing the action that maximizes the estimated Q-value from the current state (see Training Loop pseudocode, line 6).
Revised: 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).

Original: 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 ( s t , a ) ,where
Revised: 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 ( s t , a ) ,

Original:  A =
Revised: where A =

Original: {A (1) , · · · , A ( n ) } denotes the set of n actions available, π t denotes the policy of the agent at time step t , and Q π denotes the Q-estimate for policy π .
Revised: {A (1) , · · · , A ( n ) } denotes the set of n actions available, π t denotes the policy of the agent at time step t , and Q π denotes the Q-estimate for policy π .

Original: Taking the greedy ( arg max ) actions with respect to Q π t implements a new and improved policy.
Revised: Taking the greedy ( arg max ) actions with respect to Q π t implements a new and improved policy.

Original: 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).
Revised: 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).

Original: During training, We maintain a buffer D of complete episodes experienced by the agent.
Revised: During training, We maintain a buffer D of complete episodes experienced by the agent.

Original: To compute Q π t ( s t , a ) at time step t in the real-world we will rollout a simulated trajectory s
Revised: To compute Q π t ( s t , a ) at time step t in the real-world we will rollout a simulated trajectory s

Original: 1 = s t , a 1 = a , r 1 , s 2 , a 2 , r 2 , · · · , s T , a T , r T , s T + by predicting, at each simulation time step u : Termination: [Equation]
Revised: 1 = s t , a 1 = a , r 1 , s 2 , a 2 , r 2 , · · · , s T , a T , r T , s T + by predicting, at each simulation time step u : reward — r u ∼ LLM ( D r , s u , a u ) ; termination — b u ∼ LLM ( D b , s u , a u ) ; observation — s u +1 ∼ LLM

Original: [Equation] Next State: [Equation] [Equation]
Revised: ( D s , s u , a u ) ; action — a u +1 ∼ LLM ( D a , s u ) .

Original: where termination b u decides whether the simulated trajectory ends at step u .
Revised: Termination b u decides whether the simulated trajectory ends at step u .

Original: The contents of the prompt buffers D b , D r , and D s are chosen so as to maximize the relevance of the prompt contents to the current inference.
Revised: The contents of the prompt buffers D b , D r , and D s are chosen so as to maximize the relevance of the prompt contents to the current inference.

Original: D b contains ( s k , a k , b k ) tuples sampled randomly from the D , constraining the action a k tobe equalto a u , the action for which the LLM must infer termination.
Revised: D b contains ( s k , a k , b k ) tuples sampled randomly from the D , constraining the action a k to equal a u , the action for which the LLM must infer termination.

Original: D r contains ( s k , a k , r k ) tuples, again constraining a k = a u but also constraining b k =
Revised: D r contains ( s k , a k , r k ) tuples, again constraining a k = a u but also constraining b k =

Original: 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.
Revised: 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.

Original: We do not need to make next-state predictions for terminal time-steps (when the LLM has predicted b u is true).

Original: For non-terminal time-steps, the prompt includes ( s k , a k s k +1 ) tuples with a k = a u and b k = false.
Revised: For next-state predictions, the prompt includes ( s k , a k s k +1 ) tuples with a k = a u and b k = false.

Original: We also found that there was some benefit to ensuring that each modelling prompt contains a balance of certain kinds of time-steps.
Revised: We also found that there was some benefit to ensuring that each modelling prompt contains a balance of certain kinds of time-steps.

Original: For termination prediction, we balance terminal and non-terminal timesteps.
Revised: For termination prediction, we balance terminal and non-terminal timesteps.

Original: 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, regardless of ( s u , a u ) .
Revised: 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.

Original: Similarly, for reward prediction, we balance the number of time-steps corresponding to each reward value stored in D .
Revised: Similarly, for reward prediction, we balance the number of time-steps corresponding to each reward value stored in D .

Original: In order to balance two collections of unequal size, we take the smaller and duplicate a random subset of its members until the sizes are equal.
Revised: In order to balance two collections of unequal size, we take the smaller and duplicate randomly chosen members until the sizes are equal.

Original: In contrast to the other predictions, we condition the rollout policy on trajectory subsequences, not individual time-steps.
Revised: In contrast to the other predictions, we condition the rollout policy on trajectory subsequences, not individual time-steps.

Original: Prompting with sequences better enables the foundation model to apprehend the logic behind a policy.
Revised: Prompting with sequences better enables the foundation model to apprehend the logic behind a policy.

Original: Trajectory subsequences consist of ( s k , a k ) pairs, randomly clipped from the c most recent trajectories.
Revised: Trajectory subsequences consist of ( s k , a k ) pairs, randomly clipped from the c most recent trajectories.

Original: More recent trajectories will, in general demonstrate higher performance, since they come from policies that have benefited from more rounds of improvement.
Revised: More recent trajectories will, in general demonstrate higher performance, since they come from policies that have benefited from more rounds of improvement.

Original: Finally, the Q-value estimate is simply the discounted sum of rewards for the simulated episode.
Revised: Finally, the Q-value estimate is simply the discounted sum of rewards for the simulated episode.

Original: Given this description of Q-value estimation, we now return to the concept of policy improvement. 
Revised: Given this description of Q-value estimation, we now return to the concept of policy improvement. [Figure] [Table]

Original:  Policy-Improvement The arg max (line 6 of Algorithm 1) drives policy improvement in ICPI.
Revised: i [Table] Policy-Improvement The arg max (line 6 of Algorithm 1) drives policy improvement in ICPI.

Original: Critically this is not simply a one-step improvement as with Trajectory Transformer (Janner et al., 2021) but a mechanism that builds improvement on top of improvement.
Revised: Critically this is not simply a one-step improvement as with Trajectory Transformer (Janner et al., 2021) but a mechanism that builds improvement on top of improvement.

Original: This occurs through a cycle in which the arg max improves behavior.
Revised: This occurs through a cycle in which the arg max improves behavior.

Original: The improved behavior is stored in the buffer D , and then used to condition the rollout policy.
Revised: The improved behavior is stored in the buffer D , and then used to condition the rollout policy.

Original: This improves the returns generated by the LLM during planning rollouts.
Revised: This improves the returns generated by the LLM during planning rollouts.

Original: These improved rollouts improve the Q-estimates for each action.
Revised: These improved rollouts improve the Q-estimates for each action.

Original: Completing the cycle, this improves the actions chosen by the arg max .
Revised: Completing the cycle, this improves the actions chosen by the arg max .

Original: Because this process feeds into itself, it can drive improvement without bound until optimality is achieved.
Revised: Because this process feeds into itself, it can drive improvement without bound until optimality is achieved.

Original: Note that this process takes advantage of properties specific to in-context learning.
Revised: Note that this process takes advantage of properties specific to in-context learning.

Original: 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.
Revised: 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.

Original: Given this assumption, the rollout policy will improve with the improvement of the mixture of policies from which its prompt-trajectories are drawn.
Revised: Given this assumption, the rollout policy will improve with the improvement of the mixture of policies from which its prompt-trajectories are drawn.

Original: This results in a kind of rapid policy improvement that works without any use of gradients.
Revised: This results in a kind of rapid policy improvement that works without any use of gradients.

Original: 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).
Revised: 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).

Original: Since our primary results use the Codex language model (see Table 2), we use Python code to represent these values (examples are available in Table 1).
Revised: Since our primary results use the Codex language model (see Table 2), we use Python code to represent these values (examples are available in Table 1).

Original: 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.
Revised: In our experiments, we discovered that the LLM world-model was unable to reliably predict rewards, terminations, and nextstates on some of the more difficult environments.

Original: 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., 2022b).
Revised: 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., 2022b).

Original: For example, for the chain domain, the hint includes an explicit comparison ( == or !
Revised: For example, for the chain domain, the hint includes an explicit comparison ( == or !

Original: = )
Revised: = )

Original: of the current state with the goal state.
Revised: of the current state with the goal state.

Original: Note that while hints are provided hints in the initial context, the LLM must infer the hint content in rollouts generated from this context.
Revised: Note that while hints are provided hints in the initial context, the LLM must infer the hint content in rollouts generated from this context.

Original: We use a consistent idiom for rewards and terminations, namely assert reward =
Revised: We use a consistent idiom for rewards and terminations, namely assert reward =

Original: = x and assert done or assert not done .
Revised: = x and assert done or assert not done .

Original: Some decisions had to be made when representing states and actions.
Revised: Some decisions had to be made when representing states and actions.

Original: In general, we strove to use simple, idiomatic, concise Python.
Revised: In general, we strove to use simple, idiomatic, concise Python.

Original: On the more challenging environments, we did search over several options for the choice of hint.
Revised: On the more challenging environments, we did search over several options for the choice of hint.

Original: We anticipate that in the future, stronger foundation models will be increasingly robust to these decisions.
Revised: We anticipate that in the future, stronger foundation models will be increasingly robust to these decisions.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: 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 2)—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).
Revised: 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 2)—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).

Original: We next describe the six domains and their associated prompt formats, and then describe the experimental methodology and results.
Revised: We next describe the six domains and their associated prompt formats, and then describe the experimental methodology and results.

Original: 4. D OMAINS AND PROMPT FORMAT
Revised: 4. D OMAINS AND PROMPT FORMAT

Original: Chain.
Revised: Chain.

Original: In this environment, the agent occupies an 8-state chain.
Revised: In this environment, the agent occupies an 8-state chain.

Original: The agent has three actions: Left , right , and try goal .
Revised: The agent has three actions: Left , right , and try goal .

Original: 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.
Revised: 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.

Original: Episodes also terminate after 8 time-steps.
Revised: Episodes also terminate after 8 time-steps.

Original: States are represented as numbers from 0 to 7, as in assert state =
Revised: States are represented as numbers from 0 to 7, as in assert state =

Original: = n , with the appropriate integer substituted for n .
Revised: = n , with the appropriate integer substituted for n .

Original: The actions are represented as functions left() , right() , and try_goal() .
Revised: The actions are represented as functions left() , right() , and try_goal() .

Original: For the hint, we simply indicate whether or not the current state matches the goal state, 4.
Revised: For the hint, we simply indicate whether or not the current state matches the goal state, 4.

Original: Distractor Chain.
Revised: Distractor Chain.

Original: 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 , . . .
Revised: 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 , . . .

Original: The agent must therefore learn to ignore the distractor integer and base its inferrences on the information contained in the first integer.
Revised: The agent must therefore learn to ignore the distractor integer and base its inferrences on the information contained in the first integer.

Original: Aside from the addition of this distractor integer to the observation, all text representations and hints are identical to the chain environment.
Revised: Aside from the addition of this distractor integer to the observation, all text representations and hints are identical to the chain environment.

Original: Maze.
Revised: Maze.

Original: The agent navigates a small 3 × 3 gridworld with obstacles.
Revised: The agent navigates a small 3 × 3 gridworld with obstacles.

Original: The agent can move up , down , left , or right .
Revised: The agent can move up , down , left , or right .

Original: 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.
Revised: 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.

Original: This environment tests our algorithms capacity to handle 2-dimensional movement and obstacles, as well as a 4-action state-space.
Revised: This environment tests our algorithms capacity to handle 2-dimensional movement and obstacles, as well as a 4-action state-space.

Original: We represent the states as namedtuples — C(x, y) , with integers substituted for x and y .
Revised: We represent the states as namedtuples — C(x, y) , with integers substituted for x and y .

Original: Similar to chain , the hint indicates whether or not the state corresponds to the goal state.
Revised: Similar to chain , the hint indicates whether or not the state corresponds to the goal state.

Original: Mini Catch.
Revised: Mini Catch.

Original: The agent operates a paddle to catch a falling ball.
Revised: The agent operates a paddle to catch a falling ball.

Original: The ball falls from a height ofunits, descending one unit per time step.
Revised: The ball falls from a height ofunits, descending one unit per time step.

Original: The paddle can stay in place (not move), or move left or right along the bottom of the 4-unit wide plane.
Revised: The paddle can stay in place (not move), or move left or right along the bottom of the 4-unit wide plane.

Original: The agent receives a reward of 1 for catching the ball and 0 for other time-steps.
Revised: The agent receives a reward of 1 for catching the ball and 0 for other time-steps.

Original: The episode ends when the ball’s height reaches the paddle regardless of whether or not the paddle catches the ball.
Revised: The episode ends when the ball’s height reaches the paddle regardless of whether or not the paddle catches the ball.

Original: We chose this environment specifically to challenge the action-inference/rollout-policy component of our algorithm.
Revised: We chose this environment specifically to challenge the action-inference/rollout-policy component of our algorithm.

Original: 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.
Revised: 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.

Original: Successful trajectories that include movement away from the ball thus make a good rollout policies more challenging to learn (i [Table] e., elicit from the LLM via prompts) [Table]
Revised: 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).

Original: 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() .
Revised: 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() .

Original: 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.
Revised: 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.

Original: Table 1 provides an example.
Revised: Table 1 provides an example.

Original: We also include the text ball.descend() to account for the change in the ball’s position between states.
Revised: We also include the text ball.descend() to account for the change in the ball’s position between states.

Original: Mini Invaders.
Revised: Mini Invaders.

Original: The agent operates a ship that shoots down aliens which descend from the top of the screen.
Revised: The agent operates a ship that shoots down aliens which descend from the top of the screen.

Original: At the beginning of an episode, two aliens spawn at a random location in two of four columns.
Revised: At the beginning of an episode, two aliens spawn at a random location in two of four columns.

Original: 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).
Revised: 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).

Original: The agent can move left , right , or shoot .
Revised: The agent can move left , right , or shoot .

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We represent the ship by its namedtuple coordinate ( C(x, y) ) and the aliens as a list of these namedtuples.
Revised: We represent the ship by its namedtuple coordinate ( C(x, y) ) and the aliens as a list of these namedtuples.

Original: When an alien is shot down, we substitute None for the tuple, as in aliens =
Revised: When an alien is shot down, we substitute None for the tuple, as in aliens =

Original: =
Revised: =

Original: [C(x, y), None] .
Revised: [C(x, y), None] .

Original: We add the text for a in aliens: a.descend() in order to account for the change in the alien’s position between states.
Revised: We add the text for a in aliens: a.descend() in order to account for the change in the alien’s position between states.

Original: Point-Mass.
Revised: Point-Mass.

Original: A point-mass spawns at a random position on a continuous line between − 6 and +6 with a velocity of 0.
Revised: A point-mass spawns at a random position on a continuous line between − 6 and +6 with a velocity of 0.

Original: The agent can either accelerate or decelerate the point-mass.
Revised: The agent can either accelerate or decelerate the point-mass.

Original: The episode terminates with a reward of 1 once the point-mass is between − 2 and +2 and its velocity is 0 once again.
Revised: The episode terminates with a reward of 1 once the point-mass is between − 2 and +2 and its velocity is 0 once again.

Original: The episode also terminates after 8 time-steps.
Revised: The episode also terminates after 8 time-steps.

Original: This domain tests the algorithm’s ability to handle continuous states.
Revised: This domain tests the algorithm’s ability to handle continuous states.

Original: States are represented as assert pos =
Revised: States are represented as assert pos =

Original: = p and vel =
Revised: = p and vel =

Original: = v , substituting floats rounded to two decimals for p and v .
Revised: = v , substituting floats rounded to two decimals for p and v .

Original: The actions are accel(pos, vel) and decel(pos, vel) .
Revised: The actions are accel(pos, vel) and decel(pos, vel) .

Original: The hint indicates whether the success conditions are met, namely the relationship of pos to − 2 and +2 and the whether or not vel
Revised: The hint indicates whether the success conditions are met, namely the relationship of pos to − 2 and +2 and the whether or not vel

Original: == 0 .
Revised: == 0 .

Original: The hint includes identification of the aliens’ and the ship’s x -positions as well as a comparison between them.
Revised: The hint includes identification of the aliens’ and the ship’s x -positions as well as a comparison between them.

Original: 4. E XPERIMENT M ETHODOLOGY AND R ESULTS
Revised: 4. E XPERIMENT M ETHODOLOGY AND R ESULTS

Original: Methodology and evaluation.
Revised: Methodology and evaluation.

Original: 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.
Revised: 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.

Original: For all experiments c = 8 (the number of most recent successful trajectories to include in the prompt).
Revised: For all experiments c = 8 (the number of most recent successful trajectories to include in the prompt).

Original: We did not have time for hyperparameter search and chose this number based on intuition, however the c =baseline demonstrates results when this hyperparameter is doubled.
Revised: We did not have time for hyperparameter search and chose this number based on intuition, however the c =baseline demonstrates results when this hyperparameter is doubled.

Original: All results use 4 seeds.
Revised: All results use 4 seeds.

Original: For both versions of Codex, we used the OpenAI Beta under the API Terms of Use.
Revised: For both versions of Codex, we used the OpenAI Beta under the API Terms of Use.

Original: For GPT-J (Wang & Komatsuzaki, 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.
Revised: For GPT-J (Wang & Komatsuzaki, 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.

Original: All language models use a sampling temperature of 0.1.
Revised: All language models use a sampling temperature of 0.1.

Original: Comparison of ICPI with baseline algorithms.
Revised: Comparison of ICPI with baseline algorithms.

Original: We compare ICPI with three baselines (Fig. 2).
Revised: We compare ICPI with three baselines (Fig. 2).

Original: The “No arg max ” baseline learns a good policy through random exploration and then imitates this policy.
Revised: The “No arg max ” baseline learns a good policy through random exploration and then imitates this policy.

Original: This baseline assumes access a “success threshold” for each domain — an undiscounted cumulative return greater than which a trajectory is considered successful.
Revised: This baseline assumes access a “success threshold” for each domain — an undiscounted cumulative return greater than which a trajectory is considered successful.

Original: The action selection mechanism emulates ICPI’s rollout policy: prompting the LLM with a set of trajectories and eliciting an action as output.
Revised: The action selection mechanism emulates ICPI’s rollout policy: prompting the LLM with a set of trajectories and eliciting an action as output.

Original: For this baseline, we only include trajectories in the prompt whose cumulative return exceeds the success threshold.
Revised: For this baseline, we only include trajectories in the prompt whose cumulative return exceeds the success threshold.

Original: Thus the policy improves as the number of successful trajectories in the prompt increases over time.
Revised: Thus the policy improves as the number of successful trajectories in the prompt increases over time.

Original: Note that at the start of learning, the agent will have experienced too few successful trajectories to effectively populate the policy prompt.
Revised: Note that at the start of learning, the agent will have experienced too few successful trajectories to effectively populate the policy prompt.

Original: In order to facilitate exploration, we act randomly until the agent experiences 3 successes.
Revised: In order to facilitate exploration, we act randomly until the agent experiences 3 successes.

Original: “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 .
Revised: “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 .

Original: “Matching Model” is a baseline which uses the trajectory history instead of an LLM to perform modelling.
Revised: “Matching Model” is a baseline which uses the trajectory history instead of an LLM to perform modelling.

Original: 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.
Revised: 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.

Original: If a match is found, the model’s outputs the historical value (e.g. the reward associated with the state-action pair found in the buffer).
Revised: If a match is found, the model’s outputs the historical value (e.g. the reward associated with the state-action pair found in the buffer).

Original: If no match is found, the modelling rollout is terminated.
Revised: If no match is found, the modelling rollout is terminated.

Original: Recall that ICPI breaks ties randomly during action selection so this will often lead to random action selection.
Revised: Recall that ICPI breaks ties randomly during action selection so this will often lead to random action selection.

Original: As our results demonstrate, only ICPI learns good policies on all domains.
Revised: As our results demonstrate, only ICPI learns good policies on all domains.

Original: We attribute this advantage to ICPI’s ability to generalize from its context to unseen states and state/action pairs (unlike “Tabular
Revised: We attribute this advantage to ICPI’s ability to generalize from its context to unseen states and state/action pairs (unlike “Tabular

Original: Q” and “Matching Model”).
Revised: Q” and “Matching Model”).

Original: Unlike “No arg max ” ICPI is able to learn progressively, improving the policy before experiencing good trajectories.
Revised: Unlike “No arg max ” ICPI is able to learn progressively, improving the policy before experiencing good trajectories.

Original: Ablation of ICPI components.
Revised: Ablation of ICPI components.

Original: With these experiments, we ablate those components of the algorithm which are not, in principle, essential to learning (Fig. 3).
Revised: With these experiments, we ablate those components of the algorithm which are not, in principle, essential to learning (Fig. 3).

Original: “No Hints” ablates the hints described in the Prompt-Format paragraph.
Revised: “No Hints” ablates the hints described in the Prompt-Format paragraph.

Original: “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).
Revised: “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).

Original: The “No Constraints” baseline removes the constraints on these time-steps described in the same paragraph.
Revised: The “No Constraints” baseline removes the constraints on these time-steps described in the same paragraph.

Original: For example, D r is allowed to contain a mixture of terminal and non-terminal time-steps (regardless of the model’s termination prediction).
Revised: For example, D r is allowed to contain a mixture of terminal and non-terminal time-steps (regardless of the model’s termination prediction).

Original: Finally, “ c = 16 ” prompts the rollout policy with the last 16 trajectories (instead of the last 8, as in ICPI).
Revised: Finally, “ c = 16 ” prompts the rollout policy with the last 16 trajectories (instead of the last 8, as in ICPI).

Original: We find that while some ablations match ICPI’s performance in several domains, none match its performance on all six.
Revised: We find that while some ablations match ICPI’s performance in several domains, none match its performance on all six.

Original: Comparison of Different Language Models.
Revised: Comparison of Different Language Models.

Original: 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).
Revised: 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).

Original: See Tablefor details about these models.
Revised: See Tablefor details about these models.

Original: Both code-davinci-002 and code-cushman-001 are variations of [Table] including “The Pile” (Gao et al., 2020) [Table]
Revised: Both code-davinci-002 and code-cushman-001 are variations of the Codex language model.

Original: “PushShift.io Reddit” (Baumgartner et al

Original: [Table]

Original: [Table]

Original: the Codex language model.

Original: The exact number of parameters in these models is proprietary according to OpenAI, but Chen et al.
Revised: The exact number of parameters in these models is proprietary according to OpenAI, but Chen et al.

Original: (2021b) describes Codex as fine-tuned from GPT-3 Brown et al.
Revised: (2021b) describes Codex as fine-tuned from GPT-3 Brown et al.

Original: (2020), which contains 185 billion parameters.
Revised: (2020), which contains 185 billion parameters.

Original: As for the distinction between the variations, the OpenAI website describes code-cushman-001 as “almost as capable as Davinci Codex, but slightly faster.”
Revised: As for the distinction between the variations, the OpenAI website describes code-cushman-001 as “almost as capable as Davinci Codex, but slightly faster.”

Original: We found that none of the smaller models were capable of demonstrating learning on any domain but the simplest, chain .
Revised: We found that none of the smaller models were capable of demonstrating learning on any domain but the simplest, chain .

Original: 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.
Revised: Examining the trajectories generated by agents trained using these models, we [Table] 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.

Original: Since these smaller models were not trained on identical data, we are unable to isolate the role of size in these results.
Revised: Since these smaller models were not trained on identical data [Table] we are unable to isolate the role of size in these results.

Original: However, the failure of all of these smaller models to learn suggests that size has some role to play in performance.
Revised: However, the failure of all of these smaller models to learn suggests that size has some role to play in performance.

Original: We conjecture that larger models developed in the future may demonstrate comparable improvements in performance over our Codex model.
Revised: We conjecture that larger models developed in the future may demonstrate comparable improvements in performance over our Codex model.

Original: C ONCLUSION
Revised: Limitations

Original: Our main contribution is a gradient-free policy iteration algorithm for using foundation models to solve RL tasks, in which the entire locus of learning is improved prompt content.
Revised: In principle [Table] ICPI can work on any control task with a discrete state space. All that the method requires is a sequence prediction model (in the paper we use Codex) capable of inferring state transitions and action probabilities given a history of behavior.

Original: At the core of the algorithm is the use of the foundation models as both aworld model and policy to compute Q-values via rollouts. In experiments we showed that the algorithm works in six illustrative domains imposing different challenges formodel and policy elicitation, confirmed the benefit of the LLM-rollout-based policy improvement, and showed how domain knowledge hints can be incorporated into prompt formats to improve learning.
Revised: Such a model will ensure that the rollouts (described in Computing Q-values) are unbiased monte-carlo estimates of value for the current policy. Given such estimates, the algorithm will implement policy iteration, a method which is proven to converge. The domains in our paper were limited by the abilities of Codex to accurately predict transitions and actions in more complex domains.

Original: While the empirical results are quite modest (and constrained by access to the LLMs), we believe the approach provides an important new way to use LLMs that  willincrease in effectiveness as the models themselves become better (as we already saw in our comparison of small and large LLMs).
Revised: As sequence models mature, we anticipate that more domains will become tractable for ICPI.  C ONCLUSION Our main contribution is a gradient-free policy iteration algorithm, using foundation models to solve RL tasks, in which the entire locus of learning is improved prompt content.

Revised: The algorithm uses a foundation models as both a world model and policy to compute Q-values via rollouts.

Original: Although we presented the algorithm here as text-based, the approach is quite general and could be applied to any foundation models that works through prompting, including multi-modal models like Reed et al.
Revised: Although we presented the algorithm here as text-based, the approach is quite general and could be applied to any foundation models that works through prompting, including multi-modal models like Reed et al.

Original:  and Seo et al.
Revised: (2022) and Seo et al.

Original: In addition to applying to more challenging domains as high quality LLM access improves, interesting future directions include combining the approach with prompt-format discovery methods, and combining the approach with other state-of-the-art RL algorithms.
Revised: 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 quite modest, we believe the approach provides an important new way to use LLMs that will increase in effectiveness as the models themselves become better, as our size comparison study suggests.