Agent-R: Training Language Model Agents to Reflect via Iterative Self-Training

Large Language Models (LLMs) agents are increasingly pivotal for addressing complex tasks in interactive environments. Existing work mainly focuses on enhancing performance through behavior cloning from stronger experts, yet such approaches often falter in real-world applications, mainly due to the inability to recover from errors. However, step-level critique data is difficult and expensive to collect. Automating and dynamically constructing self-critique datasets is thus crucial to empowering models with intelligent agent capabilities. In this work, we propose an iterative self-training framework, Agent-R, that enables language Agent to Reflect on the fly. Unlike traditional methods that reward or penalize actions based on correctness, Agent-R leverages MCTS to construct training data that recover correct trajectories from erroneous ones. A key challenge of agent reflection lies in the necessity for timely revision rather than waiting until the end of a rollout. To address this, we introduce a model-guided critique construction mechanism: the actor model identifies the first error step (within its current capability) in a failed trajectory. Starting from it, we splice it with the adjacent correct path, which shares the same parent node in the tree. This strategy enables the model to learn reflection based on its current policy, therefore yielding better learning efficiency. To further explore the scalability of this self-improvement paradigm, we investigate iterative refinement of both error correction capabilities and dataset construction. Our findings demonstrate that Agent-R continuously improves the model's ability to recover from errors and enables timely error correction. Experiments on three interactive environments show that Agent-R effectively equips agents to correct erroneous actions while avoiding loops, achieving superior performance compared to baseline methods (+5.59%).

Host: Hey everyone, and welcome back to the podcast! I'm your host, Leo, and I'm super excited for today's episode. We're diving into some seriously fascinating stuff, a topic that's right at the forefront of AI research. It's all about how we can make AI agents smarter and more self-aware, kinda like giving them the ability to learn from their mistakes in real-time.

Guest: Hey Leo, thanks for having me. I'm thrilled to be here, and yeah, this is definitely a hot topic. We’re moving beyond just having AI follow instructions, now we're trying to give them that human-like ability to reflect and adapt, it’s pretty cool. Think about it, we don't always get it right the first time, and it's often our ability to course-correct that leads to success. So, the idea that we can translate that into AI is both challenging and incredibly promising.

Host: Absolutely! And it's not just about avoiding simple errors; it's about enabling AI to handle complex, multi-step tasks where a small mistake early on can snowball into a total failure. It’s like in a video game, where you might make a bad move early, and it affects your whole game, the difference is, our AI agent should be able to see where it went wrong and then fix its path right there, instead of waiting till the end. It's about making them resilient and adaptive, not just efficient.

Guest: Exactly, that's where the concept of 'reflection' comes in. It's a crucial thing for us, you know, thinking about what we’ve done, what went well, what went poorly, and then adjusting accordingly. And the challenge with AI is that, unlike humans, they often don’t have that natural feedback loop. Traditional methods of training AI agents mostly focus on teaching from perfect examples. So when they encounter a slightly different scenario where they are doing wrong, well, they don't always know how to get back on track, they just continue down the wrong trajectory.

Host: Yeah, that makes total sense. So, it's like, we’ve been showing them the ‘right’ way, but not teaching them how to deal with the ‘wrong’ way. It’s like giving someone a map but not teaching them what to do when they go off-road. That's a big gap. And I think this is particularly important in interactive environments where the agent is constantly making decisions and getting feedback from the environment. It's a dynamic process, not a static one, right? So just learning from perfect data is just not enough.

Guest: You nailed it. And what we’re seeing is that relying on just this ‘perfect’ training data leads to a situation where these AI agents struggle to proactively correct themselves, which leads to all sorts of cascading errors. Imagine a robot in a warehouse making a wrong turn; if it doesn't have the ability to realise its mistake and correct it, it might end up in the wrong area, and then messing up the orders, it just cascades. That's why the research into automated self-correction is so important. We need to create a system where AI can, in real-time, evaluate its actions, see where it’s veering off course, and then correct its trajectory, just like we do as humans.

