ATLAS : Agent Tuning via Learning Critical Steps (2025)

Abstract: 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.

## introduction an intelligent agent can perceive its environment, process information, and take actions to achieve specific goals or objectives @xcite . traditional ai agents, especially empowered by monte carlo tree search (mcts) or reinforcement learning, have shown great potential in complex and difficult tasks @xcite . however, these methods are task-specific and suffer from weak generalization across different task domains or environments. on the other hand, the recent rise of large language models (llms) @xcite @xcite , finetuning on crtical steps ... ... a1 a3 an example trajectory in env 1 a2 a4 example trajectory in env 2 ... a1 a3 a4 example trajectory in env 3 puzzle games science experiments house working web shopping weather prediction a2 a1 a1 a3 a4 an example trajectory in env 1 a2 a3 a4 an example trajectory in env 2 critical step selection multiple environments a1 a2 a3 a4 an example trajectory in env 3 ... ... ... 2023a,b; jiang et al., 2023) provide powerful foundation models with rich prior knowledge and reasoning capabilities, leading to remarkable generalization performance across various downstream applications @xcite . their unprecedented performance in different domains, including instruction following, reasoning, planning, and tool usage, makes llms ideal agents in multi-task settings @xcite @xcite @xcite @xcite . despite their effectiveness, transferring capabilities across tasks remains challenging due to overfitting to expert trajectories and training inefficiency. most existing agent tuning methods rely on imitation of every step of expert trajectories @xcite @xcite @xcite . for instance, fireact @xcite generates diverse expert trajectories to fine-tune llms with enhanced planning abilities, while agenttuning @xcite combines expert trajectories from various interactive environments with general instruction datasets to improve agents' generalization capability. however, fine-tuning on the entire expert trajectories can introduce expert bias, causing the model to overfit to specific behaviors and diminishing its ability to generalize to unseen environments @xcite . while experts in a few critical steps indeed provide key supervised information to improve the agent, the agent may already excel in other non-critical steps. additionally, fitting the full trajectories of one task may lead to declines in others due to the potential distribution gap and negative transfer between tasks @xcite . moreover, imitating expert trajectories on redundant steps or sub-optimal/replaceable actions incurs excessive training costs. to address the above challenges, we propose at-las, agent tuning via learning critical steps, a novel approach that identifies the critical steps within expert trajectories and only finetunes llms on those selected steps, which is illustrated in fig -ure 1 . by steering the training's focus to a few critical steps, we not only reduce the backpropagation cost for training but also lower the overfitting risk to whole trajectories. moreover, imitation on partial trajectories mitigates the expert bias, and encourages generalization across environments and tasks. in atlas, an oracle llm (gpt4o by default) is utilized as the selector to identify the critical steps semantically based on four criteria: key observations, plan formulation, recalling prior information, and pivotal actions. these critical steps are essential to the success of downstream tasks, but they are usually challenging to be generated by the base llms. our training focuses on finetuning base models on these selected steps. our experiments provide extensive comparisons between atlas and other agent-tuning strategies or llm agents. as indicated by the results in figure 2 , atlas-finetuned agent maintains and improves the generalization capabilities of base llms in diverse environments, excelling on both held-in and held-out tasks. our main contributions can be summarized as: ## preliminaries as a multi-task learning framework, we define a set of n environments as e. for any particular environment @xmath0, the agent's dynamics is formalized as a partially observable markov decision process (pomdp) (sutton, 2018) m e,i ≜ (s, a, o, t, r). in this formalization, s denotes the state space, a denotes the action space, o denotes the observation space, t : s × @xmath1, and r : s × @xmath2. upon receiving an instruction @xmath3, the agent takes action a t ∼ π θ (•|o t ; e, i) according to its policy π θ , where θ represents the policy parameters. for llm agents, π θ is an llm, and its output is not limited to actions and depends on the prompt. the environment then transitions to a new state s t+1 ∈ s based on the action, and the agent receives a new observation o t+1 ∈ o. this interaction continues as the agent engages with the cr itical step selection step-l evel fine-tuning cr itical obser vation -steps where llm agent correctly interpret or understand critical information in an observation plan creation -steps where llm agent create sub-goals based on prior observations and final objective self cor rection -steps where llm agent accurately recalls and evaluates previous actions or decisions, particularly after failure attempts cr itical action -steps where llm agent takes important actions that drive the entire process closer to the final objective o: goal: find two cellphones and put them in the bed a: check the desk first, where phones are usually kept. action: go to desk 1 plan creation o: on the desk 1, you see ... (no cellphone) action: open drawer 1 o: the drawer 1 is open. in it, you see nothing. ... non critical (follow plan) ... let's check the next drawer. action: go to drawer 2 the drawer 4 is open. in it, you see nothing. a: checked all the drawers, look for other places like cabinets. action: go to cabinet 1 self correction non critical (capable) o: in cabinet 1, you see a cellphone 2, and a houseplant 1. ... a: action: take cellphone 2 from cabinet 1 critical observation ... o: you go to bed 1, you see ... o: a: action: put cellphone 3 in/on bed 1 critical action base llm selector llm agent critical step loss finetuning critical step definition expert trajectories a: a: critical steps environment until the task is completed or the maximum step limit is reached. we adopt the react framework @xcite for the llm agent: at each step t, the agent first generates reasoning thought h t with a think prompt prompt h . define the trajectory up to step-t as the action is then drawn from the same model with another prompt prompt a , i.e., the return of trajectory τ = τ 1:t , g(τ ) = t @xmath4, reflects the agent's performance in task i within environment e. we collect a dataset d of expert trajectories for all environments in e and tasks in i. ## methodology in this section, we introduce our atlas, which identifies critical steps of the expert trajectories and finetunes llm agents solely on them. ## conclusion in this work, we propose a novel method, atlas, for fine-tuning llm agents by focusing on critical steps in expert trajectories. using an oracle llm as the selector, we first construct a dataset of critical steps, which highlights the most impactful actions in the trajectories. we then optimize the llms by computing and minimizing the loss exclusively on these critical steps. experimental results demonstrate that by leveraging only 30% of the steps, our method achieves superior performance in training environments compared to finetuning on 100% of the trajectories. moreover, the improved generalization ability of our finetuned agents is also reflected in enhanced performance on held-out tasks. these findings underscore the importance of critical step training in maintaining and enhancing the strengths of llm agents while significantly reducing training costs. this approach offers a cost-effective solution for fine-tuning llms across diverse tasks and environments.