How to Train Your Agent: Building Reliable Agents with RL

Talk transcript

1. 0:00

   [on-hold music] Um, hey, everyone.

2. 0:16

   Glad you're all here. This is the Reasoning and Reinforcement Learning track, uh, on the afternoon of the last day of the AI Engineer World's Fair. Glad you're all here.

3. 0:24

   Glad you're sharing it with us. Today, what I'm gonna talk about is, uh, a very specific case study, um, that we did. Uh, this case study, I'm gonna talk about lessons learned very concretely, um, what did and didn't work, how we were able to build an agent that worked well with reinforcement learning.

4. 0:38

   Uh, all of this-- uh, everything that I'm talking about in this presentation, this is an open source code base that we built. Um, we wanted to share these learnings, and I'll, I'll, I'll share that link with you at the end as well, um, for those of you who want to replicate what we did.

5. 0:51

   Um, so what is the project we're gonna be talking about? It's a project called ART·E. It is a natural language, uh, assistant that helps you answer questions from your email inbox.

6. 1:01

   So I'll give you an example of what we're talking about here. Um, let's say you want to ask, you know... In this case, our example question is, "When is Sherry's move to Portland targeted for?"

7. 1:10

   So you would ask this question to the assistant. It then goes and it searches your inbox. It's got several tools. So it has, like, a search tool, it has a read email tool, and then it can actually answer the final question.

8. 1:19

   You can kind of see if you, if you look here, what's going on behind the scenes. This is important, so you get a sense of kind of how this agent works and as we're talking through how we built it, how we made it work.

9. 1:27

   Um, hopefully that, that helps, uh, make the conversation very grounded in a specific task. So anyway, you see the agent, it's, it's, you know, searching for certain keywords. It get those messages back.

10. 1:37

    It's then reading one of them and, and answering the question. That's, that's what it does. Okay. So, um, question, you know, once, once we've decided this is kind of the, the task we're trying to solve, why would you re-use reinforcement learning for this specifically?

11. 1:51

    Um, and, uh, and the answer is, like, to start with, you shouldn't. In fact, to start off with, we did not. Um, so the first version of this agent, once we decided we wanted to build this, we did-- we didn't use any reinforcement learning at all.

12. 2:02

    We purely built this on prompted models. And this is the first lesson from this talk, uh, that I wanna share is, I would generally always recommend starting with getting the best performance you can with a prompted model before going to any training, including reinforcement learning.

13. 2:15

    There's a few different reasons to do that, uh, three specifically. Um, the first one is just, like, working out the bugs in your environment, right? Um, you know, maybe your tools aren't implemented properly.

14. 2:25

    Maybe they don't have access to the data you think they do. Um, we find this happens a lot, and it's a lot less frustrating to debug that, uh, you know, separately from debugging your, your training loop.

15. 2:34

    So you wanna make sure that, like, you can get at least some kind of performance, um, before you start training. Um, and then second of all, you may find as you're trying to improve the performance on, uh, on using these prompted models that, uh, you can get it working really well, and that's great.

16. 2:47

    So that means you don't need to train anything, um, and that saves you a lot of time. Um, there's a third reason as well, uh, that I'll share, which is, uh, basically once you've gone to that effort and you've done your best to get the best quality prompted baselines you possibly can, um, then, uh, if you find

17. 3:03

    that those baselines are not able to get you where you need to go and you're able to surpass them with re- reinforcement learning, it feels great. You get to gloat and be like, "Yes, I was able to beat the, the frontier models on my task."

18. 3:12

    Um, this, this... I, I, I highly recommend it. Feels good. You can, you can, like, post on X about it. There's, there's nice, you know, graphs and stuff. So this is, this is what it looks like when everything goes right.

19. 3:22

    Um, so this is an example of a training run for this ART·E model that I'm gonna be talking about. Uh, you can see that there's these, these lines for each of the prompted model baselines that we've got.

20. 3:31

    Um, so we've got o3, o4-mini, and then Gemini and, and four point one. And you can see, uh, those ones, you know, they have certain level of performance, and then you can see this, this, uh, sort of moving line, um, that's going on.

21. 3:43

    This is the model that we trained. And you can see it actually starts out significantly worse than these other models, uh, from, from the start. That's 'cause we started from a Qwen 2.5, the, the fourteen billion parameter one.

22. 3:53

    It's a relatively small model, relatively weak model, um, and so it was doing much worse than these initially. But you can see as training progresses, um, you know, initially at the beginning it-it's sort of...

