PaperlayerView exact revision on arXiv

Source-linked AI summary

ATLaS: Agent Tuning via Learning Critical Steps

Zhixun Chen, Ming Li, Yuxuan Huang, Yali Du, Meng Fang, Tianyi Zhou

arXiv:2503.02197v2cs.CLcs.AIcs.LG

TL;DR

Full-trajectory agent tuning can introduce expert bias and weaken generalization, while also fitting redundant or replaceable actions. ATLAS selects critical steps for fine-tuning; using 30% of the steps, it outperforms 100%-step baselines, particularly in multi-task learning, and improves held-out generalization.

  • Problem

    Fine-tuning on entire expert trajectories can cause expert bias, overfitting, negative transfer across tasks, and excess training cost on redundant actions.

    Source: §1 · p. 2

  • Method

    ATLAS uses an oracle LLM to select critical steps from expert trajectories and fine-tunes the agent with training loss computed only on those steps.

    Source: §1 · p. 2

  • Results

    Using 30% critical steps outperforms baseline agents fine-tuned on 100% of steps, especially in multi-task learning, and ATLAS performs better on held-out tasks.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    ATLAS maintains and improves base LLM generalization across environments, including on held-out tasks.

    Source: §1 · p. 2, §6 · p. 9

    Critical-step selection depends heavily on powerful closed-source models and focuses primarily on semantic aspects.

    Source: §7 · p. 9

Abstract

from arXiv · show

Large Language Model (LLM) agents have demonstrated remarkable generalization capabilities across multi-domain tasks. Existing agent tuning approaches typically employ supervised finetuning on entire expert trajectories. However, behavior-cloning of full trajectories can introduce expert bias and weaken generalization to states not covered by the expert data. Additionally, critical steps, such as planning, complex reasoning for intermediate subtasks, and strategic decision-making, are essential to success in agent tasks, so learning these steps is the key to improving LLM agents. For more effective and efficient agent tuning, we propose ATLaS that identifies the critical steps in expert trajectories and finetunes LLMs solely on these steps with reduced costs. By steering the training's focus to a few critical steps, our method mitigates the risk of overfitting entire trajectories and promotes generalization across different environments and tasks. In extensive experiments, an LLM finetuned on only 30% critical steps selected by ATLaS outperforms the LLM finetuned on all steps and recent open-source LLM agents. ATLaS maintains and improves base LLM skills as generalist agents interacting with diverse environments.

1 Introduction

Full-trajectory tuning can introduce expert bias, overfitting, negative transfer, and unnecessary training costs. ATLAS selects critical steps for tuning; experiments report stronger results with fewer steps and improved generalization on held-in and held-out tasks.

1 Introduction

  • Imitating entire expert trajectories can introduce expert bias, overfit specific behaviors, hurt other tasks through negative transfer, and spend compute on redundant or replaceable actions.The cited discussion also notes a potential distribution gap between tasks.Source: §1 · p. 2
  • ATLAS selects critical steps from expert trajectories and fine-tunes only on those steps, reducing training cost and the risk of overfitting whole trajectories.The selector identifies steps such as key observations, plan formulation, recalling prior information, and pivotal actions.Source: §1 · p. 2
  • Using 30% critical steps outperforms 100%-step baseline agents, especially in multi-task learning, and ATLAS generalizes better on held-in and held-out tasks.Figure 2 compares ATLAS with full-trajectory tuning across three base LLMs and reports consistent outperformance.Source: §1 · p. 2, Figure 2 · p. 2

2 Preliminaries

The preliminaries model agents interacting with environments and tasks as POMDPs, then describe the agent’s observation, reasoning, and action cycle.

2 Preliminaries

  • Each environment-task pair is formalized as a POMDP with states, actions, observations, transitions, and rewards.The transition function maps a state-action pair to a state, and the reward function maps it to a reward.Source: §2 · p. 2
  • At each step, the agent samples an action from its policy using the current observation, environment, and task instruction, after which the environment returns a new observation.The interaction continues until task completion or the maximum step limit.Source: §2 · pp. 2–3
  • Under ReAct, the LLM first generates a reasoning thought and then draws an action from the same model using a separate prompt.Expert trajectories are collected across environments and tasks, with performance represented using discounted rewards.Source: §2 · p. 3

3 Methodology

ATLAS defines critical steps by their contribution to trajectory success, uses a selector to find them in expert demonstrations, and trains the agent with loss computed only on those steps.

3 Methodology

  • ATLAS identifies critical steps in expert trajectories and fine-tunes LLM agents solely on those steps.The methodology’s training objective focuses on the selected steps rather than every trajectory step.Source: §3 · p. 3, §3.3 · p. 4