Host: So, it sounds like what we need is not just to train them on successful paths, but to equip them with the ability to analyze why certain paths failed, and then re-route themselves. It's not enough to just 'know' what the right answer is, they need to understand the 'why' behind it all. It's like they need their own internal coach that can point out mistakes and guide them back on track. And I guess this is why the usual way that we train models through imitation just doesn't work that well for this problem, as they can only copy the ‘right’ path and cannot understand the ‘wrong’ paths.

Guest: Exactly. Think of it this way: if you're learning a new skill, like playing an instrument, you'll make mistakes. But it's the ability to recognize those mistakes, and then analyze them – like, why did that note sound off? – and then adjust your technique that leads to improvement. This is the same principle we are trying to instill in AI. And, well, the old methods of providing direct feedback like 'good action' or 'bad action' from reward functions don't work that well for complex interactive tasks where errors usually only become apparent later in the interaction, after a chain of events has taken place. It's hard to define the reward on a step-by-step level when you don't know whether an action will be good or bad in the long-term picture. Plus, these interactions can get so complicated that it's hard to make a function that can effectively critique each step.

Host: Okay, so it's not just about getting the final result right but also about learning from the process itself, including all the little stumbles along the way. And it's not a simple reward or penalty, it's a nuanced understanding of what went wrong and how to do it better. So, how do researchers tackle this issue of giving AI this reflective capacity? Because it sounds pretty challenging and like a real big problem, not something that can be easily overcome.

Guest: Yeah, it's definitely a complex problem. And the biggest hurdle is the lack of good 'reflection' data. Imagine trying to teach a model to recognize errors if you don't have examples of what those errors look like, and more importantly, what to do when it encounters these errors. It's like trying to teach someone to drive without ever showing them a car crash. The real-world interactive tasks are difficult, and these types of data are notoriously hard and expensive to gather as you need to get people to identify at what step, and why that is an error. So, in order to really help models learn self-correction, we have to come up with ways to automate this process so that models can learn by themselves, and this is really what we are talking about today.

Host: Okay, so we’re talking about an automated system that can generate self-critique data for the model, basically a way to have the model criticize itself and then learn to improve. It's no longer about us giving the model all the perfect examples, but about the model learning to self-correct on the fly, based on the experience that it is gathering in its own interactions. And that brings us to, I guess, the focus of today, the new approach that you are going to tell us about today.

Guest: That’s right, Leo. Today we're gonna be discussing a new framework called 'Agent-R', and it's all about enabling language agents to reflect on the fly using an iterative self-training approach. This approach moves away from traditional reward-based systems. Instead, it focuses on constructing training samples that help the agent recover from bad trajectories by identifying where they went wrong, and splicing in a corrected path. This allows them to not only avoid failures but also to understand the dynamics that cause the failures in the first place. So, instead of just being taught what to do, they start to understand why they should be doing it. This is quite a key distinction from the previous methods.

Host: That sounds incredibly innovative. So, it's not just about punishing bad actions or rewarding good ones, it’s about creating this whole dynamic system where the agent is constantly analysing itself, learning from its mistakes in real-time, and then correcting its path as it goes, instead of waiting to get a penalty at the end. So how does this 'Agent-R' actually work in practice? I mean, this is much more than just the usual backpropagation training that we are all familiar with.

Guest: Alright, let’s dive into the nuts and bolts of Agent-R. The framework is structured into two main phases. The first phase is called 'Model-Guided Reflection Trajectory Generation.' In this phase, we use Monte Carlo Tree Search, or MCTS for short, to explore potential action paths in the environment. MCTS is this great algorithm for decision-making, and in our setting it's used to systematically search through action paths so that we can generate a varied set of trajectories, some good, some bad. This process helps ensure that our model doesn’t just see one type of solution, and that we are getting a good amount of data for both the correct and the incorrect actions. What's really key here is that we don’t just passively observe the agent's actions, but we're actively using MCTS to generate both successful and unsuccessful trajectories so that we can have data about both.