23. 4:02

    Maybe, maybe there's-- it's learning the right way to do tool calls. There's some very sharp bump as it figures out the basic stuff, and then a more gradual climb until eventually it's able to significantly outperform, uh, any of the prompted models on this task.

24. 4:14

    And this is sort of what you're... You know, in the ideal case, when everything works, this, this is what you're looking for. This is what, what you're hoping to achieve.

25. 4:21

    Um, this is another view actually of that same data we were just looking at. Um, I, I like-- I, I wanted to highlight it in this way because, uh, it's important to realize...

26. 4:30

    So on the last graph, it looked like the, the lines sort of asymptote out pretty close together. That's because they're getting near a hundred percent. But the last, um...

27. 4:37

    You can see, for example, with our best prompted model here, o3, uh, it's ninety percent accuracy. And with our RL model, we're able to get up to ninety-six percent.

28. 4:46

    And so one way to think about that is, like, sixty percent of the errors that o3 was making, um, are, are actually solved with our model, um, which is, which is quite a large, uh...

29. 4:56

    You know, we find that that's actually can be very, very important for the user experience of someone using one of these. Um, if you're getting, you know, just half as many errors, uh, that, that can make the product much stronger.

30. 5:06

    Um, so this is, this is where we got to on accuracy. There's a couple other metrics that we f-find are often very, very important. Um, and, you know, the trade-off between these does, does-- is very task dependent, but, but they matter in many cases.

31. 5:19

    Um, cost obviously is a big one. So for, for this email agentic harness that we had, we benchmarked the cost on o3, o4-mini, and our model. So if you wanted to do like a thousand searches using o3, that's gonna cost fifty-five dollars, um, which is a lot.

32. 5:35

    Uh, I think for most use cases, that probably would be cost prohibitive just from a unit economics point of view. Um, on o4-mini, we're down to eight dollars, but that's still quite expensive.

33. 5:43

    And then we drop another order of magnitude by moving to the smaller Qwen 2.5 14B. Again, this is just driven by being, it being a much smaller model, so it's, it's much cheaper to run.

34. 5:52

    Um, but we're still able to get very good performance because we've specialized it on our task. Um, beyond cost and the accuracy, um, the third metric that often comes up is latency.

35. 6:02

    Uh, particularly if you're doing, I mean, certainly anything with voice, but if there's any real-time human interaction with the task, latency is going to matter a lot. Um, and we were able to find on, on this task, we were able to get significantly better latency.

36. 6:14

    There's a number of different ways, which I'll go into in more detail later, that we were able to achieve this. Um, you know, one was just, again, moving to a smaller model helps.

37. 6:21

    There's just less, less loading from memory, less matrix multiplies. It's just you're able to get tokens out faster. Um, we were also able to train this model to have fewer turns going back and forth with the database, with the actual email, um, the list of emails.

38. 6:34

    Uh, we, we were able to train it to be more efficient with its queries. Um, and I'll go into that in a moment. And so that, that leads to lower latency.

39. 6:41

    Um, there's actually a third thing which we didn't apply here, but can help a lot with these smaller things, which is called speculative decoding. That's something you can do on large or small models.

40. 6:48

    It generally works better on smaller task-specific models because you get higher, um, acceptance rates on, on your speculator. But basically, um, there's, there's lots of reasons why smaller models work better.

41. 6:57

    Um, okay. So then the next question, uh, for those of you who haven't done this yet is like, okay, what is the effort required to do this to actually achieve these results?

42. 7:06

    Um, if you'd asked me this question a year ago, I would say, "Hey, you should really only be doing this if, you know, you're this big company and willing to put, you know, months of, of work into a project."

43. 7:15

    I think that's changing. I honestly do. Um, in this case, uh, so this, this training run, it cost us about $80 in GPU time. It did take about a week of engineering time to build this and, and caveat that was with an engineer who is familiar with this domain and, and had quite a lot of experience, uh,

44. 7:30

    you know, with machine learning and RL. Um, but I actually expect as, as we figure out the right patterns here collectively as an industry, this will keep dropping. Um, and I expect that, uh, you know, the, the sort of payback period to get, uh, a return on investment from these specialized models is actually gonna continue falling as

45. 7:45

    well. Um, and, uh, you know, part of, part of the reason I wanted to give this talk is to sort of distribute that knowl- the, the knowledge we learned and hopefully move faster towards that, that world where this is just sort of like a thing everyone knows how to do and is very easy and very fast.

