Scaling Off-Policy Reinforcement Learning with Batch and Weight Normalization

TL;DR

This work addresses the challenge of improving sample efficiency in off-policy reinforcement learning by scaling CrossQ to higher update-to-data ratios. It introduces weight normalization (WN) combined with batch normalization (BN) to stabilize training, leveraging the scale-invariance properties of BN to keep the effective learning rate constant as $UTD$ grows. Empirically, CrossQ+WN achieves competitive or superior performance across 25 continuous-control tasks from the DeepMind Control Suite and MyoSuite, often with fewer parameters and without drastic interventions like network resets. The approach preserves simplicity while enhancing robustness and scalability, with potential applicability to broader RL settings and future work exploring discrete or vision-based domains and theoretical convergence guarantees.

Abstract

Reinforcement learning has achieved significant milestones, but sample efficiency remains a bottleneck for real-world applications. Recently, CrossQ has demonstrated state-of-the-art sample efficiency with a low update-to-data (UTD) ratio of 1. In this work, we explore CrossQ's scaling behavior with higher UTD ratios. We identify challenges in the training dynamics, which are emphasized by higher UTD ratios. To address these, we integrate weight normalization into the CrossQ framework, a solution that stabilizes training, has been shown to prevent potential loss of plasticity and keeps the effective learning rate constant. Our proposed approach reliably scales with increasing UTD ratios, achieving competitive performance across 25 challenging continuous control tasks on the DeepMind Control Suite and Myosuite benchmarks, notably the complex dog and humanoid environments. This work eliminates the need for drastic interventions, such as network resets, and offers a simple yet robust pathway for improving sample efficiency and scalability in model-free reinforcement learning.

Figures (7)

• Figure 1: CrossQ + wnutd$\mathbin{=}2$ outperforms simbautd$\mathbin{=}2$ and broutd$\mathbin{=}2$. In comparison, our proposed CrossQ + wn is a simple algorithm that, unlike bro, does not require extra exploration policies or full parameter resets. We present results for $25$ complex continuous control tasks from the dmc and MyoSuite benchmarking suites. $1.0$ marks the maximum score achievable on the respective benchmarks (dmcreturn up to $1000$ / MyoSuite up to $100\%$success rate). We present iqm and $90\%$ stratified bootstrap confidence intervals aggregated over multiple environments and $10$ seeds each.
• Figure 2: Comparing performance against wall clock time, measured in environment steps per second (so larger is better) on a single RTX 4090 workstation, we observe the CrossQ + WN outperforms simba and bro across all environments. We present results for $25$ complex continuous control tasks from the dmc and MyoSuite benchmarking suites. $1.0$ marks the maximum score achievable on the respective benchmarks (dmcreturn up to $1000$ / MyoSuite up to $100\%$success rate).
• Figure 3: Growing parameter norms hinder learning. The performance benefits of CrossQ fail to scale to more complex, higher dimensional tasks such as humanoid locomotion and muscular manipulation. Investigating this, we find that the critic parameter norms increase significantly with increasing utd ratios. As a result, the effective learning rate (ELR) drops and the number of dead neurons increases. Regularizing the critic parameters with weight norm (WN), we successfully mitigate the parameter norm growth and therefore maintain a more consistent ELR, leading to better performance on these more complex tasks. Uncertainty quantification depicts the 90% stratified bootstrap confidence intervals.
• Figure 4: CrossQ wnutd scaling behavior. We plot the iqm return and $90\%$ confidence intervals for different utd ratios. Results are aggregated over 15 dmc environments and $10$ random seeds each according to agarwal2021iqm. The sample efficiency scales reliably with increasing utd ratios and remains stable even when there are no more performance gains, which is a crucial property.
• Figure 5: An ablation study comparing CrossQ + wn against a soft L2 penality on the weights, as well as other design decisions such as target networks. The results show that the hard constraint outperforms the soft approach across a range of regularization scales and tasks. Uncertainty quantification depicts the $90\%$ stratified bootstrap confidence intervals.

Theorems & Definitions (2)

• Theorem 1: van2017l2
• Theorem 2: van2017l2