Host: Okay, so MCTS is like the exploration tool, finding different possible paths in the environment. That makes sense, because you would need to get examples of all kinds of paths, not just the perfect ones, if you want the agent to be able to learn how to correct its errors. And this is really different from the traditional approach of training just on successful trajectories, as you've mentioned, right? So, instead of just following one set path, the agent is learning by navigating different possible solutions, some of which will succeed, and some of which will fail, is that it?

Guest: Exactly. But it's not just about generating random paths. We’re using MCTS, which operates by building a search tree, and simulating possible outcomes, to figure out which actions are more likely to lead to success. MCTS works in four main steps: selection, expansion, simulation, and backpropagation. In short, MCTS essentially uses simulations to decide which node to explore next, ensuring that we don’t just blindly search but try to pick paths that are most likely to lead to the goal. And we're also making sure to keep a balance between exploring uncharted territory and exploiting known good moves, so we can see a large variety of data. This will help us get both the bad and good trajectories that we can use to teach the agent to correct its own errors.

Host: Okay, I’m starting to see the big picture here. So, you are not just having the agent randomly wander around; it’s an intelligent exploration process, using MCTS, to find both the good and bad paths through the task, or the environment. But how do you then decide what is good or bad? I mean it is not always so black and white, because there might be some things that are good in the short term but then are not so great later on, and vice versa.

Guest: That's a critical question, Leo. And it's where our concept of 'reflection trajectories' comes in. We define four types of trajectories: 'initial trajectories,' which are just the starting points, and then 'bad trajectories,' 'good trajectories,' and finally, the 'revision trajectories.' We can think of it this way: the initial trajectory is the beginning of the task; then, a 'bad trajectory' is where the agent makes sub-optimal choices leading to a poor outcome, whereas a 'good trajectory' is when the agent makes correct choices that lead to a successful outcome. And the 'revision trajectory' is our secret sauce – it's what we get when we take a bad trajectory and splice it together with a good one. It’s about showing the agent how to transition from an incorrect path to a correct one.

Host: Okay, so the 'revision trajectory' is where the learning happens; it's not about the agent just avoiding the 'bad' path or solely following the 'good' path, but also understanding how you can switch between them. And this is the part that really helps it to reflect on its errors. It's like providing a detour sign on a road, instead of just showing the road or the wall. But how do you actually combine these ‘bad’ and ‘good’ trajectories to make a revision trajectory? Is that also part of this first phase, this 'Model-Guided Reflection Trajectory Generation'?

Guest: Yes, it is. This is where the model-guided reflection part comes in, and it’s a super critical component. The key challenge we were addressing is that of timely revision, which is different from just correcting an error at the end of a run. You see, if you only correct an error at the very end of a failed run, it doesn't help the agent learn how to recognize the error early on, and it doesn't prevent it from falling into that error again in the future. So, instead of waiting until the end of a bad trajectory, we’re having the agent identify its mistake based on its current capabilities, and then splicing a good trajectory from that exact point.

Host: So it's like the agent is not just making a random detour, but strategically identifying its first error and then switching to a better path from there. I guess that’s much more in line with how human thinking works when we reflect, right? We don't just restart everything every time we make a mistake; instead, we pinpoint where we went wrong and try to correct it right there and then. But how does the agent know when it has made a mistake? How do you get the model to identify this transition point?

Guest: That’s the million-dollar question. We actually have the language agent itself evaluate each action within its own self-generated bad trajectories to identify errors. It's really about seeing the error through the lens of the agent's current policy. The agent is basically acting as its own critic. When it recognizes an incorrect action, that’s where we truncate the bad trajectory and begin splicing the good trajectory, we are basically making a turn from the bad path to a good path, which will help the agent see that there was a mistake, and how it can recover. And to do this, we use a specially designed prompt that allows the agent to judge if an action is good, bad, or uncertain given the history of the task, so the prompt basically asks the agent to think about it. This step is essential, because it makes sure that the corrections are based on the agent’s learned dynamics and allows the agent to learn from its own errors. It also lays the groundwork for scalability, as the agent learns to refine the revision itself as it becomes better.