46. 7:59

    Um, so that's, that's what we'll be talking about for the rest of time is, is some more of the lessons we learned. Um, okay. So, uh, when you are using RL to train an agent or really using RL for anything else, um, I find that consistently with different problems we look at, there are, there are sort of

47. 8:15

    two hard problems that come up every single time. All right? Um, and the two hard problems are, first of all, figuring out a realistic environment, right? So if you're training an agent, you need to be training it with realistic data, with realistic inputs and outputs, tools available, everything like that to how it's gonna be used in production.

48. 8:32

    Um, because if you don't, then it's gonna be optimizing for the wrong thing and, and you won't get the results you want when you deploy it. And then the second thing, which sometimes is hard, um, sometimes isn't, this one is a little bit task dependent, is getting the right reward function.

49. 8:44

    So reward function, that just means you have to be able to know when your agent's gone through and say, in this case, given an answer to my email, you have to have some way of knowing, did it do a good job or a bad job?

50. 8:54

    All right? That's the reward function. It decides it, it, it's how you decide if it's good or it's bad. Um, some-- Depending on the domain, sometimes that's really easy.

51. 9:01

    We have, uh, I don't know if Nathan's here. He's gonna be talking next, but, um, you know, he and his team put together this thing called RLVR, which in some verifiable domains, it's actually very easy, um, to do a reward.

52. 9:11

    Oftentimes, uh, not all domains are like that. Oftentimes it is kind of hard, um, and so it's, it's somewhat task dependent. I'm gonna go through how we solve these problems specifically with ART-E.

53. 9:21

    Okay, first one, realistic environment. So for our ART-E task, what is the environment we need? What is the environment this agent's gonna be operating in? Well, it needs these tools available.

54. 9:29

    It needs to be able to go and query an email inbox. It needs to be able to, like, get emails back, um, and, and that look realistic. These emails, you know, the inbox should be large because that's what most email inboxes are like.

55. 9:40

    Um, the emails in it should be diverse, and they have to be-- look kind of like real emails. Um, so this could be kinda hard 'cause you can't just go ask like 1,000 people to, you know, give you, uh, their, their, their personal emails to train on.

56. 9:52

    Um, luckily, in this case, we were able to solve this with the help of a company that has contributed a lot, um, to just the open data ecosystem, uh, generally.

57. 9:59

    It's, it's like a quite an iconic company. Uh, perhaps I would call it a historic company. Um, I'm of course talking about Enron. Um- [laughing]

58. 10:07

    I'm hearing some laughter. So anyway, Enron was a, uh, th-they were a financialized energy company in the '90s and 2000s, committed massive fraud, ended up getting shut down by the Department of Justice.

59. 10:18

    As part of this, uh, um, you know, process, uh, the, the, the court case that they were going through, um, a dump of like 500,000 of their emails was released to the public as part of the discovery process.

60. 10:28

    Um, so that's, that's, that's great for things like this, and that's what we used as our environment, uh, for the email inboxes. All right, so now we've got realistic email inboxes, um, with tens of thousands of emails that are real emails back and forth.

61. 10:39

    Now we have to design our reward function. So as our agent is going and as our agent is, um, you know, we're asking it questions, and then it's giving us answers, we have to know, is the answer correct or not, so we can reward it when it gets the answer right, and it can learn to do that

62. 10:52

    better. There's different ways, and this part is very task dependent. Um, the way that we went about it in this case, um, was we basically turned it into a more of a verifiable problem.

63. 11:04

    And the way we did that was we actually took our email inbox. We sort of inverted the problem. We, um, we grabbed batches of 20 emails at a time, uh, from the inbox and gave them to Gemini 2.5 Pro and said, "Hey, given this set of emails, give us a few questions that a user might realistically ans-

64. 11:20

    ask that the answers are found in this email," right? And so Gemini generated the questions, it generated the answers, and then of course, the, the, the source emails that came from.

65. 11:29

    Um, and there were some extra steps on top of that. A lot of the questions it came up with looked a little bit unrealistic. We had a separate filtering step where we're like, "Okay, let's find the subset of these that, that actually look like questions that, you know, I would maybe ask."

66. 11:40

    And we end up with a list of a few thousand questions, um, along with their verified answers. Um, and so at this point, it becomes much more of a, of a sort of verified thing.

67. 11:49

    The, the reward function becomes much easier 'cause we know what the correct answer should be. And so the way we can tell if our agent did a good job is we give our agent the question, we let it go and search the email inbox and try and find the right emails and everything, and eventually it comes back

