What is an RL environment? w/ Nous Research's Roger Jin

Hey guys. Uh, so I'm Roger from Nous Research. We train open source language models. Um, today we're gonna be talking about, uh, some of our RL infrastructure. Um, so I'll first give a motivation for like reinforcement learning in general, or like how, how I like to think about reinforcement learning.

Um, and I'll talk about, a little bit about like the RL infrastructure at Nous. Um, and specifically I'll focus on, um, our environment abstraction.

Um, so start with like supervised learning. Um, yeah, you know, as you guys kind of like know the deal, you have some model, you have some dataset, some loss function. Um, and you wanna like, take some steps on your model.

You wanna make small updates to your model to, um, to optimize that loss. So it kinda like looks like this, right? Um, like you have some like kind of like nice loss function, well-behaved, like smooth, continuous. Um, and, um, and like a, a step of training and supervised learning kind of looks like you take this like, um, like this ball, um, and you like, you, well, you press play on this picture, and you kinda like pause it.

And then the, the ball moves over like a little bit, and then, uh, that's like one step of training. Um, but like does this work for like any arbitrary like loss function you can write down? Um, and the answer is like no.

Like consider accuracy, right? It takes on n to s-- n plus one like discrete values if you have like a dataset size n. Like zero over n to like dot, dot, dot, like, um, n over n, right? Um, what would like the same plot look like if we use like accuracy instead?

I guess like error rate or something, um, if you wanna make it like a classification problem. Um, so that, that's gonna look like a piecewise constant function of like any, any weight that you'd take from your model. Um, right.

So, um, the ball doesn't like, doesn't roll anymore. Um, and, uh, so that, that's like, that's like one limitation of supervised learning. Um, like, and, um, like where does this like sort of show up? Well, like if you're doing language modeling, um, every time you're picking a token, like, um, you can't backprop through that picking operation anymore pe-because like, um, the next token would be like kind of like a, um, the ch-- the index of the chosen token is like a piecewise constant function of like any weight.

So, um, you introduce like the, this piecewise constant node into your compute graph at that point. Um, and then like that means you can't really backprop through like multiple, um, picking operations for, for these tokens. Um, so what that means is like under like greedy sampling, supervised learning can't really help you with objectives on multiple tokens.

If you're, uh, if you're interested in like rewarding a language model for like how funny like a joke is or something, you can't like directly approach that with supervised learning.

Um, okay, so s-these are some limitations of supervised learning. It only works with like nice loss functions, right? Uh, it doesn't play well with objectives over like multi-step trajectories that don't decompose well into like these single step objectives.

Um, and, um, the main medium that the modeler has to communicate with the model is data. So we need like expert samples of what is good. Um, and there's this like, uh, kind of deficit which is like n-- there's no way to demonstrate what is bad in supervised learning.

You kind of just, you have to tell the model very well, like what is good. Um, and, uh, there are, there are like kind of surprising results like, uh, like weak to strong generalization. But like in general, like the model outperforming the expert labelers would be like kind of very surprising.

Okay, so moving on to RL. Um, in RL, we have this like different setup where we have this like agent in an environment, and it can take actions and influence the environment, um, and then, uh, get rewards. And, um, the rewards can sort of like be anything.

They're like up to us, and, um, they, they're usually kind of like weirdly behaved functions. Like they can be discrete, they can be like points in a game, they can be like the accuracy. Um, so, um, and, um, that, that's fully fine in RL and normal.

And the goal is to learn this like policy which controls this agent, uh, to maximize the sum of like discounted rewards over time.

Um, so like going back to the language modeling world, like, um, how does this, how does like RL directly map onto language modeling or like m-language modeling map onto RL? So the states are like the text prefixes. So the initial state is like the prompt.

Um, the actions are the next tokens. Um, that means like a language model is like exactly what a policy is, right? A policy maps a state to a probability distribution over next actions. A language model maps a text prefix to a probability distribution over next tokens.

Um, so all we have to do really is like specify our favorite reward function, like how funny like our joke is, whatever. Um, and then we can optimize our model against this with like, with like all the fancy algorithms from RL.

Um, so this is just like some write-up of like the RL objective. Um, and it kind of shows like, um, the, the, like the, the trick here is like there's this expectation, right? You're taking an expectation over the policy.

Um, and so the thing on the inside, that, that thing might be like discrete. That's like reward, whatever. Like you might have some like accuracy thing as your reward. But we're kind of like, um, we kind of like describe the reward using this like high-level language of like how humans might assign rewards.

But then we're kind of like compiling this back into like a differentiable loss function with this expectation thing of like smoothing out, uh, smoothing. Um, so that's like kind of the trick here. So in that sense, like rewards are sort of this like abstraction.