Host: Okay, so it’s not just a static rule but a dynamic evaluation process, where the agent itself, based on its understanding at that moment, identifies the turning point where it deviates from a good path. And this allows it to learn much more efficiently. I also like the idea of having a clear ‘revision signal’ that marks this transition, like you are explicitly telling the model ‘look here, this is a correction.’. This whole approach makes it much more interactive and responsive. So, this is the first phase, this clever way of generating the revision trajectories. What happens next in the second phase? Is this where the training starts?

Guest: You got it, Leo. The second phase is called 'Iterative Self-Training with Revision Trajectories.' In this phase, we are actually training the language agent using the self-generated revision trajectories. The goal is to make sure that the agent not only develops its self-reflection capabilities but also, as it is iterating, it's learning to make better decisions based on its own experiences. This is where that iterative process truly comes into play. And instead of just relying on those revision trajectories alone, we mix them with good trajectories and start training the agent. And over time, we gradually make the definition of the 'good trajectory' stricter, so it learns to go more toward the optimal trajectory. The system iteratively learns, from weaker to stronger behaviours, and it's a beautiful example of how an AI agent can learn from its own errors and also refine its own data generation process.

Host: So, instead of just being given the answers, the agent is learning from a mix of successful paths and corrected paths. And it is continuously refining the definition of what success looks like. It’s a continuous loop of error identification, correction, and learning. So, it’s not just correcting its behaviour but also getting better at identifying its errors in the first place, and how to recover from them. And then, over each iteration, the model becomes better at both the task itself and the reflection process. This is a really interesting idea. But I am wondering, why not just train on 'optimal' paths, I mean, wouldn’t that be better since those are technically the best paths?

Guest: trajectories can have some noisy actions, which can trap the agent into a dead loop. So what we're trying to get the agent to do is to not just do the actions, but to also reason about them. We have to help it to get out of the loops and see the errors so that it can refine its behaviour, rather than just trying to repeat the same actions over and over.

Host: Okay, so it's not just about reaching the destination but understanding the road, including all the detours and dead ends. And that’s what those revision trajectories are doing. So, the revision trajectories are not just about fixing errors but also about enriching the agent's overall understanding and capability by providing more information about where the agent is failing and why. And this also means you are not just training on one type of path, but you are letting the model explore all different parts of the task space, making it better at dealing with different scenarios. So, this approach, is it just a theoretical concept, or did you actually test it out?

Guest: Absolutely, Leo, it's not just theory. We put Agent-R through its paces across three very diverse and challenging interactive environments, and the results were pretty compelling. We used WebShop, ScienceWorld, and TextCraft, each with its unique challenges for language agents. So, WebShop is this virtual online shopping environment where agents have to navigate web pages, click buttons, and search for items. ScienceWorld is this text-based environment that tests agents' scientific reasoning abilities. And TextCraft is a virtual environment based on Minecraft, where agents have to craft items from raw materials. We chose these three because they represent a broad range of challenges that an agent has to be able to deal with, from web interaction to scientific reasoning and then virtual material crafting.

Host: Wow, that's quite a diverse set of environments. So, it wasn't just tested in one kind of domain. It's like testing an athlete not just in one sport, but across many different ones. It’s definitely a robust test for Agent-R. And, I guess, the idea was to see whether Agent-R could actually perform well in all kinds of situations? But before we go into the results, maybe you can briefly explain the baselines that you tested this framework against, because you mentioned that you compared against several methods, correct?

Guest: Yes, we did. We compared Agent-R against a range of both closed-source models, like GPT-3.5 Turbo, GPT-4 Turbo, and Claude 3, and open-source models, like Llama2-Chat, which is one of the popular open-source LLMs. Then we also compared against agents that were trained on expert trajectories like AgentLM and Agent-FLAN, and also against ETO, which is an agent trained using contrastive learning methods. And we also had a ‘direct revision’ baseline, which is similar to our method but without the clever way to determine the transition point from the bad to the good trajectory, basically, the model corrects only at the end of the bad trajectory. And by including all these baselines, we were able to really get a sense of how well Agent-R does compared to a broad array of other models, both new and traditional.