How We Built an Agent That Improves Itself — Zubin Aysola, Weights & Biases transcript

Okay. Hello. Thank you for showing up. Uh my name is Zuben Isola. Uh I work at Weights and Biases. I'm building the Arya agent which is something that we released to general availability on Monday. You might have seen Tim's talk uh one of the pe my peers at Weights and Biases who I work with uh on the main stage yesterday. So I'm going to give you a little bit of a deep dive into how we evaluate and build the Arya agent uh and in particular uh the evaluation framework that we use and how we use Arya to sort of self-reinforce itself in the loop of doing research for itself uh and talk about some weights and biases things.

So yeah, this is a rip from some of the slides that I presented at Nurups, but uh I think there's some interesting things about like using uh agentic harnesses. I I imagine most people in the audience or at this conference are all excited to build agent harnesses or maybe use cloud code with skills in a variety of different packages, but there's, you know, pretty strong differences in the way that you package these AI agents if you use them with different tool calls, etc. on a variety of different benchmarks.

And so inspired by that, we're doing some work at Weights and Biases to build our own agent harness to do research for you in the weights and biases platform. Um, you know, and then there's probably some other things that I don't need to convince people of as to like why we want software agents and why a software harness as opposed to like just letting the, you know, LM run free and call tools or whatever else. Uh, but that's not for here.

Uh, but one big problem presents itself, which is basically the thing that I go to bed thinking about every day, which is, you know, benchmarks, evaluations, the agents, and how you configure them are all covariant. And so like if you're trying to apply principled evaluations to any sort of system that dynamically changes, you need to have good measurement to see how the system actually performs. And you really want to see how the system performs in both your production environment and your offline environment.

Uh you know if you're RLP build from a long time ago, you might think about like the sim tore gap uh which is like basically how do we translate simulation environments into the real world environments etc. Uh and so there's sort of two things that I'll talk about uh and then give a live demo of which is weights and biases weave for agents which is basically in my opinion at least the best observability platform for both production and offline tracing of agents.

So I use it predominantly from the offline side where I build simulation environments run the Arya agent over them and then track them in weave and my entire team does that as well. So I can see all of our performance in offline metrics. And then the other half of my team works on the production side of things, deploying the agent, logging the same things in the exact same format so that I can rip those production traces into our environments and then hill climb on them or resolve our errors.

And it's a pretty nice flywheel. And then obviously the other thing that I'm here to talk about is Arya, the weights and devices agent itself, which is basically the thing that we use to now build itself because it's sophisticated enough that it can actually do that offline hill climbing uh by itself. So I'm going to just jump into a super quick live demo and then we'll talk through some other slides uh about how we built the offline evaluation framework.

So if you hop into the weights and biases platform, this is a demo project that I called Arya researches Arya here. Uh and I'm just going to chat to Arya and ask it to basically do auto research for itself. And so this is a relatively long prompt that I had it generate in actually one of these other sessions. Um but if we look at it, what it's going to do is it's going to take my codebase that is logged as a weights and biases artifact uh which we can see over here.

It's going to launch training jobs uh against this codebase which is our offline evaluation framework. It's going to review production traces that we have, add new tasks to hill climb against uh and do stuff like that. And so, you know, this thing's going to cook for a little bit. The evaluations take a little bit of time to run. And so, we'll see Arya uh basically try to write a new variant for itself. The other side of this equation, right, is the production side of things.

And so, this is a sanitized view of internal customer traces or internal traces on Arya. So we'll actually see very quickly uh that Arya will actually pop up here the exact conversation that we're running um right here. But in this trace I'm going to use this as an opportunity to take one of these production traces that are logged in weave and put them into my offline evaluation framework so that I can hill climb on them.

