1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities — part 1/4
Cleaned-text code points [0, 15741). Bibliography omitted here; source.md includes it.
# 1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities
###### Abstract
Scaling up self-supervised learning has driven breakthroughs in language and vision, yet comparable progress has remained elusive in reinforcement learning (RL). In this paper, we study building blocks for self-supervised RL that unlock substantial improvements in scalability, with network depth serving as a critical factor. Whereas most RL papers in recent years have relied on shallow architectures (around 2 – 5 layers), we demonstrate that increasing the depth up to 1024 layers can significantly boost performance. Our experiments are conducted in an unsupervised goal-conditioned setting, where no demonstrations or rewards are provided, so an agent must explore (from scratch) and learn how to maximize the likelihood of reaching commanded goals. Evaluated on simulated locomotion and manipulation tasks, our approach increases performance on the self-supervised contrastive RL algorithm by 2 × 2\times – 50 × 50\times , outperforming other goal-conditioned baselines. Increasing the model depth not only increases success rates but also qualitatively changes the behaviors learned. The project webpage and code can be found here: https://wang-kevin3290.github.io/scaling-crl/ .
## 1 Introduction
While scaling model size has been an effective recipe in many areas of machine learning, its role and impact in reinforcement learning (RL) remain unclear. The typical model size for state-based RL tasks is between 2 to 5 layers ( Raffin et al., 2021 ; Huang et al., 2022 ) . In contrast, it is not uncommon to use very deep networks in other domain areas; Llama 3 ( Dubey et al., 2024 ) and Stable Diffusion 3 ( Esser et al., 2024 ) have hundreds of layers. In fields such as vision ( Radford et al., 2021 ; Zhai et al., 2021 ; Dehghani et al., 2023 ) and language ( Srivastava et al., 2023 ) , models often only acquire the ability to solve certain tasks once they are larger than a critical scale. In the RL setting, many researchers have searched for similar emergent phenomena ( Srivastava et al., 2023 ) , but these papers typically report only small marginal benefits and typically only on tasks where small models already achieve some degree of success ( Nauman et al., 2024b ; Lee et al., 2024 ; Farebrother et al., 2024 ) . A key open question in RL today is whether it is possible to achieve similar jumps in performance by scaling RL networks.
At first glance, it makes sense why training very large RL networks should be difficult: the RL problem provides very few bits of feedback (e.g., only a sparse reward after a long sequence of observations), so the ratio of feedback to parameters is very small. The conventional wisdom ( LeCun, 2016 ) , reflected in many recent models ( Radford, 2018 ; Chen et al., 2020 ; Goyal et al., 2019 ) , has been that large AI systems must be trained primarily in a self-supervised fashion and that RL should only be used to finetune these models. Indeed, many of the recent breakthroughs in other fields have been primarily achieved with self-supervised methods, whether in computer vision ( Caron et al., 2021 ; Radford et al., 2021 ; Liu et al., 2024 ) , NLP ( Srivastava et al., 2023 ) , or multimodal learning ( Zong et al., 2024 ) . Thus, if we hope to scale reinforcement learning methods, self-supervision will likely be a key ingredient.
In this paper, we will study building blocks for scaling reinforcement learning. Our first step is to rethink the conventional wisdom above: “reinforcement learning” and “self-supervised learning” are not diametric learning rules, but rather can be married together into self-supervised RL systems that explore and learn policies without reference to a reward function or demonstrations ( Eysenbach et al., 2021 ; Eysenbach et al., 2022 ; Lee et al., 2022 ) . In this work, we use one of the simplest self-supervised RL algorithms, contrastive RL (CRL) ( Eysenbach et al., 2022 ) . The second step is to recognize the importance of increasing data availability. We will do this by building on recent GPU-accelerated RL frameworks ( Makoviychuk et al., 2021 ; Rutherford et al., 2023 ; Rudin et al., 2022 ; Bortkiewicz et al., 2024 ) . The third step is to increase network depth, using networks that are up to 100 × 100\times deeper than those typically found in prior work. Stabilizing the training of such networks will require incorporating architectural techniques from prior work, including residual connections ( He et al., 2015 ) , layer normalization ( Ba et al., 2016 ) , and Swish activation ( Ramachandran et al., 2018 ) . Our experiments will also study the relative importance of batch size and network width.
The primary contribution of this work is to show that a method that integrates these building blocks into a single RL approach exhibits strong scalability:
• Empirical Scalability: We observe a significant performance increase, more than 20 × 20\times in half of the environments and outperforming other standard goal-conditioned baselines. These performance gains correspond to qualitatively distinct policies that emerge as the scale increases.
• Scaling Depth in Network Architecture: While many prior RL works have primarily focused on increasing network width, they often report limited or even negative returns when expanding depth ( Lee et al., 2024 ; Nauman et al., 2024b ) . In contrast, our approach unlocks the ability to scale along the axis of depth, yielding performance improvements that surpass those from scaling width alone (see Sec. 4 ).
• Empirical Analysis : We conduct an extensive analysis of the key components in our scaling approach, uncovering critical factors and offering new insights.
We anticipate that future research may build on this foundation by uncovering additional building blocks.
## 2 Related Work
Natural Language Processing (NLP) and Computer Vision (CV) have recently converged in adopting similar architectures (i.e. transformers) and shared learning paradigms (i.e self-supervised learning), which together have enabled transformative capabilities of large-scale models ( Vaswani et al., 2017 ; Srivastava et al., 2023 ; Zhai et al., 2021 ; Dehghani et al., 2023 ; Wei et al., 2022 ) . In contrast, achieving similar advancements in reinforcement learning (RL) remains challenging. Several studies have explored the obstacles to scaling large RL models, including parameter underutilization ( Obando-Ceron et al., 2024 ) , plasticity and capacity loss ( Lyle et al., 2024 ; Lyle et al., 2022 ) , data sparsity ( Andrychowicz et al., 2017 ; LeCun, 2016 ) , and training instabilities ( Ota et al., 2021 ; Henderson et al., 2018 ; Van Hasselt et al., 2018 ; Nauman et al., 2024a ) . As a result, current efforts to scale RL models are largely restricted to specific problem domains, such as imitation learning ( Tuyls et al., 2024 ) , multi-agent games ( Neumann and Gros, 2022 ) , language-guided RL ( Driess et al., 2023 ; Ahn et al., 2022 ) , and discrete action spaces ( Obando-Ceron et al., 2024 ; Schwarzer et al., 2023 ) .
Recent approaches suggest several promising directions, including new architectural paradigms ( Obando-Ceron et al., 2024 ) , distributed training approaches ( Ota et al., 2021 ; Espeholt et al., 2018 ) , distributional RL ( Kumar et al., 2023 ) , and distillation ( Team et al., 2023 ) . Compared to these approaches, our method makes a simple extension to an existing self-supervised RL algorithm. The most recent works in this vein include Lee et al. (2024) and Nauman et al. (2024b) , which leverage residual connections to facilitate the training of wider networks. These efforts primarily focus on network width, noting limited gains from additional depth, thus both works use architectures with only four MLP layers. In our method, we find that scaling width indeed improves performance ( Section 4.4 ); however, our approach also enables scaling along depth, proving to be more powerful than width alone.
One notable effort to train deeper networks is described by Farebrother et al. (2024) , who cast value-based RL into a classification problem by discretizing the TD objective into a categorical cross-entropy loss. This approach draws on the conjecture that classification-based methods can be more robust and stable and thus may exhibit better scaling properties than their regressive counterparts ( Torgo and Gama, 1996 ; Farebrother et al., 2024 ) . The CRL algorithm that we use effectively uses a cross-entropy loss as well ( Eysenbach et al., 2022 ) . Its InfoNCE objective is a generalization of the cross-entropy loss, thereby performing RL tasks by effectively classifying whether current states and actions belong to the same or different trajectory that leads toward a goal state. In this vein, our work serves as a second piece of evidence that classification, much like cross-entropy’s role in the scaling success in NLP, could be a potential building block in RL.
## 3 Preliminaries
This section introduces notation and definitions for goal-conditioned RL and contrastive RL. Our focus is on online RL, where a replay buffer stores the most recent trajectories, and the critic is trained in a self-supervised manner.
#### Goal-Conditioned Reinforcement Learning
We define a goal-conditioned MDP as tuple ℳ g = ( 𝒮 , 𝒜 , p 0 , p , p g , r g , γ ) \mathcal{M}_{g}=(\mathcal{S},\mathcal{A},p_{0},p,p_{g},r_{g},\gamma) , where the agent interacts with the environment to reach arbitrary goals ( Kaelbling, 1993 ; Andrychowicz et al., 2017 ; Blier et al., 2021 ) . At every time step t t , the agent observes state s t ∈ 𝒮 s_{t}\in\mathcal{S} and performs a corresponding action a t ∈ 𝒜 a_{t}\in\mathcal{A} . The agent starts interaction in states sampled from p 0 ( s 0 ) p_{0}(s_{0}) , and the interaction dynamics are defined by the transition probability distribution p ( s t + 1 ∣ s t , a t ) p(s_{t+1}\mid s_{t},a_{t}) . Goals g ∈ 𝒢 g\in\mathcal{G} are defined in a goal space 𝒢 \mathcal{G} , which is related to 𝒮 \mathcal{S} via a mapping f : 𝒮 → 𝒢 f:\mathcal{S}\to\mathcal{G} . For example, 𝒢 \mathcal{G} may correspond to a subset of state dimensions. The prior distribution over goals is defined by p g ( g ) p_{g}(g) . The reward function is defined as the probability density of reaching the goal in the next time step r g ( s t , a t ) ≜ ( 1 − γ ) p ( s t + 1 = g ∣ s t , a t ) r_{g}(s_{t},a_{t})\triangleq(1-\gamma)p(s_{t+1}=g\mid s_{t},a_{t}) , with discount factor γ \gamma .
In this setting, the goal-conditioned policy π ( a ∣ s , g ) \pi(a\mid s,g) receives both the current observation of the environment as well as a goal. We define the discounted state visitation distribution as p γ π ( ⋅ ∣ ⋅ , g ) ( s ) ≜ ( 1 − γ ) ∑ t = 0 ∞ γ t p t π ( ⋅ ∣ ⋅ , g ) ( s ) p^{\pi\left(\cdot\mid\cdot,g\right)}_{\gamma}(s)\triangleq(1-\gamma)\sum_{t=0}^{\infty}\gamma^{t}p^{\pi\left(\cdot\mid\cdot,g\right)}_{t}(s) , where p t π ( s ) p_{t}^{\pi}(s) is the probability that policy π \pi visits s s after exactly t t steps, when conditioned with g g . This last expression is precisely the Q Q -function of the policy π ( ⋅ ∣ ⋅ , g ) \pi(\cdot\mid\cdot,g) for the reward r g r_{g} : Q g π ( s , a ) ≜ p γ π ( ⋅ ∣ ⋅ , g ) ( g ∣ s , a ) Q^{\pi}_{g}(s,a)\triangleq p^{\pi\left(\cdot\mid\cdot,g\right)}_{\gamma}(g\mid s,a) . The objective is to maximize the expected reward: max π 𝔼 p 0 ( s 0 ) , p g ( g ) , π ( ⋅ ∣ ⋅ , g ) [ ∑ t = 0 ∞ γ t r g ( s t , a t ) ] . \max_{\pi}\mathbb{E}_{p_{0}(s_{0}),p_{g}\left(g\right),\pi\left(\cdot\mid\cdot,g\right)}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{g}\left(s_{t},a_{t}\right)\right]. (1)
#### Contrastive Reinforcement Learning.
Our experiments will use the contrastive RL algorithm ( Eysenbach et al., 2022 ) to solve goal-conditioned problems. Contrastive RL is an actor-critic method; we will use f ϕ , ψ ( s , a , g ) f_{\phi,\psi}(s,a,g) to denote the critic and π θ ( a ∣ s , g ) \pi_{\theta}(a\mid s,g) to denote the policy. The critic is parametrized with two neural networks that return state, action pair embedding ϕ ( s , a ) \phi(s,a) and goal embedding ψ ( g ) \psi(g) . The critic’s output is defined as the l 2 l^{2} -norm between these embeddings: f ϕ , ψ ( s , a , g ) = ‖ ϕ ( s , a ) − ψ ( g ) ‖ 2 f_{\phi,\psi}(s,a,g)=\|\phi(s,a)-\psi(g)\|_{2} . The critic is trained with the InfoNCE objective ( Sohn, 2016 ) as in previous works ( Eysenbach et al., 2022 ; Eysenbach et al., 2021 ; Zheng et al., 2023 ; Zheng et al., 2024 ; Myers et al., 2024 ; Bortkiewicz et al., 2024 ) . Training is conducted on batches ℬ \mathcal{B} , where s i , a i , g i {\color[rgb]{0,0.55,0}s_{i},a_{i},g_{i}} represent the state, action, and goal (future state) sampled from the same trajectory, while g j {\color[rgb]{0.75,0,0}g_{j}} represents a goal sampled from a different, random trajectory. The objective function is defined as: min ϕ , ψ 𝔼 ℬ [ − ∑ i = 1 | ℬ | log ( e f ϕ , ψ ( s i , a i , g i ) ∑ j = 1 K e f ϕ , ψ ( s i , a i , g j ) ) ] . \min_{\phi,\psi}\mathbb{E}_{\mathcal{B}}\left[-\sum\nolimits_{{\color[rgb]{0,0.55,0}i=1}}^{|\mathcal{B}|}\log\biggl({\frac{e^{f_{\phi,\psi}({\color[rgb]{0,0.55,0}s_{i}},{\color[rgb]{0,0.55,0}a_{i}},{\color[rgb]{0,0.55,0}g_{i}})}}{\sum\nolimits_{{\color[rgb]{0.75,0,0}j=1}}^{K}e^{f_{\phi,\psi}({\color[rgb]{0,0.55,0}s_{i}},{\color[rgb]{0,0.55,0}a_{i}},{\color[rgb]{0.75,0,0}g_{j}})}}}\biggr)\right].
The policy π θ ( a ∣ s , g ) \pi_{\theta}(a\mid s,g) is trained to maximize the critic: max π θ 𝔼 p 0 ( s 0 ) , p ( s t + 1 ∣ s t , a t ) , p g ( g ) , π θ ( a ∣ s , g ) [ f ϕ , ψ ( s , a , g ) ] . \max_{\pi_{\theta}}\mathbb{E}_{\begin{subarray}{c}p_{0}(s_{0}),p(s_{t+1}\mid s_{t},a_{t}),\\ p_{g}\left(g\right),\pi_{\theta}(a\mid s,g)\end{subarray}}\left[f_{\phi,\psi}(s,a,g)\right].
#### Residual Connections
We incorporate residual connections ( He et al., 2015 ) into our architecture, following their successful use in RL ( Farebrother et al., 2024 ; Lee et al., 2024 ; Nauman et al., 2024b ) . A residual block transforms a given representation 𝐡 i \mathbf{h}_{i} by adding a learned residual function F i ( 𝐡 i ) F_{i}(\mathbf{h}_{i}) to the original representation. Mathematically, this is expressed as: 𝐡 i + 1 = 𝐡 i + F i ( 𝐡 i ) \mathbf{h}_{i+1}=\mathbf{h}_{i}+F_{i}\left(\mathbf{h}_{i}\right) where 𝐡 i + 1 \mathbf{h}_{i+1} is the output representation, 𝐡 i \mathbf{h}_{i} is the input representation, and F i ( 𝐡 i ) F_{i}(\mathbf{h}_{i}) is a transformation learned through the network (e.g., using one or more layers). The addition ensures that the network learns modifications to the input rather than entirely new transformations, helping to preserve useful features from earlier layers. Residual connections improve gradient propagation by introducing shortcut paths ( He et al., 2016 ; Veit et al., 2016 ) , enabling more effective training of deep models.
## 4 Experiments
### 4.1 Experimental Setup
#### Environments.
All RL experiments use the JaxGCRL codebase ( Bortkiewicz et al., 2024 ) , which facilitates fast online GCRL experiments based on Brax ( Freeman et al., 2021 ) and MJX ( Todorov et al., 2012 ) environments. The specific environments used are a range of locomotion, navigation, and robotic manipulation tasks, for details see Appendix B . We use a sparse reward setting, with r = 1 r=1 only when the agent is in the goal proximity. For evaluation, we measure the number of time steps (out of 1000) that the agent is near the goal. When reporting an algorithm’s performance as a single number, we compute the average score over the last five epochs of training.
#### Architectural Components