3.1 Definition of Critical Steps

  • Critical steps require precise decisions and execution and substantially affect return, while non-critical steps allow changes without compromising the trajectory.The paper categorizes critical steps as Plan Creation, Critical Observation, Critical Action, and Self Correction.Source: §3.1 · p. 3

3.2 Selection of Critical Steps

  • A selector summarizes a trajectory’s high-level plan and sub-goals, then identifies and categorizes critical steps, selecting no more than a specified percentage of the trajectory.The paper motivates semantic, context-aware selection because task-critical steps can differ in meaning.Source: §3.2 · pp. 3–4

3.3 Agent Tuning on Critical Steps

  • Training retains the full trajectory as input but computes teacher-forced loss only on critical-step tokens, maximizing their likelihood under the model.The objective is defined over the critical-step dataset, with θ denoting the trainable LLM parameters.Source: §3.2 · p. 4, §3.3 · p. 4

4 Experiments

Across diverse held-in and held-out tasks, ATLAS fine-tunes on selected critical steps and improves agent performance. Experiments compare this approach with full-trajectory, non-critical, random, and value-based selection, and test different selection ratios and selector models.

4.1 Environments; 4.2 Implementation Details; 4.3 Baselines

  • Experiments span held-in and held-out environments, fine-tune Llama-3.1-8B on GPT-4o-filtered AgentTraj-L steps, and compare closed- and open-source models, tuned agents, unfiltered data, and perplexity-selected steps.The environment set includes tool-use and game tasks; the perplexity baseline selects the top 30% of steps by base-model perplexity.Source: §4.1 · pp. 4–5, §4.2 · p. 5, §4.3 · p. 5

4.4 Main Results

  • With only 30% of expert-trajectory tokens, ATLAS beats baselines on most held-in tasks and improves average performance by over 5% versus unfiltered tuning; critical-step training also enhances held-out generalization.It achieves comparable held-out performance without the large general-instruction data mix used by other approaches.Source: §4.4 · pp. 5–6, Table 1 · p. 6

4.5.2 Different Selection Ratio

  • A 30% critical-step selection ratio yields the best agent performance on both held-in and held-out tasks.The authors report only a marginal increase in selected critical steps above 30% and do not evaluate higher ratios.Source: §4.5.2 · p. 6

4.5.1 Critical Steps vs. Non-Critical Steps

  • Critical-step tuning improves held-in and held-out performance, whereas tuning on the same number of non-critical steps harms performance across tasks.Table 2 compares critical steps, non-critical steps, and complete trajectories.Source: §4.5.1 · p. 6, Table 2 · p. 7

4.5.3 Critical Steps vs. Random Steps

  • Randomly selected steps improve performance as their ratio increases, but remain inferior to ATLAS-selected critical steps at all three tested ratios.The comparison tests 10%, 20%, and 30% selection; the gap is particularly notable at higher ratios.Source: §4.5.3 · pp. 6–7, Table 3 · p. 7

4.5.4 Value Function defined Critical Steps

  • On BabyAI and Weather, ATLAS performs comparably to estimated value-function selection, whose rollout estimates can be inaccurate and computationally expensive.The value-function estimate uses repeated rollouts; some environments and long-horizon tasks make this approach infeasible or impractical.Source: Table 4 · p. 7, §4.5.4 · p. 7

4.5.5 Choices of Selector LLMs

  • GPT-4o-selected steps fine-tune agents better than steps selected by Llama3.1-70B-Instruct on Alfworld, BabyAI, and Weather.The comparison indicates that selector capability affects the quality of the critical-step dataset.Source: §4.5.5 · pp. 7–8, Table 5 · p. 8

4.6 Critical Step Verification

  • On 100 BabyAI and Maze tasks, initializing from critical steps yields notable gains over starting at the beginning of the task.The authors infer that the base model can generally generate non-critical steps autonomously, supporting a training focus on critical steps.Source: Table 6 · p. 8, §4.6 · p. 8

5 Related Work

Prior work explores LLM agents’ planning and interaction capabilities, alongside methods that fine-tune agents using diverse trajectories, datasets, and feedback.

5.1 LLM Agents

5.2 Tuning Agents

6 Conclusion

ATLAS selects critical steps from expert trajectories and trains only on those steps; using 30% of the steps outperforms full-trajectory fine-tuning in training environments and improves held-out task performance.

6 Conclusion

  • Using only 30% of expert-trajectory steps, ATLAS achieves superior performance in training environments compared with fine-tuning on 100% of trajectories.The conclusion also reports improved performance on held-out tasks.Source: §6 · p. 9
  • ATLAS uses an oracle LLM to select critical steps, constructs a critical-step dataset, and computes the training loss exclusively on those steps.Source: §6 · p. 9
  • The authors conclude that critical-step training improves generalization while reducing training costs and maintaining or enhancing LLM-agent strengths.Source: §6 · p. 9

