Advancing Large Language Model Agent via Iterative Contrastive Trajectory Optimization

Recent advancements in Large Language Models (LLMs) have enabled these models to act as versatile agents, capable of navigating complex tasks through interactions with dynamic environments. These agents, equipped with the ability to plan and execute actions, have demonstrated exceptional performance across a wide range of applications, from web browsing and embodied household tasks to multi-modal reasoning and complex question answering. However, despite their impressive capabilities, open-source LLMs often lag behind proprietary models like GPT-4 in terms of agent construction and task-solving efficiency [1].

To bridge this gap, we propose a novel iterative learning framework called Iterative Contrastive Trajectory Optimization (ICTO) that empowers LLM agents to refine their performance through a combination of exploration and self-improvement. Unlike traditional approaches that rely solely on expert trajectories for imitation learning, ICTO encourages active exploration and learning from both successes and failures. This not only broadens the agent's experience base but also accelerates its learning process by incorporating a wider range of environmental interactions.

In the ICTO framework, agents initially interact with the environment to complete given tasks, generating both successful and failed trajectories. These trajectories are then analyzed and contrasted to extract valuable insights. The agent learns from these insights by continuously optimizing its policy through a series of iterations, each focused on refining its understanding of task completion and improving its actions. Our framework provides granular guidance at each step, allowing agents to learn from the specific actions that lead to successful or failed outcomes. Through iterative optimization, ICTO aims to refine the agent's actions and decision-making processes, ultimately enhancing its overall performance and adaptability in diverse environments.

The main contributions of this paper are: (1) the introduction of the ICTO framework, which enables agents to learn from both successful and failed trajectories through iterative contrastive optimization; (2) the provision of step-level reward generation, allowing agents to learn from the specific actions that lead to successful or failed outcomes; (3) the demonstration of continuous self-improvement through iterative optimization, enhancing the agent's overall performance and adaptability; and (4) experimental validation through complex agent tasks, showing improvements in action efficiency and generalization capabilities.

II.

Related works

Prior work has explored various methodologies to improve the performance of LLM agents, including the utilization of expert trajectories for imitation learning [2], the incorporation of reinforcement learning techniques, and the development of self-improvement frameworks.

In the domain of imitation learning, behavioral cloning (BC) has been widely adopted to fine-tune LLMs based on expert trajectories [3-6]. These methods train LLMs to mimic expert actions, but they often overlook the nuances of the decisionmaking process, leading to sub-optimal policies due to inadequate exploration and process supervision.

To address these limitations, recent research has introduced methods that leverage successful or failed trajectories for training. For example, Song et al. [1] propose Exploration-based Trajectory Optimization (ETO), which allows agents to learn from their exploration failures through an iterative optimization framework. Similarly, Xiong et al. [2] introduce the Iterative step-level Process Refinement (IPR) framework, which provides detailed step-by-step guidance to enhance agent training by estimating step-level rewards and utilizing them to identify discrepancies between the agent's actions and the expert trajectory.

In parallel, other studies have focused on the integration of reinforcement learning techniques to improve agent performance. For instance, Fu et al. [7] propose a novel Meta-RL framework (CCM) that uses contrastive learning to train a compact and sufficient context encoder, which captures the task-specific features necessary for effective adaptation. Yang et al. [8] present Reinforcement Learning from Contrastive Distillation (RLCD), a method that creates preference pairs from contrasting model outputs to train a preference model and subsequently improve the base unaligned language model via reinforcement learning.

Furthermore, Wang et al. [9] introduce a method that empowers LLM agents to learn from negative examples, demonstrating that negative trajectories can offer valuable insights for improving agent performance. Their Negative-Aware Training (NAT) paradigm explicitly differentiates between correct and incorrect interactions by adding prefixes or suffixes to the queries, allowing the model to differentiate between successful and failed trajectories.

III.

Methodology

A.

Problem Formulation

In developing an intelligent agent, we model the task as a Partially Observable Markow Decision Process (POMDPs), which is formally represented by a tuple (U, S, A, O, T, R). The components of this tuple are defined as follows: U:

The instruction space, representing the set of all possible tasks or commands that the agent might receive from an external source or user.

S:

The state space, which includes all possible states that the agent could occupy in a given environment. Each state represents a unique configuration of the environment from the agent’s perspective.

A:

The action space, encompassing all potential actions the agent can execute in various states. An action a_t ∈ A is taken at each time step t based on the current policy.

O:

The observation space, denoting the set of all possible observations the agent can perceive from the environment. Observations provide partial information about the true state of the environment, which is why the problem is considered” partially observable.

