1 Introduction↩︎

Large language models (LLMs) [1]–[5] have enabled increasingly capable autonomous agents for long-horizon tasks [6]–[11]. Yet robust generalization remains a central bottleneck. Even with reinforcement learning, agents trained for interactive decision making can overfit to their training environments and generalize poorly to out-of-distribution (OOD) tasks [12], [13]. This challenge is especially acute in long-horizon settings, where each action must stay grounded in the task instruction despite long interaction histories and many spurious correlations. Across tasks, local observations and available actions can remain highly similar, making it easy for the model to rely on state cues rather than the intended goal.

A central question, then, is whether current LLMs truly solve tasks by following the instruction or instead memorize how familiar tasks are usually carried out. To study this question, we use ALFWorld, ScienceWorld, and WebShop [7], [14], [15], and manually corrupt the task description, making it impossible to interpret while leaving the rest of the prompt unchanged. However, current models, including GPT-5.4, often fail to recognize that the instruction is invalid and instead continue as if it were valid, even when explicitly told that they can ask for clarification.

In many such cases, the model effectively reconstructs the original task from the corrupted input. This behavior reveals a specific failure mode: when the prompt resembles a familiar task, the model defaults to a learned policy even though the instruction no longer cleanly supports that behavior. We term this phenomenon task insensitivity.

We further ask whether the same failure mode harms OOD generalization when models face similar but distinct tasks. In controlled task splits for all three environments, agents trained with supervised fine-tuning or reinforcement learning often reuse action patterns learned during training even when the new task requires different behavior. For example, as shown in Figure 1, a model trained on heat bread and place it on countertop may choose go to microwave 1 when evaluated on cool bread and place it on the countertop [16], despite knowing that a microwave should not be used for cooling when asked directly.

To understand why this shortcut behavior develops during training, we analyze attention during action generation. We find a consistent drift: attention to the task instruction decreases, while attention to current observations increases. This pattern is consistent with a policy that increasingly predicts actions from local context rather than from the task itself.

Taken together, these diagnostics point to a training gap: standard next-action supervision provides no explicit pressure to preserve strong dependence on the task instruction. To address this gap, we introduce Task-Perturbed NLL Optimization, a lightweight contrastive regularizer that encourages the original action to become less likely when the task description is replaced with a different task.

Our contributions are as follows:

• We provide empirical evidence that LLM agents suffer from task insensitivity: they often reconstruct and execute familiar tasks under corrupted or replaced instructions, even when clarification is explicitly permitted.

• We use task perturbations as behavioral probes and show that the same failure mode appears under controlled OOD task shifts, where action dependence on the task instruction weakens over training and attention drifts from task tokens toward local observations.

• We introduce Task-Perturbed NLL Optimization, a lightweight contrastive regularizer that improves task sensitivity and OOD robustness in the controlled settings we study while preserving more stable task-focused attention.

2 Related Work↩︎

Agent Generalization. Recent studies of agent training dynamics and RL fine-tuning suggest that improved optimization does not automatically yield OOD robustness [12], [16]; models may continue to rely on brittle heuristics inherited from the training distribution [12], [13], [17]–[19]. Most of this literature evaluates final success under distribution shift. Our focus is more specific: we ask whether the learned action policy remains conditionally grounded in the task instruction as training progresses.

Shortcut learning. A broad literature shows that neural models often rely on shortcuts and spurious correlations that perform well in-distribution but fail under distribution shift [20]–[22]. In language models, instruction tuning improves task following [5], [23], but optimization can still drift toward easier local regularities rather than the intended instruction semantics [24], [25]. Here we focus on agentic training, and we show that current models gradually attend less to the task description and become more prone to shortcut learning.

3 Diagnostics: Task Insensitivity in Agentic Training↩︎

3.1 Evidence of Task-Level Overfitting↩︎

To isolate whether strong benchmark performance reflects instruction following or memorization, we design a minimal diagnostic: we corrupt the task description while keeping the agent prompt otherwise intact. Concretely, for each task in ALFWorld, ScienceWorld, and WebShop [7], [14], [15], we manually rewrite the task description to be lexically similar to the original but semantically ambiguous or underspecified. For ALFWorld and ScienceWorld, we directly remove descriptive words; for example, "cool the bread and put it on the countertop" can be rewritten as "cool bread countertop," which cannot be interpreted unambiguously. For WebShop, we introduce conflicting attributes such as "loose fit and slim fit". We use 300 samples for each environment for evaluation.

