[NeurIPS Best Paper] 1000 Layer Networks for Self-Supervised RL — Kevin Wang et al, Princeton

Yeah. We had a poster yesterday, and then today we'll have an oral talk. Were you just, like, mobbed?

Oh, There was a lot of people. It's like three hours straight of like, you know, like waves of people to like that were we were trying to stupid. So I've never received the best paper.

Oh, I just like woke up one day and like checked my email and then They just tell they just, like They was like, oh, like, that's pay like, I just saw your email. Oh, you were like, been awarded best paper. I'm like But maybe you know from the reviews as well.

Right? Is that okay? Yeah.

We know from the reviews that we did well, but there's a difference between, like, doing well on the reviews and getting best paper. So that part, we didn't actually know. Yeah.

Yeah.

And, yeah, I guess I led the project, like, started the project, and then well, I was very happy to collaborate with Ishaan and Nicole and Ben also.

context. So so yeah. So we're all from Princeton.

Yeah. We're and thanks to Alan for booking you guys. So this project actually started from, like, an IW seminar.

So like like, independent work research seminar that Ben was teaching. And this was, like, actually, like, like, one of my first experiences in, like, ML research. So it was really valuable to, like, get that experience.

And then Yoshana was also in that seminar and working on adjacent things, so we collaborated a lot during that seminar. And then, yeah, the project turned out to have some pretty cool results. And then later on, also, like, Halt working on sort of similar things also joined it on the project and became, like, a good collaboration.

Any other perspectives you guys wanna chime in with? I guess maybe I should just go over, like, an overview of our project. Yes.

Okay. Sorry. Yes.

So the way that I kind of view our project is that if you look at the landscape of deep learning, you know, you have NLP, like language, vision, and then RL. And, like, as Ben kind of alluded to, you know, like, in in language, in vision, we've sort of converged to these, like, paradigms of scaling to massive networks. Right?

Like, hundreds of billions of parameters, trillions of parameters, and there's been, you know, a lot a lot gained from in deep learning from from that. Right? And then but then it seems like in the third sort of branch of deep learning in deep RL, that has not yet been the case.

Like, I was very surprised, like, coming into some like, you know, Ben's class and seminar when I was looking at the networks. Oh, why were you just using, like, a simple, like, two layer MLP for, like, these frontier sort of, you know, state of the art RL algorithms? And so I was very curious.

Like, can we design RL algorithms? Can we sort of put together a recipe for RL that can allow it to scale in potentially, you know, analogous ways that language and vision might scale. And so what we did is that we know that traditional RL, like, let's say, like, value value based RL doesn't really scale.

Right? This is pretty clear from the literature. So we tried a different approach to RL called self supervised RL where instead of learning, like, a value function, we're learning representations of states, actions, and future states such that the representations along the same trajectory are pushed together, the representations along different trajectories are pushed apart.

And this is just, like, a different approach to RL that allows us to learn in a self supervised manner. So there's we can solve task reach goals without any human crafted reward signal. And so we know that self supervised learning is scalable in these different areas in deep deep learning, so can self supervised RL scale in similar ways?

When we first tried it, it actually didn't work. Like, we've made the networks deeper. The performance, like, totally degraded.

But then we also but then I separately, was like there's also some other work, like, in in our literature, like, we tried, like, residual connections, and then there's other a few other architectural components that we had to put into the recipe. And then all of a sudden, like, one day, like, I ran this experiment, and there was, like, this one environment in which there was, like like, going from, like like, doubling the depth didn't really do anything, but, like, doubling the depth again with these different components suddenly, like, skyrocketed performance in this one environment.

we know that standard RL is not super scalable. And so, like, why can this different approach or different objective RL be scalable? I think it's because we're fundamentally shifting the burden of learning from something like like q learning or, like, regressing to, like, TD errors, which we know is quite spurious and noisy and biased, to fundamentally, like, a classification problem.

We're trying to classify whether future state is along the same trajectory or along a different trajectory. And we do this with representation learning. Right?

And we know that classification, cross entry loss, and representation learning is scalable in the deep learning literature. Right? If we think about language and, like, some of the objectives there.

