CompleteP for RL: Maintaining Feature Learning When Scaling Deep Reinforcement Learning

The maximal update parameterization ($\mu P$) provides an approach to scale parameters while preserving feature learning in deep neural networks. In fixed data distribution settings it has been found to yield more stable learning dynamics, more consistent learned features, and enable optimal hyperparameters such as learning rate to transfer from small models to larger ones. However, it is unclear if these benefits readily transfer to reinforcement learning problems, where learning dynamics are coupled to the non-stationary data distribution induced by an agent's own actions. We empirically study how two regimes, the ''rich'' CompleteP and ''lazy'' Neural Tangent Kernel (NTK) parameterizations affect hyperparameter transfer, feature and policy consistency, and learning behavior as we scale reinforcement learning agents. Ultimately, we show that agents trained using CompleteP consequentially improves compute and reward efficiency compared to the NTK parameterization across multiple control tasks and variants.

Deep reinforcement learning - training neural networks through trial and error - is an emerging approach for solving difficult problems in robotics and control. One reason these networks are effective is because they often improve as they become larger. Another is that they can change their internal representations during training - a process known as "feature learning". Recent theory has shown that feature learning is not automatically preserved when networks are scaled up. In some settings, increasing the number of parameters can weaken the network’s ability to change its representations. Importantly, this behavior can be controlled without changing the architecture itself: the way parameters are initialized and connections are scaled can affect how strongly feature learning depends on model size. Most of this theory has been developed for fixed datasets, where the examples used for learning remain mostly static. Reinforcement learning adds a new complication: the data an agent learns from depends on its own past actions. As a result, when an agent learns from scratch, its training experiences can change rapidly over time. In our work, we tested whether theoretical predictions about these scaling choices hold in deep reinforcement learning. We studied the "CompleteP" parameterization, which is designed to preserve feature learning as models grow, and the "NTK" parameterization, which suppresses feature learning at large scale. We found that agents trained with CompleteP learned more efficiently and showed better hyperparameter transfer: optimization settings that worked for smaller models could be reused more reliably when scaling to larger models. This ultimately reduces the computational overhead needed to perform expensive hyperparameter sweeps to train large robotic agents.