We show some examples in Figure 2. Critically, we explicitly instruct the model to ask a clarification question if the task is unclear; the full prompt template is provided in the appendix.

Under this setup, the model has two options: it can request clarification, or it can treat the corrupted description as sufficient and proceed to action. We define two metrics: Inquiry denotes the fraction of cases in which the model asks a clarification question, and Hit denotes the fraction of cases in which the model takes the same action it would take under the original, uncorrupted task.

As shown in Table 1, models usually choose to act on the corrupted task description even when they are explicitly allowed to ask for clarification. Rather than flagging the corrupted instruction as unclear, they often infer a well-formed task and attempt to solve it. A high Hit rate does not, by itself, prove that the model has semantically reconstructed the original task; however, it does show that the corrupted prompt often elicits behavior aligned with the original task.

To more directly estimate whether the model semantically reconstructs the original task intent, we use GPT-5.4 as a judge model to determine whether a generated response appears to be attempting the original task. We randomly select 100 samples and compare the GPT-5.4 judgments with human annotations. GPT-5.4 agrees with the human judgments on 91% of the samples. As shown in Figure 3, corrupted tasks frequently elicit responses consistent with the original intent. Taken together, these results indicate that when faced with an ambiguous instruction, models default to familiar policies rather than reasoning from the provided text—a phenomenon we term task insensitivity.

3.2 Training-Time Dynamics Associated with Overfitting↩︎

The corrupted-task diagnostic above shows that task insensitivity is present across models even without deliberate training. We next ask how this failure mode evolves when agents are explicitly trained on a restricted set of tasks. To study this, we train agents via supervised fine-tuning and reinforcement learning, and track diagnostic metrics across training checkpoints.

Figure 4 shows the probability that the model’s predicted action under a corrupted task description remains aligned with the original task. This probability increases over training, which suggests that the model becomes increasingly likely to recover familiar task patterns from incomplete or ambiguous instructions.

In addition, the training responses always restate the task description at the beginning of the response. We therefore evaluate the similarity between the restated task and the original task when the model is asked to solve the corrupted task. Figure 5 shows this similarity, measured by bge-base-en. This similarity increases over training, which suggests that the model becomes increasingly likely to recover familiar task patterns from corrupted instructions.

The corruption-based diagnostic establishes that agents can ignore or reconstruct underspecified instructions. We next ask whether the same failure mode matters for realistic generalization: when the task is valid but differs from training, does the model still over-rely on familiar action patterns?

We next move to a more realistic OOD setting in which the model is trained on one subset of task types and evaluated on another. In ALFWorld, ScienceWorld, and WebShop, we split each environment into disjoint training and evaluation task groups along semantic categories (details in the appendix). For example, in ALFWorld we train on “heat and place” tasks and evaluate on “cool and place” tasks. After training, the model often reuses training-set heuristics in contexts where they are inappropriate: an agent trained on “heat \(A\) and place it on \(B\)” may still go to the microwave when asked to “cool \(A\) and place it on \(B\)”. Crucially, this is not a knowledge failure—when asked directly whether a microwave should be used for cooling, the same model answers correctly.

To test this account more directly, we conduct a controlled task replacement experiment. For these training prompts, we replace the original task description with similar OOD task descriptions, and manually filter out cases where the original action would still solve the new task (retaining only pairs that genuinely require a different action, with 200 samples per dataset). As shown in Figure 6, we track action consistency across training checkpoints and find that the model becomes increasingly likely to predict the same action even when the task changes.

These results show that training can also increase overfitting to the seen task distribution: when faced with similar but different tasks, the model may still reuse the action pattern associated with the training task even after the task has changed.

4 Observed Attention Drift During Training↩︎