And so I'm just going to do that live. We'll ask you know for a relatively simple task. Um this is going to be a little interesting because it's going to be slightly self-referential but I'll go here. I'll pop up in a new Arya tab and I'll just ask it hey take this trace log it into our offline evals I don't think spelling matters too much and then run the candidate and production agent on it and we'll have it cook for that as well and we'll see those evaluations get logged into weave as they're running um so that's basically the tight version of what I do basically every day which is look at production traces in our weights and biases weave platform run evaluations in our codebase which is our offline client evaluation sandbox uh and then deploy new versions of the agent uh and work with my team to do that.

And if we look at our production project, this is, you know, a bunch of traces that we have over the last seven weeks of our nightly uh like CI jobs that run where we evaluate the agent in its production format and other candidate variants that we cut. And so we can see our relative performance over time. Uh like you know, for a while the CI broke, we had Arya have to fix itself last night, but you know, we're getting about like 66% performance on some of the tasks, etc.

And so we have a variety of tasks um that we train the agent to do. And so now I'm just going to jump into a material version of the presentation to go a little bit deeper into the theory behind building a simulation environment for offline hill climbing for basically building a software agent where traditionally you might expect that you are really robust if you're doing reinforcement learning but I mostly prompt engineer these days because the sophisticated models are relatively good at performing tasks at weights and biases and so we're really working on building skills for the software agent more than doing reinforcement learning.

But I think applying the same methodology to how you build the agent. So the same robustness of how you simulate environments uh is really helpful for our work. So we benchmark a bite-wise identical version of the agent in our production environment and our simulated environment. There's a lot of text on these slides that I and Claude sort of slop out. Uh but you know we'll just go for the graphics. Um so I think the exciting part here is that both using the logging framework from weave as well as just in the way that we've designed the system our research and production code are exactly the same.

And in fact, there's like a 4-hour sync job that happens between production to our research environment so that we don't get any drift when you, you know, researchers are cutting new variants of the agent, new skills, etc., and doing their hill climbing. And so, we have this nice little tight loop that basically mirror each other on two sides of the stack from our deployment layer and our offline benchmarking layer. Uh, and then, you know, we expose internally as a team a bunch of run eval commands to generate score trajectories.

And so that's what we were talking about with this like large project where we have a bunch of trajectory examples where if I click into one of these, you know, we can see like our relative scores of a variety of different metrics when it loads. Um, but essentially what we get is these super robust traces of the agent rolling out over time. And there's a lot of signal to reap from this. And I think the thesis behind this is generate a ton of traces and then decide what you're going to do with said ton of traces as you try to figure out how the agent performs. you're measuring emergent properties and then trying to align the agent in particular directions.

And so the two ways that that happens is I look manually at the tasks or like we would do in that first example, I would ask Arya to review the roll out that it generated for itself or review some other rollout and decide what went wrong or what went well and try to reinforce that behavior through prompting or anything else. Uh and then you know the actual method behind the agent is relatively simple. We want to test on a variety of the different models both models offered through corewave inference models offered through you know the foundation model players etc.

And so we define you know a relatively agnostic software stack for how we treat compaction and how we prepare context and how we assemble UI payloads etc. And the again mental pattern behind this is to really try to make it very very simple to have lots of mutations of the exact same configuration. You want to basically YAML define different configurations of the agent to sort of get multiple parallel uh variants and then test them all and see what happens because ultimately you can have insight into how to make the problem better.

But if we take an adage from old reinforcement learning training or just simple model training, it's just better to run more experiments than fewer. And so that's the thesis behind the agent harness itself too. And so then you know we talk about like the sandbox stack behind the agent. And I think that you know there's a lot of mentions of code mode here at this conference and other things like that. And one thing that we really want to see happen because we want to build the best product for our customers is to just have a sandbox environment that they can do anything they wanted.

So for example prior to yesterday I wasn't entirely certain that Arya was going to be able to do a bunch of parallel executions of itself. And so I asked it both while I was writing this presentation and preparing for this to just create the sandbox environment uh for itself to run a bunch of parallel executions of its own research loop. Right. Uh, and so like having that unconstrained environment is really useful for getting the agent to do emergent things.

And so, you know, I I'm not really sure what the six phrases per record uh means, but I think from a theory standpoint, we want to build then a good pattern for generating these simulation environments. And so we have a relatively agnostic DAG that you might think about in a traditional machine learning context, which is like we take a configuration, which are these YAML files. We hydrate them. So we load the live data that we need to.

We set up an environment and the environments and weights and biases are relatively expensive because we have lots of production data that we might want to test against or other things like that. These are full machine learning training logs. We might have to simulate uh GPU executions if you're doing auto research etc. That's a relatively intensive process and so you might want to parallelize that. Then you would rehydrate because oftentimes there's runtime configurations that you can't encode in a YAML specification and so you hot patch that data back into the config config.

Uh and then you run the agent and running the agent is trivial. You run it, as I mentioned at the beginning, in the bite-wise identical version that you have in production, and then you score it. And scoring is where I spend a lot of my time, which is basically trying to think about the robustness of how you measure results. Uh, and I imagine most people here are pretty good at thinking about things like that. And I think it's really worth spending the time once you have a configuration of how to run evaluations, thinking about the robustness of our emails. uh a Slack message that I probably can't show live on the screen that I got from one of my teammates today made me very happy where he basically just spent all of yesterday thinking about the health of our evaluations and the drift between our evaluations and production.

Uh and I think that kind of mental model of like you know why are things working well why are things not working well? What are the gaps that we see in these two patterns is really useful. Um in our case weights and biases Arya scores itself in two patterns which is normatively which gives us basically did we pass a task or not and then relativistically where we can set styles based on one variant where it asks questions to the user and one variant where it doesn't ask questions to the user and we can sort of see which one behaves better uh with a relative scoring.

So a pretty traditional formulation from a reinforcement learning standpoint. And then of course if you're running this in parallel you got to tear it down because you want to start it again and you don't want to clobber your teammates's work etc. Uh and so we'll see what happens in this project because a bunch of my teammates are running eval right now and so we'll see how Arya performs. Uh and then you know more traditionally on tasks I think the eval tasks are just YAML specifications that we define as a starting condition of an environment with a bunch of user configurations as well as then an ending condition that we sort of want to get to.

Uh in our case tasks are flows from users. So we simulate them in three ways. Uh we do just simple text which is like hey here's an instruction to Arya. You might imagine that the task that it generates in this chat over here um where it's logging a specific trace from production uh into our local framework actually looks like you know just a simple question that I asked it basically which was that question of Arya doing auto research itself um we can just ask it to continue um as it does that and that task is just a single question to it there's other variants of the tasks that exist as well where we simulate a language model that has a persona from a user and we ask it, hey, pretend to be this user, ask certain questions in a particular order so that we can simulate multi-turn environments and interactions with our agent.

Uh, and then you just run a bunch of these. So, we have 886 tasks. We categorize them by levels. We expose them to our product team so that they can decide whether or not the tasks are good enough or they reflect things that we care about from our benchmarks, etc. And then we just run them a bunch of times. Yeah. And then, you know, the ultimate thing here is that the trajectory is like the meat and the data that I live behind.

And so I look at our production traces that exist in our traffic and our offline traces that exist. And I really try to exploit the behavior patterns from both of those. And so we build tooling ourselves to understand our traces. Uh we have like this behavior trace project etc that we run in both prod and offline. And yeah that's the eval flywheel. So you know fundamentally every production miss or every production goodness as well because I think it's useful to hill climb in a positive direction uh becomes a task for our agent framework.

And so we use that to sort of define the agent and make it better. And these are, you know, our real scores straight from weave. So we have a bunch of different categories that we measure ourselves on. I think the agent is really good at conceptually guiding you. We want to make it better at doing some error analysis for projects. You know, we try to make these as difficult as possible for the agent to really test the metal of how this thing performs.

And so I'll hop back to Arya now. There's sort of five minutes left um with sort of a demo of what it's actually doing. So right now it's going to run a bunch of code as it's trying to rip this production trace and log it into the framework. If I hop into this uh live demo that I just triggered with that nice long prompt, you know, Arya is going to tell me the research runs that it did. So we can click through some of these links.

Uh I asked it to write a report using weights and biases reports from the uh run that it just had. If we scrub this, we can see that this is the most recent production trace that it just logged. It wrote itself a new task, ran it, and then scored it. And so this is like the production variant getting benchmarked on our specific task. If we look at the report that it generated, this is the demo that I asked it to do. Um, so you know, it ran a real womb agent trace, turned it into a WBF regression task.

WBAF is the factory that I build, which is the weights and biases agent factory. That's the principle behind our offline benchmarking. It identified that the problem was that we weren't calling weave.log, one of our SDK calls properly in the sandbox. It replicates the source trace, which we can click into here, uh, and then runs a bunch of variants of the agent to see how it makes itself better. And so, you know, it added a hill climb target uh which was basically like here's what to do to fix it.

Ran the trace and then we can see if it gives us a little explanation of exactly what happened at the end. Uh how did the prod versus candidate variant perform and you know this is now what I live in. And so I think the flywheel here is exceptionally compelling for me where instead of going back to my cloud code and writing offline benchmarks and trying to think about what I want the agent to do, I'm just going to live in this platform instead and hop between my production tracing project where I'm looking at traces roll in as they come through over time where AR is, you know, building itself and then hop back into my offline evaluation project to sort of see what my team is doing, what I'm doing, and what changes I want to make. what change did we make to the prod skill or to the agent variant etc.

And you know to sort of like watch this thing roll out. I think a big challenge that we have is it's really easy to go like auto mode for some of these tasks where you want to just see the agent do everything. And I haven't written a line of code in maybe eight months because I just tell Claude to write all my code for me. And that's a really nice pattern, but it doesn't absolve yourself of the thought behind how we want to make improvements.

And I think using these tools to improve themselves is very valuable because you get to spend more of your time in the gray of actually trying to think how we make this system better, right? And so if you're building an agent, I implore you to use weights and biases to both track that agent etc. Uh we have a really excellent model tracking platform as well. So some of the other demos that we've shown here are Arya training machine learning models on H200's running on core infrastructure and doing auto research for Karpathy's nano chat etc or full-scale production projects.

But I think the exciting thing in this case for me is that I don't need to leave this platform in order to do my job now right. I have full observability from a production standpoint. I have full observability from an offline standpoint. And if we look at, you know, the things that Arya is saying, it's running lots of evaluations on my behalf to improve itself. And that auto mode that I just mentioned is really exciting because now I get to spend all of my time thinking about how I can make this system reinforce itself better.

The tight little guard rails that I can put around it to make it a more useful tool for all of you and then inherently for me. Um, but yeah, that's a that's a quick demo. We can we can look at the production call trace that we have here from Arya. This is the evaluation that it just ran. So if we look into like the actual trace, we'll see that, you know, AR did a bunch of work. It ran this predict and score column. We have a bunch of tool calls that ran.

Uh and you know, if we hop back into the chat, we'll see the actual variant that it ran. Um it looks like the candidate variant just had a tight little prompt that we injected into the system prompt or in one of the skills to sort of solve that exact SDK error uh and mitigate it. And that kind of replication from production to simulation to agent then defining the improvement pattern is the thing that I really would like to leave everybody with.

So that's coreweave Arya. Uh welcome to weights and biases. I hope that was an interesting talk in sort of how we decide to build uh our offline evaluation metrics, how we reap data from production and the flywheel of trying to build these agents and the things that I think about on a day-to-day basis. Uh and yeah, let's get started. Thank you all so much. >> [music]