Mental Modelling of Reinforcement Learning Agents by Language Models
Authors: Wenhao Lu
Xufeng Zhao
Josua Spisak
Jae Hee Lee
Stefan Wermter
Venue: Accepted by TMLR ()

Abstract: Can emergent language models faithfully model the intelligence of decision-making agents? Though modern language models already exhibit some reasoning ability, and theoretically can potentially express any probable distribution over tokens, it remains underexplored how the world knowledge these pre-trained models have memorized can be utilised to comprehend an agent's behaviour in the physical world. This paper empirically examines, for the first time, how well large language models (LLMs) can build a mental model of reinforcement learning (RL) agents, termed agent mental modelling, by reasoning about an agent's behaviour and its effect on states from agent interaction history. This research attempts to unveil the potential of leveraging LLMs for elucidating RL agent behaviour, addressing a key challenge in explainable RL. To this end, we propose specific evaluation metrics and test them on selected RL task datasets of varying complexity, reporting findings on agent mental model establishment. Our results disclose that LLMs are not yet capable of fully realising the mental modelling of agents through inference alone without further innovations. This work thus provides new insights into the capabilities and limitations of modern LLMs, highlighting that while they show promise in understanding agents with a longer history context, preexisting beliefs within LLMs about behavioural optimum and state complexity limit their ability to fully comprehend an agent's behaviour and action effects.

Published in Transactions on Machine Learning Research (12/2024)
Mental Modelling of Reinforcement Learning Agents by Lan- guage Models
Wenhao Lu University of Hamburg
wenhao. wenhao..u@ruii-hmmbuu.. de
Xufeng Zhao University of Hambury
aajeng..hhoo&hunittammbuuy
Josua Spisak University of Hamburg
josua.spisakkuuuiihoommtt de
Jae Hee Lee University of Hamburg
jee. hew..e..eee@uiihammbbbyya.. de
Stefan Wermter University of Hamburg
stefee..vermteeruuni-mmoee de
Reviewed on OpenReview: https: openreview. net/ forum? id- JN76WWoPee
Abstract
Can emergent language models faithfully model the intelligence of decision-making agents? Though modern language models already exhibit some reasoning ability, and theoretically can potentially express any probable distribution over tokens, it remains underexplored how the world knowledge these pre-trained models have memorized can be utilised to compre- hend an agent's behaviour in the physical world. This paper empirically examines, for the first time, how well large language models (LLMss can build a mental model of reinforco- ment learning (RL) agents, termed agent mental modelling, by reasoning about an agent's 's behaviour and its effect states from agent interaction history. This research attempts to unveil the potential of leveraging LLMs for elucidating RI agent behaviour, addressing a key challenge in explainable RL. To this end, we propose specific evaluation metrics and Lest them on selected RL task datasets of varying complexity, reporting findings on agent mental model establishment. Our results disclose that LLMs are not yet capable of fully re- alising the mental modelling of agents through inference alone without further innovations. This work thus provides new insights into the capabilities and limitations modern LLMs, highlighting that while they show promise in understanding agents with longer history context, preexisting beliefs within LLMs about behavioural optimum and state complexity limit their ability to fully comprehend an agent's behaviour and action effects.
1 Introduction
Large language models (LLMs) perform surprisingly well in some types of reasoning due to their common- sense knowledge (Li et al., 2022b), including math, symbolic, and spatial reasoning (Kojima et al., 2022; Yamada et al., 2023; Momennejad et al., 2023; Zhao et al., 2024b). Still, most reasoning experiments focus on human-wrilten lext corpora (Cobbe et al., 2021; Lu et al., 2022), rather than real or simulated sequential and temporal data, such as interactions of reinforcement learning (RL) agents with physical simulators. The latter scenario unveils the potential of leveraging LLMs for clucidating RL agent behaviour, with which we may further facilitate human understanding of such behaviour long-standing challenge in explainable RL. (Milani et al., 2024; Lu et al., 2024). Leveraging LLMss for this purpose is tempting as they can provide

---

Published in Transactions on Machine Learning Research (12/2024)
explanatory reasoning over sequence of actions in human-readable language, and this is possible due to their known ability to in-context learn from input-ouiput pains (Garg et al., 2022; Min et al., 2022; Li et al., 2023). This capability allows LLMs to integrate both interaction history and task-specific knowledge for more nuanced explanations, unlike most existing explainability techniques for RL (Milani et al., 2024), which have lailed to consider the dynamics of agent-environment interactions.
?
non -rontransparent
unroll
approximated
reasoning
Trajectory
Mental Model
Figure 1: A conception of LLMs approximating agent's 's mental model where language models reason over non- non-transparent agent rationale into more interpretable artefacts for end-users.
There is an ongoing debate about whether the next-token prediction paradigm of modern LLMs can model human-like intelligence (Merrill & Sabharwal, 2023; Bachmann & Nagarajan, 2024). While next-token pre- dictors can theoretically express any conceivable token distribution, it remains underexplored how the world knowledge these models have memorised during the pre-training phase (Roberts et al., 2020) can be utilised to comprehend an agent's behaviour in the real or simulated physical world. In this work, we conduct the first empirical study to examine whether LLMs can build a mental model (Johnson-Laird, 1983; Bansal et al., 2019) of agents (Figure 1), termed agent mental modelling, by reasoning about an agent's behaviour and the consequences from its interaction history. Understanding LLMs' ability to interpret agent behaviour could guide the development of agent-oriented LLMs that plan and generate sequences of embodied actions. Though recent studies (Li et al., 2022a; Huang et al., 2023) show that LLMs can aid in planning for em- bodied tasks, they merely demonstrate limited understanding of the physical world, often constrained by basic, low-level control mechanisms. Further, this agent underslanding could also inform the use of LLMs communication mediators between black-bor box agents and various stakeholders (see Sec. 5.3 for a discussion).
Understanding RL agent behaviour is more complex for LLMs than solving traditional reasoning tasks, which often involve the procedure of plugging different values into equations (Razeghi el al., 2022). In our study, we formalize the process of agent mental modelling, requiring LLMs to not only comprehend the actions taken by the agent but also perceive the resulting state changes (details in Sec. 3).
The contributions of this paper include: 1) we shed light on evaluating LLMs' ability to build mental model of RL agents, including both agent behaviour and environmental dynamics, conducting quantitative and qualitative analyses of their capabilities and limitations; 2) we present empirical evaluation results in RL tasks, offering a well-designed testbed for this research with proposed evaluation metrics, and discuss the broader implications of enabling agent mental modelling.
2 Related Work
Mental Modelling for Explaining Agents. With rapid advancements in RL techniques, understanding the behaviour of black-box RL agents has become increasingly critical yet challenging (Qing et al., 2022; Milani et al., 2024). Various methods have been proposed to improve human comprehension of agent deci- sions, including visual explanations that highlight salient state features (Greydanus el al., 2018; Iyer et al., 2018), building surrogate models like decision trees to approximate original complex agents (Bastani et al., 2018), and learning interpretable agents through causal reasoning (Lu et al., 2024). From psychological perspective, the explainability of RL agents relates mental modelling, i.c., understanding how learning systems function (Johnson-Laird,, 1983; Bansal et al., 2019). Previous studies often focus isolated aspects
2

---

Published in Transactions on Machine Learning Research (12/2024)
of agent behaviour, such state saliency or reward importance (e.g., a salient state leading to specific action or a reward motivating a behaviour (Juozapaitis et. al., 2019)). Moreover, most techniques fail to engage with the agent's learning process or interactions with the environment. This work explores the po- tential of LLMs to facilitate mental modelling of agents by reasoning over agent-environment context, with the anticipation of providing more comprehensive explanations beyy-o SLate/reward imporlance succcee--schh as revealing behavioural patterns and biases in human-readable text. LLMs are uniquely suited for this task due to their strong in-context learning capabilities (Brown et al., 2020), allowing them to interpret an agent's interaction history and provide richer insights into its behaviour than existing explainable RL techniques. Given their emergent capabilities (Wei et al., 2022a) (e.g., to mediate between humans and machine learning models), this paper takes a concrete step toward leveraging LLMs for more intuitive mental modelling of RL agents.
In-Context Learning. LLMs have exhibited strong performance inferring answers to queries upon being given input-output pairs without gradient updates, capability known as in-context learning (Brown et al., 2020; Garg et al., 2022; Min et al., 2022; Li et. al., 2023). In this study, focus on evaluating LLMs understanding of agents within in-context learning but applied to sequential macisinn-making settings (Xu et al., 2022). Here, the context is in the form of state-action-reward tuples instead of input-output tuples. Closely related to our work is in-context reinforcement learning, where pre- transformer architecture- based models are fine-tuned to predict actions for query states in a task, given history interactions (Laskin et al., 2022; Lee et al., 2023; Lin et al., 2023; Wang et al., 2024). Unlike this line of work, we aim to evaluate LLMs' capability of building a mental model of RL agents via in-context learning, instead of optimising LLMs.
Internal World Models. LLMs can also be grounded to specifie task such reasoning in the physical world or fine-tuned for enhanced embodied experiences (Liu et al., 2022: Xiang et al., 2023). However, because our focus is on the off-the-sheit performance of LLMs, we avoid this by creating collection of interactions of RL agents with physics engines (o.g., MuJoCo (Todorov et al., 2012)). This results in a more challenging dataset benchmarking that does not explicitly query the LLMs for physics understanding, instead testing their inherent capability to understand the dynamics and rationale behind an RL agent's actions. This allows of LLMs, which may offer capabilities for planning, predicting, and reasoning, as seen in works embodied task planning (Ahn et al., 2022; Driess et al., 2023).
look into the inherent internal world model (Lake et al., 2017; Amos et al., 2018)
3 LLM-Xavier Evaluation Framework
Our work studies the capability of LLMSs to understand and interpret RL agents, i.e., agent mental modelling in the context of Markov Decision Process (MDP) M (Puterman, 2014), including policies म : :S + A and transition function T S x A S, See Figure? for an overview of the LLM-Xavier' evaluation framework.
S, where S represents the state space and A represents the action space.
3.1 In-Context Prompting
The evaluation is carried out in the context of an RL task T which can be viewed as the instantiation of an MDP M. For each T, we compile a dataset of interactions between the agent and the task environment, consisting of traversed stato-action- tuples, denoted as εr == {(s,, a,, ri)ssz where L indicates the task episode length. Further, the subset of the interaction history with a time window (history size) H ending time t is denoted as E, = (ss, 1r)) i.c., capturing the most recent H tuples up to time 1 < 1.
The in-context learning prompts we constructed consist of task-specific background information, agent be- haviour history, and evaluation question prompts (sec Appendix B for example instantiated prompts):
a) A system-level prompt outlining the MDP components of the enviromment in which the agent oper- ates, including the state and action space, along with a brief task description.
Inspired by Xavier from Men who can read minds, signify its ability to model the mental states of RT. agents.
3

---

and Published in Transactions on Machine Learning Research (12/2024)
System Prompts
Task Prompts
Argue or against action
choice
Answer
Predict next action the agent
would take Predict next state
unroll
t
ground truth
Evaluation
Trajectory
Figure 2: An overview of the LLM-Xavier workflow for oflline evaluating LLMs' understanding of RL agents. The evaluation involves LLMs in-context analysing RL agent's 's past rollouts to argue about current actions, predict future actions, infer state changes resulting from the agent's actions (details in Sec. 3.2).
b) Specific prompts tailored individual evaluation purposes (Sec. 3.2), adapted based on whether the RL sctting involves discrete continuous state/action space.
c) With subsets of interaction history ε, leading up to the current time as the in-context history, we prompt LLMS respond to various masked-out queries Tquery, corresponding different evaluation questions, via inference over y LLMM(IIqquay&&... &..H).
3.2 Evaluation Metrics
Evaluating the extent which LLMs can develop a mental model requires examining their understanding of both the dynamics (mechanics) of environments that RL agents interact with and the rationale behind the agent's chosen actions. Mental modelling implies comprehensive grasp of agents, which is difficult to fully capture. However, we believe that predictability is key component of understanding. To systematically evaluate these aspects, we designed a series of targeted evaluation questions focused on predictability.
Actions Understanding. To assess LLMs' comprehennsio of the behaviour RL agents, we evaluate their ability accurately predict the internal strategies of agents, including
1) predicting next action y = d given Equery = St+1s
2) deducing last (previous) action y dc+1 given Equery
3) judging the neat action that is given
0 if agree with given action d++1
given Tquery = 814
if disagree with given action d:41
The third evaluation question assesses LLMs' reasoning by having them judge the rationale behind the agent's actions, rather than directly explaining why specific action was taken.
Dynamics Understanding. To assess the awareness LLMs to infer state transitions caused by agent actions, the evaluation of dynamics understanding includes
(1) predicting next state = Se42 given Equery = (8t+1, ar+1), and
deducing last (previous) state y 8.41 given Equery (ae+1, 8t42)

---

Published in Transactions on Machine Learning Research (12/2024)
3.3 Tost-processing of LLMs' Predictions in Practice
The evaluation prompts are tailored to the RL tasks' action or state space. For discrete actions, LLMs predict single integer within the action range. For continuous actions, explore two options: (1) predicting which bin (from a manually divided set of 10) the next action will fall into, and (2) directly predicting the absolute action value within the valid range for each action dimension. For continuous state prediction, we adopt predicting relative changes (e.g., increase, decrease, unchange) instead of exact state values. This approach assesses the LLMs' ability to sense state transitions (As), e.g., changes in physical properties in physics tasks. Detailed evaluation prompts are provided in Appendix B.
We extract predictions by post-processing the generations y f [LM(() Ei,tr) with regular expressions and compute performance by comparing them to the ground truth from the dataset. seeurayy/matchingg rale for predicting states and actions is calculated as the number of correct predictions divided by the query length E.. Note that if LLMs predict absolute action values, both predicted and ground truth values are quantised into 10 bins, and matching is based on bin alignment. Detailed prediction post-processing is discussed in Appendix C.
4 Experimental Setup
We empirically evaluate contemporary open-source and proprictary LLMs on their understanding of the agent's ' mental model, including Llama3-8B2, Llama3-70B, GPT-3.5 and GPT-4o" models". All language models are prompted with the Chain-of-Thought (CoT) strategy (Wei et al., 2022b), explicitly encouraged to provide reasoning with explanations before jumping to the answer.
Offline RL Datasets. To benchmark LLMs' ability to build a mental model of an agent's behaviour, we selected a variety of tasks featuring different state spaces, action spaces, and reward spaces, resulting in a dataset comprising seven lasks (Brockman et al., 2016) with approximately 2000 query samples, represented as (st, ar rt) tuples. Four of the seven tasks are classic physical control tasks of increasing complexity, while the other three are from the Fetch environment (Plappert et al., 2018), which includes a 7-DoF arm with a two-fingeres parallel gripper. Brief descriptions of each task are provided below, with corresponding visuals in Figure 3. See Appendix A.2 for detailed descriptions and Table 3 in Appendix A.1 for task statistics.
MountainCar Task: control a car placed at the bottom of a sinusoidal valley by applying accelerations in either direction to reach the goal state on top of the right hill.
Acrobot Task: control two-link robotic chain to swing its free end above specified height by applying torques at the joint.
Pendulum Task: apply torque to swing a pendulum from a random position to an upright stance, balancing it above a fixed point.
LunarLander Task: control rocket's engine to land designated pad, using discrete actions (engine on or off), with the challenge of precise landing.
FetchPickAndPlacs Task: a 7-DoF robot moves a block to a target position on a table or in mid-air using gripper.
FetchPush Task: 7-DoF robot pushes block to a target position table using locked gripper.
FetchSlide Task: a 7-DoF robot hits a puck to slide it to a larget, position on a slippery table, outside its workspace.
com/ /lama3/
3 www.rrrrmmoooociiimmmmo www.oo.con.onnnaeeooo www.aaiioom////////o///ee/ mwwece///ooo///////oooo/osssssss tuubb---tto
4 hhttps:/ http:////paaooooooopoooo wpw.mmooopnnicomomomoo
ceceaai.oo//ooss /ocss/ models gpt-44
tnxernet, get -t-err,, and gt-
5

---

Published in Transactions on Machine Learning Research (12/2024)
MountainCar
Acrobot
LunarLander
Pendulum
FetchPickAndPlace
FetchPush
FetchSlide
Figure 3: Visual representations of the
tasks utilised in the evaluation experiments.
GPT-3.5
GPT-4o
Llama3-8b
Llama3-70b
Next Action Prediction
Action Prediction
Next State Prediction
Last State Prediction
Matching Rate (%)
90
60
MountainCar
80
45
70
45
30
60
80
30
15
0123
5
10
20
0123
10
20
0123
10
20
0123
5
10
20
75
40
30
60
Acrobot
20
45
10
30
30
0123 5
20
0123
20
0123
10
0123
10
20
60
60
45
60
45
45
Pendulum
40
30
30
20
15
15
0
5
20
0123
0123
0123
64
64
16
Lunarl LandLander
12
56
56
48
48
3
5
20 0123
0123
10
0123
10
History Size
Figure 4: Comparative plots of LLMs' performance across various tasks with different history sizes (utilising indexed history in prompts, i.c., E,h presented with indices prefixes). The Pendulum task evaluates continuous action prediction. First two columns show agent understanding; last two columns show dynamics understanding. A description of these tasks can be found in Appendix A.2.

---

Published in Transactions on Machine Learning Research (12/2024)
5 Results and Discussion
5.1 LLMs can utilise agent history to build mental model
The experimental results indicate that LLMs can predict agent behaviours to a certain extent, for example in MountainCar, where they achieve over 75% accuracy, far surpassing the random guess baseline (1/3 chance for three action choices). However, performance declines with more challenging tasks like Acrobot and FetchPickAndPlace, as illustrated in Figure 4 (or Table 7 in Appendix F,5 for averaged accuracy) and Figure 32, which feature larger state and action spaces. We hypothesise that complex tasks require more specialized knowledge, whereas common-sense knowledge about cars and hills aids LLMs predictions in the MountainCar task.
Longer histories enhance understanding but can degrade with excess. We study the impact of the size of history provided in the context. As expected, is shown in Figure 4, providing longer history generally improves LLMs' understanding of agent behaviours. However, the benefits of including more history saturate and may even degrade, as seen in action prediction across different models, particularly Llama3- 70b. This indicates that current LLMs, despite their long context length, struggle to handle excessive data in context. In this case, more data may hinder the ability to model the agent's 's behaviour, which is in contrast with a typical learning scenario where model performance rapidly increases as learning samples increase.
The issue of performance decline due to an excessively long history becomes more pronounced for dynamics predictions, as evidenced in the MountainCar results (refer to Figure 26 for extensive details). However, as task complexity increases, the detrimental effects of redundant history may diminish (as observed in Acrobot results in Figure 27 and Fetch-series results in Figure 32), primarily because of the challenges posed by complex state and action spaces.
On the whole, across different tasks, smallor models show significant variations in optimal history sizes for both behaviour and dynamics predictions (see Figure 4 or Table 8 in Appendix F.6 for extracted values). In contrast, larger models like GPT-3.5 are more consistent, achieving the highest accuracy with a history size of 5 (GPT-40 with size of 1). This indicates that larger models handle contextual information more consistently in inference to extent, while smaller models fail to utilise history statistics effectively. To further explore beyond basic accuracy, we offer a statistical analysis on how gradually increasing history size affects previous predictions Humprove/worrenp in Appendix F.7.
Regressing on absolute action values is easier than predicting action bins. Surprisingly, LLMs perform better at predicting absolute action values than at predicting the bins into which the estimated action falls (refer to Appendix B.5 for differences in prompts). At most, LLama3-8b can allocate the numbers into categories with a mere 10.87% accuracy for the Pendulum task (GPT-3.5 achieves 39.19%), but performs better in predicting numeric values with accuracy of up to 47.73% (GPT-3.5 scores 56.82%). A detailed comparison of the averaged accuracy across LLMs is depicted in Figure 5. We hypothesise that predicting bins requires additional math ability to categorise values using context information. Refer to Figure 39 and Appendix F.6.1 for the illustrative discrepancy.
5.2 LLMs' dynamics understanding has the potential to be further improved
Inferring the dynamics in simulated world for different tasks can be challenging many aspects, such as reasoning on high-dimension state, computing physies consequences, and on.
To investigate LLMs' potential of understanding dynamics, first, we investigate the impact of providing dynamics principles, which turns out to improve both behaviour and dynamics prediction when the dynamics context is informed to LLMs (see Figure 36 for details).
Further, we explicitly examined prediction performance across state components for each dimension. As depicted in Figure 6, LLMs find it relatively casier to sense car position (element 0) than velocity (clement 1) for the MountainCar task; in contrast, for the Acrobot task, LLMs exhibit nearly uniform prediction accuracy across all state elements due to the difficulty in sensing state changes (see Appendix F.2 for
7

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5
GPT-40
Llama3-Bb
Llama3-8b No Bins
Llama3-70b
Liama3-70b No Bins
GPT-4o
Pendulum Task
No History Indexed
History Indexed
Average Matching Rate (%)
40
30
20
10
NA Pred.
LA Pred.
NA Pred.
LA Pred.
Evaluation Metrics
Figure 5: Comparison of models' performance in predicting absolute action values and action bins for the Pendulum task with and without indexed history in prompls. "NA. Pred." and "LA. Pred." stand for next and last action prediction, respectively, and are used throughout the paper. Hatching indicates numeric prediction accuracy (*No Bins").
details). We hypothesise that LLMs more proficient in linear regression, noted in Zhang al. (2023), and the dynamics equation in MountainCar is almost linear, whereas it is non-lincar in Acrobot.
Interestingly, the small model (Llama3-8h 8b) is comparable to or even outperforms a larger model like GPT- 3.5 in predicting individual state elements in some tasks, such as Acrobot. This suggests that while small models have inferior predictive ability in actions, their understanding of action effects may not be significantly influenced by the model size, but more likely by state complexity (e-g., predicting y coordinate is easier the lunar lander is more likely descent in most steps). Refer to Appendix F.1 and F. for more illustrative results.
GPT-3.5 Next State Prediction
GPT-4o Next State Prediction
Liama3-8b
Next State Prediction
Llama3-70b Next State Prediction
80
100
State Element 0
Accuracy (9)
75
State Element
State Element
State Element
State Element 0
State Element
State Element
State Element
25
5
20
20
10
20
GPT-3.5
Last State Prediction
GPT-4o
Last State Prediction
Llama3-8b
Last State Prediction
Llama3-70b
Last State Prediction
Accuracy (%)
50
75
40
30
20
0123
5
0123 5
20
0123
10
20
History Size
History Size
History Size
History Size
Figure 6: Dynamics of LLMs' performance on predicting individual state element for the MountainCar task (with indexed history in prompts).
State span is statistically influential in LLMs' dynamics understanding. Based on the observations above, we examined the factors that statistically influence LLMs' perception of state changes. Figure demonstrates decent lincar relationship between the span of state elements and LLMs' ability detect these changes in tasks like MountainCar, Acrobot, and LunarLander, especially with GPT-3.5. For Acrobot, observe low slope with moderate fit, while LunarLander shows positive slope with moderate
8

---

Published in Transactions on Machine Learning Research (12/2024)
fit. Our hypothesis posits that in given task, smaller state span makes sensing changes difficult due to the minimal differences in state values. Conversely, a larger span tends to improve prediction accuracy it amplifies the discrepancies between state values, making changes more discernible. Further details and discussions are provided in Appendix F.3.
Linear Fit GPT-3.5
Linear Fit - Llama3-8b
Linear Fit Llama3-70b
90
Tasks
Tasks
Tasks
80
80
MC (R2=1.00)
MG (R*=1.00)
MC (R7=1
70
AC (R*=0.58)
AC (R*-0.14)
AC (R²=0.17)
Dummy Accuracy (8)
Pen (#740.40)
Pen (R7-0.35) 33)
Pen (R?=0.00) 00)
60
(R7=0.25)
LL. (R7=0.40)
60
LL (R7=0.21)
50
FetchPNP (8*=0.72)
40
FetchPNP (8*=0.14)
50
FetchPNP (R*=0.47)
FetchPush (R.00.00)
FetchPush (R--0.02)
FetchPush (R2-0.02)
40
FatchSlide (R'=0.01)
30
FetchSidde (R²=0.01)
FerchSlide (R (0006)))
30
20
20
10
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Dummy Span
Figure 7: The R² values (also reflected by the linewidth) and slopes of the fitted lincar lines for all models across various lasks. A standardised dummy state span (0-1) and accuracy scale (0-100) are utilised to account for differing state spans among tasks in the analysis. "MC" stands for MountainCar, "AC" for Acrobot, "Pen" for Podulum, "LIL" for LunarLander, "FetchPNP" for FetchPickAndPlace
5.3 Understanding error occurs from various aspects
With the anticipation that LLMs' explanatory reasoning (clicited via CoT) can benefit the human un- derstanding of agent behaviour, in addition to the existing quantitative results, we further examined the reasoning error types across LLMs by manually reviewing their judgments on the rationale of actions taken. Table shows an examination of the MountainCar task, highlighting that LLama3-8b displays the most errors. Mcanwhile, GPT-3.5, despite having superior task comprehension (c.g., referring momentum strategies), is less effective at retaining task descriptions in memory compared to Tlama3-70b. Notably, GPT-4o's responses are of higher quality overall. While GPT-4o is occasionally inconsistent with the op- timal momentum strategy (i.c., initially accelerating left), it displayed a significantly better grasp of this strategy than other models. Furthermore, it avoids fundamental mistakes such as misinterpreting numbers or the presented information. Detailed error type reports are available in Appendix G.1.
Table 1: Error counts in LLMs' responses for MountainCar task over 50 steps, with various error types. Note: Multiple error types occur within single step.
Error
GPT-3.5
Llama3-8b
Llama3-70b
GPT-4o
Types
(1) Task Understanding
9
30
16
(2) Logic
5
19
(3) History Understanding
3
18
(4) Physical Understanding
1
2
0
(5) Mathematical Understanding
2
25
2
0
(6) Missing Information
9
10
13
6
Confirmation bias in LLMs: a likely error from initial assumptions of R.L agent's optimality.
In the manual review, queried LLMss to judge possible next action (see corresponding prompt in

---

Published in Transactions on Machine Learning Research (12/2024)
Appendix B.4) given the history of the last three actions and states, history size determined to be optimal for most models (see Table 8). The provided next action was sometimes correct (if it was the agent's 's action) and sometimes incorrect, ensuring LLMs made context-based conclusions rather than merely agreeing or disagrecing with the prompt. We evaluated whether the LLMs' judgments were correct according to a human reviewer, independent of the RL agent's action correctness. An automatic evaluation compared LLMs' decisions to the RL agent's actions (as described in Sec. 3.2).
The manual evaluation did differ from the automatic evaluation, shown in Table 2. The table's percentages refer to the proportion of LLMs responses deemed correct. This difference stems from considering different action ground truth since the RL agent occasionally acts illogically, loading to the human reviewer deeming those actions incorrect, while automatic evaluation considers them correct from the RL agent's perspective.
Hypothetically, this difference may arise from LLMs forming mental models of actions they believe an opti- mal RL agent should take. This reveals another type of undorstanding error, where LLMs' responses may be influenced by preexisting beliefs or assumptions about optimal behaviour. As result, LLMs sometimes propose more optimal actions than the RL agent, particularly in larger models (evidenced by the increased accuracy in human evaluation), where these beliefs in optimality might be more strongly reinforced. Explor- ing whether such beliefs exist in LLMs could be promising research direction, potentially benefiting from advanced explainability techniques for LLMs Zhao et al. (2024a).
Despite these differences, the variation is minimal and not statistically significant (paired t-test, p-vauue 0.45 0.05), suggesting that the RL agent's behaviour is near-optimal or acceptable to the human evaluator. Thus, LLMs' action judgments may offer (less biased) explanations for RL agent actions; when agent actions are sub-optimal, LLMs may help explain why those actions are deemed inappropriate.
The accuracy of models evaluated manually or automatically for 50 steps in the MountainCar task with the metric judging nexl action.
Table 2:
Model
Manual
Automatic
GPT-3.5
60%
67%
Llama3-8b 3-bb
40%
52%
Llama3-70b
67%
65%
GPT-4o 85%
81%
Pre-training biases may affect understanding despite history length. Interestingly, we found that LLMs prediction errors may be influenced by pro-training biases towards cautious actions. Figure 33 Appendix F.7 shows that increasod history size does not affect late-episode steps where RL agents consistently accelerate right to finish. Instead, LLMS mentally "suggest." conservative actions, like accelerating left. to avoid overshooting, regardless of history context size. We hypothesise that these LLMs prediction errors may not be solely duc to difficulty in processing extensive histories but also stem from inherent biases about agent behaviour and environment. context, consistent with our earlier observations on confirmation bias.
5.4
Data format influences understanding
Prompting format generally has an impact on LLMs' reasoning performance. In the context of agent un- derstanding, we do an ablation study to investigate the robustness of prompts on the history format and provided information. We find that:
1) Excluding the sequential indices from the history context in prompts for LLMs generally negatively impacts their performance in most tasks, indicating that LLMs still struggle to process raw data and indexing helps The resulting performance variations are reported in Figure 8.
2) Task description, despite not being directly relevant to numerical value regression as in statistics, is essential for better understanding of both agent behaviour and dynamics, which brings the promise of utilising LLMs digest additional information beyond mere numerical regression when mental modelling agents. The ablation results can be found in Appendix F.8.2.
10

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5
GPT-3.5 History Indexed
Liama3 db Acrobot
Liama3-se History Indexed
Llama3-70b
Liama 70b History Indoxed
MountainCar
Pendulum
LunarLander
50
64
80
Average Matching Rote (%6
70
88
18
68
15
60
400
00
12
8.2
40
30
24
24
20
18
10
NA Pred. LA Pred. NS Pred. LS Prod.
Pred LA Pred, NS Fred. LS Prod.
NA Pred, LA Predi NS Prod. LS Pred.
NA Predi LA Pred. Prod. Pred.
Evaluation Metrics
Figure 8: Performance comparison of language models with and without indexed history in prompts various lasks. Bars with hatching indicate accuracy with indexed history in prompts. NS Pred. Next State Prediction, T.S Pred. = Last State Prediction.
6 Conclusion
This work studies underexplored aspect of next-token predictors, with a focus whether LLMs build mental model of agents. We proposed specific prompts to evaluate this capability. Quantitative evaluation results disclose that LLMs can establish agent mental models to some exlent only since their understanding of state changes may diminish with increasing task complexity (e.g., high-dim spaces); their interpretation of agent behaviours may tumble for tasks with contimuuus actions. Analysis of evaluation prompts reveals that their content and structure, such as history size, lask instructions, and data format are crucial for the effective establishment, indicating areas for future improvement. further review of LLMs error responses (elicited via CoT prompting) highlights qualitative differences in LLMs' understanding performance, with models like GPT-4o showing superior comprehension and fewer errors compared to the small Llama3 model. Additionally, it reveals that LLM responses may be influenced by preexisting beliefs about optimal behaviour, which can be a hurdle or a catalyzer for understanding. These findings suggest the potential and limitations of in-context mental modelling of agents within MDP frameworks and underscore the possible role LLMs as communication mediators between black-bo box agents and stakeholders, pointing to future research avenues.
Potential Direction
It remains unclear whether LLMs can benefit from thousands of agent trajectories compared to the lim- ited number of examples studied in this paper. We hypothesise that large amounts of demonstrations (state-action-reward tuples) in the prompt could enhance the capacity that LLMs have already developed. Additionally, fine-tuning LLMs with demonstrations (Lin et al., 2023; Wang et al., 2024) from specific do- mains may further improve their understanding capacity in these domains. Further analysis on this aspect is left for future work.
At this stage, LLMs lack several key aspects required for true understanding, often due to reasoning errors and limitations in predictability. Relying on LLMs for mental modelling without oversight could be problematic. Therefore, our key takeaway is that rather than using a fully automated, LLM-driven evaluation system for agent mental modelling, domain experts expected to be involved upfront to audit LLMs' reasoning and ensure reasonable predictability and reliability before presenting the results to non-expert practitioners. A few concrete suggestions for utilising LLMs in agent mental modelling are provided in Appendix H.
We recognise that the issue hallucination may exist. To increase the robustness and reliability of using LLMs for explaining an agent's behaviour, a detailed analysis of this behaviour necessary before being deployed a setting where they directly interact with humans, Also, our evaluation results underscore the need for developing methods to mitigate hallueinations.
11

---

Published in Transactions on Machine Learning Research (12/2024)
Our evaluation framework can be adapted for non-deterministic MDPs with minor adjustments. In discrete state spaces (e.g., FrozenTake (Brockman et al., 2016)), LLMs can output distribution over possible next state values, with multiple queries needed to approximate the distribution, evaluated using metrics like KL- divergence. Alternatively, LLMs could rank next state values, compared to the ground truth probabilities derived from simulator. continuous stale spaces, stales can be quantised into bins, and LLMs can output a distribution or ranking over sampled values from each bin. For multi-dimennional states, distributions or rankings can be generated for each dimension separately. This approach also applies predicting actions, where the ground truth differs by algorithm: for PPO (Schulman et al., 2017), it would be the mean and variance of the Gaussian policy, and for DQN (Hasselt et al., 2016), the Q-values for discrete actions. In continuous action spaces, LLMs can output mean and deviation for a Gaussian distribution, which can be compared to the learned (Gaussian) distribution ion of the PPO agent.
Our study provides a macro-level analysis by examining the average model performance over multiple RL. datascts of varying types. However, the capability of LLMs build a mental model of agents may vary across different datasets. While our analysis discusses this aspect, it is important to explore ways of standardising this type of benchmarking for language models, which may evolve LLMs become more intelligent. A long-term goal of this rescarch is to facilitate human understanding of more intelligent agents in critical domains, and we see this work as foundational step towards developing progressively more agent-oriented language models with realistic world models in mind.
Our experiments are limited to uni-modal RL tasks (i.e., using proprioceptive states), but extending them to tasks (e.g., incorporating vision, auditory, and touch feedback) straightforward. Multi-modal inputs can provide LLMs with richer environmental information than state vectors, and we hypothesise that these additional signals may enhance LLMs' agent mental modelling.
Future-Proofing the Approach for LLM-Basse Agents
While our current framework focuses on evaluating LLM's capability of mental modelling traditional RL agents, it remains equally relevant as RL agents evolve into LLM-based counterparts that can provide human-readable justifications for their decisions. One might argue that these self-generated explanations reduce the need for external evaluation. However, these internal justifications may not reflect truly objective mental model as they can be shaped by the agent's underlying optimisation processes, biases, and emergent behaviours. To address this, when evaluating LLM-based agents, we can extend our framework to compare their self-reported rationales against the mental model formed by an independent evaluator LLM.
Furthermore, as LLM-driven agents rapidly adapt policies in response dialogue histories, external instruc- tions, or internal reasoning steps, we can integrate temporal tracking of these shifts in real-time. This enables the evaluator LLM to detect subtle policy shifts and predict. future actions based on evolving rationales. As LLM-based agents expand into multi-domain environments with non-traditional state, action, and reward structurcs-sude as language- baaed based tasks-our framework can similarly evolve, challenging the evaluator LLM to generalise beyond conventional RL dynamics. Finally, as we leveraged human input for error-type annotation and analysis, future- pooorengoofing involves integrating semi-automated tools to handle intricate LLM- agent interactions at scale. Human evaluators may focus on higher-level interpretation and strategic insights, while automated metrics track the consistency and coherence of evolving mental models over time.
Broader Impact Statement
We do not anticipate any immediate ethical or societal implications from our research. However, since we explore LLMs applications for enhancing human understanding of agents, it is important to be cautious about the potential for fabricated or inaccurate claims in LLMs' explanatory responses, which may arise from misinformation and hallucinations inherent the LLMs employed. It is recommended use our proposed evaluation prompts and task dataset with care and mindfulness.
Acknowledgements
This research was funded by the Federal Ministry for Economic Affairs and Climate Action (BMWK) under the Federal Aviation Research Programme (TuFO), Projekt. VeriKAS (20X1905)
12

---

Published in Transactions on Machine Learning Research (12/2024)
References
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, et al. Do as i can and not as i say: Grounding language in robotic affordances. In ar preprinl S0o:2204.01691, 2022.
Brandon Amos, Laurent Dinh, Serkan Cabi, Thomas Rothörl, Alistair Muldal, Tom Erez, Yuval Tassa, Nando de Freitas, and Misha Denil. Learning awareness models. In International Conference on Learning Representations, 2018. URL https://opennreview..t http://wwwwe..ttttttoorrrrddddde
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. ar preprint 0rXiv:243.069633 2024.
Gagan Bansal, Besmira Nushi, Ece Kamar, Walter S. Lasecki, Daniel S. Weld, and Eric Horvitz. Beyond accuracy: The role of mental models in human-ai team performance. Proceedings of the AAAI Conference on Human Compulalion and Crowdsourcing, 7(1):2 Oct. 2019. doi: 10.1609 hhcomp.v7il.b. URL
https://ojs.. aaai. wwoooog///oo/eoeo /in//iidexx..e pptp://oon//arrriccl//r/ee/////s
Osbert Bastani, Yewen Pu, and Armando Solar-Tezama. Verifiable reinforcement learning via. policy extrac- tion. Advances in newral information processing systems, 31, 2018.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jomas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016. URL
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural informalion processing systems, 33:1877 13:1877-19011 1901, 2020.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arv preprint Xrx::22100. 10. 14168, 2021.
Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multimodal language model. In International Conference on Machine Learning, pp. 8469-8448 PMLR, 2023.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in- context? case study of simple function classes. Advances in Neural Information Processing Systems, 35: 11583-30593 2022.
Samuel Greydanus, Amurag Koul, Jonathan Dodge, and Alan Fern. Visualizing and understanding atari agents. In International Conference on Machine Learning (ICML), pp. 1792 1801. PMLR, 2018.
Hado van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. Tn Proceedings of the Thirtieth AAAI Conference on Artificial Intelligence, AAAIIG,, pp. 2094 2100. AAAI Press, 2016.
Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tomp- son, Igor Mordatch, Yevgen Chebotar, et al. Inner monologue: Embodied reasoning through planning with language models. In Conference on Robot Learning, pp. 1769 1782 PMLR, 2023.
Rahul Tyer, Yuezhang T.i, Huao Li, Michael Lewis, Ramitha Sundar, and Katia Sycara. Transparency and explanation in deep reinforcement learning neural networks. In Proceedings of the 2018 AAAI/ACM Conference on AI, Elhics, and Society, pp. 144 150, 2018.
Philip Nicholas Johnson-Laird. Mental models: Towards cognitive science of language, inference, and consciousmess. Harvard University Press, Cambridge, MA, 1983.
Zoe Juozapaitis, Anurag Koul, Alan Fern, Martin Erwig, and Finale Doshi Velez. Explainable Reinforcement learning via Reward Decomposition. In IJCAT/ECAT Workshop Eaplainable Artificial Intelligence, 2019.
13

---

Published in Transactions on Machine Learning Research (12/2024)
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are vero-shot reasoners. Advances in neural information processing systems, 35:22199-2222 2022.
Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40:e253, 2017.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arx:::2200. 14215, 2022.
Teven Le Scao and Alexander M Rush. How many data points prompt worth? In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human 2021.
Language Technologies, pp.
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofrr Nachum, and Emma Brun- skill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2023.
Shuang Li, Xavier Puig, Chris Paxton, Yilun Du, Clinton Wang, Linxi Fan, Tao Chen, De-An Huang, Ekin Akyürek, Anima Anandkumar, et al. Pre-trained language models for interactive decision-making, Advances in Newral Information Processing Systems, 35:31199 31212, 2022a.
Xiang Lorraine Li, Adhiguna Kuncoro, Jordan Hoffmann, Cyprien de Masson d' A'tuume Phil Blunsom, and Aida Nematzadeh. systematic investigation commonsense knowledge in large language models. Proceedings of the 2022 Conference on Empirical Methods Natural Language Processing, pp. 11838-
11855, 2022b.
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers algorithms: Generalization and stability in in-context learning. In International Conference Machine Learning, pp. 19565 19594. PMLR, 2023.
Timothy P Lillierap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.0200 2015.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In The Twelfth International Conference on Learning Representations, 2023.
Ruibo Liu, Jason Wei, Shixiang Shane Gu, Te-Yen Wu, Soroush Vosoughi, Claire Cui, Denny Zhou, and Andrew M Dai. Mind's eye: Grounded language model reasoning through simulation. arXiv preprint arXiv:221000335 2022.
Pan Lu, Swaroop Mishra, Tanglin Xia, Liang Qiu, Kai-Wer Chang, Song-Chun Chun Zhu, Oyvind Tafjord, Peter Clark, and Ashwin Kalyan. Learn to explain: Multimodal reasoning via thought chains for science question answering. Advances in Neural Information Processing Systems, 35:2507-225 2521, 2022.
Wenhao Lu, Xufeng Zhao, Thilo Fryen, Jae Hee Lee, Mengdi Li, Sven Magg, and Stefan Wermter. Causal state distillation for explainable reinforcement learning. In Causal Learning and Reasoning, pp. 106 142. PMLR, 2024.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations, 2023.
Stephanic Milani, Nicholay Topin, Manucla Veloso, and Fei Fang. Explainable reinforcement learning: A survey and comparative review. ACM Comput. Surv., 56(7), apr 2024. ISSN 0360-0300. doi: 10.1145/ 3616864. URL https://doi. cereii.org/11.111 11113/061145/33168644
14

---

Published in Transactions on Machine Learning Research (12/2024)
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannanch Hajishirzi, and Luke Zettle- moyer. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Lenguage Processing, pp. 11048 11064, 2022.
Swaroop Mishra, Daniel Khashabi, Chitta Baral, Yejin Choi, and Hannaneh Hajishirzi. Reframing in- structional prompts to gptk's language. In 60th Annual Meeting of the Association for Computational Linguistics, ACL 2022, pp. 589 612. Association for Computational Linguistics (ACL), 2022.
Ida Momennejad, Hosein Hasanbeig, Felipe Vieira Frujeri, Hiteshi Sharma, Nebojsa Jojic, Hamid Palangi, Robert Ness, and Jonathan Larson. Evaluating cognitive maps and planning in large language models with cogeval. Advances in Neural Information Processing Systems, 36, 2023.
Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, Vikash Kumar, and Wojeiech Zaremba. Multi- goal reinforcement learning: Challenging robotics environments and request for research, 2018.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Yumpeng Qing, Shunyu Liu, Jie Song, Huiqiong Wang, and Mingli Song. A survey on explainable reinforce- ment learning: Concepts, algorithmss, challenges. arXiv preprint (Xxi:2221.06666 2022.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot numerical reasoning. Findings of the Associalion for Computational Linguistics: EMNLP 2022, 2022.
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of language model? Tn Proceedings of the 2020 Conference on Empirical Methods in Natural Longuage Processing (EMNLP), pp. 5418 5426, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimiza- tion algorithms. arXiv preprint arvv: 1xX:x:1170000000 707.06347, 2017.
Emannel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. Tn 2012 IEEE/RS) International Conference on Intelligent Robots and Systems, pp. 5026 5033, 2012. doi: 10.1109/IROS.I /tt.//OO........... 201266386109.
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers learn temporal difference methods for in-context reinforcement learning. arXiv preprint arXiv:2405.1386. 2024.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgcaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large language models. Transactions on Ma- chine Learning Research, 2022a. ISSN 2835-8856. Survey Certification.
URL https: /ttttps:///oae https:: //wp:///eee/./eeeee///eomee http:///pprrieeeeee/o//o//////aa wwt..e/ooorrurdddyyyssssuuuo
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fci Xia, Ed Chi, Quoc V Le, Denny Zhou, al. Chain-of-thoughl prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837, 2022b.
Jiannan Xiang, Tianhua Tao, Yi Gu, Tianmin Shu, Zirui Wang, Zichao Yang, and Zhiting Hu. Language models meet world models: Embodied experiences enhance language models. Advances in neural infor- mation processing systems, 36, 2023.
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine learnóng, pp.
PMLR, 2022.
Yutaro Yamada, Yihan Bao, Andrew Kyle Lampinen, Jungo Kasai, and Ilker Yildirim. Evaluating spatial understanding of large language models. Transactions on Machine Learning Research, 2023.
15

---

Published in Transactions on Machine Learning Research (12/2024)
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. In RO-FoMo: Robustness of Few-shol and Zero-shol Learning in Large Foundalion Models, 2023.
Haiyan Zhao, Hanjie Chen, Fan Yang, Ninghao Liu, Huiqi Deng, Hengyi Cai, Shuaiqiang Wang, Dawei Yin, and Mengnan Du. Explainability for large language models: A survey. ACM Transactions Intelligent Systems and Technology, 15(2):1-38, 2024a.
Xufeng Zhao, Mengdi Li, Wenhao Lu, Cornelius Weber, Jae Hee Lee, Kun Chu, and Stefan Wermter. Enhancing Zero-Shot Chain-of- Thhought-Thhoght Reasoning in Large Language Models through Logic. In 2024 Joint International Conference Compulational Lánguistics, language Resources and Evaluation (LREC- COLING 2024), May 2024b. URL https://arts wwtg/bs/////aeeeeeeeee 13339.
16

---

Published in Transactions on Machine Learning Research (12/2024)
A Statistics of Our Offline-RI RL Datasets
A.1 Data Collection
The datasel of interaction histories (episodes) is collected by running RL agents in each task. Unlike Liu el al. (2022), whose physics alignment dataset contains text-based physical reasoning questions resembling physics textbooks, our dataset comprises interactions of RL agents with various physics engines (environments). For each task, episodic histories are collected by rumning single-task RL algorithms (Lillicrap et al., 2015; Hasselt et al., 2016; Schulman et al., 2017) to solve that task. An overview of the task dataset statistics is provided in Table
Table 3: A statistical overview of the task dataset tested in the experiment.
# of episodes
Length per episode
Tasks
State Space
State Dim
Action Space
Action Dim
MountainCar
continuous
discrete
choices) (3 choices)
Acrobot
3
100
continuous
discrete
LunarLander
3
250
continuous
8
discrete
(4 choices)
Pendulum
3
~ 50
continuous
continuous
DnvertedDoublleeenduall
continuous
continuous
FetchPickkAddPin
10
~ 10
continuous
25
contimuous
FetchPush
10
10
continuous
25
continuous
FetchSlide
10
25
continuous
continuous
A.2 A Full Task Description
Below, we provide complete description of the MountainCar task (at the last of the section), including its MDP components. For the remaining tasks, only the task descriptions are provided. Most of the texts are credited to https:/ggymnasium farama.org/-
Acrobot Task Description. The Acrobot environeent is based on Sutton's work in "Generalization in Reinforcement Learning: Successful Ezamples Using Sparse Coarse Coding" and Sutton and Barto's 's book. The system consists of two links connected linearly to form chain, with one end of the chain fixed. The joint between the two links is actuated. The goal is to apply torques on the actuated joint to swing the free end of the ouler-link above given height. while slarting from the initial stale of honging downwords.
Pendulum Task Description. - The inwerted pendulum swingup problem is based on the classic problem in control theory. The system consists of a pendulum attached at one end to a fixed point, and the other end being free. The pendulum slarts a random position and the goal is to apply lorque on the free end to swing it into an upright position, with its center of gravity right above the fixed point.
LunarLander Task Description. - This environment is classic rocket trajectory optimization problem. According to Pontryagin's mazimum principle, il is optimal to fire the engine al full throttle turn. il off. This is the reason why this environment has discrete actions: engine on or off. The landing pad is always at coordinates (0,0). The coordinates are the first two numbers in the state vector. Landing outside of the landing pad is possible. Fuel is infinite, so an agent can learn to fly and then land on its first attempt.
The task the environment is for a manipulator move block to largel position on top of lable or in mid-uir. The robot is 7-DoF Fetch Mobile Munipulator with two-fingered parallel gripper (i.e., end effector). The robot is controlled by small displacements of the gripper in Cartesian coordinates and the inverse kinematics are computed internally by the MuJoCo frumework. The gripper can be opened or closed in order to perform the graspping operation of pick and place. The lask is also continuing which means that the robol has to maintain the block in the targel position for an indefinite period of time.
RetchPickAndPar Task Description.
FetchSlide Task Description. The task in the environment is for monipulator to hit puck in order to reach target position top of long and slippery table. The table has low friction coefficient in
17