To understand the mechanism behind the shortcut behavior reported in Section 3, we study attention patterns during action generation. We decompose the agent prompt into five regions: Overall task, Response style, Current observation, available actions and Others. As shown in Figure 7, during training, attention to the overall task region decreases, while attention to the current observation region increases (see Appendix for the full per-region breakdown). The observed pattern suggests a relative reweighting rather than a complete collapse of task use. In most models, task tokens still receive substantial attention throughout training, but the gap between task-related attention and observation-related attention narrows over time. This is precisely the regime in which errors such as confusing “heat” with “cool” become plausible: the model is not blind to the task, but it may no longer weight task information strongly enough when local context supports a familiar policy. An analogous trend also appears under reinforcement-learning training: Appendix Figure 16 shows that attention to the task instruction also declines across GRPO checkpoints. In these environments, “cooling something” and “heating something” can induce nearly identical local states before the heating or cooling action is taken; the task instruction may be the only meaningful difference. If the model underweights the task instruction, it can easily choose the wrong tool, such as a microwave for a cooling task.

We view the observed attention drift as a plausible optimization bias rather than a formal property of transformers. In long-horizon agent training, the task instruction is nearly constant within a trajectory, while the current observation and recent history change at every step. Because supervision is applied at every decision point, these dynamic signals can continue to provide fresh action-discriminative evidence, especially after the model has already captured the coarse task identity.

A simple toy decomposition makes this intuition more concrete. Consider the simplified predictor \[\hat{y}_t = w_s x_s + w_d x_t + b,\] where \(x_s\) is a trajectory-level task signal shared across the episode and \(x_t\) is a step-varying state signal. Under the squared loss \[L = \frac{1}{2N}\sum_{t=1}^{N}{(\hat{y}_t - y_t)}^2,\] the gradients satisfy \[\frac{\partial L}{\partial w_s} = x_s \bar{\delta}, \qquad \frac{\partial L}{\partial w_d} = \frac{1}{N}\sum_{t=1}^{N}\delta_t x_t,\] where \(\delta_t = \hat{y}_t - y_t\) and \(\bar{\delta}=\frac{1}{N}\sum_{t=1}^{N}\delta_t\). The main message is qualitative. The update associated with the static task signal depends on the average residual over the trajectory, while the update associated with the dynamic state signal depends on step-specific error-feature interactions. Once the model has already captured coarse task-level information, \(\bar{\delta}\) may become small, whereas \[\frac{1}{N}\sum_{t=1}^{N}\delta_t x_t\] can still remain informative because \(x_t\) varies across steps. This creates an optimization bias toward features that continue to explain the remaining step-level errors.

This toy view does not analyze transformer attention directly, and it does not imply that task tokens become useless. Rather, it offers a simple explanatory picture: during training, the model may gradually rely more on current observations and recent history for next-action prediction, which is consistent with the empirically observed decline in relative attention to task tokens.

In summary, the observed attention drift aligns with the behavioral evidence from previous sections: as training progresses, models become more vulnerable to relying on local, state-level cues when solving similar-but-different tasks. Although this analysis does not establish a complete causal account of internal decision-making, it offers a plausible optimization picture that directly motivates our intervention: if task insensitivity is associated with an optimization bias toward dynamic local-context features, then training should explicitly reward conditioning on the task instruction.

5 Method: Task-Perturbed NLL Optimization↩︎

The analysis above suggests that, in the settings we study, the issue is not a lack of access to task information, but rather a lack of training pressure to preserve task sensitivity. If local state and history are sufficient to explain the training action, standard next-action training provides little incentive to maintain a strong dependence on the task tokens. Our goal is therefore to add a direct signal that tests whether the predicted action still depends on the task. If the task instruction is replaced with a different task, the original ground-truth action should become substantially less likely. We turn this intuition into a contrastive regularizer.

5.1 Formulation↩︎

We augment standard supervised fine-tuning (SFT) with a task-perturbation regularizer. For each training instance, we construct a perturbed prompt by replacing the original task description with a similar but distinct task from the same environment. In practice, we sample replacement tasks from the top-\(10\) most lexically similar task descriptions in the training set (measured by edit distance). These nearby replacements create hard counterfactuals: the surrounding context remains plausible, but the correct action should typically change.