68. 12:01

    with an answer. And then we can just use an LLM as judge, a very simple one, and say like, "Hey, you know, here's the question. Here's the, the golden answer that, that we believe is right.

69. 12:09

    Here's the answer we got from our, from our model. Is it right or not?" Um, we did have to do a little bit of, uh, iteration there, making sure that the judge was well calibrated on like what, you know, what, what counts as, as correct or not.

70. 12:20

    But by and large, this worked pretty well, um, and was able to make this more of a verified task. Um, so, so that's how, that's how we solved the, the reward function problem, was by having that, you know, turning this into something where we had more of a golden dataset.

71. 12:33

    Um, okay. So once you've solved that problem, those problems, once you have your environment, once you have, um, your, your reward function defined, then basically you just kind of have to run a loop over and over and over again where you have your agent go through and it tries to, um, solve the problem, and then you figure

72. 12:48

    out if it's good or it's bad. Um, and then you just, uh, you know, reward if it is g- if it's good and punish if it's bad, and that's it.

73. 12:56

    Um, and, uh, you do this over and over and over again, and then hopefully, if you've got everything set up right, um, it learns what good looks like, it learns what bad looks like, um, and it starts doing it right.

74. 13:07

    Um, and then again, this is, this is the curve we saw earlier where, where you can see it, it, it starts getting better over time. Um- Okay, a few other, like, interesting learnings from this project.

75. 13:18

    Um, one thing is we found that there's actually... You can throw a lot of stuff into your reward function, uh, beyond just the primary thing you're trying to solve for.

76. 13:27

    And so we actually ended up... We- there were like sort of eight different little things that we gave extra credit for. Um, and I'm gonna share two of them here.

77. 13:34

    So the first one here is, um, is we were trying to have it optimized for the number of turns, how many times back and forth, how many times it had to query the email inbox before it came up with the right answer, right?

78. 13:45

    So because the most important thing, of course, is, is getting the answer right, but between two answers that both get it right, we would rather it took fewer turns back and forth because that's fewer tokens, that's lower latency, lower costs.

79. 13:55

    Um, it's just like a, a, a more efficient agent. So, um, so you can see here on this first graph that early on, it... while it was getting its feet wet and figuring out what worked, it, it ended up spiking up to over six turns w- on average.

80. 14:07

    So it would go back and forth a bunch of times with the email inbox and, and try and find the right thing. But then once it was able to, like, w- figure out how to use the tools efficiently, figure out, like, you know, the, the right way to construct keywords and find the right email, it was able

81. 14:18

    to get very efficient and actually fast, uh, better than any of our prompted models, uh, on this metric of, um, using fewer turns. And again, this was just 'cause we gave it a little bit of extra...

82. 14:27

    It was, it was a very small amount relative to the reward for getting it right, but a little bit of extra credit on, on using for fewer turns, and it was able to, to use that, um, to, to optimize against that.

83. 14:38

    Um, another extra reward function we gave it is, um, to try and discourage it from hallucinating answers. So, um, obviously, the best thing is to get the right answer.

84. 14:47

    If you can't find the right answer, it's much better to say, "Hey, I don't know," than to make up an answer in, in a situation like this. So we basically penalized it if, um, if the reward model said, "Hey, you got the answer wrong," and, but it had tried to give an answer, give an answer, that was

85. 15:01

    like a much lower reward than if it just said, "Hey, I don't know. I can't solve this problem." And as you can see, that worked quite well, um, compared to any of the prompted models, including o3.

86. 15:09

    We ended up with a significantly lower hallucination rate, uh, because that was part of our reward function. Um, again, these are, these are things that are just sort of like extra credit, but, um, we found that, like, you can throw in a bunch of these and, and it can ot- jointly optimize all of them at the same

87. 15:23

    time, which is super powerful. Okay. I wanna talk a little bit about reward hacking. Um, it's, it's something that comes up a lot when you're trying to do this, and it's kind of a fun thing to talk about.

88. 15:32

    Um, this is an iconic video some of you might have seen. Uh, this was released by OpenAI, uh, almost a decade ago at this point of, um... They were, they were trying to, uh...

89. 15:40

    They had this environment where you were trying to, uh, get this boat to complete a race. And instead of learning to complete compu- uh, complete the race, it learned that, "Oh, if I just go in this, like, little circle that's not even part of the racetrack, I can, like, just get a bunch of points."

90. 15:53

    Um, and so it just started doing that over and over and over again instead of, like, actually following. Um, this is something that comes up a lot if you're doing reinforcement learning, and it's basically just the difference between, um, uh, the difference between what you actually want the model to do and what you can measure, um, like,