T:S×A→S:

The transition function, which describes the probability of moving from one state s_t ∈ S to another state s_t +1 ∈ S after taking an action a_t ∈ A. This function captures the dynamics of the environment.

R:S×A→[0,1]:

The reward function, defining the immediate reward r_t ∈[0,1] the agent receives after taking action a_t in state s_t. The reward function quantifies the desirability of actions in specific states, guiding the agent towards preferred behaviors.

The objective of the LLM agent is to learn a policy π_θ (a_t | u, s_t) parameterized by θ, which maps a given task instruction u ∈ U,current state s_t ∈ S, and observation o_t ∈ O to a probability distribution over possible actions A. The agent seeks to maximize the expected cumulative reward over time, which is mathematically formulated as: 1 $J (π_{θ}) = E_{π_{θ}} [\sum_{t = 0}^{\infty} γ^{t} R (s_{t}, a_{t})]$ where γ ∈[0,1] is a discount factor that balances the importance of immediate rewards versus future rewards.

B.

Iterative Contrastive Trajectory Optimization (ICTO) Framework

Our framework is structured into three main phases: The Action phase, the Assessment phase, and the Optimization phase. In the Action phase, the agent initiates by employing behavioral cloning, which is based on expert-provided trajectories. Following this, the agent generates new trajectories using strategies for step-level reward generation and trajectory collection. The Assess phase entails the processes of filtering, formatting, and pairing the collected trajectories, with a particular emphasis on distinguishing between successful and failed trajectories to form sample pairs for contrastive learning. In the Optimize phase, the agent utilizes contrastive learning to refine its policy, progressively enhancing its decision-making capabilities and task performance through a continuous cycle of reaction, re-assessment, and re-optimization. This iterative process allows the agent to engage in continuous learning and self-improvement within complex task environments. Figure 1 illustrates the proposed ICTO Framework.

1)

Action Phase

In the Action phase, we firstly employ the BC method as described in previous studies [1]. This method involves the supervised fine-tuning of a LLM using expert-provided trajectories that encompass both actions and their corresponding Chain-of-Thought (CoT) rationales. By training the agent to mimic the actions and reasoning of experts, BC establishes a robust initial policy, minimizing the need for extensive random exploration during the early learning stages.

Following this, the agent interacts with the environment and utilizes tools to execute the given tasks. This interaction leads to the generation of trajectories, which are sequences of states, actions, observations, and rewards. Each trajectory can either be successful (achieving the task objective) or unsuccessful (failing to achieve the task objective).

Let D^exp = {(u,a₁,o₁,r₁, …, a_n, o_n, r_n)} denote the set of collected trajectories, where:

u ∈ U is the task instruction.

a_t ∈ A is the action taken at step t.

o ∈ O is the observation received at step t.

r_t ∈ [0,1] is the reward obtained at step t.

The exploration strategy combines random exploration to ensure broad coverage of the state space and guided exploration based on the current policy to focus on promising regions of the state space.

2)

Assessment Phase

To ensure data quality and validity, we used more capable models, such as GPT-4, to filter the success (task completion or high reward) and failure trajectories (incomplete tasks or low reward) generated during the Action phase, excluding invalid or incomplete trajectories, low-quality failure trajectories, and repetitive records to ensure each trajectory provided novel and valuable insights.

Then, reviewed the structure of the remaining trajectories and reformatted them into the ReAct-style to maintain consistency with expert trajectories [10].

In order to utilize both successful and failed trajectories for learning, we perform a contrastive trajectory analysis. This process involves pairing each failed trajectory T_fail with a corresponding successful trajectory T_succ that accomplishes the same task under similar conditions. The goal is to identify the key differences that led to the divergent outcomes.

For each trajectory pair ⟨T_fail,T_succ⟩, we compute step-level rewards based on the difference in cumulative rewards between the successful and failed trajectories up to each step t. The step-level reward R_t for the action taken at step t is calculated as: 2 $R_{t} = \sum_{i = 1}^{t} (r_{i}^{succ} - r_{i}^{fail})$

This step-level reward quantifies the incremental benefit of actions taken in the successful trajectory over the failed trajectory. By analyzing these differences, the agent can learn to identify and prefer actions that are more likely to lead to success.

3)

Optimization Phase