The core desideratum is simple: if the task changes, the original response should become less likely. A naive way to enforce this would be to directly enlarge the gap between the negative log-likelihood under the original prompt and that under the perturbed prompt. However, such an objective can over-penalize the perturbed case by driving \(NLL_{\text{perturb}}\) (the NLL of the response under the perturbed task) unnecessarily high. Instead, we ask only that the current model maintain a sufficient amount of separation between the two prompts, calibrated by a frozen reference model.

Specifically, for each instance \(i\), the reference model computes the token-level negative log-likelihood of the response under the original prompt and the perturbed prompt, denoted by \(NLL^{(i)}_{\text{vanilla,ref}}\) and \(NLL^{(i)}_{\text{perturb,ref}}\), respectively. We then define the reference ratio \[\rho^{(i)}_{\text{ref}} = \frac{NLL^{(i)}_{\text{perturb,ref}}}{NLL^{(i)}_{\text{vanilla,ref}} + \epsilon},\] where \(\epsilon > 0\) is a small constant for numerical stability. This ratio measures how much less compatible the original response becomes after the task is replaced, according to the reference model.

During training, the current model computes the corresponding ratio \[\rho^{(i)}_{\text{cur}} = \frac{NLL^{(i)}_{\text{perturb}}}{NLL^{(i)}_{\text{vanilla}} + \epsilon}.\] Rather than encouraging \(\rho^{(i)}_{\text{cur}}\) to grow without bound, we penalize it only when the current separation is weaker than the reference-calibrated one: \[\mathcal{L}^{(i)}_{\text{floor}} = \max\left( 0,\, \log \rho^{(i)}_{\text{ref}} - \log \rho^{(i)}_{\text{cur}} \right).\] The full training objective is then \[\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{SFT}} + \lambda \mathcal{L}_{\text{floor}},\] where \(\mathcal{L}_{\text{SFT}}\) is the standard supervised fine-tuning loss. This objective preserves the desired effect of task perturbation, namely making the original response less likely when the task changes, while removing the incentive to keep increasing \(NLL_{\text{perturb}}\) once a sufficient reference-calibrated separation has been reached.

5.1.0.1 Extension to RL.

We apply the same idea in reinforcement learning after rollout and advantage computation. We select \(k\%\) of samples with negative advantage, replace their task descriptions with perturbed tasks, and reuse the same advantage scores in the subsequent logit computation and gradient update. This discourages the model from assigning high probability to the original action once the task has been changed. We do not use a margin here because the advantage already controls the optimization strength.

6 Experiments↩︎

6.1 Experimental Setup↩︎

We evaluate our method on embodied-agent benchmarks with controlled train/test task splits designed to measure OOD generalization. Our primary environments are ALFWorld, ScienceWorld, and WebShop [7], [14], [15]. We train each model three times and report the mean performance across the three runs. Further details are provided in the appendix.

We consider two training settings. In the first, we apply Task-Perturbed NLL Optimization on top of supervised fine-tuning (SFT). In the second, we incorporate the same regularizer into reinforcement-learning-based agent training with GRPO [26]. Unless otherwise stated, all models use the same backbone, prompt template, decoding configuration, and environment interfaces within each comparison group. We compare against three baselines: vanilla SFT, contrastive instruction tuning (CoIN) [27], and task augmentation. CoIN augments tasks and uses contrastive learning to align similar tasks and separate dissimilar ones. We use GPT-5.4 to rewrite task descriptions for CoIN training. For task augmentation, we randomly replace the original task with one rewritten version, allowing the task wording to vary across steps for training. For SFT, we use \(\lambda=0.3\); for RL, we use \(k=30\%\).

6.2 Main Results↩︎

Figure 8 and Tables 2 and 4 summarize the main quantitative results on ALFWorld, ScienceWorld, and WebShop. Our method improves OOD performance in most settings. The gains are most consistent under GRPO and in task pairs where local state is highly similar across tasks but the correct action depends critically on the task instruction. This is precisely the regime in which shortcut learning should be most harmful in our controlled task-split evaluations. We study the effect of \(\lambda\) in Table 3. The RL results are reported in more detail in Appendix Table 4, which shows that the same intervention also improves both in-domain and OOD performance under GRPO for Qwen3-4B and Qwen3-8B.

Overall, the results suggest that explicit task-sensitivity regularization can improve OOD generalization in these controlled settings when the task instruction provides information that cannot be recovered reliably from state and history alone.

