SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathbb{C}(\mathbf{x},\mathbf{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  $(\mathbf{C}(\mathbf{x},\mathbf{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
Figure 3: Comparison of ICPI with ablations. The  $y$ -axis depicts regret (normalized between 0 and 1), computed relative to an optimal return with a discount-factor of 0.8. The  $x$ -axis depicts time-steps during training. Error bars are standard errors from four seeds.
$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,  $\mathcal{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,  $\mathcal{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.
Figure 4: Comparison of different language models used to implement ICPI. The  $y$ -axis depicts regret (normalized between 0 and 1), computed relative to an optimal return with a discount-factor of 0.8. The  $x$ -axis depicts time-steps of training. Error bars are standard errors from four seeds.
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 domain, 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.
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.
Table 1: Pretrained Large Language Models (LLMs) Used in Experiments
Model	Parameters	Training data
GPT-J (B. Wang et al. 2021)	6 billion	“The Pile” (Leo Gao et al. 2020), an 825GB English corpus incl. Wikipedia, GitHub, academic pubs
InCoder (Fried et al. 2022)	6.7 billion	159 GB of open-source StackOverflow code
OPT-30B (Zhang et al. 2022)	30 billion	180B tokens of predominantly English data including “The Pile” (Leo Gao et al. 2020) and “PushShift.io Reddit” (Baumgartner et al. 2020)
Codex (M. Chen et al. 2021a)	185 billion	179 GB of GitHub code
[1] 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.
[2] 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).
[3] 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-AAbstract.html (visited on 05/16/2022).
[4] 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).
[5] 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).
[6] 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/1457c0d6bfb4967418bf8ac142f64a-AAbstract.html (visited on 05/14/2022).
[7] Tom Brown et al. "Language models are few-shot learners". In: Advances in neural information processing systems 33 (2020), pp. 1877-1901.
[8] 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).
[9] 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).
[10] 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).
[11] Mark Chen et al. "Evaluating large language models trained on code". In: arXiv preprint arXiv:2107.03374 (2021).
[12] 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).
[32] 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).
[33] John Schulman et al. "Proximal policy optimization algorithms". In: arXiv preprint arXiv:1707.06347 (2017).
[34] 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).
[35] Noah Shinn et al. Reflection: Language Agents with Verbal Reinforcement Learning. 2023. arXiv: 2303.11366 [cs.AI].
[36] 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).
[37] 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).
[38] 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).
[39] Alex Tamkin et al. "Understanding the capabilities, limitations, and societal impact of large language models". In: arXiv preprint arXiv:2102.02503 (2021).
[40] 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).
[41] Ben Wang et al. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. May 2021. URL: https://github.com/kingoflolz/mesh-transformer-jax.
[42] 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).
[43] 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).
[44] 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.emnlp-demos.6 (visited on 01/22/2022).
[45] 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).
[46] Susan Zhang et al. OPT: Open Pre-trained Transformer Language Models. _eprint: 2205.01068. 2022.
A Example prompts for each domain
Table 2: This table provides example prompts for each domain, showcasing the text format and hints. Hints are highlighted in grey.
Chain assert state  $= = 6$  and state  $! = 4$  state  $=$  left() assert reward  $= = 0$  assert not done   
Distractor assert state  $= = (6,3)$  and state  $! = (4,3)$  state  $=$  left() assert reward  $= = 0$  assert not done   
Maze assert state  $= = C(i = 2,j = 1)$  and state  $! = C(i = 1,j = 0)$  state, reward  $= 1$  left() assert reward  $= = 0$  assert not done   
Mini Catch assert paddle  $= = C(2,0)$  and ball  $= = C(0,4)$  and paddle.x  $= = 2$  and ball.x  $= = 0$  and paddle.x > ball.x and ball.y  $= = 4$  reward  $=$  paddle.left() ball descend () assert reward  $= = 0$  assert not done   
Mini Invaders assert ship  $= = C(2,0)$  and aliens  $= = [C(3,5),C(1,5)]$  and (ship.x, aliens[O].x, aliens[1].x)  $= = (2,3,1)$  and ship.x < aliens[O].x and ship.x > aliens[1].x ship.left() assert reward  $= = 0$  for a in aliens: a descended () assert not done   
Point-Mass assert pos  $= = -3.45$  and vel  $= = 0.00$  and pos  $< - 2$  and vel  $= = 0|$  pos, vel  $=$  decel(pos, vel) assert reward  $= = 0$  assert not done
B Tabular Q-Learning Hyperparameter Sweeps
Tabular Q learning rate
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
This figure presents learning curves for Tabular Q-Learning with a count-based exploration bonus. Actions are chosen using the following formulation:  $a \gets \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
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:
Number of Hidden Layers 1,2
Hidden Size 256, 512, 1024
Actor Learning Rate 0.001, 0.002, 0.005
Critic Learning Rate 0.0001, 0.0005, 0.001
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.
Figure 5: This figure presents learning curves for the random prediction ablation described above on the Chain domain. The Y-axis represents regret and the X-axis represents timesteps in the environment. Each colored line refers to a different prediction which was randomized. The red line is the original ICPI learning curve, included for reference.
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.
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:
assert state  $= = 5$  and state  $! = 4$  # the current, starting state   
state  $=$  left() # a new state is the result of the left() action   
assert not done   
assert reward  $= = 0$    
assert state  $= = 4$  and state  $= = 4$  # The new state is state 4, the goal state   
# Here the rollout policy unexpectedly chooses to move right   
# instead of trying the goal.   
state  $=$  right()   
assert not done   
assert reward  $= = 0$    
assert state  $= = 5$  and state  $! = 4$
# Again the rollout policy outputs a suboptimal action: reward = try_goal(state) assert done # try_goal always terminates an episode with 0 reward assert reward  $= = 0$
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  $\gamma$  under the optimal policy (one timestep away from a terminal state with reward 1).
The estimate for action try_goal() is accurate in state 5:
assert state == 5 and state != 4 # the current, starting state  
reward = try_goal(state)  
assert done  
assert reward == 0
However the rollout also underestimates the value of action right():
assert state == 5 and state != 4 # the current, starting state
state = right()
assert not done
assert reward == 0
assert state == 6 and state != 4
reward = try_goal(state) # left() would be the optimal policy, not try_goal()
assert done
assert reward == 0
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.
assert state == [2, 3] and state != [4, 3]
state = left()
assert not done
assert reward == 0
assert state == [1, 2] and state != [4, 2]
state = left() # The optimal action is right()
assert not done
assert reward == 0
assert state == [0, 0] and state != [4, 0]
state = left() # The optimal action is right()
assert not done
assert reward == 0
assert state == [0, 0] and state != [4, 0]
state = left() # The optimal action is right()
assert not done
assert reward == 0
assert state == [0, 0] and state != [4, 0]
state = left() # The optimal action is right()
assert not done
assert reward == 0
The rollout for the try_goal() action is correct: