How language models extrapolate outside the training data: A Case study in Textualized Gridworld

## Abstract

Abstract Language models have demonstrated impressive capabilities across various natural language processing tasks, yet they struggle with planning tasks requiring multi-step simulations. Inspired by human cognitive processes, this paper investigates the optimal planning power of language models that can construct a cognitive map of a given environment. Our experiments demonstrate that cognitive map significantly enhances the performance of both optimal and reachable planning generation ability in the Gridworld path planning task. We observe that our method showcases two key characteristics similar to human cognition: generalization of its planning ability to extrapolated environments and rapid adaptation with limited training data. We hope our findings in the Gridworld task provide insights into modeling human cognitive processes in language models, potentially leading to the development of more advanced and robust systems that better resemble human cognition.

## 1 Introduction

Figure: (a) Overall process
Refer to caption: /html/2406.15275/assets/x1.png

Figure: (a) Exploration-based Optimality($\downarrow$) Extrapolability($\uparrow$)
Refer to caption: /html/2406.15275/assets/x3.png

Language models have recently demonstrated remarkable proficiency in a variety of complex tasks (Brown et al., ; Touvron et al., ; Chowdhery et al., ; Chen et al., ), from natural language understanding to code generation, primarily through the training objective of next-token prediction. This training paradigm enables language models to excel at general planning tasks by leveraging their extensive learned knowledge and pattern recognition capabilities (Ahn et al., ; Liang et al., ; Song et al., ). However, it is also true that they often falter in scenarios that require robust, long-horizon planning tasks (Dziri et al., ). It is in contrast that humans naturally employ model-based planning, internally construct models to simulate outcomes and guide optimal decision-making, as extensively documented in cognitive science literature  (Daw et al., ). This distinction aligns with the dual-process theory of reasoning: System 1 processes are fast, automatic, and pattern-based, akin to model-free planning, while System 2 processes are slower, deliberative, and involve explicit reasoning (Daniel, ).

Existing methods suggests different ways to model System 2 process. Exploration-based planning methods such as Tree of Thought (ToT) (Yao et al., ) use external simulation via multi-turn inference to enhance the exploration ability of language models. While demonstrating strong exploration capabilities, they fall short in achieving optimal planning because they conduct explicit rollouts rather than incorporating the simulation process within the action decisions. On the other hand, the Chain of Thought (CoT) (Wei et al., ) methodology is also widely used to model the System 2 process of language models, injecting an optimal path into the reasoning process. Despite its impressive enhancement in reasoning and planning, existing CoT methods focus on generating a solution from start to goal, which does not align with the true notion of "simulation," i.e., analyzing real-world systems and predicting outcomes. Naive CoT only provides intermediate steps from the start to the goal, which is inappropriate as a simulation process.