Figure 9 shows that under our method, attention to the task description is better preserved over training. Figure 10 shows that this increase is accompanied by a slower rise in conditional consistency across checkpoints. Since higher consistency means that the model repeats the same action after the task is replaced, the flatter trend indicates that our method suppresses the growth of task-insensitive behavior and better preserves sensitivity to task changes. Taken together, these patterns support the interpretation that the intervention is effective not because it adds a qualitatively new capability, but because it directly reinforces dependence on the task instruction.

7 Conclusion↩︎

We identify task insensitivity as a concrete and measurable source of OOD failure in language agents. Across three benchmark environments, we show that agents can reconstruct corrupted tasks, reuse familiar action patterns under controlled task shifts, and gradually drift away from task-focused attention during training. We further show that making task sensitivity an explicit training target through Task-Perturbed NLL Optimization improves OOD generalization in the controlled settings we study. Taken together, these results suggest that preserving action-level dependence on task instructions is a promising direction for improving agent robustness beyond standard next-action training.

Limitations↩︎

Our empirical study covers only three agent settings: ALFWorld, ScienceWorld, and WebShop. These environments span embodied tasks, scientific interaction, and tool use, but they still represent a narrow slice of agentic behavior. Our attention analysis is mechanistically suggestive rather than causal: the observed drift away from task tokens is consistent with our optimization account, but it does not by itself establish how the model internally computes its decisions.

LLM usage↩︎

We use LLMs to polish the paper and help create figures.

8 Additional Experimental Details↩︎

This appendix is organized into four parts. We first summarize additional experimental details, then provide supplementary theoretical analysis, implementation details for Task-Perturbed NLL Optimization, and finally collect prompts together with additional visualizations.

8.1 Dataset Split↩︎

We conduct experiments in three environments: ALFWorld, ScienceWorld, and WebShop. For all environments, we set the maximum number of interaction steps to 30. For ALFWorld and ScienceWorld, we follow the OOD splits of [13]. In ALFWorld, we designate Cool & Place and Pick Two & Place as held-out task types. In ScienceWorld, we reserve the final task type of each topic for OOD evaluation. For WebShop, we treat product categories as task types and randomly designate 30% of the categories as OOD, using the remaining categories for in-domain training and evaluation.

8.2 Training Details↩︎

We train for five epochs in SFT and 200 steps in GRPO. The learning rate is \(1\times10^{-5}\) for SFT and \(1\times10^{-6}\) for GRPO. All experiments are conducted on eight NVIDIA H20 GPUs. The batch size is 16 for SFT and 32 for GRPO, with a group size of 8 for GRPO. To penalize outputs that do not follow the required format, we apply a reward penalty of \(-0.1\). The KL regularization coefficient is set to 0.01.

9 Additional Theoretical Analysis↩︎

9.1 Gradient Analysis of Static Task Signals↩︎

This appendix provides a simplified analysis intended only to build intuition. It does not model the full optimization dynamics of autoregressive transformers, nor does it prove that task attention must decrease in practice. Instead, it isolates a basic optimization asymmetry between a trajectory-level task signal and step-varying state signals, offering intuition for why task information can become relatively under-emphasized during training.

Consider a toy predictor for the action representation at decision step \(t\): \[\hat{y}_t = w_s x_s + w_d x_t + b,\] where \(x_s\) is a static task feature shared across the trajectory, \(x_t\) is a dynamic state feature that varies with \(t\), and \(b\) is a bias term. We optimize the squared-error objective \[L = \frac{1}{2N}\sum_{t=1}^{N}{\left(\hat{y}_t - y_t\right)}^2.\] Let \(\delta_t = \hat{y}_t - y_t\) denote the residual. Then the gradient on the static-feature weight is \[\frac{\partial L}{\partial w_s} = \frac{1}{N}\sum_{t=1}^{N}\delta_t x_s = x_s \cdot \bar{\delta},\] where \(\bar{\delta} = \frac{1}{N}\sum_{t=1}^{N}\delta_t\) is the average residual. By contrast, the update on \(w_d\) is proportional to \(\frac{1}{N}\sum_{t=1}^{N}\delta_t x_t\), which continues to reflect step-specific variation through the changing feature \(x_t\). The same average residual \(\bar{\delta}\) also drives the bias update, so once coarse trajectory-level error has been largely absorbed by \(b\) and the static term, the remaining loss is more readily reduced by fitting the dynamic feature. In this sense, the static task signal can receive weaker continuing action-discriminative updates than the dynamic state signal. Although highly simplified, this toy view is intended only as explanatory intuition. It does not claim that transformer attention directly equals feature importance, nor that task attention must literally collapse. Rather, it highlights a training bias in which step-varying features may remain more useful for reducing residual error at individual decision points. This is consistent with our empirical observation that, over training, the model increasingly relies on the current observation and recent history, while the relative influence of task tokens declines.

9.2 Local Attention-Reallocation Analysis↩︎

This appendix gives a local first-order analysis of a simplified two-source softmax attention module. The goal is not to prove that attention drift must occur in a full Transformer. In full models, queries, keys, values, residual connections, MLP blocks, and layer normalization co-evolve during training. Instead, the analysis isolates a simple local mechanism: if a dynamic local-context representation provides a better first-order descent direction than a static task representation, gradient descent on the attention logits reallocates probability mass toward the dynamic source.

9.2.0.1 Setup.

Consider one action-prediction step. We approximate the prompt as containing two sources of information: a static task source with value \(v_{\mathrm{task}}\), and a dynamic local-context source with value \(v_{\mathrm{dyn}}\). The dynamic source may summarize the current observation, recent interaction history, and other step-specific prompt regions. Let \(s_{\mathrm{task}}\) and \(s_{\mathrm{dyn}}\) be their attention logits, and define the relative dynamic logit \[r = s_{\mathrm{dyn}} - s_{\mathrm{task}}.\] The corresponding two-source softmax weights are \[\alpha_{\mathrm{dyn}} = \sigma(r), \qquad \alpha_{\mathrm{task}} = 1-\sigma(r),\] where \(\sigma(\cdot)\) denotes the logistic sigmoid. The attention output used for action prediction is \[o = \alpha_{\mathrm{task}} v_{\mathrm{task}} + \alpha_{\mathrm{dyn}} v_{\mathrm{dyn}}.\] Let \(\ell(o,y)\) be the action-prediction loss, and define the output-space gradient \[g = \nabla_o \ell(o,y).\]

9.2.0.2 Proposition.

Assume the attention weights are non-degenerate, \[0 < \alpha_{\mathrm{task}} < 1, \qquad 0 < \alpha_{\mathrm{dyn}} < 1.\] If the dynamic source has a local predictive advantage over the task source, \[g^\top (v_{\mathrm{dyn}} - v_{\mathrm{task}}) < 0,\] then an infinitesimal gradient descent step on \(r=s_{\mathrm{dyn}}-s_{\mathrm{task}}\) increases \(\alpha_{\mathrm{dyn}}\) and decreases \(\alpha_{\mathrm{task}}\).

9.2.0.3 Proof.

Since \(\alpha_{\mathrm{dyn}}=\sigma(r)\) and \(\alpha_{\mathrm{task}}=1-\sigma(r)\), we have \[\frac{\partial \alpha_{\mathrm{dyn}}}{\partial r} = \alpha_{\mathrm{dyn}}\alpha_{\mathrm{task}}, \qquad \frac{\partial \alpha_{\mathrm{task}}}{\partial r} = -\alpha_{\mathrm{dyn}}\alpha_{\mathrm{task}}.\] Therefore, \[\frac{\partial o}{\partial r} = \alpha_{\mathrm{dyn}}\alpha_{\mathrm{task}} \left( v_{\mathrm{dyn}} - v_{\mathrm{task}} \right).\] By the chain rule, \[\frac{\partial \ell}{\partial r} = g^\top \frac{\partial o}{\partial r} = \alpha_{\mathrm{dyn}}\alpha_{\mathrm{task}} g^\top \left( v_{\mathrm{dyn}} - v_{\mathrm{task}} \right).\] The non-degeneracy assumption implies \(\alpha_{\mathrm{dyn}}\alpha_{\mathrm{task}}>0\). Hence, under the local predictive advantage condition, \[\frac{\partial \ell}{\partial r}<0.\] For a sufficiently small gradient descent step with learning rate \(\eta>0\), \[r^+ = r-\eta\frac{\partial \ell}{\partial r} > r.\] Since \(\sigma(r)\) is strictly increasing, \[\alpha_{\mathrm{dyn}}^+ = \sigma(r^+) > \sigma(r) = \alpha_{\mathrm{dyn}}.\] Because \(\alpha_{\mathrm{task}}=1-\alpha_{\mathrm{dyn}}\), it follows that \[\alpha_{\mathrm{task}}^+ < \alpha_{\mathrm{task}}.\] This proves the claim. \(\square\)

9.2.0.4 Expected reallocation over decision steps.

Now consider decision steps \(t=1,\ldots,N\), each with relative logit \(r_t\), attention weights \(\alpha_{\mathrm{dyn},t}\) and \(\alpha_{\mathrm{task},t}\), output gradient \(g_t\), and source values \(v_{\mathrm{dyn},t}\) and \(v_{\mathrm{task},t}\). A first-order gradient descent step gives \[\Delta r_t = -\eta \alpha_{\mathrm{dyn},t}\alpha_{\mathrm{task},t} g_t^\top \left( v_{\mathrm{dyn},t}-v_{\mathrm{task},t} \right).\] Therefore, if \[\mathbb{E}_t \left[ \alpha_{\mathrm{dyn},t}\alpha_{\mathrm{task},t} g_t^\top \left( v_{\mathrm{dyn},t}-v_{\mathrm{task},t} \right) \right] <0,\] then \[\mathbb{E}_t[\Delta r_t] > 0.\] Thus, when the dynamic local-context source has an average first-order predictive advantage, gradient descent increases its expected relative attention logit.

9.2.0.5 Interpretation.

The condition \(g^\top(v_{\mathrm{dyn}}-v_{\mathrm{task}})<0\) means that moving the attention output from the static task representation toward the dynamic local-context representation reduces the loss to first order. Under this condition, the softmax attention update increases the relative logit of the dynamic source and decreases relative task attention.

This local mechanism complements the static-versus-dynamic gradient intuition in the main text. Static task tokens can identify the broad task, but once coarse task-level information has been learned, dynamic local-context tokens may continue to explain residual step-level action errors. Whenever this happens, gradient descent can reallocate attention toward the dynamic source. In the simplified two-source model, this appears as attention drift away from task tokens. In full Transformers, this should be interpreted only as a local mechanism consistent with the empirical drift we observe, not as a proof that attention drift must occur.

10 Implementation Details of Task-Perturbed NLL Optimization↩︎

Response masking. The response template begins with a task-restatement field (the initial Task:… slot), which explicitly repeats the original instruction. If this field were included in \(NLL_{\text{perturb}}\), the perturbed prompt would be penalized for not matching the copied task string rather than for predicting the action itself. Therefore, when computing the task-perturbation regularizer, we mask tokens in this field by setting their labels to IGNORE_INDEX. This masking is applied only to the regularizer; the main SFT loss is computed on the full response.

Reference-ratio computation. For each training instance, we use a frozen reference model to compute \[\rho^{(i)}_{\text{ref}} = \frac{NLL^{(i)}_{\text{perturb,ref}}}{NLL^{(i)}_{\text{vanilla,ref}} + \epsilon},\] where \(NLL^{(i)}_{\text{vanilla,ref}}\) and \(NLL^{(i)}_{\text{perturb,ref}}\) are the negative log-likelihoods of the response under the original and perturbed prompts, respectively, and \(\epsilon > 0\) is a small constant for numerical stability. This reference ratio serves as a sample-wise target for how strongly the current model should separate the original and perturbed prompts.

Training objective. For the current model, we compute \[\rho^{(i)}_{\text{cur}} = \frac{NLL^{(i)}_{\text{perturb}}}{NLL^{(i)}_{\text{vanilla}} + \epsilon}.\] We then apply a hinge-style lower-bound regularizer: \[\mathcal{L}_{\text{floor}} = \frac{1}{B}\sum_{i=1}^{B} \max\left( 0,\, \log \rho^{(i)}_{\text{ref}} - \log \rho^{(i)}_{\text{cur}} \right).\] The final training objective is \[\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{SFT}} + \lambda \mathcal{L}_{\text{floor}},\] where \(\mathcal{L}_{\text{SFT}}\) is the standard SFT loss and \(\lambda\) is the regularization weight. This formulation preserves the desired effect of task perturbation while preventing optimization from unnecessarily inflating \(NLL_{\text{perturb}}\) once the reference-calibrated separation is reached.

11 Prompts and Additional Visualizations↩︎

11.1 Agent Prompts↩︎

You are an expert agent operating in the ALFRED Embodied Environment. Your task is to: {task_description}
Prior to this step, you have already taken {step_count} step(s). Below are the most recent {history_length} observations and the corresponding actions you took: {action_history}
You are now at step {current_step} and your current observation is: {current_observation}
The admissible actions in the current situation are: [{admissible_actions}].

Now it's your turn to take an action.
You should first reason step-by-step based on the current situation. This reasoning process MUST be enclosed within <thinking></thinking> tags.
Once you've finished your reasoning, you should choose an admissible action for the current step and present it within <action></action> tags.

You are an expert agent operating in the ScienceWorld environment, which is a text-based virtual environment centered around accomplishing tasks from the elementary science curriculum.
Your current task is: {task_description}

Prior to this step, you have already taken {step_count} step(s). Below are the most recent {history_length} observations and the corresponding actions you took: {action_history}
You are now at step {current_step} and your current observation is: {current_observation}
Here are the actions you may take:
{action_descriptions}

Current available actions:
{available_actions}

Now it's your turn to take an action. You should first reason step-by-step about the current situation. This reasoning process MUST be enclosed within <thinking></thinking> tags.
Once you've finished your reasoning, you should choose an appropriate action for the current step and present it within <action></action> tags.

You are an expert autonomous agent operating in the WebShop e-commerce environment.
Your task is to: {task_description}.
Prior to this step, you have already taken {step_count} step(s). Below are the most recent {history_length} observations and the corresponding actions you took: {action_history}
You have a total budget of {max_steps} steps for this task. You are now at step {current_step}, so you have {remaining_steps} step(s) remaining including the current step. If you use too many steps without completing the task, the episode will fail, so avoid unnecessary searching or paging when you already have enough information to make progress.
You are now at step {current_step} and your current observation is: {current_observation}.
The admissible actions in the current situation are:
[
{available_actions}
].

Now it's your turn to take one action for the current step.
You should first reason step-by-step about the current situation, then think carefully about which admissible action best advances the shopping goal. This reasoning process MUST be enclosed within <thinking></thinking> tags.
Once you've finished your reasoning, you should choose an admissible action for the current step and present it within <action></action> tags.

If we explicitly allow the model to ask for clarification, the prompt becomes:

You are an expert agent operating in the ALFRED Embodied Environment. Your task is to: {task_description}
If the task is unclear, you may ask a clarification question within <inquiry></inquiry> tags.
Prior to this step, you have already taken {step_count} step(s). Below are the most recent {history_length} observations and the corresponding actions you took: {action_history}
You are now at step {current_step} and your current observation is: {current_observation}
The admissible actions in the current situation are: [{admissible_actions}].

Now it's your turn to take an action.
You should first reason step-by-step based on the current situation. This reasoning process MUST be enclosed within <thinking></thinking> tags.
Once you've finished your reasoning, you should choose an admissible action for the current step and present it within <action></action> tags.
If the task is unclear, you may ask a clarification question within <inquiry></inquiry> tags.

11.2 Judge Prompt↩︎

Determine whether the generation is attempting to solve the original task.
Return valid JSON only.

[ORIGINAL TASK]
{task_description}

[GENERATION]
{model_response}

Judge whether the generation is trying to solve the original task.

Return JSON with EXACT keys:
{
  "reconstruction_label": "exact|partial|wrong|uncertain",
  "reason_short": "<one sentence>",
  "matched_evidence": "<short evidence span>",
  "conflict_evidence": "<short conflict span or empty>",
  "confidence": <float 0.0..1.0>
}

11.3 Per-Region Attention Breakdown↩︎

As a complement to Figure 7 in Section 4, we report attention allocation for each prompt region individually across training checkpoints.

References↩︎