7 Limitation

The authors identify dependence on powerful closed-source models and semantic-only selection as areas for improvement in ATLAS’s critical-step selection process.

7 Limitation

  • ATLAS’s critical-step selection depends heavily on powerful closed-source models, motivating approaches that reduce selection time and computational cost.Source: §7 · p. 9
  • Selection focuses primarily on semantics; combining semantic criteria with other metrics could improve selection accuracy and reduce fine-tuning tokens.Source: §7 · p. 9
  • The authors report improvements in efficiency and generalization for multi-environment fine-tuning while noting that aspects of the method remain open to improvement.Source: §7 · p. 9

A Environment Details

The appendix describes environments spanning household interaction, navigation, tool use, planning, crafting, and scientific tasks, with task-specific evaluation measures and step limits.

A Environment Details

  • ALFWorld tests household navigation and commonsense interaction, while BabyAI adapts grid-world instruction tasks and MAZE evaluates directional movement toward a goal.ALFWorld and MAZE use success rate; BabyAI uses task-specific rewards, with task-length caps of 30, 15, and 20 rounds, respectively.Source: A Environment Details · p. 12
  • SciWorld evaluates interactive scientific reasoning, and TextCraft evaluates hierarchical item crafting with success-based rewards.SciWorld trials are capped at 30 iterations; TextCraft tasks range from 1 to 4 procedural steps and allow up to 20 rounds.Source: A Environment Details · p. 13
  • Weather and Movie provide API-backed information-retrieval tasks, evaluated by success rate with per-task iteration caps.Weather offers 18 commands and a 10-iteration cap; Movie has 16 actions and a 12-iteration cap.Source: A Environment Details · pp. 12–13
  • WebShop evaluates product procurement and Wordle evaluates lexical deduction, both using success rate under capped task rounds.WebShop sequences are limited to 10 rounds and Wordle trials to 8.Source: A Environment Details · p. 13
  • Academia and TodoList assess tool-mediated information and agenda tasks, while Sheet evaluates spreadsheet edits against a reference template.Academia and TodoList use success rate; Sheet uses a 0–1 reward based on structural and content congruence.Source: A Environment Details · pp. 13–14
  • Jericho uses modular subtasks to keep interactive exploration tractable, while PDDL tests multi-step symbolic planning through translated textual states and actions.Jericho subtasks are limited to 15 steps; PDDL progress is measured by normalized matching to goal states.Source: A Environment Details · p. 14

B.1 Training Configuration

The appendix specifies training settings and describes perplexity-based and rollout-based approaches for selecting critical steps.

B.1 Training Configuration

  • Training uses Adam with a 2e-5 learning rate, cosine schedule, 3 epochs, 0.03 warmup, batch size 128, and maximum length 8192.The ATLAS fine-tuning run takes approximately 8 hours on 4 NVIDIA A100 80G GPUs.Source: §B.1 · p. 14

B.2 Perplexity Selection

  • Perplexity selection scores each step by exponentiated average negative token log-likelihood and selects the top 30% within each trajectory.Higher perplexity indicates a step that is more challenging for the model to generate.Source: §B.2 · p. 14

B.3 Value Estimation

  • Monte Carlo rollouts estimate action value from step t; a consecutive-step value difference exceeding 0.1 marks the corresponding expert action trainable.The discount factor γ is set to 0.99.Source: §B.3 · pp. 14–15

C Additional Results

Table 7 presents additional results for other backbones, random selection, and different critical-step selection ratios.

C Additional Results

  • Table 7 reports detailed results for other backbone models, random selection, and different critical-step selection ratios.The table caption says the best performance in each section is highlighted in bold.Source: C Additional Results · p. 15

D Prompts Details

The appendix describes how critical steps are defined and categorized, what the selection prompts request, and which task-inference prompts are reused.

D.1 Prompts of Critical Step Selection

  • A critical step is a key action or decision that, if performed correctly, significantly increases the likelihood of task success.Source: §D.1 · p. 15
  • The prompt categories cover plan creation, critical observation, decisive action, and self-correction after failure or suboptimal outcomes.Source: §3.2 · p. 4, §D.1 · p. 15
  • The task instructions request a plan, up to a specified maximum of critical action steps, and an explanation of each choice.Source: §D.1 · p. 15
  • The selection prompt asks for a high-level plan, chosen critical steps, and reasons linking each step to a category and subsequent mistake avoidance.Source: §D.1 · p. 15
  • Table 7 covers other backbone models, random selection, and different critical-step selection ratios; its caption says best performance per section is bolded.Source: Table 7 · p. 16

D.2 Prompts for Task Inference

  • For task inference, the authors adopt AgentGym's setup and reuse its prompts.Source: §D.2 · p. 16

Related papers

Loading 2503.02197v2…