Like we end up with a smooth like differentiable loss at the end. Um, but like we have this like language of rewards to communicate intent to the model.

Um, yeah. Um, and like you can like do some math to like optimize the model, um, policy gradient. Okay, so, um, uh, the, like, you know, the, the math that's im-- or the, the algorithm that's implied by like policy gradient is, is sort of like reinforce.

Um, where you can like take rollouts. Um, you can take samples from your model. Um, and, um, and, uh, kind of like estimate this, uh, this gradient in a policy gradient with the, with the, the mean of like these, uh, uh, reward returns from the, the rollouts.

Um, and, um, this is like a, you know, quite a simple algorithm. If, if you kind of like look at, if you kind of stare at like this part, um- It sort of looks like just like weighted supervised fine-tuning, right?

Like you have this like log of like the, the probability of a token and then like some like weight on it. Um, and reinforce, this is like the return, but it doesn't have to be that. Like, you can replace it with like some advantage, um, like GRPO does that, right?

It's like the, the reward norm, uh, the, the, the, uh, the mean, uh, like the, the normalized rewards. Um, and, uh,

and, uh, this like, uh, this is kind of like nice, right? Uh, so in RL, like we can work with like nearly arbitrary rewards, so discrete, non-differentiable, whatever. Um, you can have objectives over these like multi-step trajectories. Uh, you can communicate intent to the model like in this like additional way, and you can also demonstrate like what is bad in addition to like what is good simply through like negative rewards.

Um, the model outperforming expert labelers is, is possible. The model learning like self-correction is like expected. Um, and, um, yeah, we've seen like kind of cool emergent behaviors of like, you know, like o1, o3, r1 kind of like, kind of these, these like thinking models that kind of learn their own approach to, to solving problems.

Okay. Um, so we, we're now like kind of moving into this era of experience. So like since, you know, twenty twenty or so, like maybe even before that, we've been scaling up data, um, pre-training data, SFT. Um, the goal being to get to, to get this like very diverse, large pool of, uh, of, of, of good inputs, high-quality inputs to train our model.

Um, and now like we're kind of motivated to do the same thing for RL environments. And, um, like previously, I guess like the format of the data wasn't so important because like there was, there would be these like, you know, ShareGPT or OpenAI format or whatever.

But, and it was kind of annoying that there was like multiple formats, but it was kind of easy to interconvert between them. Um, but these environments are sort of like written in code, and it's like much harder to kind of like share these environments.

So, um, in open source, we kind of like wanna build up our collection of like, just like how we've built up this like huge pool of like data. We want to build up this like huge pool of environments.

Um, and so like there just needs to be some kind of like standard for open source developers to, um, to all like work together to like build up this, you know, to scale environments up to like millions and millions of environments.

Um, so like, so we wanna pose the question here, like what is the abstraction that will allow open source to scale up to a million environments? What would that look like? Um, okay. So, um, I guess like first, like we can think about like where, where do environments like live in the larger picture of like some training infrastructure.

Um, so at Nous we have this like distributed training, like infrastructure. It's, it's pretty standard. Um, we've got these like three microservices. So the trainer is one microservice, handles like backprop. Um, inference, um, is another microservice, right? It's like an OpenAI API, um, that, that serves the model, um, the current policy.

Um, and we also have this like environment manager/like rollout service, which, which handles all the environments.

Um, and each of these environments, if you have like multiple of them, is it's, it's also its own microservice. So, um, that means that each environment can have its own like kind of code execution environment or like, like, you know, Python environment, or like it can run on its own hardware.

If for some reason you want to like run an environment like on your phone or something, like you could, you could just do that. Um, and these, these environments run like async from the trainer, right? So, um, so they don't like really care about the trainer.

They're just like, their goal is just like constantly produce rollouts, constantly produce like data and never, never stop. Um, if you want to like resurrect some kind of ancient environment from like twenty nineteen, um, like yeah, you, you can just like make like a Docker container or something, and, um, and you can, you can run like this like ancient environment.

Okay. And, um, the, the environment interface, like I guess the minimal environment interface that, that we tr- we kind of expose is rather simple. So there's just like these like two things. There's like a lot of other stuff that you can do, like, you know, like logging and whatever.

Um, but like at a minimum, like these two things. So get item and collect trajectories. So get item is like kind of like, you know, just a data loader, kind of like a function. Um, so you can get like a row from your dataset, or if you like are working with like larger batches, you can kind of get, get like a batch from your dataset.

Um, if, if you have like some fancy curriculum you wanna do, um, get item kind of can manage that curriculum for you. Uh, and then collect trajectories takes in any kind of item that's returned by get item and then maps it to, uh, a group of, uh, scored rollouts.

