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

Adam Lee, M Ganesh Kumar, Blake Bordelon, Cengiz Pehlevan
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:64359-64419, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-lee26ai, title = {{C}omplete{P} for {RL}: Maintaining Feature Learning When Scaling Deep Reinforcement Learning}, author = {Lee, Adam and Kumar, M Ganesh and Bordelon, Blake and Pehlevan, Cengiz}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {64359--64419}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/lee26ai/lee26ai.pdf}, url = {https://proceedings.mlr.press/v306/lee26ai.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T CompleteP for RL: Maintaining Feature Learning When Scaling Deep Reinforcement Learning %A Adam Lee %A M Ganesh Kumar %A Blake Bordelon %A Cengiz Pehlevan %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-lee26ai %I PMLR %P 64359--64419 %U https://proceedings.mlr.press/v306/lee26ai.html %V 306 %X 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.
APA
Lee, A., Kumar, M.G., Bordelon, B. & Pehlevan, C.. (2026). CompleteP for RL: Maintaining Feature Learning When Scaling Deep Reinforcement Learning. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:64359-64419 Available from https://proceedings.mlr.press/v306/lee26ai.html.

Related Material