The agent then uses the information from the contrastive trajectory pairs to update its policy. We apply Direct Preference Optimization (DPO) to optimize the policy. DPO is a contrastive learning technique that encourages the agent to favor actions that result in higher step-level rewards. The loss function for DPO is defined as: 3 $L_{DPO} = E_{T_{fail},}_{T_{succ} \sim D^{e x p}} [\log \frac{π_{θ} (a_{t}^{succ} ∣ u, s_{t})}{π_{θ} (a_{t}^{fail} ∣ u, s_{t})}]$ where π_θ (a_t | u, s_t) represents the probability of taking action a_t given the instruction u and state s_t. The objective of this loss function is to maximize the likelihood of actions taken in successful trajectories while minimizing the likelihood of actions taken in failed trajectories.

The optimization process involves adjusting the policy parameters θ to increase the preference for actions that are more likely to lead to success, effectively learning from both positive and negative experiences.

The ICTO method is inherently iterative. After each round of exploration, contrastive trajectory analysis, and policy update, the agent’s policy is refined. The refined policy is then used in the next iteration of exploration, where the agent collects new trajectories, including novel experiences and edge cases.

C.

Implementation Details

1)

Initialization

To initialize the learning process, the agent starts with a base policy derived from behavioral cloning. Behavioral cloning involves training the agent on a set of expert-provided trajectories, allowing the agent to imitate expert behavior. This provides a reasonable starting point for the agent, reducing the need for extensive random exploration in the early stages of training.

2)

Reward Model

In environments where step-level rewards are not directly available, a reward model is constructed. This model estimates the rewards based on observed state-action pairs. The reward model, parameterized by a neural network with parameters ϕ, is trained using the collected trajectories to predict rewards as follows: 4 $R_{ϕ} (s_{t}, a_{t}) \approx r_{t}$

The reward model training objective minimizes the mean squared error between the predicted and actual rewards: 5 $L_{reward} = E_{(s_{t}, a_{t}, r_{t}) \sim D^{e x p}} [{(R_{ϕ} (s_{t}, a_{t}) - r_{t})}^{2}]$

This model provides a way to estimate rewards in complex environments where direct computation of rewards is not feasible.

IV.

EXPERIMENTS

A.

Experimental Setup

1)

Datasets and Environments

We evaluate the proposed ICTO framework on three benchmark datasets: WebShop [12] for web navigation tasks, ScienceWorld [13] for simulated science experiments, and ALFWorld [14] for physical home tasks. Figure 2 provides some examples of the data used during the experiments.

The experiment environment is summarized in Table 1, our experiments were conducted on an Intel Core i9-10900K CPU and an NVIDIA Tesla V100 PCIe 32GB GPU. The LLM agent is trained using the Llama2-7B Chat model [11] as a basis. To enhance the agent’s capabilities, we Intel Core i9-10900K CPU and an NVIDIA Tesla To enhance the agent's capabilities, we implemented a 2-epoch fine-tuning process with a batch size of 64 and a cosine learning rate scheduler, where 3% of the total steps are used for the warm-up phase. The maximum learning rate is set to 5 × 10⁻⁵. For optimization, we used the AdamW optimizer. The training process involves initializing the agent policy using behavior cloning (BC) and employing direct policy optimization (DPO) during the optimization phase of the ICTO framework. Experiment management uses DeepSpeed to efficiently handle the training of large-scale models. Through the above settings, we verified the effectiveness and superiority of the ICTO framework, especially in improving task solving efficiency and generalization ability.

TABLE I.

Experimental environment

Component	Details
CPU	Intel Core i9-10900K
GPU	NVIDIA Tesla V100 PCIe 32GB
LLM Agent Model	Llama2-7B Chat
Optimizer	AdamW Optimizer
Experiment Management Tool	DeepSpeed

2)

Baselines

We compare ICTO against several baseline models to benchmark its performance: Supervised Fine-Tuning (SFT):

Utilizes expert trajectories for behavioral cloning.

ETO:

Leverages exploration failures for iterative optimization.

IPR:

Offers detailed step-by-step guidance to refine agent training.

RLCD:

Enhances the base language model by generating preference pairs from contrasting model outputs.

NAT:

Differentiates between correct and incorrect interactions by modifying queries with prefixes or suffixes.

3)

Evaluation Metrics

we employ several key metrics:

Average Reward This metric quantifies the mean cumulative reward achieved by the agent across all episodes, offering insight into overall performance and learning efficiency. The average reward R̄ can be expressed as: 6 $\bar{R} = \frac{1}{N} \sum_{i = 1}^{N} E [\sum_{t = 1}^{T_{i}} γ^{t} R_{t}^{(i)}]$ where N represents the total number of episodes, T_i is the total time steps in episode i, γ ∈[0,1] is a discount factor that balances immediate and future rewards, and $R_{t}^{i}$ denotes the reward received at time step t in episode i.