Um, so, uh, I, I guess it's kind of like, uh, good that we, uh, uh, or okay, so I'll, I'll try to like motivate why we, why we have this like collect trajectories function. Um, so like I guess another like common kind of abstraction is like we, that we split up the rollout part from the scoring part or the inference part from the scoring part.

But collect trajectories is, is a fusion of both of these. It handles both inference and scoring, and we deliberately chose that because like, what happens when you try to like do like multi-turn or like, like multi-agent, um, with like a separate scorer function?

Um, then, then things kind of get like weird. Um, so we have this like fused, um, inference and scoring function, collect trajectories, where we're, we're kind of agnostic to how many like inf-- how many times like inference and scoring happen.

Um, so, uh, multi-step, multi-agent are kind of like automatically supported by, by this method. Okay. And, um, the atomic unit of training data that's produced by this function is a, is a group. Um, so this kind of encapsulates like what the data preparation would look like for like DPO, GRPO, um, SFT with like, you know, group size one.

Um, it also allows you to customize like what a group even means, right? So GRPO is like, has a setup where like a group is like, um, you have like multiple rollouts for the same problem, but it doesn't like kind of have to be that way.

Um, if you wanted to, like for example, ta- sample a batch of like similar problems, for example, I have this example here of like same like chess position, but it's just like black to move in one and white to move in the other.

Um, on the left it's like, uh, mate in one for black. On the right it's like mate in 18 for white, right? So, um, so if you want to like group like similar problems together instead of like having the same problem and like do multiple rollouts, like you can totally do that with, with this.

So we, we don't lock you into like the GRPO official like definition of what a group is.

Um, another like kind of weird quirky thing about, uh, our design is that, uh, at least for text, the thing that's returned by each of these environments is like the literal tokens. So it, it's not like, it's not text, it's not like messages, it's, it's the tokens.

Um, and, um, this allows you to deal with like weird little quirks about like inference. Like, for example, like maybe your, your inference API, um, gives you like an EOS token even if like your, uh, your max context length is like exceeded.

So it inserts this like kind of fake EOS. The model didn't wanna stop, but it was like forced to stop. Um, so you kinda like wanna strip that out when you're training on it so the model doesn't like stop in the middle of its like generation.

Um, say, say you wanted to like make a, uh, like a waifu group chat environment or something, right? So, um, and, um, you have like the, the roles being like the usernames or something, right? Um, and that, that requires a custom chat template.

Um, and that's, that's totally fine. We can support a per environment chat template in, in this framework. Um, another thing you can do like with in, in like the token world is like the, the trainer is now like agnostic to like chat versus instruct model.

Um, and what that means is like you can, you can do all like the cool like R one zero kind of style like experiments. Um, but it also means that you can mix together instruction, uh, instruct, and base model fine-tuning like all at once, right?

So, um, uh, so for example, you can do like if you were training this like code assistant model, or you can do chat code assistant style RL. Um, um, but you can also mix that with like base model style like autocompletion, um, RL.

Um, and another thing you can do if you have like to-- the environment return tokens is you can override the gRPC advantages. So the, the gRPC advantages are like just like kind of token level weights on, on the SFT, right?

Um, and, um, if you had some reason to like, um, some way to like kind of get token level rewards, for example, you have like some stack trace from your, um, from your like interpreter, um, and you see like exactly which lines are, are like involved in the error.

Um, the, these environments can then override the advantages like computed by the trainer. And you can have like, um, token level rewards or like, you know, weird like token level like re-weighings of the advantages.

Okay. Um, and, um, we wanted to make this kind of like extensible. So, um, this is just an example of like some, uh, kind of like extension that I, I thought of. Um, like suppose that you wanted different environments to support different like attention masking schemes.

So like if you had like super long context for some environment, maybe you would want, wanna use like sliding window attention for that one. Um, if, if you had some kind of like, uh, some kind of like smarter way to do like sliding window, like, like, you know, if you maybe just like, if you had some like protein, um, question answer environment where you put in the protein sequence, you gotta like try to ask a question about it.

Maybe you wanna only do attention, um, where like the contact points of the, uh, of the tertiary structure are. Um, so like maybe you wanna do something where you like, um, your, your environment passes the, uh, inputs to like a flux attention block mask getter, um, within the, this like score data group that's returned.

Yeah. So, so these environments are, are kind of like extensible. You can come up with your own like weird wacky things that you want to like give the trainer. Okay. Um, and, um, so all this was like joint work with, uh, Dakota, Shannon, Jay Chen, and, and Technium, and, uh, thanks to Alessio and, uh, Swyx for hosting.