---

Published in Transactions on Machine Learning Research (12/2024)
order to make it slippery for the puck to slide and be able to reach the target position which is outside of the robot's workspace. The robot. is 7-DoF Fetch Mobile Manipulator with two-fingered porallel gripper (i.e., end effector). The robot is controlled by small displacements of the gripper in Cartesian coordinates and the inverse kinematics are computed internally by the MuJoCo framework. The gripper locked in a closed confguuation since the puck doesn't need to be graspped. The lask also contónaing which means that the robot has maintain the puck in the target position for an indefinite period of time.
FetchPush Task Description. - The task in the environment is for manipulator to move block to target position top of table by pushing with its gripper. The robot is 7-DoF' Fetch Mobile Manipulator with two-fingered parallel gripper (i.e., end effector). The robot is controlled by small displacements of the gripper in Carlesian coortinates and the inverse kinematics are computed internally by the MuJoCo fromework. The gripper locked in closed configuration in order to perform the push task. The task is also contiruing which means that the robot has maintain the block the target position for an indefinite
period of time.
MountainCar Task Prompt
task_ _ess_deacription
The Mountain Car MDP is deterministic MDP that consists of car placed stochastically at the bottom a sinusoidal valley, with the only possible actions being the accelerations that can be applied to the car in either direction. The goal of the MDP is to strategically accelerate the car to reach the goal state on top the right hill.
observation_ space
The observation is ndarray with shape (2,) where the elements correspond the following: position of the car along the x-axis (range from -1.2 to 0.6), velocity of the car (range from -0.07 to 0.07)
action_etd _aaton_ppace
There are 3 discrete deterministic actions, Accelerate the left
1: Do not accelerate
2: Accelerate the right
reward_ seward_space
The goal is reach the flag placed top of the right hill quickly possible, as such the agent is penalised with a reward of -1 for each timestep.
transition disiooi_dddmcaca
Given action, the mountain follows the following transition dynamics, velocity_te t+1 velocity_t + (action * force cos (3 position_t) gravity position_tws t+1 position_t velocity_t+3
where force = 0.001 and gravity 0025.. The collisions at either end are inelastic with the velocity set to upon collision with the wall.
init_state
The position of the car is assigned uniform random value in [-0.6, -0.4]. The starting velocity of the car is always assigned to 0.
termination
The episode ends if the position the is greater than equal 0.5 (the goal position on top of the right hill).
Prompt Examples
The structured input template used for querying LLMs consists of system prompt containing the task description, MDP components, and prompt with specific evaluation questions, shown in B. and B.2,
18

---

Published in Transactions on Machine Learning Research (12/2024)
respectively. An example prompt for predicting the next action for tasks with discrete action space is depicted in B.3, and for judging agent actions in B.4. The prompts for each evaluation metrie may vary slightly depending on the task type (i.e., state and action space as illustrated in Table 4), detailed in B.5.
B.1 System Prompt
System Prompt
Below a description of the (task name] task.
Task description:
jeeeek_deerrppton
Observation space:
(observation_spatt
Action space:
[action_space}
Reward space:
(reward_space)
Transition dynamics:
Initial state:
(init_state)
Termination:
(termination)
B.2
Offline Evaluation Prompt
Offline Evaluation Prompt
Given snippet of an episode (generated by reinforcement learning agent optimally trained for solving the given task) of
the states:
{states)
the corresponding actions taken by the RL agent,
(actions)
and the rewards received:
(rewards)
Your task is to analyse the sequence of states, actions, and rewards to address the question: {question}
19

---

Published in Transactions on Machine Learning Research (12/2024)
B.3 Example Next Action Prediction Prompt
Next Discrete Action Prediction Prompt
next step (i} (indexed from 0), agent transitioned state s(i} {state). Based your observation and understanding of the agent's 's behaviour, you predict the action a(i} (an integer from the given range) the RL agent will most likely take at step {i}?
Please first provide a compact reasoning before your answer to the action choice. Think step by step and use the following template in your provided answer:
1. [Reasoning]:
2. [Prediction]:
3. [Formatting]:
Return a list with the following example format, python final action choice
action_choice [0]
Please choose only action, even if multiple actions seem possible.
B.4 Example Action Judgment Prompt
Action Judgment Prompt
In next step {i} (indexed from 0), the agent transitioned to the state s(i} = (state}. In this state s(i), the agent chose to take action a(i} (action). Based on your observation and understanding of the agent's 's behaviour up to this point, critically evaluate the rationale behind the agent's choice action a(i} in state s(i}. You're encouraged scrutinize the correctness of the action a(i), especially if your analysis suggests that the action might be flaved or suboptimal.
After your evaluation, determine whether to accept or reject the agent's a action a(i}. Think step by step and use the following template in your provided
- answer:
1. [Reasoning]:
2. [Juttffiattio]: Critique whether the action is correct given the historical context and the state s(i}, and whether there are any reasons to doubt its correctness.
3. [Formatting]:
Final vote is [True] (i.e., argue for the agent's action) or [False] (i.e., argue against agent's action) Return list with the following example format, python
final vote is to argue against the agent's ' action final_ vote = [False]
Please choose either True or False, even in cases of uncertainty or multiple possibilities.
20

---

Published in Transactions on Machine Learning Research (12/2024)
Table 4: Different state spaces and action spaces of MDPs (tasks) considered in the experiments.
continuous action
discrete action
continuous state
discrete state
X
B.5 Evaluation Prompts in Practice
The evaluation prompts (parts b, c in Sec. 3.1) are adapted based on the nature of the RL tasks, specifically the type of action or state space (discrete or continuous). For tasks with discrete action spaces, LLMs are prompted output single integer within the action range. For tasks with continuous actions, we evaluate two options:
Predicting bins: The action range is manually divided into 10 bins, and LLMss are queried to predict which bin the RL agent's 'ss next action will fall into.
Predicting absolute numbers: LLMs are queried directly output the exact action value within the valid action range for cach dimension of the action space.
For tasks involving continuous state prediction, we adopt predicting relative changes (e.g., increase, de- crease, unchange) instead of exact state values. This approach assesses the LLMs' ability to sense state transitions (As), e.g., changes physical properties in physics tasks.
C Post-processing LLMs' Predictions
We evaluate LLMs using metrics that require predicting states and actions. We extract LLMs' responses through pattern matching and compute evaluation results by comparing them with the ground truth state- action pairs from the episodes on which the LLMs are evaluated.
For predicting discrete actions, compute the matching rate of the LLMs predieted actions with the ground truth. For predicting continuous actions, if LLMs are prompted to predict bins, compute the matching rate as we did for discrete actions, with the ground truth represented by the bin index to which it belongs. However, if LLMs are queried to directly predict absolute action values, we quantize both the predicted and ground truth values into bins (by dividing the original action range into 10 bins) and then measure whether they fall into the same bin.
For predicting continuous states, we evaluate if LLMs correctly predict the change in state, As, categorising increases as 1, decreases as O, and unchanged as 2. We then compute the accuracy classification score for their prodictions. We also record the accuracy of predicting changes in individual state elements, As;-
D Pseudo-code for Performing Evaluation Metrics
D.1 Pseudo-code for Predicting Next Action
Algorithm 1 presents an example pseudo-code for the next action prediction tasks.
E
Types of LLMs' Understanding Failures
Table 5 lists error types we located by analysing LLMs' responses to evaluation questions for the MountainCar
task. The MountainCar task is an RL task in which the agent has to correctly steer car to move it onto a mountain with the actions: "Left Acceleration", "Right Acceleration", and "No Acceleration".
21

---

Published in Transactions on Machine Learning Research (12/2024)
Algorithm 1: Offline Evaluation of L.Mss Agent Understanding
Data: Load offine RL dataset ÊT for a task T
Input: Load the LLM model LLM
Result: Evaluation results
Initialize aclion matching counter N <- 0
Set the history size H ;
Sct the maximum time steps Tmaa for t + to maz do
if 1. H then
Extract the last transitions Et, 841,
// Prepare input for LLM including current state and history ELH Predict de+1 + (8
if â a441 then L Increment counter N N+++1;
if Task gool is achieved then break:
L
Compute evaluation results using the counter N and the query length, e.g., Tmaz H
Table 5: List of error types LLMs' responses
No.
Error Types
Task Understanding
(
Logic
History Understanding
(4) Physical Understanding
Mathematical Understanding
Missing Information
Examples of each error type from various LLMs responses:
Task Understanding
Llama3-8b replies-
"However, one could argue that the agent's velocity is already quite high ([0.0075)), and acceleroting further might not be necessary Additionally, the agent's position is still relatively far from the goal ([0.5729) 0.5), it may need to adjust ils stralegy soon."
counterproductive if it leads to overshooting the goal position.
A Analysis: In this example, the model makes multiple mistakes, it misunderstands the task as it thinks overshooting the goal would be bad, and it also wrongly interprets the state -0.5729 as 0.5729.
Logic
Llama3-70b replies-..
"Given the current state s20, it's unclear why the agent would choose to accelerate to the right, the position is still below the goal and the velocity is not high enough to guarantee reaching the goal."
Analysis: Here the reasoning logic does not make sense, If the goal is getting to the right, acceler- ating to the right should make sense especially the velocity is not high enough yet-
Llama3-8h replies
"This pattern suggests that the agent might be trying control the car's movement by alternating between accelerating left (action 0) and not accelerating (action 1) to navigate the valley towards the goal position on the right hill."
22

---

Published in Transactions on Machine Learning Research (12/2024)
Analysis: LLM argues that accelerating left will navigate the car toward the right.
History Understanding
GPT-3.5 replies, given history of actions: a20 = 2, 021 0, a22 = 0, 423 = 0
"The agent has been moving towards the right hill, and in the previous steps, has been accelerating to the right (action 2) to gain momenlun."
Analysis: The LLM claims the agent has been accelerating to the right using action 2 when the last three actions were action 0.
Physical Understanding
Llama3-70b replies-
"Given the agent's behaviour of accelerating to the right when the car is far from the right hill and not accelerating when it is closer, the choice of action 1 (no acceleration) in slate s18 seems reasonable. This action allows the car to conserve momentum and potentially reach the goal position more efficiently.."
Analysis: Here the LLM wrongly believes that not accelerating could allow the car to conserve momentum, however moving upwards gravity will cause the car to lose momentum should it stop accelerating.
GPT-3.5 replies
"The agent seems to be attempting to move towards the right hill by repeatedly accelerating to the left."
Analysis: LLM doesn't realize that accelerating the left won't move the car to the right.
Mathematical Understanding
Llama3-8b replies, given the state history $25 [-4.0838 x 10-1 1.7447 10-3], 926 [-4.0849 x 10-1 -1.020 10_ and the new state 827 = [-4.104 10-1 - 10-3]
"Looking al the sequence of states provided, the car is moving to the right (position is increasing) while the velocity is decreasing."
The LLM does not realize that the position is decreasing, moving to the left it wrongly interprets the numbers.
Analysis:
Missing Information
Llama3-8b replies
"The action of not accelerating might delay the agent's arrival at the goal position, especially when it is very close to the goal. It. is crucial for the agent to maintain its momentun and continae accelerating towards the goal to ménimize the time taken to reach the flag."
The car needs to accelerate to the left to get to a position from which it can build enough momentum towards the right overcome the right hill. The LLM missing the information about the environment that would allow it to understand this behaviour.
Analysis:
GPT-4o replies-
"Given the car's current state, taking action '0' (accelerate the left) might not be the most optimal choice. The car is already moving left with very small velocity, and further accelerating to the left might not significantly help in building the necessary momentum to move rightward. Instead, it might be more beneficial to start accelerating to the right (action 22-) to begin the process of climbing the right hill."
Analysis: The LLM is missing the exact information about how much velocity is needed at which position to accurately judge when the agent should switch back to accelerating rightward.
23

---

Published in Transactions on Machine Learning Research (12/2024)
F Additional Results: Evaluating LLMs' Understanding Across Various Tasks
F.1 State Element Prediction Accuracy with Increased History
In the task of predicting (full) states, we also plot the prediction accuracy for individual state clements and how they vary with increased history size for different tasks: Figure 9 for the Pendulum lask, Figure 10 for the Acrobot task, and Figure 11 for the LunarLander task.
GPT-3.5
Liaras3-00
LlarreG-70ee
66
State Cimment
State Clement
Stute E Exmmet
State Exmmet
State E Element
State Element
45
40
State Element0 1
25
GPT-3355 State Prediction
Llame3-BB State Prediction
Llama3-706
60
50
50
45
86
45
35
30
1.2
16
20
10
Figure 9: Dynamics LLMs' performance on predicting individual state element for the Pendulum task (with indexed history prompts).
GPT-3.5 Next State Preciction
Llama3-B0 3- Next State Prediction
Llama3-70b Next State Prediction
State El
65
State Elemen10 10
Shele Elerent
Skele Elanaet
45
45
8
Siate Eloment
40
Blale Element
40
State Element 5
88
Elmment
36
State
20
20
Slalc Eoment
10
10
20
10
20
GPT-3.5 Last State Prediction
Liarma3-8b Lasi State Prediction
Lluma3-706 - Last State Prediction
70
65
16
56
50
45
8
45
86
0.3
10
20
0 1.2.2.
10
20
1.2.3
10
20
History Sizs
Hastoy Size
History
Figure 10: Dynamics of LLMs' performance on predicting individual state element for the Acrobot task (with indexed history in prompts).
24

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5 Nexi State Prediction
Liamad-Bb Nee State Predicton
Limma3-70b Nexi State Prediction
100
80
Stlle Z
State Eeeent
State Bement 10
State E Fammrt
8
Slate Earrnt
60
State Elemett
Blult E
Btale Eenest
70
State Eleme112 12
State Element
60
60
State Element 0
Steie Elamerl 2
GPT-3.5 Leat Prediction
Llama3-8b
Prediction
Llama3-700 Last Prediction
100
100
100
00
80
30
20
History Size
Hiltory Bize
History Size
Dynamics of LLMs' performance on predicting individual state element for the LunarLander task (with indexed history in prompts).
Figure 11:
F.2 Average State Element Prediction Accuracy
In addition to reporting the dynamics of prediction accuracy for individual state elements, we report the averaged prediction accuracy for state elements in the MountainCar task (Figure 12), the Pendulum task (Figure 13), the Acrobot task (Figure 14), the LunarLander task (Figure 15), the FotchPickAndPlacee task (Figure 16), the FetchPush task (Figure 17), and the FetchSlide task (Figure 18). Refer to Table for the meanings of the state element indices in Fetch-series tasks.
We find that LLMs are slightly more sensitive to changes in angular velocity than angle, shown by the Pendulum and Acrobot results. In addition, LLMs show greater sensitivity to changes in position than linear velocity, as demonstrated by MountainCar and LunarLander results.
25

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-35
History
Hiskey
Next State Prediction
Averaged State Elemen Accurery
State Elamants
Last State Prediction
Roaraged Store Element Aucurecy (%)
Podton
Wiedl
Figure 12: LLMs' averaged performance in predicting individual state element for the MountainCar task.
The red dotted line represents the accuracy (33.33%) of randomly guessing state changes for each element (applicable to the following figures).
GPT-3-5
GPT-35 Hisoory indecad
Liemad-bb
Llamdd-Bb History Indexed
Uama337066
Liame3-706 History Indexed
Next State Prediction
Averaged Shote Elemen Accurecy
Shvte Elaments
Last State Prediction
Ausrapeed State Eleement Azcurscy (%)
Angular
Figure 13: LLMSs averaged performance in predicting individual state element for the Pendulum task.
26

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5
GPT-3.5
Liama--b
Llaned-BL History
Ulera3-706
Llamad-706 Hiskey
Next State Prediction
Averaged State Clement Acourery
State Elemonts
Last State Prediction
Avorgge: States Element Acsurety (N)
argh
arggl
arggl
2
vildy
Arga vilddy
Figure 14: LLMs' averaged performance in predicting individual state element for the Acrobot lask.
GPT-3.5
Liama3-8b
Liama3 History
Liama3-706
Llama3-
Prediction
Avergged Stotee Element Acturery
State Elementss
Last State Prediction
Aesrogee Sjate Elemmett Ausuracy (%6
poalion
wiliily
y ywbbdy
aegiee Elementss
inpdlr whkcly
kg
Figure 15: LLMs' averaged performance in predicting individual state element for the Lunar Landea task.
27

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5
GPT-3.5
Liama3-Bb
Liamad db History
Laana3-706
History Indexed
Next State Prediction
Averaged State Elemnent Aature Acturecy
State Elements
Last State Prediction
Avengged State Element Acturey (5)
Figure 16: LLMss averaged performance in predicting individual state element for the FetchPickAndPlace task. Refer to Table 6 for the meanings of the state clement indicos Fetch-scries tasks.
GPT-3.5 indexed
Lisma3-db
Liama3 Bb Indexed
Llama-To
Liam33 Indexed
Prediction
Averaged State Elemert Aupuracy
State Elements
Last State Prediction
Aeoragee Statts Element Acsurity (5)
Figure 17: LLMs' averaged performance in predicting individual state element for the FetchPush Push task.
28

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3-5
GPT-330 History indeced
Liema3-bb
Liamad db Hiletor Indewed
Uamaa3700
Liama3-70b History Indexed
Next State Prediction
Averagged Stote Element Accurncy
15
12
Elements
Last State Prediction
Averaged State Elemer Auorary
15
State Elermnnts
Figure 18: LLMss averaged performance in predicting individual state element for the FetchSlide task.
Table 6: List of state elements in Fetch-scries tasks
Indices
State Elements Meanings
0
End effector x position
1
End effector position
2
End effector position
3
Block position
4
Block y position
5
Block position
6
Relative block position with respect to gripper position
7
Relative block position with respect to gripper y position Relative block position with respect to gripper 2 position
8
9
Joint displacement the right gripper finger Joint displacement of the left gripper finger
10
11
Global x rotation of the block in XY Z Éuler frame rotation Global y rotation of the block in XYZ Euler frame rotation
12
13
Global rotation of the block in XYZ Fuler frame rotation Relative block velocity in direction with respect to the gripper Relative block velocity in direction with respect to the gripper Relative block velocity in 2 direction with respect to the gripper
14
15
16
17
Block angular velocity along the x axis Block angular velocity along the y axis Block angular velocity along the z axis
18
19
20
End effector velocity in direction
21
End effector velocity in direction
22
End eflector velocity in z direction
23
Right gripper finger velocity
24
Left gripper finger velocity
29

---

Published in Transactions on Machine Learning Research (12/2024)
F.3 Influential Factors in State Element Prediction
Figure 19 shows the linear relationship between stale span and prediction accuracy, while Figure 20 shows the corresponding residuals-fitted plot. Figure 21 illustrates the relationship between state variance and prodiction accuracy. Detailed state span (min-max) statistics for cach state element across all tasks are detailed in Figure 22. For each task in Figure 19 and Figure 21, we used the raw state span during analysis, as LLMS process the unscaled state values. The linear analysis was conducted between the raw state span and prediction accuracy. However, to compare relationships across tasks, we used a 0-1 dummy scale for plotting purposes, shown in Figure 7.
The linear relationship analysis for different models and tasks reveals key insights. For GPT-3.5 on the MountainCar task, the positive slope indicates that accuracy improves with span, and the perfect r² of shows a flawless fit, meaning the model's predictions align perfectly with the data. For GPT-3.5 on the Acrobot task, the moderate slope and r² of 0.58 suggest a decent correlation with some unexplained variability. In contrast, the Llama3-70b model on the Pendulum task shows a low or near-zero slope and an y² of O, indicating the model does not effectively explain the variability in the data for this task.
State element variance also influences dynamics prediction, being slightly more significant than state span in tasks like Pendulum and LunarLander, as shown in Figure 21. This offers a different perspective on how LLMs' dynamics understanding can be affected and generalized. Additionally, the linear relationship is often less significant for smaller models, suggesting that larger models are more sensitive to these changes. In general, from a statistical point of view, larger spans and greater state variance facilitate the LLMs' perception of changes due to larger value discrepancies.
GPT-3.5 Liner
GPT-3.5AC
GPT-3.5 Linee
GPT-3.5 FeecheNP
Fetthuuh Litrerr
GPT-3.5
45
Lineer Fi
Linoar Fii
Fil oay
DC Heaar Fii
Linear Fi
Acammy
60
50
35
60
30
Linear
40
0
40
40
0.0
0.5
1.0
5
10
15
10
1.0
0.4
Liarma3-Bo Eb MC
LiamaB Bb AC
Llamma-88 8p Pan
Liema3-8u 4b LL
Blama3-3b FatchPNP Linear A
Liame3-8b Fetee P Push Linear A 0.02
LiamaB-88 Bb FatchEllde Linear R 0.01
Linear Rf: 1.00
Linear Rf: 0.14
Linerr R 0,33
Liner R 0.40
50
100
en
UC
80
...ull
Ausrecy
42
45
60
40
GC
50
4c
40
20
40
20
20
36
Linear Fi:
36
Linser Fit
20
0.0
0.5
10
10
0.4
1emmc3-700 MC
Iamme3-70hh AC
Ulmma3-70b Pen Lincar R?: 0.00
Llams3-70h 700 LL Linear R
Llamm3-70e FetchiPPP Lincar R 0.47
Llmms3-70s FetchPPush Linear &:
1emma3-700 Feechölide Linear A: 0.08
Lincar RP:
Linear R 0.17
Lincar Fi
45.0
47.0
Lincar Fh:
100
......
Lincar FIt
80
.or.. Linear Fit
Bc
-.... Linear Fit
Linear F1
Acounit
42.5
46.5
6G
60
10
40.0
46.0
37.5
Lnner'1 Fit
20
20
0.0
0.5
1.0
5
10
15
10
15
0.5
1.0
1.5
0.2
0.4
2
Spon
Spen
Span
Span
Span
Span
Spon
Figure 19: The linear relationship between state span and state element prediction accuracy across all LLMs and tasks. Each blue dot represents the average prediction accuracy for the corresponding state element. "MC" stands for MountainCar, "AC" for Acrobot, "Pen" for Pedulum, "LL" " for LunarLander, "FetchPNP" for Fetch Pick AndPaaee.
30

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5 MC
AC
GPT-335 Pen
GPT-3.5LLL
GPT-3355 FetoiPNPP
GPT-336 FanchPush
GPT-3.5 Fatchölde
0.04
40
Reckboale
0.02
20
20
0.00
20
-0.02
-5
-0.04
-10
25
50
75
50
40
45
50
75
30
40
50
45
44
45
50
56
Liama3-00
Lame3-00 Pen
Liams3-bb LL
Liamad-8b FetchFNP
Llamaã-do PetchPush
Ultmm3-tio FelchtSiide
5.0
0.02
2.5
20
10
20
0.00
0.0
0
-0.02
-2.5
-0.04
-10
-20
42.5 45.0 47.5
25
30
35
Liams--7b MC
Liamad-- AC
Liama 70b Pcn
Liama3- 70b LL
Lloma3 70b FotchPPP
Liamaa 70e FetchPushh
Llamad /0b FelchSlide
Reeidalt
0.04
0.02
0.5
20
22
40
20
0.0
-0.02
-10
9C Predicled Vauus
25
45.0 Predicted Values
62.5
46 Predicied Values
46
50
75
30 Prodicted Valuea
40
60
47.5 50.0 52.5 Predicted Values
20 Predicted Values
40
Predicted Valuea
Figure 20: The residuals versus fitted values across all LLMs and tasks. Randomly distributed residuals around zero indicate that the linear models are appropriate for the task data.
GPT-3-5 MC
GPT-3 AC Lirean ear R?, 0.68
GPT-3.5 Pen Li Lineer R?, 0.45
GPT-3.5 LL
GPT-3.5 FetchPNP
GPT-3.5 FetchPush Linear R?, 0.00
GPT-3.5 FetchSide
Linear A?
1.00
Linear R 0.70
Linsar R?, 0.71
Linear R 0.01
100
Linear
Linear
Linear
80
40
80
Acourty
40
30
40
00
40
60
60
25
Linear Fit
36
Linear Fit
30
Br Linear Fit
1100 L Lineer Fit
20
0.2
0.4
2
0.5
1.5
1.5
0.2
G.4
0.0
0.1
0.00
0.25
0.50
0.00
0.25
0.52
Liams--b R?:
Llmma3-ib Linear RR:
Liama3-00 eb Uneer
Liama3 FetchPPP Linear
Laam3-BB Bb Fetch Push Linear
Liama3-8b Felchitide
Unner
50
Linear Fi
Linear Fil
80
Linear Fi
Actomna
10
60
40
40
38
20
20
20
Linear
Linear
Linear
20
0.5
0.1
0.00
0.00
0.25
LismaB-70a MC
Liama3- 70b AC
Llamaa-706 Pen
Liama3-70h LL Linear 0.85
Llamma370B FatchPPPP Linsar 0.50
L Liama3-700 70th FatchPush Linear R²: 0.04
Lioma3: 70h FeschSiid
Linear R²: 1.00
Linear R 0.14
Linear R 0.00
Linear R²: R00000
47.0
100
45.0
80
Linwer
80
Ascuncy
00
42.5
40.0
46.0
00
80
40
40
20
37.5
..... Linear Fit
46.5
20
Lincar Fit
20
0.4
2 Variance
0.5 Varlance
0.2 Verlance
0,1 Verlance
0.00
0.25
0.50
Variance
Varianca
Variance
Figure 21: The linear relationship between state variance and state element prediction accuracy across all LLMs and tasks.
31

---

Published in Transactions on Machine Learning Research (12/2024)
MountannCar Min and Mix Valuee for Each State Element
Acrobot Min and Max Values for Lach State Element
025
Vax
Max
Valie
-0.0
-10
States Dimensions
States Dimensions
Pendulum
Min Values Each State Clement
LunorLanderr
and Max Values for Each Clement
Mir
Min
6.
1.0
255
0.0
800
States
FetchPick/onPPlse
Min and Max Values for Each State Element
IetthhPush
Min and Max: Values for Each State Element
20
Mn
Max
1.5
Max
0.
-1.
States Dimeraions
FalchSSice - Min and Max Valuess for Each State Element
Max
States Dimensions
Figure 22: Minimum and maximum values of state elements across all tasks.
F.4 Average Action Element Prediction Accuracy
It is observed that all LLMs fall short of scoring reasonable accuracy in predicting 4-dim actions for the Fetch-series lasks compared to the four classic control lasks. Therefore we report the averaged prediction accuracy for individual action elements in the FetchPickAndPlace (Figure 23), the FutchPush (Figure 24), and the FetchSlide tasks (Figure 25).
32

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-35
Hiletry
History
Next Action Prediction
Anstaqued Action Elemet Acouracy
Last Action Prediction
Arergged Aution Element Accornay
10
Action
Figure 23: LLMs' averaged performance in predicting individual action element for the FetchPickAndPlacr lask. "0" represents x displacement, "1" displacement, and "2" z displacement, "3" gripper displacement.
The red dotted line represents the accuracy (10%) randomly guessing action bins (total 10) for each element, applicable to the following figures.
GPT-3.5
Liama3-8b
Liama3 db History
Liama3-70e
Action Prediction
20.0
Averaged Addon Dlemen Aocuray
17.6
125
76
60
Action Elements
Last Action Prediction
Averoged Acurascy
Figure 24: LLMs' averaged performance in predicting individual action element for the FetchPush task. "0" represents x displacement, "1" y displacement, and "2" z displacement, "3" gripper displacement.
33

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5
GPT-3.5 Hisooyy indoxed
Lisma3-a 8b
Gi Liama3 Bb History Indexed
Uama3-706
Liama3 70b History Indexed
Next Action Prediction
20.0
Avwraged Aclon Element Accunncy
15.
129
160
60
Astion Elements
Last Prediction
Rveraged Aation Element
15
Elements
Figure 25: LLMs' averaged performance in prodicting individual action clement for the FetchSlide task. "0" represents displacement, "1" displacement, and "2" displacement, "3" gripper displacement.
F.5 Average Comparison of Model Predictions
Table displays the average accuracy of LLMs' predictions regarding the agent's behaviour and the resulting state changes in Mountainear, Acrobot, and Pendulum tasks.
MountainCar
Arrobot
Pendulum
GPT-3.5
Llama3-8b
Llama3-70b
GPT-3.5
Llama3-8b
Llama3-70b
GPT-3.5
Llama3-8h
Llama3-70h
Pred.
74.60%
39.10%
18%
43.94%
46.29%
15%
17.08%
3.72%
11.77%
81.48%
68.63%
87.06%
46.36%
44.95%
64.73%
17.49%
3.51%
12.42%
LA Pred.
76.73%
61.83%
78.87%
39.25%
47.24%
55.32%
14.28%
1.58%
14.02%
80.06%
73.90%
76.85%
44.62%
42.35%
55.40%
20.63%
1.89%
13.86%
33.43%
30.81%
37.04%
0.30%
0.26%
0.13%
9.52%
8.34%
7.61%
NS Pred.
37.41%
33.65%
40.68%
0.00%
0.42%
0.43%
7.89%
6.65%
5.49%
31.97%
22.12%
29.32%
1.14%
2.95%
1. 1.69%
6.46%
10.54%
10.22%
LS Pred.
32.41%
22.45%
35.25%
0.61%
2.32%
2.87%
5.41%
8.27%
7.64%
Table 7: Comparison of model predictions with and w/o indexed history. Light grey cells show results with indexed history. NA Pred. = Next Action Prediction; LA Pred. = Last Action Prediction; NS Pred. = Next State Prediction; Pred. Last State Prediction.
F.6 Dynamic Performance of All Evaluation Metrics
The dynamics of LLMs' understanding performance (with and without indexed history in prompts) with increasing history size for the MountainCar task (Figure 26), the Acrobot task (Figure 27), the Pendulum task (Figure 28 and Figure 29), and the LunarLander task (Figure 30).
Among all results, is observed that models' understandigg agent behaviour improves significantly with smaller history sizes but does not increase further with larger histories. In some cases, like with Llama3-70b,
34

---

Published in Transactions on Machine Learning Research (12/2024)
it may degrade. Generally, model performance in action prediction tends increase and then likely saturate as the history size grows. The optimal history size for each language model is detailed in Table 8.
In complex tasks like Acrobot, history size has less impact on model performance in state prediction. We hypothesise that this is due to the complex relationships in the interaction data, where adding more history does not enhance the LLMs' understanding of the environment dynamics. For moderately complex tasks (e.g., Pendulum), model performance initially increases with small history size, consistent with our earlier finding for predicting actions. This is demonstrated in the third column of Figure 28.
GPT-3.5 Nexd Action Prediction
GPT-3.5 Last Action Prediction
GPT-3.5 Next State Prediction
GPT-3.5 Last State Prediction
90
45.0
80
42.5
Salthing Rorn (%)
85
40.0
80
36.0
75
32.5
25
30.0
65
20
0.1
20
0.1
20
0.1
10
20
012
10
20
Llama3-8b
Next Action Prediction
Llama3-8b Last Action Prediction
Liama3-8b
Next Stato Prodiction
Llama3-88
State Prodiction
Hutury
80
40
Meicing Rain (%)
80
36
25
80
85
25
10
0123
20
012
10
20
01233
10
20
0123
10
20
Liaama-770-
Llama3-70b
Prediction
Llame3-70b Next State Prediction
Llama3-70b Last Prediction
Hutery
50
46
Hutory
Mencing Rate (%)
80
70
36
30
26
20
0123
20
0.1
20
012
10
20
0
10
20
History Size
Hastory Size
History Size
History Size
The dynamics of LLMs' understanding performance with increasing history size for the Moun- tainCar task.
Figure 26:
Model
MountainCar
Acrobot
Pendulum
LunarLander
GPTT-355
5
Llama3-8b
20
5
Llama3-70b
0
10
GPT-4o
1
Table 8: Optimal history size for each language model, evaluated across four metrics for individual tasks.
35

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-335 Next Action Prediction
GPT-3.3 - Aation Prediction
GPT-3.5
Next State Prediction
GPT-335 Last State Prediction
1.50
Hiltay inceedd
1.25
Hirtyr deod
50
1.09
15
15
0.75
40
0.53
8
35
0.25
0:0
C128
10
0.123
5
10
0123
0122
10
Llama3-80 Next Action Prediction
Ulmma--b- - Last Action Prediction
Llama3-8h Next State Prodiction
Llama3-ab
Last State Pradiction
2.5
Hiltry Inccod
56
1.5
35
0.5
c1233
01223
0123
0123
Llama3 3-b
Next Action Prediction
Llama3-706
Last Action Prediction
Ulmma3-7bb Next State Prediction
Larma3-70b
Last, State Prediction
75
Hadry Inceued
(8)
6
8
8
50
8
45
0.5
40
0.3
35
History Sizo
History Bizo
History Size
History Size
Figure 27: The dynamics of LLMs' understanding performance with increasing history size for the Acrobot task.
GPT-3.5 Next Action Prediction
GPT-3.5 Last Action Prediction
GPT-3.5 - Next State Prediction
GPT-3.5 Last State Prediction
Hkatayy Incead
30
Batry ddead
153)
20
20
15
10
10
011
20
0.123
0
1
Llarma3-B Sb
Next Action Prediction
LIamm3--ae Sb Last Action Prediction
Liama3-88 Bb
Next State Prediction
Liamatl ab
Last State Prediction
History Inceand
20
20
10
0120
Liamme3-70 Next Action Prediction
Llama3-70p
Lest Action Prediction
Llama3-70b
Nexi State Prediction
Llarme--0b
State Prediction
18
17.5
Haleyy treuwve
15.0
40
16
12.5
ts
14
160
12
7.6
10
5.
2.5
Hissory Site
History Stze
History Size
Hottoy Size
Figure 28: The dynamics of LLMs' understanding performance with increasing history size for the Pendu- lum task with discretized actions in evaluation prompls.
36

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-3.5 Next Action Prediction
GPT-3.5 Last Action Prodiction
GPT-3.5 Next State Prediction
GPT-3.5 Last State Prodiction
Histay Inceed
Hirty deod
koadingg
-0
30
18
20
012
20
0.12
0.123 6
0.1
10
Liomme-So- Next Action Prediction
Llama3-80 8-80--8----- - Leel Action Prediction
Llema3-86 Next State Prodiction
Liama3-ab
Last State Prodiction
Hilery tnccoee
20
10
10
0122
21233
Liams3- 70b
Next Action Prediction
Llama3-70p
Last Action Prediction
Larma3-70b
Next State Prediction
Llarra3-70b Last State Prediction
17.5
Haary Incaudd
15.0
(N)
8
60
12.5
104
18
30
7.5
8
20
50
2.5
History Size
History Slbo
History Size
History Size
Figure 29: The dynamics of LLMs' understanding performance with increasing history size for the Pendu- lum task with continuous actions in evaluation prompts.
GPT-3.5 Next Aztin Prediction
GPT-3.5 - Action
Action Prediction
GPT-3.5 Next State Prediction
GPT-3.5 3.5 Last State Prediction
70
15
Hi Inceadd lattry deead
45
40
2
10
0.22
012.3
0123
Liama3-Bb Bb Next Action Prodidiion
Llama3-8t- Last Action Prediction
Llama3-8b
Next State Prediction
Liama3- Bb
Last State Prediction
Hitkry Inceadd
60
50
6
Llamm33770
Next Action Prediction
Llama3-70b Last Action Prediction
Llamu3-70b- Next State Prediction
Llama3-70b Last State Prediction
23
Huleyy Ircewdd
60
67.5
52.5
80
80.0
45.0
History Size
History Size
History Size
History Size
Figure 30: The dynamics of LLMs' understanding performance with increasing history size for the Lu- narLander task.
37

---

Published in Transactions on Machine Learning Research (12/2024)
F.6.1 Comparative performance of models on predicting discretized actions
Contimuing from the plot of LLMs' performance on the Pendulum task with continuous actions (third row of Figure 4 in the main text), Figure 31 presents a comparative plot of LLMs' performance on the Pendulum task with discretized actions. The results show that LLMs are generally less effective classifying action bins compared to regressing on absolute action values.
Next Action Prediction
Lest Action Prediction
Next Prediction
State Prediction
35
17.5
GPT-3.5
Maidiing Rone (%)
30
Naithing Rato (80
Rate
15
Salle Maldiag Rate (5p
15.0
12.5
Liame3-700
Uame3-Bb
10
10.0
15
7.5
10
0.0
0123
5
20
012
10
0.123
5
History
History
History Size
History
Figure 31: A comparative plot of LLMs' performance the Pendulum task with discretized actions, following the plot of predicting continuous actions (third row of Figure 4 in the main text).
F.6.2 Comparative performance of models on Fetch-series tasks
Contimuing from Figure 4 in the main text, Figure 32 presents the performance of LLMs the Fetch tasks, characterized by continuous actions (d4 4) and high-dimensional state space (ds = 25). The results show that LLMs struggle with state predictions, achieving zero accuracy with very large state elements.
Action Prediction
Last Prediction
State Prediction
Last Prediction
0.34
0.04
Mastieg
0.02
0.02
0.00
0.00
-0.09
-0.04
04
Action Prediction
Last Prediction
State Prediction
Last Prediction
1.5
0.34
1.0
1.0
0.02
0.5
86
0.00
0.0
5.0
-0.02
02
-3.5
6.04
Action Prediction
Last Prediction
State Prediction
Laei Stare Prediction
0.04
0.04
Madcing Bata (S)
0.02
0.00
History Bise
History Size
History Bize
History Size
Figure 32: Comparative plots of LLMs' performance on the Fetch-series task, following the plots in Figure 4 in the main text).
38

---

Published in Transactions on Machine Learning Research (12/2024)
F.7 Relative Improvement and Worsening Rates with Increased History Size
When increasing history size in prompls for LLMs, observed that previously incorrect predictions at certain time steps may become correct, and vice versa (see Figure 33 for an example with the MountainCar task). We define the relative improvement rate for a history size H > 0 compared to fixed history size O cooreet in and the relative worsening rate for the same H as #[corect-+icooredll >H.
#incorrecto
#oorrecto
As shown in Figure 33, incrcasing history size does not impact steps near the episode end, where RL agents consistently accelerate to the right to cross the finish line. In contrast, LLMs, likely influenced by pre-training on human preferences, suggest more cautious actions, such as accelerating left to avoid overshooting, This safcty bias, ingrained during pre-training, significantly affects their action predictions regardless of history length. Thus, LLM prediction failures are not solely due difficulty processing extensive histories; their inherent beliefs about agent behaviour and environment context also impair prediction accuracy in certain tasks.
Further, by examining wise mismatches in the Acrobot task, it is observed that LLMs often justify their predictions mainly on the rewards received in previous steps, neglecting history patterns, especially when the history is short.
Figure 34 and Figure 35 illustrate how these improvement/ wereeemeen//wwseeeeeee /worsening rates change as history size increases. Analysing the impact of longer history correcting previous incorrect predictions when H is fixed at 0 shows that increasing history size improves LLMs' understanding of agent behaviour initially, but further increases may degrade it Conversely, increasing history size has negligible influence on dynamics predictions.
MountainCar Next Action Prediction
History Size
Mismalch Count
3
4
20
3
2
20
40
60
80
100
Step Indices
Figure 33: Step-wise mismatches with increasing history size for the GPT-335 model on the Mountain Car task (the same mismatch pattern is observed across different episodes). Each bar (corresponding to different history size) cach step index indicates mismatch between the LLMs' prediction (based that history size) and the action laken by the RL agent al that step.
At certain steps (e.g., around 40), the model with h 0 fails, while the model with h > succeeds. In contrast, at other steps (c.g., at the last steps), increasing history size minimally affects the LLMs² predictions.
39

---

Published in Transactions on Machine Learning Research (12/2024)
GPT-316
Liamed-50 --
Limma-7700
MountainCar Next/ctiio Prediction
MountainCar
Last Acton Prediction
MountainCar
Next State Prediction
MountainCar
Last State Prediction
improvmton
0.8
0.5
0.5
0.4
0.6
0.4
03
02
0.2
0.1
1...3
1.33
12.3
12.3
Acnobot
Next Prodiction
Acrobot
Last Action Prodiction
Acrobot
Next State Prediction
Acrobat
Last Prediction
0.5
0.100
impovvmentt
0.125
0.3
0.000
0.028
0.008
0.000
0.000
1.2
1.2
1
10
1
10
Pendulum
Next Adtion Prediction
Pendulum - Last Action Predidion
Pendulum Next State Prediction
Penduuuu Last State Prodiction
0.
improveent
0.0
0.15
0.4
0.3
04
0.10
0.10
0.2
0.00
0.06
01
0.00
Action Prediction
Lunart
Last Prediction
Lunael LandeLander
State Prediction
Lunarl Prediction
0.10
0.5
Rekatve Impovarent
0.5
0.4
0.06
0.
02
123
10
1223
5
10
12.
10
20
125
10
20
History
History
Hatry
Figure 34: The dynamics of relative improvement with increasing history size across all LLMs and tasks (with indexed history in prompts). The first two columns represent agent understanding; the last two columns represent dynamics understanding.
F.8
Ablation Study
F.8.1
Comparison of models without using task dynamics
Figure 36 illustrates the performance variation when dynamics equations are excluded from the prompts.
F.8.2
Comparison of models without using task instructions
Akin to prior works by Mishra et al. (2022); Le Scao & Rush (2021), which show that task framing in prompt influences language models, we observe a similar effect. When removing task instruction from evaluation prompts, models' understanding performance across the majority of evaluation metrics is significantly de- grading, as demonstrated context. (i.e., sequence of numerical values) remaining unchanged. We hypothesise that LLMs ability to mental model agents is enhanced by more informative context.
MountainCar (Figure 37) and Acrobot (Figure 38) tasks; despite the history
F.8.3 Comparison of models: action bins vs. absolute values prediction
Figure 39 presents the evaluation results of LLMs on Pendulum tasks, comparing predictions of action bins (the first three rows) with predictions of absolute action values (the last three rows).
40

---

Published in Transactions on Machine Learning Research (12/2024)
Mourtaan Nood Action Prodictin
Meurral nCur
Lead Action Prudicion
HourmrnnCar Nood State Prudiction
Mounta Lund Stata Prudicion
0.30
200
065
200
Acrobot heod Aot Prodiction
Acrocct Last Action Prodicioon
Acrobot - Noxt State Prodiction
Acrobot Lost Stane Prodiction
0.5
12
06
6.5
12.3
1.2
Paxcuuuu
Nox Action Prodiction
Pondulurt Last Action Procictio
Pondulum Noxt Blate Prodiction
Pendulum Last Blate Prodict on
Dhibb Wateee))
02
0.
123
1:33
1:3
LunorLanddrr
NewtArtion Prediction
LunorLancer
Last Action Frediction
aurart.ander
Next State Prediction
Lunorl.ander
Last State Frediciion
64
123
122
123
Figure 35: The dynamics of relative worsening with increasing history size across all LLMs and tasks (with indered history in prompts).
Next Action Predidin
Last Action Preciction
Neot State "eeckdan
Last State Precidion
0m131
73
C125
19
0.1.3
0123
0.133
35
22
0.133
01:33
0123
AextAction Picdition
Last Aation Prodiction
Next State Prediction
Last State Prockdion
Limad--0 Dyra.
8
20
50
35
0120
Figure 36: Comparative plots of LLMs' performance on MountainCar with different history sizes (with indexed history in prompts). The suffix of model names "No Dyna." indicates not using dynamics equations in prompts.
41

---

Published in Transactions on Machine Learning Research (12/2024)
Most
Action Prediction
Noxt State Prodiction
Prediction
82
27.6
960
28
2541
01233
2123
D1223
0122
Next Prediction
Last Action Prediction
Next State Prodiction
State Prediction
Llmol
Liared
35
03
30
20
0125
10
20
2
0.2
90
0 2
Nexl Action Preticton
Last Action Prediction
Neat Stade Pred diion
Lusi State Prediction
Lberad--T
20
45
Lberad--T
42
55
58
23
0123
0123
20
01233
10
0.123
10
Figure 37: Comparative plots of LLMs' performance on MountainCar with different history sizes (with indered hislory in prompts). The sulfix of model names "No Inst." indicates not using task description in prompts.
42

---

Published in Transactions on Machine Learning Research (12/2024)
Prediction
Last Action Prediction
Next State Prediction
Prediction
1.0
2.0
2.5
01233
0.123
Nest Action Prodiction
Last Action Prodiction
Next State Prediction
Lost Stane Prediction
Llmrl
Liamedk
2.0
1.5
30
2.0
0123
0:1.3
0:225
5
15
012
Next, Action Prediction
Lusd Ausion Predicion
Next Stute Prediction
Last Stark Predictio
Libraa--TE
Liwradt-TE
(R)
2.0
35
1.
30
2.5
8.0
0123
0.223
26
0:223
12
20
0122
12
Figure 38: Comparative plots of LLMs' performance on Acrobot with different history sizes (with indezed history in prompts). The suflix of model names "No Inst." indicates not using task description in prompts.
43

---

Published in Transactions on Machine Learning Research (12/2024)
Noxt Action Prediction
Prediction
Next Prodiction
Lost State Prediction
10.
30
10
Milding Rud 500
12.5
90.2
20
7.5
15
5.
2.5
0.2
Nest Aution Prediction
Last Action Prediction
Next State Prediction
Last Stute Prediction
Limmdd db
30
Lianad da No irst.
16
Hatring Ruts 18)
15
20
0:123
Next Prodiction
Prodiction
Lirrao No Inst
Macoing Rate (h)
25
2010
17.6
15.0
12.6
0-
2
10
6
10
012
20
Noxt Action Prediction
Action Prediction
Next State Prediction
Last State Prediction
16.
10
15.2
Nolding Rub N0
40
50
41.2
7.5
30
5.
25
70
0.2
Next Aution Prediction
Last Action Prediction
Next State Prediction
Last State Prediction
Ulmed db
20
Lianad da No Inst.
Harring Serriag Runo
15
20
8
88
0.123
Prediction
Prediction
Utmra No
50
Maszking Rate mh3
36
0123
10
20
013
10
0.2.3
0.123 3
20
Figure 39: Comparative plots of LLMs' performance Pendulum with different history sizes (with indexed history in prompts). First. three rows: predicting action bins; Last three rows: predicling absolule action values.
44

---

Published in Transactions on Machine Learning Research (12/2024)
G LLMs Erroneous Responses in MountainCar Task
task. Even the RL agent occasionally makes such mistakes.
Explanations of Various Error Types in LLMs Reasoning. A manual review of the MountainCan task across three GLMMs-GT--355 Llama3-8b, and 1lama3-70b-revealle significant differences in their explanations that were not necessarily anticipated from the quantitative analysis. Table 5 provides an overview of the types and Table for counts of errors found in cach model. During the evaluation, a single response could contain multiple error types. Despile Llama3-8b producing the shortest responses, it also had the highest error count.
Aside from the errors, GPT-3.5 demonstrated bet- ter understanding of the task, often referring to the need to accelerate left to gain momentum for climbing the right hill. This was rarely mentioned by Llama3- 70b and never by Llama3-8b, indicating GPT-3.5's superior task comprehension and explanatory ability. Llama3-700 however, had an advantage in maintain- ing coherence, as it was less likely to contradict its arguments, unlike GPT-3.5, which occasionally ar- gued against an action before ultimately support- ing it. Both GPT-3.55 and Llama3-8h also displayed misunderstandings of the actions, such as incorrectly defining "action 0 (no acceleration)". This suggests common-sense bias toward interpreting as no ac- tion. Llama3-70b was betler at relaining the task description in memory.
(1) The first type of error, understanding the task, appeared frequently when the LLMs had to evaluate a proposed action, such as no accelcration the MountainCar task. All three models lended to be concerned about overshooting the goal of reaching position of >= 0.5. However, in this task, overshoot- ing is irrelevant since the goal is to surpass 0.5. Similar replies across models suggest this mistake stems from a shared common- sense notion. Additionally, Llama3-8b often failed to recognise the presence of hill on the left side.
G.1 A Compact Analysis of Error Types
Table I shows quantitative analysis of the frequency of different error types committed by the LLMs for the MountainCar task. The evaluation highlighted various types of errors (see Table the Appendix), with Llama3-8h displaying the most errors despite its shorter responses. A common error among all mod- cls misinterpreting the goal of the task, reflect- ing a shared common sense misundersianding Logi- cal errors, particularly in oscillation movements, were prevalent in GPT-3.5 and Llama3-70b, while Llama3- 8b frequently produced paradoxical replics. Misun- derstanding the task history and physical principles was rare but present. Mathematical errors, especially disregarding the minus sign, occasionally impacted reasoning. Notably, GPT-3.5 demonstrated a bet- ter task understanding by referring to momentum strategies in the task, an insight less frequently or never mentioned by Llama3-700 70b and Llama3-8b re- spectively. Llama3-700 did have one other advantage over other models as it was less often confused by its argument and excelled in maintaining task descrip- tions. Despite occasional errors in defining actions, GPT-3.5's superior comprehension of the task con- tributed to its higher-quality explanations.
(2) Logical mistakes noted in GPT-3.5 and Llama3-70b when the LLMs justified moving left without recognising the need for oscilla- tion to gain momentum, leading to paradox- ical replies. These types of errors were more prevalent in Llama3-bb.
(3) Misunderstanding the history refers to the occasional misinterpretation or incorrect rep- ctition of the history provided to the LLMs.
(4) Physical misudderstanding, though rare, in- volved incorrect responses regarding the ef- fects acceleration on velocity and similar cases.
(5) Mathematical errors commonly involved the LLMS disregarding the minus sign, leading them to believe that -0.5 is closer to 0.5 than 0.3. Although these mistakes led to awkward reasoning, they seldom significantly worsened the final decision.
Guideline on Utilising LLMs for Agent Mental Modelling
(6) A common and human-like error involved judging when switch directions to either gain or use momentum in the MountainCar
LLMs are not yet capable of robust mental modelling. Below are some concrete suggestions for LLM and RL
practitioners:
45

---

Published in Transactions on Machine Learning Research (12/2024)
Exercise caution. In general, LLMs are prone to hallucination, and their reasoning should only be trusted after expert auditing to ensure a high rating of reliability/trusee /ttostttttoooiinss..... It is crucial to explore ethical metrics to quantify Lrustworthiness.
Consider multimodal input. While our evalua- tion used only textual imput, practitioners may sider incorporating multimodal environmental data (e.g,, visual or sensory inputs), which could poten- tially enhance LLMs' understanding of behaviour and dynamics, leading to more trustworthy reasoning.
Incorporate domain-spccific RL data. LLMs could benefit from action- and dynamics-related data trajectories) during pre- training post-training stages to improve their abil- ity to model RL agents. Practitioners may consider this practice when LLMs fall short of their expecta- tion in mental modelling tasks in their domain.
(e.g., sequential/tempont
Explore new learning paradigms. Theoreticians may consider developing new learning paradigms that enables LLMs to learn and align with causo-effect lationships in input-output data, which could help LLMs perform agent mental modelling in a more trustworthy manner for a wider range of complex and unseen tasks.
46


Reviews:
[{"review_id": "KWknM2G8tY", "reviewer": "Reviewer_au9G", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\nThis work addresses a relatively unexplored area of agent mental modeling using LLMs.\nThe study is conducted on a diverse set of RL tasks, ranging from classic control problems to more complex robotic manipulation scenarios. \n\n**Weaknesses:**\n\nConceptually, would there be a need for agent mental modeling if the RL agent itself is based on an LLM that already generates human-readable text for its actions? Addressing this aspect seems to be overlooked in the paper.\n\nPredicting the next state or action does not necessarily reflect a true understanding of the agent mental model or underlying dynamics of the environment. A model might achieve accurate predictions for the wrong reasons. \nFor the mental modeling to be meaningful, it is crucial not only to predict states/actions, but also to provide interpretable and usable reasoning. Without understanding the reasons behind predictions, it seems like we are back at square one.\n\nOne limitation of this work is primarily focusing on pre-existing tasks which limits the evaluation scope and could impact the results as LLMs could have gotten biased during pretraining with some knowledge about these tasks or optimal behavior. It would be interesting to see a similar analysis on new sets of tasks where LLMs could potentially exhibit less bias.\n\nThe introduction does not motivate why LLMs would be particularly suited for elucidating RL agent behavior beyond their ability to generate human-readable text.\n\nError type (1), \"task understanding,\" is rather broad. How is task understanding evaluated in this context?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates the capacity of large language models (LLMs) to construct mental models of RL agents, which is formed as “agent mental modeling”. To take advantage of LLMs for this agent behavior modeling is a novel direction. The study looks at how well LLMs can predict an agent’s actions and their effects on the environment states within a markov decision process setup.\n\nThe authors prompt LLAMA3 8B, 70B instruct and GPT 3.5 with histories of actions and states, to predict agents’ actions and environment states as a proxy for understanding environment dynamics. This is done on a range of tasks including MountainCar, Acrobot, LunarLander, Pendulum, FetchPush, FetchSlide and FetchPickAndPlace.\n\nThe paper provides insights into LLMs’ capabilities and limitations in building a mental model of agents. They show that LLMs exhibit some promise in constructing mental models of agents. However, their performance is notably and inconsistently affected by factors such as history length, task complexity, pretraining bias and input data format.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "In section 3, it would be helpful to change “last” to previous” as “last” could be interpreted as the final step or action, potentially confusing the reader.\n\nIn section 3, could you clarify “Judging the next action that is given”? Why $y=0$ when there is an agreement with action $a_{t+1}$?\n\nThe paper could benefit from further addressing how predictions of actions/states contribute to understanding of an agent's mental model and environment dynamics.\n\nIn Figure 4, it would be helpful to put task names on the right side of the rows.\n\nFor the LunarLander task, do you have any intuition as to why LLAMA3 8B outperforms the 70B model significantly? \n\nThe term \"matching rate\" should be clearly defined in the paper. It can be characterized as the prediction accuracy.\n\nFigure 5 text could be bigger. There seems to be enough space.\n\nTable 1 provides a nice analysis. To improve its readability, there is sufficient space to include both the error type names and numbers in the first column.", "broader_impact_concerns": "There are not any particular concerns on the ethical implications of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "MLnd9bGiVN", "reviewer": "Reviewer_fLL5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Positive:\n- relatively extensive empirical study\n- a potentially high impact and (to the best of my knowledge) original work helping to understand how LLMs can model RL agents\n\n\nNegative:\n\n- not too many technical insights following from this paper\n- limited technical descriptions\n- limited literature study, not so much in reference count, but more in terms of the explanation and discussion.\n- lack of detail in some of the explanations, for instance, the bias issues.\n- The large list of results makes it difficult to assess their correctness and as the analyses are rather short and do not stay close to the data, they take a relatively big leap.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates the capability of LLMs to reason about an agent’s behaviour and its effect on states via in-context learning. A large empirical study is conducted, including various evaluation metrics. The results show that LLMs are not yet capable of fully mentally modelling agents, and the results highlight some bias issues, impact of parameters such as history size, task instruction, and data format. There is also a comparison to human evaluations for comparison purposes.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The literature and technical descriptions should be extended where possible, as the paper does not give much context and technical understanding currently.\n\nIn the results 5.1, it is stated that the LLMs can accurately predict agent behaviours, “surpassing the random guess baseline”. This is a relatively low standard for accuracy. This statement can be adapted.\n\nThe confirmation bias in LLMs section seems somewhat speculative. Maybe the authors can try to reduce the gap between the hypothesis and the data (which is only performance data). In other words, there are a few hypotheses, but there are currently no tests for them, any of which may improve the paper.", "broader_impact_concerns": "The algorithm used does not pose any specific ethical concern other than those already implied by LLMs and data analysis techniques.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rusNEBmbGf", "reviewer": "Reviewer_Eu2S", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper studies an important problem for RL practitioners, which is motivated by the emergence of LLMs. Given LLMs can build powerful internal models of the world, it is logical to wonder whether they can help in the mental modelling of RL agents. This work makes a concrete step in this direction.\n- The authors propose a detailed evaluation framework with a clear template. It does not only consider actions understanding but also dynamics understanding. Given the LLM is asked to provide information about its reasoning process, this allows to get a glimpse into the mental model of the LL  with respect to the RL Task at hand.\n- The experimental study is detailed and involves several tasks (from easy to very hard ones). The task selection is interesting because an LLM could in principle use common-sense knowledge (built from vast knowledge bases) to reason about the tasks. The authors analyze various aspects (impact of history length, data format, data span and data variance,  impact of indexing, error types etc.). The empirical evaluation involved significant manual/human work, especially as far as the error type analysis is concerned.\n- The classification of error types is interesting and sheds light on the internal reasoning inconsistencies of LLMs, and how these can affect RL agents.\n\nWeaknesses\n- I found it surprising and even confusing that the authors do not present extensive empirical results with GPT-4o, as they do with the other three LLMs. Table 1 already hints at the ability of GPT-4o to significantly reduce the total number of errors. It is well understood in the LLM literature that recent models (usually but not always with more parameters) can outperform older models. I was thus not able to understand why the authors performed such an extensive study using the three LLMs but omitting GPT-4o (with the exception of the experiment on error types, where GPT-4o was included). To me this is not simply a question of not including the most recent LLMs; the bigger question is whether the more recent LLMs would corroborate the current findings, or whether they would show much better performance. Is it for instance possible that the recent state=of-the-art LL:s could be much more successful in the mental modeling of RL agents, given they may have much better models of the world, common-sense facts/reasoning, etc.?\n- The evaluation framework and, in particular, the action understanding and dynamics understanding presented in this work seem to make more sense for deterministic MDPs.\n    - If a given action at a given state results in a distribution of next states (as opposed to a unique next state), then the LLM may fail to predict the next state, even if it outputs a valid state, i.e., a state appearing in the distribution over next states.\n    - If there are multiple actions resulting in a given next state from a given start state, then the LLM may output any of them and still be correct.\n- How to deal with the above point is not obvious at all. One idea is to restrict our interest to the space of deterministic MDPs with deterministic policies. If we want to cover general MDPs, we may want to output a distribution over states or actions, and possibly compute the KL-divergence between the ground truth and the LLM output? There are multiple possibilities here, and it is not clear to me what exactly the authors had in mind.\n- I am not so sure that this work has a clear takeaway message to its readers. The work claims that their extensive empirical study shows that LLMs are not yet fully capable of realizing the mental modelling of agents, This is for sure an important finding. But then what is the author's suggestions? How exactly should LLMs be used? Perhaps they should not be used at all? Or be used with a great deal of caution? My concern is that there is not a very clear guideline or general takeaway message for LLM and RL practitioners and theoreticians. At the end of  the day, the work must male concrete suggestions based on its findings.\n- I was not personally clear how Figures 19 was produced. My understanding is that each point in the plot corresponds to one state variable; state variables with a bigger space then get a higher span value of the x-axis. But how exactly are these values normalized to the 0-1 scale for each task. Is this done linearly? Some details there would have helped. Also, Furthermore, if the span is the real problem, would it have helped to simple rescale all state variables in each task, so that they all have the same high span? In that case, perhaps the variance would have been the main issue?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper deals with the following questions: can modern LLMs, trained on vast amounts of world knowledge and possessing some already proven reasoning ability, be utilized to mentally model RL agents, and understand their actions as well as the environment they operate in? In this direction, the paper introduces the concept of agent mental modelling, which consists in building a mental model of the RL agent with the help of LLMs. The authors propose a specific evaluation framework that involves both actions as well as dynamics understanding, and detail a concrete template to interact with the LLM in-context and without re-optimizing or finetuning. The authors subsequently conduct an extensive experimental study on a variety of RL tasks of varying difficulty, where they investigate several aspects related to agent mental modelling. The evaluation studies in depth various practical parameters as well, such as history size, data format, data span and data variance, impact of indexing, error types etc. The authors reach the conclusion that LLMs are not yet fully capable of full agent mental modelling; in particular, despite their promise, pre-existing beliefs within LLMs (or general problems like hallucinations) prevent them from fully comprehending the agents' actions and behavior.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would be eager to change my score, but I first want to understand the following;\n- Why did the authors not experiment with GPT-4o (except for the error type experiment)? It is possible that the more recent, state-of-the-art LLMs would have shown superior mental modelling capabilities. To me, the bigger problem is as follows: the authors state that \"LLMs are not yet capable of fully realizing the mental modelling of agents through inference alone without further innovations.\"  Is it possible that this will not be as big of a problem with more recent LLMs such as GPT-4o? At least this is what the current Table 1 hints at?\n- Is the evaluation framework + template introduced in this work valid for any MDP? It seems to me that the target of this work is primarily deterministic MDPs with deterministic policies. I cannot immediately see how the proposed evaluation framework could be used for arbitrary MDPs and arbitrary RL policies. Have the authors thought about experimenting with such MDPs?\n- If LLMs are not yet capable of a robust mental modelling, then what is the takeaway message for practitioners? Use LLMs with a lot of caution? This work demonstrates that LLMs can very often be wrong, and can even produce a variety of distinct error types. Should practitioners then use them, and how exactly?\n\nEven though it is not as critical, I would appreciate it if the authors could provide clarifications to the following:\n- Can they elaborate more on the span, as I detailed above?\n- I was not able to follow Appendix F.7. For step index 100, as the history size increases, the mismatch count also increases (from 0 to 7). But then how do the authors conclude that for larger steps, increasing history size minimally affects the LLMs’ predictions? I misinterpreted that figure, so additional clarifications could help.", "broader_impact_concerns": "No concern beyond what I mentioned previously.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies an important problem for RL practitioners.
The authors propose a detailed evaluation framework with a clear template. It does not only consider actions understanding but also dynamics understanding. Given the LLM is asked to provide information about its reasoning process, this allows to get a glimpse into the mental model of the LLM with respect to the RL Task at hand.
The experimental study is detailed and involves several tasks (from easy to very hard ones). The task selection is interesting because an LLM could in principle use common-sense knowledge (built from vast knowledge bases) to reason about the tasks. The authors analyze various aspects (impact of history length, data format, data span and data variance, impact of indexing, error types etc.). The empirical evaluation involved significant manual/human work, especially as far as the error type analysis is concerned.

The paper will benefit from a section discussing the potential for future-proofing this approach, especially as more agents become LLM-based and can independently provide human-readable rationales before taking actions. This would add value to the discussion.
Decision: Accept (poster)