So in some sense, we're kind of blurring the the lines. We're doing reinforcement learning. It's still an actor critic reinforcement learning algorithm.

It's like a goal condition reinforcement algorithm. But the objective the burden of, like, learning of the of of solving that RL task shifts to something that's more similar to objectives that you might see in language and vision that we know have scaled so much. And so I think yeah.

I think that's, like, one of the fundamental insights that we've seen is that it seems like by approaching RL in this different approach, we're able to, like, get so much more out of we were able to scale our networks, like, significantly beyond what was, like, standard use in Aural. Can I jump in?

Yeah. So, actually, if you look at a lot of our tasks there, particularly sort of like robotics tasks. So this is personally, I I'd be very curious about how a work like this could impact, like, the robotics field.

Like, my understanding of robotics is that a lot of robotics are now there's kind of multi a few different approaches. Like, one approach is we wanna train the robots using imitation learning. So we try to collect, like, an insane amount of data.

We have a ton of cubers human supervision, and we try to scale up this data, and we're, like, learning with imitation learning. Like but on the other hand, potential like, perhaps there's another approach, which is, like, for example, like, goal condition reinforcement learning, where we can actually train robotic agents and train RL agents to solve meaningful tasks with absolutely no human supervision. No demonstration.

It's much more scalable. Yeah. So, yeah, so this could serve as an alternate approach.

And perhaps instead of, like, scaling data, like, scaling manual, like, human supervision, which which is, you know, not super scalable, if there are ways to sort of make goal condition reinforcement learning scalable, and, like, we can just scale the architecture or we can scale Because you're focused on all your objectives. Yeah. Right.

With with certain different objectives. I think that could be very exciting to see how to see how that can affect a field like robotics, for example. Yeah.

I don't know if there's any others. Yeah. I guess, like, to your original question of, like, the trade offs.

Like, one of the trade offs one of the limitations that we say is, like, obviously, if you make the networks bigger, the it it will takes longer to run. Right? So if you, like, double the depth at some level of depth, you you it might take, like, twice as much to, like, take make a forward pass through the network.

Right? However, this is not so, like, within our paper, like, for most environments, we are able to, like, saturate like, get to, like, almost perfect performance within just, you know we don't even need to get to, like, a thousand layers. Like, maybe just 64 layers, for example, is sufficient.

And in this regime, like like, the the latency of the network is not necessarily actually even the not necessarily, like, a a significant bottleneck. Like, you can imagine there's a lot of tasks in which especially in RL, that, like, collecting data might be the bottleneck. Right?

And making forward passes through our network may not be the bottleneck. And so in our environment, we in our research, we specifically use the JAX g c r l environment, which is a JAX based GPU accelerated environment. So we can collect, like, thousands of, like, environment trajectories, like, in parallel at the same time so that we're able to, like, make, like Oh, this is built in.

Right. This is built in so that we can collect, you know, like, like, a thousand trajectories at the same time along all these environments, and so makes that make sure that, like, we have enough data to, like, saturate the learning from things. Wow.

think the data is crucial like here. Yeah. I guess even to build on that, I like drawing analogies to like successes in other areas of deep learning.

Like, for example, in large language models, the reason why we're able to scale to such large networks is that we found a paradigm in which we can leverage the entire Internet scale of data to learn. Right? And so data in RL traditionally has been hard to come by.

But now with these, like, GPU accelerated environments, we can collect hundreds of millions of times of the data within just a few hours. And so I think that this serves as, like, a really good test bed for us to be able to also find ways to scale up, like, network capacity and get similar kind of gains.

do you change that? Oh, I'm not saying that we're changing it. I wanna leverage insights from that to apply to our RL.

It's I I feel like you should go the other way. You think you should go the other way? Yeah.

I mean, that that would be a very interesting research direction too. But, actually, yeah, even on that point, like, one of the things I was thinking about is that the way that our our objective works is in some set it's not exactly next word prediction, but it's kind of like next state prediction. Right?

You imagine you're at some current state and you're at some current action, and we wanna predict whether or not this future state, this this certain state is a future state along the same trajectory or a different trajectory. And so in some sense, we are actually doing some sort of, like Implicit world model. Implicit, like, you know, like, in in I don't know if that's a bad word this this is.