Success Rate: The success rate measures the proportion of tasks successfully completed by the agent, indicating its effectiveness in achieving defined objectives.

Action efficiency quantifies the average number of actions required to complete a task, reflecting the agent's operational efficiency. The metric is calculated as: 7 $η = \frac{1}{N} \sum_{i = 1}^{N} \frac{1}{T_{i}} \sum_{t = 1}^{T_{i}} E [I (a_{t}^{(i)} \in A_{optimal} (s_{t}^{(i)}))]$ where $a_{t}^{i}$ is the action taken at time step t in episode $i, s_{t}^{i}$ is the corresponding state and $A_{optimal} (s_{t}^{i})$ represents the set of optimal actions for state $s_{t}^{i}$ .

Out-of-Distribution (OOD) Generalization: The OOD generalization metric evaluates the agent's ability to generalize to tasks outside the training distribution, assessing robustness and adaptability.

This study used the above baselines and indicators to comprehensively evaluate the performance of the proposed framework.

B.

Results

This study reports on ICTO performance and baselines for three benchmark tasks. Table 2 summarizes the results. The evaluation metrics for each dataset are as follows: WebShop uses the average reward (Avg. Reward) as the performance indicator, ScienceWorld is evaluated based on the success rate (Success Rate), and ALFWorld utilizes action efficiency (Action Efficiency) as its metric. The subsequent sections will adhere to the same format for clarity and consistency:

TABLE II.

Comparison of ICTO and Baseline Performances

Method	WebShop	ScienceWorld	ALFWorld
SFT	63.1	70.0%	12.5
ETO	67.4	72.3%	11.2
IPR	68.3	73.8%	10.8
RLCD	65.8	71.5%	11.5
NAT	66.5	72.0%	11.0
ICTO (ours)	70.2	75.6%	9.7

ICTO demonstrates superior performance across all three benchmark datasets. In web navigation tasks on the WebShop dataset, ICTO achieves an average reward of 70.2, a success rate of 75.6%, and an action efficiency of 9.7, outperforming all baseline methods. In ScienceWorld, ICTO achieves an average reward of 67.3, a success rate of 72.5%, and an action efficiency of 10.2, effectively solving complex reasoning and planning tasks. In ALFWorld, ICTO records an average reward of 62.1, a success rate of 74.3%, and an action efficiency of 10.5, demonstrating its proficiency in embodied household tasks.

Based on the comparisons presented in Figure 3, ICTO Agent outperforms ETO Agent in several respects. ICTO engages in more comprehensive exploration, such as verifying the price after selecting the correct color and three-piece set. ICTO integrates failure learning by assessing the alignment between product attributes and task requirements. Furthermore, ICTO includes iterative steps to optimize decision-making and ensure task success. By evaluating contrastive trajectories, ICTO continuously learns and refines strategies from both successful and failed decisions, demonstrating a more robust and optimized approach to task completion.

Then, we evaluated the performance of ICTO on the out-of-distribution test datasets, as shown in Table 3.

TABLE III.

Generalization Performance of ICTO on OOD Tasks

Method	WebShop	ScienceWorld	ALFWorld
SFT	52.3	60.0%	15.0
ETO	55.8	62.0%	14.2
IPR	57.1	63.5%	13.8
RLCD	54.2	61.0%	14.5
NAT	56.0	62.5%	14.0
ICTO (ours)	59.5	66.0%	12.5

ICTO shows strong generalization capabilities on OOD tasks. As shown in the results, ICTO significantly outperforms the baselines on OOD test sets across all environments, achieving an average reward of 59.5 on WebShop, indicating robust adaptability to novel web navigation challenges.

Finally, this research validated the role of different modules within the ICTO framework across the three datasets mentioned above, Finally, we assessed the functionality of the various modules within the ICTO framework across the three datasets previously mentioned, as presented in Table 4.

TABLE IV.

Ablation Study of ICTO Modules

Training Scheme	WebShop	ScienceWorld	ALFWorld
w/o Contrastive Learning	64.2	67.8%	11.6
w/o Behavioral Cloning	60.7	62.5%	13.1
Iteration=1	66.1	69.2%	12.8
Iteration=2	68.5	70.6%	12.3
Iteration=3	70.9	72.3%	11.7
Iteration=4	72.3	73.1%	11.0
Iteration=5	72.0	72.8%	10.5