91. 16:09

    what you're actually rewarding it for. And, and if you... almost always, if you let one of these run long enough, it will figure out some way to exploit your measure, um, and it will figure out some way to, to get a really high reward, um, without actually solving the problem, and you need to just watch for that.

92. 16:22

    So I'm gonna give a couple examples here. Um, this is a, th- this is a graph, um, from another project actually, not this one. Uh, so an engineer on our team was, uh, was working on this game called NYT Connections.

93. 16:33

    Some of you might know. You get 16 words, and you have to put them in, like, four groups of four. It's quite a challenging game, especially for these language models, 'cause it requires a lot of world knowledge and, like, you know, lateral thinking.

94. 16:44

    Anyway, um, so, so they were trying to train this model to do it, and, uh, it, it, it wasn't figuring it out, wasn't figuring it out, what if... it wasn't figuring it out.

95. 16:50

    And then boom, you, you can see here around step 40, it just, like, takes off, and it's like, "Okay, we figured out how to, how to solve this." And, and this engineer, I'm, I'm, I'm gonna, I'm gonna call out...

96. 16:57

    Where's, where's, where's Anki on our team? He's here at the conference. Yeah. He's great. You should talk to him after. But, um, he was like, "Hey, we, we, we solved it.

97. 17:03

    Like, we got NYT Connections." And, like, and it's like, okay, the graph looks good. Let's look at what it's actually doing. What it was actually doing is it had figured out there was a bug in how we wrote the verification, and if it just put every single word in every single category, it was able to get a

98. 17:16

    perfect score- [laughing] ... um, 'cause we weren't verifying that there were, in fact, only four words, uh, in each category. Um, so this is another example. This is a fun one.

99. 17:26

    So I was, I was training a model, um, to produce really good titles for Hacker News, um, titles that would get a thing up voted. So I had this reward model I'd trained on, like, existing Hacker News, um, articles and how many up votes they got, and I was, I was trying to train this model to produce

100. 17:39

     new titles. And it was working really well for a while, you can see, and, and sort of subjectively as well. I, I looked at a bunch of these, these titles generated and, and for these first, like, thousand steps or so.

101. 17:48

     It was actually learning things that I was like, "Okay, as someone who spends way too much time on Hacker News, yeah, that, that does look like a good title.

102. 17:54

     You're, you're doing a good job." And then you can see around step, um, 1,200 here, it just, like, jumps a bunch, right? And it's like, okay, um, it clearly figured something out.

103. 18:02

     I don't know what it figured out, um, but we should look at that. Um, and so, uh, what... Turns out what the model had figured out was that it could just completely ignore the content of the post and generate the same title for every single one of them, and that would, like, maximize its score. [laughing]

104. 18:18

     So it generated this title, "Google lays off 80% of workforce." [laughing] Literally, every single article, this was, this was what it labeled it as. And the reward model was like, "Yes, that is gonna get up voted on Hacker News for sure." [laughing]

105. 18:28

     Which, which it probably would, to be fair. Um,

106. 18:32

     so, so anyway, the way, the way we solved this, um, what we found is that it's, it's really important to watch out for this. Solving it a- typically involves modifying in some way your retor- your, your reward function to penalize things like that.

107. 18:44

     So in the second example I talked about, it was actually quite an easy fix once we identified it, um, which was just add an extra LLM-as-judge that looked at the title, looked at the content, and said, "Hey, is there anything in the title that's not supported by the content?"

108. 18:55

     And we added that on and, and, and it, it actually worked great. Um, the important thing here is you wanna be looking at your, your rollouts, not just blindly trusting the reward function, figuring out what's actually happening.

109. 19:05

     Um, anyway, so, uh, that's it. Um, I'm, I'm almost out of time, so I'm gonna stop. A couple of QR codes for you. Um, everything in this presentation, and there's a much longer write-up I have of this whole project.

110. 19:17

     It includes the code. It includes the artifacts, datasets along the way. Um, you can, you can check that out there. Um, one more thing is, uh, we have a Discord that's open.

111. 19:26

     We have an open source project for training reinforcement learning models. Um, we have a Discord you can go to if you're interested in this kind of thing. Um, we...

112. 19:34

     we're all in there. We answer questions. There's lots of people from the community trying to do these things. So if, uh, if you're interested in building things with this, um, feel free to join it.

113. 19:41

     And, uh, yeah, happy, happy to chat there. And, um, yes, thank you everyone. Uh, th- appreciate your time. [upbeat music]