Or, like, in language, you you do a cross entry loss to classify the next token. Right? And here, we're just doing a binary classification of, like, whether or not some next state is some Yeah.

Yeah. Yeah. It's a classification.

Yeah. Yeah. Yeah.

Yeah. Like, you push the frontier capabilities with the with with that and then you distill Distill it back. I actually, this is a good point.

Like, if you go out to our website, like, this is one of the future directions that we list at the very bottom. Oh, okay. Yeah.

We we we we would love to see if we could get similar performance. Like, we pushed the you know, like, we do achieve state of the art performance on goal condition RL and Jack's GC RL by a significant amount, and so it was very exciting to see the the, like, the the sort of frontier of the ability to train RL agents sort of pushed. And if we can do that in a way that also sort of is just as efficient as a standard, you know, networks, that would be very cool.

So, you know, like, is it Yeah.

So Yeah. So if there's ways to, like, to scale down to a smaller model or prune the model and maybe not and still retain performance, that's a very interesting research structure that we should put. Let's talk about other future directions.

What what what else is your personal passions?

last papers that I will tackle during the PhD. Personally, I would I'm very curious of like, can we like, what's the real like, can we push? I'm I'm I'm curious about, like, advancing the frontier as much as possible.

So if you actually look at our paper, we focus on scaling depth, but we notice that we see that scaling width actually also improves performance, and we also find that actually by scaling depth, we actually unlock the ability to scale along batch size as well.

like yeah. Right. So, like okay.

I guess for context, like, in traditional RL, like, value based RL, scaling batch size is not super effective. But there's we also can see there's also other work in other areas of deep learning that show that scaling batch size is only most effective when there's, like, a large enough network capacity to take advantage of the scaled batch size. And we actually find that, you know, perhaps so one hypothesis might be, like, perhaps the reason why scaled batch isn't that effective in traditional RLs because, like, we've been using these tiny networks that haven't been able to capture that.

And one of our experiments is that, like, because we are enabled sequential training of deep network, we actually were able to this is a great test bed for, you know, like, testing this hypothesis, and we find that indeed, as we scale to network capacity, we also unlock this different dimension of scaling Biopsy. And so all I about to say is that I'm very curious for someone, like, with enough compute to, like, take some of these environments, scale up batch scale up depth to the maximum capability, also scale along with also scale along batch size. And let let's like, basically, like, in the same way that in language, we're we're scaling on so many different axes, can we unlock different dimensions of scaling as well, and what capabilities and how far can we push the frontier of training these RL agents from doing that?

Before we pass it, Sean, when you say enough compute, what kind of compute budget did you have? How how does it how you know, I just wanna That's that's see what you guys got. Good question.

So we we wanted to make sure that this is we we wanted to make it such that, like, you know, it's quite accessible. So I keep the nice thing is that all of our experiments, even a thousand layer networks, can be run on one single 80 gigabyte h 100 GPU. So that's So it's dollars.

Yeah. Right. Right.

Right. So everything can be run on one GPU. But in theory, if we had, you know, like a distributor training setup and, like, can just, like, blast compute through this and really wanted to push the frontier, it'd be very interesting to see how things go.

Yep. Cool.

So I actually haven't gotten a chance to go to the conference that much. I'm actually working full time now. So Oh, dang.

Yeah. So so far, I I I actually literally just got my badge, like, a few moments before my intercession.

sort of, like, a very eye opening. I think that the general question is that people thought it was a very eye opening paper because, like, the objective is quite simple. It's quite elegant.

And for us to be able to, like, you know, like I I don't wanna say, like, overturn, but, like, sort of challenge the conventional wisdom that, like, RL is not super scalable and push it to such limits, like a thousand layers deep and see continued improve improve performance. I think the general impression that I've gotten is that, you know, this this this could be, like, a really cool like, if if we can sort of build along this direction and that, like, we can really scale along to all these different dimensions and push the frontier of the ability for RL. I'm very curious to see how that goes.

Alright. Well, thank you so much for dropping by. Congrats on the paper again.

Yeah. And good luck in your future work. Thank you.

Thanks for having us. Yeah. Yep.

Thanks.