When humans solve planning tasks like Gridworld, they typically conduct iterative simulations until reaching the goal, while avoiding deadend states. After reaching the goal, they work backward to figure out the current move to get to the goal state. This simulation does not require actual rollout but rather the construction of a "cognitive map." Based on the intuition, we train language models with datasets augmented with a cognitive map, comprising three key processes: Sampling, Propagation, and Backtracking. Namely, the model samples plausible actions that can be performed at each state. It then propagates for each action that reaches to a new state. The model explores(samples and propagates) potential actions and outcome states until the goal is reached. Once the goal is reached, it backtracks traces back from the goal state to the start state to refine the optimal path(See Section [4](#S4) for details).

We show that models fine-tuned using this approach exhibit significant characteristics similar to human cognitive map, which are:

- •
Extrapolation: The ability to solve problems in larger environments which is unseen during the training stage.
- •
Rapid Adaptation: Learning effectively with a small training dataset.

In this paper, we aim to show the strength of constructing cognitive maps in path planning Gridworld(^1^11We choose Gridworld since the world is easily scalable, so we can easily test the planning ability for extrapolated data. We further discuss the application of our method in general agent tasks in Section [7](#S7)). By leveraging the global world representations that the cognitive map gives to the language model, we aim to model the cognitive process of humans, thereby enhancing the performance of language models in complex, long-horizon planning tasks.

## 2 Related work

### 2.1 Planning in cognitive science

Dual-process theories of cognition distinguish between System 1 and System 2 reasoning, where System 1 involves fast, automatic, and intuitive thinking, and System 2 involves slow, deliberate, and analytical thought (Daniel, ). System 1 and System 2 reasoning aligns with model-free and model-based planning, respectively. Unlike model-free planning, model-based planning benefits by constructing and utilizing internal models to simulate outcomes and guide decisions (Daw et al., ). Evidence from cognitive science suggests that humans rely on model-based planning for complex decision-making tasks, using internal simulations to evaluate the consequences of different actions (Doll et al., ).

One of the key characteristics of human cognition is extrapolation. Human ability to solve problems in larger, more complex environments that were not encountered during the initial learning phase has been largely observed accross cognitive science domains (Yousefzadeh and Mollick, ; Fellini and Morellini, ; Stojic et al., ). This capability allows humans to apply learned knowledge to novel situations, demonstrating a remarkable level of generalization. Additionally, human cognition is characterized by rapid adaptation, the ability to learn effectively from a small amount of data (Lake et al., ). This enables humans to quickly adjust their strategies and behaviors based on limited experiences or feedback, which is essential for navigating dynamic and unpredictable environments. These traits underscore the flexibility and efficiency of human cognitive processes, setting a benchmark for artificial intelligence systems aiming to replicate human-like problem-solving and learning abilities.

### 2.2 Planning in language model domain

Language models are primarily trained using next-token prediction, which facilitates model-free planning by allowing them to generate text based on learned patterns and associations from vast amounts of data. This training approach enables language models to perform zero-shot planning in various domains, such as general planning and code generation, without requiring task-specific training Chen et al. (). Language models excel at pattern matching and leverage their learned world model to plan and generate coherent text in given tasks Brown et al. ().

Despite their successes, language models often struggle with planning tasks that require extensive foresight and detailed reasoning. Studies have shown that language models perform poorly on tasks requiring robust planning and complex decision-making, highlighting the limitations of model-free approaches (Xie et al., ; Valmeekam et al., ).

### 2.3 Exploration-based planning

One limitation of language model planning is the lack of opportunity to explore a wide range of world states. To address this, exploration-based methods have been designed to enable language models to explore more effectively. One such example is Tree of Thought (ToT) (Yao et al., ), which leverages language models’ knowledge to sample child nodes and simultaneously evaluate the value of the current state to structure and search the tree structure at the same time. Also one can enhance the searching by using Monte Carlo Tree Search (MCTS) algorithm (Kocsis and Szepesvári, ) to explore states and update their value (Hao et al., ; Zhou et al., ). These exploration-based planning methods enhance the model’s ability to reach the goal, with a perfect probability if we have infinite time and inference cost. However, their goal is to enhance the reachability of the plan while sacrificing the optimality of the generated plan.

### 2.4 Imitation-based planning

One way to inject optimal planning capability is by imitating the optimal demonstrations. We call it as Chain of Thought (CoT)(Wei et al., ) methodology, and it has demonstrated significant improvements in language models’ reasoning and planning abilities. They involve guiding the model through a series of chains generated by human or ground truth that simulates step-by-step reasoning, enhancing the model’s decision-making processes via few-shot demonstration or fine-tuning. They imitate ground truth reasoning demonstrations to improve the performance in tasks requiring complex reasoning (Nye et al., ; Yao et al., ).

One strength of training with CoT is that the demonstration does not have to be optimal. Methods such as Searchformer (Lehnert et al., ) experimentally showed that imitating automatic tree search planning such as A* rather than human demonstration is also helpful when planning. The searching algorithm conducts tree search via heuristic cost function, so it can cover more world model representation than ground truth demonstration. Yet they sacrifice the optimality of the plan, they show that the policy of the model can further imply more optimal solutions than the demonstrated algorithm.

#### Limitation of imitation-based planning in extrapolated data

Despite the great potential of expressiveness power of CoT, imitation-based planning does not work well, even for frontier language models such as GPT 4 (OpenAI et al., ). They are shown to struggle at solving and generalizing compositional tasks such as multiplication, Einstein’s puzzle (Dziri et al., ) or Blocksworld puzzle (Valmeekam et al., ; Stechly et al., ) in an extrapolated data even after fine-tuning or giving few-shot demonstrations. It is in contrast with features of human cognition planning.

One of the biggest reasons for the fallacy is that human policy is compact and entangled so they do not let language models see the global representation of the world model. Such entanglement often lets hallucination happen during generation (Peng et al., ). Theoretical approaches highlight the effectiveness of CoT (Merrill and Sabharwal, ; Feng et al., ), but they also note that CoT demonstration should be long enough to potentially solve planning problems that necessitate System 2 reasoning.

To sum up, exploration-based planning method lacks in generating the optimal plan, while imitation-based planning lacks in succeeding in the extrapolated data(Figure [2](#S1.F2)). We claim that imitating cognitive map deduces optimal model-based planning in extrapolated data points. By explicitly reasoning the representation of global states, we hypothesize that language models can achieve optimal planning capabilities. This approach mirrors human cognitive processes and enables language models to handle complex, long-horizon planning tasks more effectively Daw et al. ().

## 3 Planning task interacting with world model

Planning task with environment feedback in language model can be formalized as a Markov decision process with instruction space $\mathcal{U}$, state space $\mathcal{S}$, action space $\mathcal{A}$, observation space $\mathcal{O}$, metric space $\mathcal{C}$, transition function $T:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}$, and metric function $R:\mathcal{U}\times(\mathcal{S}\times\mathcal{A})^{n}\rightarrow\mathcal{C}$ where $n$ is the trajectory length. Note that for language model domain, $\mathcal{U}$, $\mathcal{A}$, and $\mathcal{O}$ are given as sequences of language tokens.

Given an instruction $u\in\mathcal{U}$, the model $\theta$ first generates the action $a_{1}\sim\pi_{\theta}(\cdot|u)\in\mathcal{A}$ according to its policy $\pi_{\theta}$. For each state $s_{t}\in S$ and its observation $o_{t}\in\mathcal{O}$, the agent generates the corresponding action in the $t+1$ step $a_{t+1}\sim\pi_{\theta}(\cdot|u,a_{1},o_{1},\ldots,a_{t},o_{t})\in\mathcal{A}$, which concludes to a new state $s_{t+1}=T(s_{t},a_{t})$ and its observation $o_{t+1}$. The interaction loop repeats until the task has terminated for some reason(succeeded, failed, number of steps exceeded maximum value, etc.), and the action trajectory is denoted as:

$$ $e=(u,a_{1},o_{1},\ldots,o_{n-1},a_{n},o_{n})\sim\pi_{\theta}(e|u),$ $$

$$ $\pi_{\theta}(e|u)=\prod_{j=1}^{n}\pi_{\theta}(a_{j}|u,a_{1},o_{1},\ldots,o_{j}),$ $$

Finally, we define the (action, state) trajectory $k=((a_{1},s_{1}),\ldots,(a_{n},s_{n}))$ accordingly. The final reward is computed based on the metric function $R(u,k)$.

Note that we can apply reasoning before deciding the action with so-called a "thinking" process. Namely, we have a thinking space $\mathcal{M}$(of language subset) which generates $m\sim\pi_{\theta}(\cdot|u)\in\mathcal{M}$,
then generate $a_{1}\sim\pi_{\theta}(\cdot|u,m)\in\mathcal{M}$ upon the generated thought.

## 4 Cognitive map of a given world model

Figure: Algorithm 1 Cognitive map construction algorithm

Our proposed process can be divided into three different stages: sampling, propagation, and backtrack. The sampling stage can be defined as the process where the model expects potential actions that can be applied for each state. The propagation stage explores potential outcomes of the actions sampled from the sampling stage. The backtrack stage is the process of tracing back from the goal state to refine and select the optimal path based on the outcomes of the simulation (See Algorithm [1](#alg1) for the pseudocode).

### 4.1 Sampling

During the sampling stage, the model samples possible actions in $S(s)\subset A$ regarding the current state $s$, which leads to a new state that was not propagated before. This sampling process is repeated until reaching the desirable goal. This iterative approach enables the cognitive map of the agent to explore the various states of the given world model.(^2^22In this paper, we treat only for Gridworld case so there are small number of action candidates to sample(there are always four moves, up, down, left, or right). We discuss possible sampling strategies that the agent can make in a bigger world in Section [7](#S7).)

### 4.2 Propagation

During the propagation stage, the model expects a transited state $T(s,a)\in\mathcal{S}$ for each action $a$ sampled from the current state $s$. For example, in the Gridworld, the model simulates different paths by considering movements in four directions (up, down, left, right) from each cell. The goal is to explore various routes to reach the target cell, considering obstacles and the grid’s boundaries. This stage is crucial for understanding the global representation of the world model by expecting the future consequences of each action before committing to the actual decision.

### 4.3 Backtrack

Once the sampling and propagation stages have concluded by reaching a desirable goal state, we need to backtrack through the simulated paths to determine the most efficient route taken to achieve the goal. This involves assessing the propagated paths and identifying the one that reaches the goal. Backtrack search identifies $T^{-1}(s,a)\in\mathcal{S}$ for each state-action pair to ensure that the selected path is valid. By refining the decisions made during the simulation stage, backtracking ultimately guides the model to make informed, strategic choices based on the simulated outcomes. Since the propagation stage deduces only one path with minimal steps, we can guarantee the optimality of the generated path.

### 4.4 Training how to generate cognitive map

To sum up, constructing the cognitive map $m$ is a sequential application of sampling($S$), propagation($T$), and backtrack($T^{-1}$). We train the language model $\theta$ via supervised learning method so that the model can successfully construct the cognitive map without any external interaction, as shown in Figure [1(b)](#S1.F1.sf2).

## 5 Experimental setup

### 5.1 Basic setup

#### Task, train, and inference detail

In this paper, we set textualized Gridworld (Brown, ) as the main task. Gridworld is a task that involves path planning from the start state to the goal state while avoiding the pit, wall, and grids outside the world. Especially, we ensure that there is only one path from start to goal for each environment(For the example of input including instruction and the world information fed to the model, refer Appendix [A.1](#A1.SS1)).

We train the model for 1 epoch with 50K samples of size 10$\times$10 at largest. After training, we test the model with 3K samples with each grid being 20$\times$20 at largest. We utilize Llama-3-8B model(^3^33[https://llama.meta.com/llama3/](https://llama.meta.com/llama3/)) throughout the whole experiments. In each turn, we set the maximum token length of the model to be 8192.
(For more details about experimental setup, refer Appendix [A.2](#A1.SS2)).

#### Optimal planning analysis: single-turn setting

We see the capability of cognitive map on generating the optimal plan without any further observations(Corresponding to 4-1 of Figure [1(a)](#S1.F1.sf1)). We evaluate the optimality of the trajectory generated by $a_{1},a_{2},\ldots,a_{n}\sim\pi_{\theta}(\cdot|u,m)\in\mathcal{A}^{n}$. This is equivalent to generating the whole plan in a single-turn, so we train the model to generate the whole plan in one turn and evaluate it.

Given the instruction $u$ and generated (action, state) trajectory $k$, we first define optimal path $k^{*}$ as an element of $\operatorname*{arg\,min}_{k\in\{k\mid r(u,k)=\text{success}\}}|k|$. Since we ensure there is only one possible $k^{*}$, we only need to check whether the generated trajectory is $k^{*}$. Now we can define $R(u,k)$ as follows:

- •
$R(u,k)=\textsc{success}$ if $k=k^{*}$
- •
$R(u,k)=\textsc{fail}$ otherwise

#### Reachable planning analysis: multi-turn setting

We also investigate whether the cognitive map can deduce a reachable plan(which doesn’t require optimal planning) in a multi-turn setting(Corresponding to 4-2 of Figure [1(a)](#S1.F1.sf1)). We evaluate the reachability of the trajectory $a_{1},a_{2},\ldots,a_{n}$ generated by $a_{t+1}\sim\pi_{\theta}(\cdot|u,m,a_{1},o_{1},\ldots,a_{t},o_{t})\in\mathcal{A}$. There are multiple interactions between the model and the environment in a muli-turn setting, so failure cases can be divided into three cases(deadend, max step, and invalid). Especially, given the instruction $u$ and generated (action, state) trajectory $k$, we define $R(u,k)$ as follows:

- •
$R(u,k)=\textsc{success}$ if $k[-1][1]=goal$
- •
$R(u,k)=\textsc{deadend}$ if $k[-1][1]\in P$
- •
$R(u,k)=\textsc{max step}$ if $|k|>max$
- •
$R(u,k)=\textsc{invalid}$ if $\exists a\in k[:][0]\mid a\notin A$

In the paper, we set the maximum steps $max=200$ for each environment. We denote $P$ as the set of deadend states, i.e. set of pits and walls in the grid, and all the states outside the grid.

For each observation, we give the information of the current state along with possible moves that do not directly result the deadend state along with its corresponding states. For example, if the current state $s_{t}$ is (11, 4) and there the possible actions are right or left, the observation $o_{t}$ is "Current:\n(11, 4)\nPossible:\n(10, 4)\nleft\n(12, 4)\nright".

### 5.2 Cognitive map design choices

Our construction of the cognitive map is straightforward; We sample all 4 possible directions("up", "down", "right", and "left") as actions, and we propagate each action iteratively until reaching the goal state. After reaching the goal, we backtrack until reaching the start state.

#### Marking deadend

In this paper, we variate two different cognitive maps when marking the deadend state. We can either verbalize all samples and mark the actions resulting deadend state with a special token(<deadend>)(denoted as "marking deadend"), or just verbalize all the actions including the deadend states(denoted as "w.o. marking").
For example, the verbalized cognitive map at state (0, 0) in Figure [1(b)](#S1.F1.sf2) is as follows:
marking deadend: "(0, 1) up (0, -1) <deadend> (1, 0) right (-1, 0) <deadend>"
w.o. marking: (0, 1) up (0, -1) down (1, 0) right (-1, 0) left".
See Appendix [B.1](#A2.SS1) for the whole construction.

#### Baselines

We name none and cot as our baseline experiments representing imitation-based planning. none implicitly learns how to conduct path planning, while cot learns a verbalized backward trace as a Chain of Thought(CoT) demonstration. For example, the cot reasoning of the Figure [1(b)](#S1.F1.sf2) is "(3, 2)up\n(3, 1)\nright\n(2, 1)\nup\n(2, 0)\nright\n(1, 0)\nright\n(0, 0)", while we do not verbalize anything for none. See Appendix [B.1](#A2.SS1) for the whole construction.

#### Backward cognitive map construction: bwd

Backward chaining is a powerful approach that simplifies complex problems by focusing on the desired outcome and systematically working back to the starting point. LAMBADA (Kazemi et al., ) shows that backward chaining helps in reasoning tasks. We adopt the intuition to see if constructing the cognitive map in a backward manner enhances the planning ability of the language model.

In this paper, we define two types of construction, fwd and bwd. For fwd, the construction from the start to the goal is identical to the procedure stated in Section [4](#S4). For bwd, we build the reversed cognitive map starting from the goal state to the start state. We have the sampling function identical to fwd($S$). For the propagation stage, we expect a reverse transition state $T^{-1}(s,a)\in\mathcal{S}$ for a given state $s$ and action $a$. For the backtracking stage, we backtrack the path from start to goal by iteratively searching $T(s,a)\in\mathcal{S}$. See Appendix [B.2](#A2.SS2) for the example.

#### Additional setups

We also experiment effects of excluding each component when designing the cognitive map. First, we see the effect of excluding the backtrack stage(denoted as "w.o. backtrack"). Also, we observe the effect when we only sample possible moves instead of all moves(denoted as "w.o. all"). See Appendix [B.1](#A2.SS1) , [B.2](#A2.SS2) for examples.

### 5.3 Comparsion of the cognitive model with exploration-based planning

We also compare the planning ability of our method with exploration-based planning. We train the language model which traverses the world through a tree search algorithm via Depth-First Search(DFS), denoted as DFS(Pseudocode of the algorithm is in Algorithm [2](#alg2)). See Table [9](#A4.T9) for the example.

## 6 Results and analysis

### 6.1 Overall result analysis

**Table 1: Optimal and reachable rate of generated plans via single- and multi-turn settings: The first two columns(none and backtrack) are the baselines for imitation-based learning, and the rest are different design choices of constructing the cognitive map. Also bwd constructs the map starting from the goal state, while fwd starts from the start state. See Appendix [B](#A2) for actual prompts.**
|  | Implicit baseline | Explicit baseline | Cognitive map |  |
| --- | --- | --- | --- | --- |
| Optimal | None | CoT | marking deadend | w.o. marking |
| bwd | 0.190 | 0.265 | 0.705 | 0.765 |
| fwd | 0.190 | 0.252 | 0.585 | 0.618 |
| Reachable | None | CoT | marking deadend | w.o. marking |
| bwd | 0.321 | 0.287 | 0.885 | 0.724 |
| fwd | 0.321 | 0.339 | 0.854 | 0.816 |

**Table 2: Qualitative comparison of reachable and optimal plan generation rate between baselines and our methods. The degree of the darkness at (x, y) coordinate of each plot denotes the performance of the corresponding model in Gridworld of size x$\times$y. The red box denotes the boundary of the training data. We provide visualization of all the experiments in Appendix [C](#A3).**
|  | Implicit baseline | Explicit baseline(best) | Cognitive map(best) |
| --- | --- | --- | --- |
|  |  |  |  |
| Optimal | none: 19.0% | bwd backtrack: 26.5% | bwd w.o. marking: 76.5% |
|  |  |  |  |
| Reachable | none: 32.1% | fwd backtrack: 33.9% | bwd marking deadend: 88.5% |

#### Cognitive map compared with baselines

Table [1](#S6.T1) shows both the optimal and reachable rate of plans generated by each experiment. As shown in the table, w.o. marking shows the best performance among the experiments both for bwd and fwd cognitive map construction in the optimal planning. (76.5% and 61.8% for bwd and fwd construction, respectively). Unlike in the optimal planning setting, marking deadend shows the best performance both for bwd and fwd cognitive map construction(88.5% and 85.4% for bwd and fwd construction, respectively). Our experiments show that cognitive maps improve performance in the Gridworld path planning task. It boosts the optimal planning performance by up to 57.5% and reachable planning by up to 56.4% with implicit fine-tuning baseline, and enhancing optimal planning by up to 50% and reachable planning by up to 54.6% with CoT fine-tuning baseline. We also qualitatively show the planning performance in both optimal and reachable settings with respect to the size of the grid in Table [2](#S6.T2).

#### Cognitive map further enhances reachability

A notable observation is the performance gap between optimal and reachable planning. Reachable planning experiments show significant improvements compared to optimal planning across all configurations, except for all backtrack in fwd cognitive map construction. For instance, w.o. backtrack from fwd map construction more than doubled its score (42.3% to 85.2%). This implies that although the cognitive map was designed to find the optimal plan, it also substantially enhances reachable planning capability.

#### bwd approach enhances the performance of the cognitive map

As shown in Table [1](#S6.T1), both the optimal and reachable planning got the highest performance with bwd cognitive map construction. This aligns with the findings in LAMBADA (Kazemi et al., ), where they show that backward chaining helps in reasoning tasks.

Figure: (a) Success rate comparison: For our method, we provide additional performance of the model checkpoints at step 125, 250, 375, 500, 625, 1250, 2500, and 3125.
Refer to caption: /html/2406.15275/assets/x32.png

Figure: Figure 4: Number of inference steps per optimal steps for succeeded instances. We compare exploration-based planning(DFS: red) with our method(bwd marking deadend: blue) in the reachable planning setting. Each scattered dot denotes the succeeded instance, and the line plot denotes the regression line of the success. Note that the regression line of our method is almost identical to a $x=y$ graph, which means that cognitive map helps generating optimal performance even for the reachable planning setting.
Refer to caption: /html/2406.15275/assets/x34.png

### 6.2 Extrapolation ability

As stated in Section [5.1](#S5.SS1), we train the model on world sizes of up to 10x10 and test it on world sizes of up to 20x20 world to investigate the extrapolation ability. Extrapolation success is defined as the ability to succeed in rollouts beyond the 10x10 boundary. As shown in Table [2](#S6.T2), the darkness outside the boundary of the red box shows that the cognitive map helps in planning on the extrapolated data. We also observe a consistent tendency for experiments with all inclusion, which have a wider coverage of success rate in the extrapolated data. These results align with the result discussed in Section [6.1](#S6.SS1). We qualitatively analyze all the results in Appendix [C.2](#A3.SS2) and [C.3](#A3.SS3).

### 6.3 Rapid adaptation

We experiment speed of the language model learning robust bwd marking deadend cognitive map construction by both watching its reachability and loss curve. As shown in Figure [3(a)](#S6.F3.sf1), the success rate of the model converges at the early stage(79.13% at step 500, 74.5% at step 625), which implies that it is easy for a model to learn how to generate a robust cognitive map. We also observe that the loss curve of the model converges to 0 much faster than other baseline methods(Figure [3(b)](#S6.F3.sf2)). It suggests that the language model rapidly learns how to construct the cognitive map.

**Table 3: Comparsion of reachability with exploration-based planning(DFS) and our best(bwd marking deadend cognitive map)**
| Cognitive map(best) | Exploration-based |
| --- | --- |
| 0.885 | 0.945 |

### 6.4 Compare with exploration-based planning

Despite the significant planning ability demonstrated by the cognitive map, we observe that there is still room for improvement, particularly in reachable planning(See Appendix [D](#A4) for detailed setting and qualitative analysis). The results show that exploration-based planning is more effective in generating reachable plans(See Table [3](#S6.T3)). We interpret the incident due to an inductive bias of exploration-based planning toward exploring via tree search algorithm. With an infinite amount of time, tree search algorithm almost perfectly ensures the reachability toward the goal.

On the other hand, we show that constructing the cognitive map has benefits in generating the plan with optimal steps(See Figure [4](#S6.F4)). We see that the exploration-based planning has a mean inference steps for reaching the goal for $n$ optimal steps as $O(n^{2})$, while our method shows the optimal performance also for the reachable plan generation.

## 7 Discussion

#### Why is cognitive map easier to learn than CoT or implicit policy?

One of the left curiosity is interpreting the result we obtained from Section [6.3](#S6.SS3), which implies that the language model rapidly learns how to construct the cognitive map. The result suggests that the language model easily learns how to construct a cognitive map. How is it possible?

In this paper, the generation of cognitive map $m\in\mathcal{M}$ given an instruction $u\in\mathcal{U}$ is modeled as a composition of following functions: sampling($S$), propagation($T$), and backtrack($T^{-1}$). Given the current state $s_{t}$, the sampling strategy selects actions $S(s_{t})\subset\mathcal{A}$. For each action $a\in S(s_{t})$, propagation generates the next state $T(s_{t},a)$. These steps iterate until we reach the goal state. Once the goal state $g$ is reached, we backtrack to the start state $T^{-1}(s_{t},a)$. Assuming that applying the transition function $T$ and its inverse $T^{-1}$ is trivial for the model, each generation step of $m$ is straightforward. Hence if conducting each function is straightforward, we can also say that the overall generation of $m\sim\pi_{\theta}(\cdot|u)\in\mathcal{M}$ is also trivial.

#### Why does cognitive map deduce robust planning?

Human demonstration of path planning is often abstract and entangled. This entanglement occurs between the human policy $\pi_{\text{human}}$ and the replication policy from observed human demonstrations $\pi_{\text{human demo}}$. $\pi_{\text{human demo}}$ tends to be more complex and interwoven since it has various hidden factors. Namely, our objective is to expect

$$ $e=(u,a_{1},o_{1},\ldots,o_{n-1},a_{N})\sim\pi_{\text{human}}(e|u),$ $$

but the human demonstration is given as

$$ $e^{\prime}=(u,a^{\prime}_{1},o^{\prime}_{1},\ldots,o^{\prime}_{n-1},a^{\prime}_{n})\sim\pi_{\text{human demo}}(e^{\prime}|u),$ $$

where $(a^{\prime}_{1},\ldots,a^{\prime}_{n})\subset(a_{1},\ldots,a_{N}),n\ll N$.

Our underlying assumption is that $\pi_{\text{human}}$ is relatively simple, and we can model it by composing sampling($S$), propagation($T$), and backtrack($T^{-1}$) functions. We saw in the Gridworld that these functions are straightforward and easily replicated. By employing these simple and well-defined functions, the model effectively addresses the problem of data path planning. This approach allows the model to generate robust planning ability, including extrapolation beyond the training data.

#### Extension to general world

Our discussion of planning is restricted to the Gridworld environment throughout this paper. Gridworld has a relatively small world size compared to general agent tasks such as those performed by web agents (Yao et al., ) or travel agents (Xie et al., ). To extend planning capabilities to these larger, more complex worlds, we may need to modify the sampling strategy. The current sampling approach may not scale well in larger environments. A potential solution is to train the model to sample the top $n$ preferable actions, which would allow the model to efficiently navigate and plan in much larger and more intricate worlds. We leave the investigation of sampling strategy as a further research direction.

## 8 Acknowledgement

We thank Seonghyeon Ye, Hyeonbin Hwang, Sejune Joo and Juyoung Suk for their invaluable advice and support throughout this research.

## Appendix A Experimental Details

### A.1 Input detail

Table [4](#A1.T4) describes a sample input of the model, describing the instruction of the Gridworld and the specific world information.

**Table 4: The prompt for all cases**
| Common Prompt | human: You are given a rectangular gridworld, where you can move up, down, left, or right as long as each of your x, y coordinates are within 0 to the x, y size of the grid. If you move up, your y coordinate increases by 1. If you move down, your y coordinate decreases by 1. If you move left, your x coordinate decreases by 1. If you move right, your x coordinate increases by 1.\n\nYou will interact with the gridworld environment to reach the goal state, while avoiding the pit and the wall. You cannot move through the wall or move outside the grid. If you fall into the pit, you lose. If you reach the goal, you win. For each of your turn, you will be given the possible moves.\n\nYou should respond your move with either one of ’up’, ’down’, ’left’, or ’right’.<br>gpt: OK<br>human: Grid is from (0, 0) to (3, 2). Goal: (3, 2)\nCurrent: (0, 0)\nThe pit is at (3, 0). The wall is at (1, 1), and (2, 2).\nCurrent:\n(0, 0)\nPossible:\n(0, 1)\nup\n(1, 0)\nright |
| --- | --- |

### A.2 Experimental setup details

We use one 8 Nvidia A100 node for both training and inference. For the training steps, we use FSDP framework (Zhao et al., ) and cosine annealing learning rate scheduler (Loshchilov and Hutter, ) for 1 epoch. We utilize bfloat16 floating-point format and a warmup ratio of 0.03. We set the weight decay as 0. We set the batch size of 2 for each GPUs, so the effective batch size is 16 per step. We train each model for $\text{50000}/\text{16}=\text{3125}$ steps. For inference, we use VLLM framework (Kwon et al., ).

While exploring the pure planning ability of the language model, we did not want the model to refuse to explore extrapolated data only because it has never seen the coordinate. To handle the bias, we adjust the starting position of the grid using uniform sampling while ensuring that the entire grid, with its x and y coordinates, fits within the range of 0 to 19. This method, illustrated in Figure [5](#A1.F5), minimizes bias related to the unseen x and y coordinates by randomizing the starting point within the defined bounds.

Figure: Figure 5: Visualization of configuring Gridworld instance for train(left) and test(right) dataset. To evaluate the extrapolation ability, we set the size of the grid as $X,Y\sim\leavevmode\nobreak\ Unif(2,10)$ for train and $X,Y\sim Unif(2,20)$ for test. Also we set the starting point of the grid as $\Delta x\sim Unif(0,20-X),\Delta y\sim Unif(0,20-Y)$ for both train and test.
Refer to caption: /html/2406.15275/assets/x35.png

## Appendix B Cognitive map description

### B.1 Cognitive map example: fwd

Table [5](#A2.T5) describes a sample of fwd cognitive map construction for each experiment.

**Table 5: fwd cognitive map example for each experiment**
| Design choice | Cognitive map example |
| --- | --- |
| none |  |
| cot | Thought:\nStep 1:\nStep 2:\nStep 3:\nStep 4:\nStep 5:\nBacktrack:\n(3, 2)up\n(3, 1)\nright\n(2, 1)\nup\n(2, 0)\nright\n(1, 0)\nright\n(0, 0) |
| w.o. all backtrack | Thought:\nStep 1:\n(0, 1)\nup\n(1, 0)\nright\nStep 2:\n(0, 2)\nup\n(2, 0)\nright\nStep 3:\n(1, 2)\nright\n(2, 1)\nup\nStep 4:\n(3, 1)\nright\nStep 5:\n(3, 2)\nup |
| w.o. all | Thought:\nStep 1:\n(0, 1)\nup\n(1, 0)\nright\nStep 2:\n(0, 2)\nup\n(2, 0)\nright\nStep 3:\n(1, 2)\nright\n(2, 1)\nup\nStep 4:\n(3, 1)\nright\nStep 5:\n(3, 2)\nup\nBacktrack:\n(3, 2)up\n(3, 1)\nright\n(2, 1)\nup\n(2, 0)\nright\n(1, 0)\nright\n(0, 0) |
| w.o. marking backtrack | Thought:\nStep 1:\n(0, 1)\nup\n(0, -1)\ndown\n(-1, 0)\nleft\n(1, 0)\nright\nStep 2:\n(0, 2)\nup\n(0, 0)\ndown\n(-1, 1)\nleft\n(1, 1)\nright\n(1, 1)\nup\n(1, -1)\ndown\n(0, 0)\nleft\n(2, 0)\nright\nStep 3:\n(0, 3)\nup\n(0, 1)\ndown\n(-1, 2)\nleft\n(1, 2)\nright\n(2, 1)\nup\n(2, -1)\ndown\n(1, 0)\nleft\n(3, 0)\nright\nStep 4:\n(1, 3)\nup\n(1, 1)\ndown\n(0, 2)\nleft\n(2, 2)\nright\n(2, 2)\nup\n(2, 0)\ndown\n(1, 1)\nleft\n(3, 1)\nright\nStep 5:\n(3, 2)\nup\n(3, 0)\ndown\n(2, 1)\nleft\n(4, 1)\nright |
| w.o. marking | Thought:\nStep 1:\n(0, 1)\nup\n(0, -1)\ndown\n(-1, 0)\nleft\n(1, 0)\nright\nStep 2:\n(0, 2)\nup\n(0, 0)\ndown\n(-1, 1)\nleft\n(1, 1)\nright\n(1, 1)\nup\n(1, -1)\ndown\n(0, 0)\nleft\n(2, 0)\nright\nStep 3:\n(0, 3)\nup\n(0, 1)\ndown\n(-1, 2)\nleft\n(1, 2)\nright\n(2, 1)\nup\n(2, -1)\ndown\n(1, 0)\nleft\n(3, 0)\nright\nStep 4:\n(1, 3)\nup\n(1, 1)\ndown\n(0, 2)\nleft\n(2, 2)\nright\n(2, 2)\nup\n(2, 0)\ndown\n(1, 1)\nleft\n(3, 1)\nright\nStep 5:\n(3, 2)\nup\n(3, 0)\ndown\n(2, 1)\nleft\n(4, 1)\nright\nBacktrack:\n(3, 2)up\n(3, 1)\nright\n(2, 1)\nup\n(2, 0)\nright\n(1, 0)\nright\n(0, 0) |
| w.o. backtrack | Thought:\nStep 1:\n(0, 1)\nup\n(0, -1)\ncut\n(-1, 0)\ncut\n(1, 0)\nright\nStep 2:\n(0, 2)\nup\n(0, 0)\ncut\n(-1, 1)\ncut\n(1, 1)\ncut\n(1, 1)\ncut\n(1, -1)\ncut\n(0, 0)\ncut\n(2, 0)\nright\nStep 3:\n(0, 3)\ncut\n(0, 1)\ncut\n(-1, 2)\ncut\n(1, 2)\nright\n(2, 1)\nup\n(2, -1)\ncut\n(1, 0)\ncut\n(3, 0)\ncut\nStep 4:\n(1, 3)\ncut\n(1, 1)\ncut\n(0, 2)\ncut\n(2, 2)\ncut\n(2, 2)\ncut\n(2, 0)\ncut\n(1, 1)\ncut\n(3, 1)\nright\nStep 5:\n(3, 2)\nup\n(3, 0)\ncut\n(2, 1)\ncut\n(4, 1)\ncut |
| marking deadend | Thought:\nStep 1:\n(0, 1)\nup\n(0, -1)\ncut\n(-1, 0)\ncut\n(1, 0)\nright\nStep 2:\n(0, 2)\nup\n(0, 0)\ncut\n(-1, 1)\ncut\n(1, 1)\ncut\n(1, 1)\ncut\n(1, -1)\ncut\n(0, 0)\ncut\n(2, 0)\nright\nStep 3:\n(0, 3)\ncut\n(0, 1)\ncut\n(-1, 2)\ncut\n(1, 2)\nright\n(2, 1)\nup\n(2, -1)\ncut\n(1, 0)\ncut\n(3, 0)\ncut\nStep 4:\n(1, 3)\ncut\n(1, 1)\ncut\n(0, 2)\ncut\n(2, 2)\ncut\n(2, 2)\ncut\n(2, 0)\ncut\n(1, 1)\ncut\n(3, 1)\nright\nStep 5:\n(3, 2)\nup\n(3, 0)\ncut\n(2, 1)\ncut\n(4, 1)\ncut\nBacktrack:\n(3, 2)up\n(3, 1)\nright\n(2, 1)\nup\n(2, 0)\nright\n(1, 0)\nright\n(0, 0) |

### B.2 Cognitive map example: bwd

Table [6](#A2.T6) describes a sample of bwd cognitive map construction for each experiment.

**Table 6: bwd cognitive map example for each experiment**
| Design choice | Cognitive map example |
| --- | --- |
| none |  |
| cot | Thought:\nStep 1:\nStep 2:\nStep 3:\nStep 4:\nStep 5:\nBacktrack:\n(0, 0)\nright\n(1, 0)\nright\n(2, 0)\nup\n(2, 1)\nright\n(3, 1)\nup\n(3, 2) |
| w.o. all backtrack | Thought:\nStep 1:\n(3, 1)\nup\nStep 2:\n(2, 1)\nright\nStep 3:\n(2, 0)\nup\nStep 4:\n(1, 0)\nright\nStep 5:\n(0, 0)\nright |
| w.o. all | Thought:\nStep 1:\n(3, 1)\nup\nStep 2:\n(2, 1)\nright\nStep 3:\n(2, 0)\nup\nStep 4:\n(1, 0)\nright\nStep 5:\n(0, 0)\nright\nBacktrack:\n(0, 0)\nright\n(1, 0)\nright\n(2, 0)\nup\n(2, 1)\nright\n(3, 1)\nup\n(3, 2) |
| w.o. marking backtrack | Thought:\nStep 1:\n(3, 3)\ndown\n(3, 1)\nup\n(2, 2)\nright\n(4, 2)\nleft\nStep 2:\n(3, 2)\ndown\n(3, 0)\nup\n(2, 1)\nright\n(4, 1)\nleft\nStep 3:\n(2, 2)\ndown\n(2, 0)\nup\n(1, 1)\nright\n(3, 1)\nleft\nStep 4:\n(2, 1)\ndown\n(2, -1)\nup\n(1, 0)\nright\n(3, 0)\nleft\nStep 5:\n(1, 1)\ndown\n(1, -1)\nup\n(0, 0)\nright\n(2, 0)\nleft |
| w.o. marking | Thought:\nStep 1:\n(3, 3)\ndown\n(3, 1)\nup\n(2, 2)\nright\n(4, 2)\nleft\nStep 2:\n(3, 2)\ndown\n(3, 0)\nup\n(2, 1)\nright\n(4, 1)\nleft\nStep 3:\n(2, 2)\ndown\n(2, 0)\nup\n(1, 1)\nright\n(3, 1)\nleft\nStep 4:\n(2, 1)\ndown\n(2, -1)\nup\n(1, 0)\nright\n(3, 0)\nleft\nStep 5:\n(1, 1)\ndown\n(1, -1)\nup\n(0, 0)\nright\n(2, 0)\nleft\nBacktrack:\n(0, 0)\nright\n(1, 0)\nright\n(2, 0)\nup\n(2, 1)\nright\n(3, 1)\nup\n(3, 2) |
| w.o. backtrack | Thought:\nStep 1:\n(3, 3)\ncut\n(3, 1)\nup\n(2, 2)\ncut\n(4, 2)\ncut\nStep 2:\n(3, 2)\ncut\n(3, 0)\ncut\n(2, 1)\nright\n(4, 1)\ncut\nStep 3:\n(2, 2)\ncut\n(2, 0)\nup\n(1, 1)\ncut\n(3, 1)\ncut\nStep 4:\n(2, 1)\ncut\n(2, -1)\ncut\n(1, 0)\nright\n(3, 0)\ncut\nStep 5:\n(1, 1)\ncut\n(1, -1)\ncut\n(0, 0)\nright\n(2, 0)\ncut |
| marking deadend | Thought:\nStep 1:\n(3, 3)\ncut\n(3, 1)\nup\n(2, 2)\ncut\n(4, 2)\ncut\nStep 2:\n(3, 2)\ncut\n(3, 0)\ncut\n(2, 1)\nright\n(4, 1)\ncut\nStep 3:\n(2, 2)\ncut\n(2, 0)\nup\n(1, 1)\ncut\n(3, 1)\ncut\nStep 4:\n(2, 1)\ncut\n(2, -1)\ncut\n(1, 0)\nright\n(3, 0)\ncut\nStep 5:\n(1, 1)\ncut\n(1, -1)\ncut\n(0, 0)\nright\n(2, 0)\ncut\nBacktrack:\n(0, 0)\nright\n(1, 0)\nright\n(2, 0)\nup\n(2, 1)\nright\n(3, 1)\nup\n(3, 2) |

**Table 7: Planning performance with all exclusion**
|  | w.o. all | with all |  |  |
| --- | --- | --- | --- | --- |
| Optimal | w.o. all backtrack | w.o. all | w.o. marking | marking deadend |
| bwd | 0.296 | 0.277 | 0.765 | 0.705 |
| fwd | 0.295 | 0.290 | 0.618 | 0.585 |
| Reachable | w.o. all backtrack | w.o. all | w.o. marking | marking deadend |
| bwd | 0.394 | 0.283 | 0.724 | 0.885 |
| fwd | 0.416 | 0.345 | 0.816 | 0.854 |

**Table 8: Planning performance with backtrack exclusion**
|  | w.o. backtrack | with backtrack |  |  |
| --- | --- | --- | --- | --- |
| Optimal | w.o. marking backtrack | w.o. backtrack | w.o. marking | marking deadend |
| bwd | 0.406 | 0.423 | 0.765 | 0.705 |
| fwd | 0.528 | 0.516 | 0.618 | 0.585 |
| Reachable | w.o. marking backtrack | w.o. backtrack | w.o. marking | marking deadend |
| bwd | 0.739 | 0.852 | 0.724 | 0.885 |
| fwd | 0.672 | 0.624 | 0.816 | 0.854 |

## Appendix C Additional results

### C.1 Additional investigations

#### Effect of all inclusion

Both cognitive map constructions w.o. all (w.o. all backtrack: 29.6% and 29.5% for bwd and fwd construction, respectively; w.o. all: 27.7% and 29.0% for bwd and fwd construction, respectively) suffer at generating the optimal plan, achieving under 30% success rate. However, every cognitive map construction with the inclusion of all shows better performance, with fwd marking deadend being the lowest among them(58.5%).

We observe a similar trend in the reachable planning test. While both baseline methods w.o. all performed at best 41.6% (fwd w.o. all backtrack), the lowest performance among every cognitive map construction was 72.4% (bwd w.o. marking). This implies that the inclusion of all significantly enhances the planning capability, leading to more successful and efficient pathfinding.

#### Effect of backtrack inclusion

Both cognitive map constructions w.o. backtrack (w.o. marking backtrack: 40.6% and 52.8% for bwd and fwd construction, respectively; w.o. backtrack: 42.3% and 51.6% for bwd and fwd construction, respectively) suffer at generating the optimal plan. However, every cognitive map construction with the inclusion of backtrack shows better performance, with fwd marking deadend being the lowest among them(58.5%).

The analysis in the reachable planning test was slightly blurry, yet there was an obvious trend. For each experiment, adding backtracking enhanced the performance of the planning in most settings(except bwd construction w.o. marking). This implies that the inclusion of backtrack slightly enhances the planning capability.

### C.2 Vizualization for optimal planning experiments

For optimal planning, we have only success or failure cases. Hence we only provide the success rate for each experiment.

#### fwd construction

See Figure [6](#A3.F6) for success rate of each experiment.

Figure: (a) none: 19%
Refer to caption: /html/2406.15275/assets/x36.png

#### bwd construction

See Figure [7](#A3.F7) for success rate of each experiment.

Figure: (a) none: 19%
Refer to caption: /html/2406.15275/assets/x44.png

### C.3 Visualization for reachable planning experiments

For reachable planning, we have one success cases and three different failure cases(deadend, max step, and invalid). Hence we provide the visualization of all 4 cases.

#### fwd construction

For success rate, see Figure [8](#A3.F8). For failure cases, see Figure [9](#A3.F9) for deadend, Figure [10](#A3.F10) for max step, and Figrue [11](#A3.F11) for invalid rate.

Figure: (a) none: 32%
Refer to caption: /html/2406.15275/assets/x52.png

Figure: (a) none: 64%
Refer to caption: /html/2406.15275/assets/x60.png

Figure: (a) none: 0%
Refer to caption: /html/2406.15275/assets/x68.png

Figure: (a) none: 3%
Refer to caption: /html/2406.15275/assets/x76.png

#### bwd construction

For success rate, see Figure [12](#A3.F12). For failure cases, see Figure [13](#A3.F13) for deadend, Figure [14](#A3.F14) for max step, and Figrue [15](#A3.F15) for invalid rate.

Figure: (a) none: 32%
Refer to caption: /html/2406.15275/assets/x84.png

Figure: (a) none: 64%
Refer to caption: /html/2406.15275/assets/x92.png

Figure: (a) none: 0%
Refer to caption: /html/2406.15275/assets/x100.png

Figure: (a) none: 3%
Refer to caption: /html/2406.15275/assets/x108.png

## Appendix D Comparison of cognitive model with exploration-based model

See Figure [16](#A4.F16) for comparing the success rate of reachable planning between the cognitive model and exploration-based planning. See Algorithm [2](#alg2) for the algorithm of DFS. Moreover, Table [9](#A4.T9) describes an example of an output of exploration-based planning.

Figure: (a) Ours: 88.5%
Refer to caption: /html/2406.15275/assets/x116.png

**Table 9: The prompt for DFS case**
| DFS Prompt | gpt: up<br>human: Current:\n(0, 1)\nPossible:\n(0, 2)\nup\n(0, 0)\ndown<br>gpt: up<br>human: Current:\n(0, 2)\nPossible:\n(0, 1)\ndown\n(1, 2)\nright<br>gpt: right<br>human: Current:\n(1, 2)\nPossible:\n(0, 2)\nleft<br>gpt: left<br>human: Current:\n(0, 2)\nPossible:\n(0, 1)\ndown\n(1, 2)\nright<br>gpt: down<br>human: Current:\n(0, 1)\nPossible:\n(0, 2)\nup\n(0, 0)\ndown<br>gpt: down<br>human: Current:\n(0, 0)\nPossible:\n(0, 1)\nup\n(1, 0)\nright<br>gpt: right<br>human: Current:\n(1, 0)\nPossible:\n(0, 0)\nleft\n(2, 0)\nright<br>gpt: right<br>human: Current:\n(2, 0)\nPossible:\n(2, 1)\nup\n(1, 0)\nleft<br>gpt: up<br>human: Current:\n(2, 1)\nPossible:\n(2, 0)\ndown\n(3, 1)\nright<br>gpt: right<br>human: Current:\n(3, 1)\nPossible:\n(3, 2)\nup\n(2, 1)\nleft<br>gpt: up |
| --- | --- |

Figure: Algorithm 2 Tree search via DFS

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: I think this paper is well written but it over presents its relations to human cognition. I will increase my score if authors could provide a way to address these weaknesses in the revision.
soundness: 3
strengths: - The paper is written clearly
- The results are strong on the task: their backward version with markings of dead ends (Table-1 last column) generalizes to larger grids
summary: The authors study LMs ability to solve path finding  in synthetic GridWorld problems. They propose fine-tuning an LM to first produce path search traces (called cognitive map) then produce the path. These traces are obtained by running two different algorithms (called forward and backward) on these grids. They compare their method to CoT prompting for extrapolation to larger grids.
weaknesses: - The **cognitive maps** have too much emphasis yet we see only a single synthetic task. How these maps related to how humans solve this task? Is there any experiment with humans to show similarity— eye gazing, asking them how they arrived at their solutions? I saw some references in the paper but did not find a strong discussion of it.
- The fine-tuning of rule based solved cognitive maps is not an ecologically valid comparison to humans, as humans are not trained with such intermediate maps but they could come up with themselves.  
- The experiments start too late on Page 7, you could move most of  3.3 to appendix.
- The method is not novel as fine-tuning with intermediate structures or CoTs  do exist in the literature. The structure of CoT and how they constructed is different but specific to this task.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Even for a proof-of-concept paper, some experimental results somehow demonstrating the scalability of the method would be better, e.g., providing results on different tasks / environments.
soundness: 2
strengths: 1. Innovative approach using cognitive maps to emulate human-like System 2 reasoning in language models.
2. Rigorous experimental design testing multiple configurations of cognitive map construction (forward, backward) and planning tasks.
3. Experimental results demonstrate significant performance improvements.
summary: The paper investigates language models' extrapolation abilities in novel environments using a textualized Gridworld path-planning task. By introducing cognitive maps inspired by dual-process theory, the study demonstrates enhanced planning capabilities in language models over traditional Chain of Thought (CoT) approaches, particularly in extrapolated, larger environments.
weaknesses: I feel like the experiments can only prove that "cognitive maps" is a good method for the navigation task in Gridworld, but cannot support the claim that it is helpful for generalization. 

The generalization ability of LLMs comes from large-scale pretraining on diverse tasks so that the models learn the underlying rules behind them and can generalize to unseen tasks. In this sense, evidences given by **training**  (or say, overfitting) a language model on a very specific domain seems not so convincing. 

Specifically, training the model with cognitive map provides extra supervision signals, so the performance increase is reasonable on such a small dataset in a restricted domain, but it is unclear whether the conclusion will still hold when scaling up.

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: - I didn't find these details in the paper: how was the training set constructed? How did you obtain the cognitive map of training samples?  What are its statistics (e.g., lengths of inputs, outputs, and plans)?
- From Figure 2, it appears that the cognitive map traverses the entire search tree. Could this cause issues with length if the grid size is large?
- In the results shown in Table 1, performance reaches 1 when the grid size is small. Does this suggest that limited diversity may be an issue, allowing the model to solve test samples simply by memorizing the training samples?
soundness: 3
strengths: - The task is clearly defined, the process is detailed, and the experiments show effectiveness of their proposed method;
- The motivation is clear and writing is explicit, with precise language, well-organized structure, and clear communication of ideas;
- The analysis is fairly thorough.
summary: This paper introduces a path planning task in a textualized Gridworld, which requires a simulation process to obtain human-like cognition. 
They show that conventional approaches such as end2end generation and COT fail to generalize to larger environment. In stead, they design a  cognitive maps based approach to mimic human thinking process to enhance model planning abilities in extrapolated environments. Experiments show that their method can enable model to better generalize to larger sizes with better results.
weaknesses: - My primary concern is that "extrapolation" or "generalizability" extends beyond simply increasing the grid size. Rather than merely expanding the grid, it would be more insightful to evaluate model performance on grids that differ in structure or complexity from the training set;
- The baselines used are relatively simple. Incorporating stronger methods, such as [1] or Tree of Thoughts for planning construction, would better support claims about the model's effectiveness. Additionally, clarifying why the proposed cognitive map approach outperforms previous planning methods would strengthen the argument;
- The test set lacks diversity, raising questions about the model's generalization capabilities and its applicability to real-world scenarios.


[1] Reasoning with Language Model is Planning with World Model

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 2
questions: - What does "simulative reasoning" mean, and how is it different from "sequential logic" as in traditional CoT?
- What evidence is there that OpenAI-o1 does "tree search" or "internal simulation" different from traditional CoT?
- How does this work relate to [2], and what accounts for differences between the authors' results compared to [2]?
- How might the authors see their methods being used at scale and/or with real-world domains? Would "cognitive map-regarded training" be applicable to general-purpose LLMs like o1 and Claude, or would it only be used in a domain-specific fine-tuning stage?
soundness: 3
strengths: - Paper is well-motivated, and the question of whether in-context planning extrapolates to new data is compelling.
- Task seems appropriate for studying how algorithms succeed/fail to extrapolate.
- Problem formulation is thorough and described well
- A number of experiment details and different analyses provided in the Appendix.
summary: This work aims to test whether LLMs are capable of planning on out-of-distribution (extrapolated) data, and proposes a new method for fine-tuning LLMs to more effectively plan and generalize to new data. The authors propose to use a textualized gridworld as the domain of study, and OOD generalization is tested by training models on (N x N) grids and testing on (A x B) grids, where A and B may be greater than N. They test two task varients, offline and online planning, where in the first a model outputs an entire plan in one shot, and in the latter a model iteratively outputs a single action, then feedback from the environment is appended to the context, then the model produces another action, and so on. Their method consists of three stages where the LLM chains actions to and from the goal. The authors test a few variants on this method and compare to two baselines, one being a simpler CoT prompt. They fine-tune LLMs with data in their cognitive maps format, and find that with 500-1000 training steps the LLMs reach peak performance. Their method vastly exceeds the performance of their baselines in OOD generalization.
weaknesses: - The authors make strong claims about "simulation" in reasoning, but do not sufficiently describe what this means or provide concrete justifications for their claims. E.g. 
	- "The model [o1]’s success is attributed to its ability to conduct ”internal simulation” before providing an answer." I haven't seen any work making this claim, and none is cited here.
	- "These observations collectively imply that cognitive maps tap into a form of simulative reasoning that is fundamentally different from the sequential logic typically employed in CoT." I'm not sure how these results support this claim, or what "simulative reasoning" or "sequential logic" mean. 
- The authors' description of OpenAI's o1 model seems off in my reading, along with the interpretation of results. To my understanding, o1 is not "a tree-searching augmented LLM capable of internal simulation before planning". The OpenAI o1 blog post [1] specifies that the model uses a (seemingly more or less traditional) CoT that is merely hidden from the user, and even gives examples of these chains. Perhaps there's something I'm missing here.
- For the authors' claim that this "emphasize[s] the need for cognitive map-regarded training", it is unclear to me what this approach would be like for real-world domains and applications.
- Missing citations, in particular [2] seems highly relevant since it tests planning and cognitive map representations in LLMs, and similarly finds that current LLMs struggle to do planning in-context. However, the models they test seem to fare much better than these authors' baselines. Another related work not cited is [3].
- The explanation of the Cognitive Map method (Section 4) wasn't very clear to me, and figures 1 & 2 didn't help me much.
	- Fig. 1 - It's not obvious to me what the takeaways are or how to map this to the main theory/methods. The only clear takeaway I see is showing extrapolation to larger grids. The grid with the shaded area (bottom-middle) seems intended to emphasize "no cognitive map / no simulative reasoning" but it looks like it could mean partial vs. full observability or something else.
	- Fig. 2 - I'm not sure whether the "cognitive map" spoken by the agent is distinct, or in a different text format, from the steps printed in the "Output" panel. The numbers mentioned in the caption aren't shown in the figure. The figure could also reference the optimal vs. reachable plan distinction mentioned in the caption.



[1] https://openai.com/index/learning-to-reason-with-llms/

[2] Momennejad et al. (2024). Evaluating cognitive maps and planning in large language models with CogEval

[3] Yamada et al. (2023). Evaluating spatial understanding of large language models.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Question:
1. How would you cognitive map method generalise to other approaches?
2. Comparison to previous work: How does this work compare to approaches like: Act-Re (which has a fine-tuning mechanism), or StateAct, (which proposes a `state-tracking')?
soundness: 3
strengths: Strengths:
1. Great analysis of extrapolation ability of LLMs.
2. Detailed experiments and results
3. Strong and impressive results on the controlled task of grid-world navigation.
summary: The paper investigates the generalisation & extrapolation of LLMs in a controlled setting and proposes a new method for better extrapolation in the controlled setting. Specifically, the authors investigate whether and how LLMs generalise in the task of path planning. To this end the paper uses a textual grid-world of varying sizes. The authors train the models on grid sizes of 10x10 and then evaluate the models on grid sizes up to 20x20. Direct Answer and CoT produce poor results barely generalising beyond 10x10. The author's method of `cognitive maps' shows strong performance and generalisation to the grid sizes up to 20x20. Interestingly zero-/ and few-shot prompting elicits very poor performance except for models such as o1 (however, still lower than the author's approach). Finally, the author's detail their experiments and results in detail.
weaknesses: Weaknesses:
1. Focus on one control task, it would be interesting to do additional experiments of the like done in `Physics of Language Models' https://physics.allen-zhu.com/
2. Question of how to apply to tasks beyond the specific control task?

Meta-review: This paper studies the generalization capabilities of data-driven models trained on a toy grid-world task. The find that a CoT approach (dubbed "cognitive maps") enables better generalization to unseen environments.

On the positive side, this paper studies an interesting (albeit toy) setting, and the experiments are generally well done. On the negative side, it is debatable how generalizable the findings are to more real-world cases. Moreover, I find the "human-like" claim unconvincing, given that there were no human experiments to support this claim. (The authors have added discussion around this point during the rebuttal, but mere discussion is not enough in my opinion).

I am therefore recommending that this paper not be accepted.

Decision: Reject