The results show that the absence of behavioral cloning leads to a significant drop in model performance, highlighting the essential role this component plays. Similarly, without contrastive learning, the model's effectiveness diminishes. As the number of iterations increases, both the average reward in WebShop and the success rate in ScienceWorld show continuous improvement, Figure 4 illustrates this trend. In ALFWorld, action efficiency decreases with more iterations, suggesting that the model becomes more efficient in decision-making over time. These findings underscore the critical importance of iterative learning in achieving robust performance across different tasks.

C.

Analysis

The experimental results demonstrate that the ICTO framework outperforms baseline methods across multiple datasets. On the WebShop dataset, ICTO achieved an average reward of 70.2, significantly surpassing SFT's 63.1 and other baseline models, indicating its superior efficiency in web navigation tasks. Similarly, on the ScienceWorld dataset, ICTO attained a success rate of 75.6%, outperforming baseline methods such as IPR and ETO, showcasing its exceptional performance in complex scientific experiment tasks. In the ALFWorld dataset, ICTO demonstrated an action efficiency of 9.7, which is lower than the higher values of 11.2 for ETO and 11.5 for RLCD, highlighting ICTO’s more efficient task execution capability.

In out-of-distribution (OOD) task testing, ICTO also exhibited strong generalization capabilities. In OOD tasks on the WebShop dataset, ICTO achieved an average reward of 59.5, notably higher than SFT’s 52.3 and RLCD’s 54.2, indicating its better adaptability to unseen tasks. Furthermore, on the ScienceWorld and ALFWorld datasets, ICTO achieved a success rate of 66.0% and an action efficiency of 12.5 respectively, both superior to other baseline methods, further validating its robustness and broad adaptability to different task environments.

Ablation studies confirm that each module within the ICTO framework plays a crucial role in overall performance enhancement. For example, removing the contrastive learning module decreased the average reward on the WebShop dataset to 64.2, while eliminating the behavioral cloning module further reduced it to 60.7. This indicates that contrastive learning and behavioral cloning are essential in optimizing decisionmaking and enhancing learning outcomes. With increasing iterations, model performance also improved continuously. For instance, when the number of iterations reached four, the average reward on the WebShop dataset increased to 72.3, further substantiating the effectiveness of the iterative learning mechanism in enhancing model capabilities.

D.

Discussion

The ICTO framework exhibits notable performance improvements over baseline methods, primarily due to its training strategy that integrates successful and failed trajectories. Compared with SFT, ICTO effectively leveraged failed exploration trajectories to improve the decisionmaking process. When compared with ETO, ICTO provided more refined step-by-step guidance, enhancing learning outcomes. Compared with IPR, ICTO’s contrastive learning mechanism offered stronger learning signals, resulting in better performance across various task environments. Additionally, ICTO surpassed RLCD and NAT, verifying the effectiveness of its contrastive learning and stepwise optimization strategies.

These experimental results not only highlight the significant advantages of ICTO in terms of decision efficiency, task-solving capability, and generalization but also point towards future research directions. Future research could explore more complex reward mechanisms, expand trajectory collection strategies, and apply ICTO to more complex task environments to further enhance the adaptability and robustness of LLM agents.

V.

Conclusions

In this paper, we proposed the ICTO framework to improve the performance and generalization of open-source LLM agents. ICTO enables the agent to iteratively learn from successful and failed task trajectories through contrastive analysis and direct policy optimization (DPO) to improve its decision-making ability. Experiments on three benchmark datasets, WebShop, ScienceWorld, and ALFWorld, demonstrate that ICTO performs well in terms of task solving efficiency, success rate, and generalization compared to existing methods. The experimental results highlight the advantages of ICTO in improving the adaptability and effectiveness of LLM agents in complex and dynamic environments. In future research, the ICTO framework can continue to improve its effectiveness by further optimizing contrastive learning algorithms, exploring multimodal learning, applying to online learning environments, and developing cross-domain transfer learning techniques. In addition, as ICTO is deployed in more application scenarios, considering its ethical and social impacts will also become an important part of research.

Langue:: Anglais

Périodicité:: 4 fois par an
Sujets de la revue:: Informatique, Informatique, autres

RSS Feed de la revue

Advancing Large Language Model Agent via Iterative Contrastive Trajectory Optimization

Chengang Jing

Xin Jing

Kun Li

Publié en ligne: 31 déc. 2024

Pages: 19 - 27

DOI: https://doi.org/10.2478/ijanmc-2024-0033

Mots clésIterative Optimization, Large Language Models, Agent

© 2024 Chengang Jing et al., published by Sciendo

This work is licensed under the Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International License.

Mots clés
Iterative Optimization, Large Language Models, Agent