Compute- and Memory-Efficient Reinforcement Learning with Latent Experience Replay

1 INTRODUCTION . Success stories of deep reinforcement learning ( RL ) from high dimensional inputs such as pixels or large spatial layouts include achieving superhuman performance on Atari games ( Mnih et al. , 2015 ; Schrittwieser et al. , 2019 ; Badia et al. , 2020 ) , grandmaster level in Starcraft II ( Vinyals et al. , 2019 ) and grasping a diverse set of objects with impressive success rates and generalization with robots in the real world ( Kalashnikov et al. , 2018 ) . Modern off-policy RL algorithms ( Mnih et al. , 2015 ; Hessel et al. , 2018 ; Hafner et al. , 2019 ; 2020 ; Srinivas et al. , 2020 ; Kostrikov et al. , 2020 ; Laskin et al. , 2020 ) have improved the sample-efficiency of agents that process high-dimensional pixel inputs with convolutional neural networks ( CNNs ; LeCun et al . 1998 ) using the past experiential data that is typically stored as raw observations form in a replay buffer ( Lin , 1992 ) . However , these methods demand high memory and computational bandwidth , which makes deep RL inaccessible in several scenarios , such as learning with much lighter on-device computation ( e.g . mobile phones or other light-weight edge devices ) . For compute- and memory-efficient deep learning , several strategies , such as network pruning ( Han et al. , 2015 ; Frankle & Carbin , 2019 ) , quantization ( Han et al. , 2015 ; Iandola et al. , 2016 ) and freezing ( Yosinski et al. , 2014 ; Raghu et al. , 2017 ) have been proposed in supervised learning and unsupervised learning for various purposes ( see Section 2 for more details ) . In computer vision , Raghu et al . ( 2017 ) showed that the computational cost of updating CNNs can be reduced by freezing lower layers earlier in training , and Han et al . ( 2015 ) introduced a deep compression , which reduces the memory requirement of neural networks by producing a sparse network . In natural language processing , several approaches ( Tay et al. , 2019 ; Sun et al. , 2020 ) have studied improving the computational efficiency of Transformers ( Vaswani et al. , 2017 ) . In deep RL , however , developing compute- and memory-efficient techniques has received relatively little attention despite their serious impact on the practicality of RL algorithms . In this paper , we propose Latent Vector Experience Replay ( LeVER ) , a simple technique to reduce computational overhead and memory requirements that is compatible with various off-policy RL algorithms ( Haarnoja et al. , 2018 ; Hessel et al. , 2018 ; Srinivas et al. , 2020 ) . Our main idea is to freeze the lower layers of CNN encoders of RL agents early in training , which enables two key capabilities : ( a ) compute-efficiency : reducing the computational overhead of gradient updates in CNNs ; ( b ) memory-efficiency : saving memory by storing the low-dimensional latent vectors to experience replay instead of high-dimensional images . Additionally , we leverage the memory-efficiency of LeVER to adaptively increase the replay capacity , resulting in improved sample-efficiency of offpolicy RL algorithms in constrained-memory settings . LeVER achieves these improvements without sacrificing the performance of RL agents due to early convergence of CNN encoders . To summarize , the main contributions of this paper are as follows : • We present LeVER , a compute- and memory-efficient technique that can be used in conjunction with most modern off-policy RL algorithms ( Haarnoja et al. , 2018 ; Hessel et al. , 2018 ) . • We show that LeVER significantly reduces computation while matching the original performance of existing RL algorithms on both continuous control tasks from DeepMind Control Suite ( Tassa et al. , 2018 ) and discrete control tasks from Atari games ( Bellemare et al. , 2013 ) . • We show that LeVER improves the sample-efficiency of RL agents in constrained-memory settings by enabling an increased replay buffer capacity . • Finally , we show that LeVER is useful for computation-efficient transfer learning , highlighting the generality and transferability of encoder features . 2 RELATED WORK . Off-policy deep reinforcement learning . The most sample-efficient RL agents often use off-policy RL algorithms , a recipe for improving the agent ’ s policy from experiences that may have been recorded with a different policy ( Sutton & Barto , 2018 ) . Off-policy RL algorithms are typically based on Q-Learning ( Watkins & Dayan , 1992 ) which estimates the optimal value functions for the task at hand , while actor-critic based off-policy methods ( Lillicrap et al. , 2016 ; Schulman et al. , 2017 ; Haarnoja et al. , 2018 ) are also commonly used . In this paper we will consider Deep QNetworks ( DQN ; Mnih et al . 2015 ) , which combine the function approximation capability of deep convolutional neural networks ( CNNs ; LeCun et al . 1998 ) with Q-Learning along with the usage of the experience replay buffer ( Lin , 1992 ) as well as off-policy actor-critic methods ( Lillicrap et al. , 2016 ; Haarnoja et al. , 2018 ) , which have been proposed for continuous control tasks . Taking into account the learning ability of humans and practical limitations of wall clock time for deploying RL algorithms in the real world , particularly those that learn from raw high dimensional inputs such as pixels ( Kalashnikov et al. , 2018 ) , the sample-inefficiency of off-policy RL algorithms has been a research topic of wide interest and importance ( Lake et al. , 2017 ; Kaiser et al. , 2020 ) . To address this , several improvements in pixel-based off-policy RL have been proposed recently : algorithmic improvements such as Rainbow ( Hessel et al. , 2018 ) and its data-efficient versions ( van Hasselt et al. , 2019 ) ; using ensemble approaches based on bootstrapping ( Osband et al. , 2016 ; Lee et al. , 2020 ) ; combining RL algorithms with auxiliary predictive , reconstruction and contrastive losses ( Jaderberg et al. , 2017 ; Higgins et al. , 2017 ; Oord et al. , 2018 ; Yarats et al. , 2019 ; Srinivas et al. , 2020 ; Stooke et al. , 2020 ) ; using world models for auxiliary losses and/or synthetic rollouts ( Sutton , 1991 ; Ha & Schmidhuber , 2018 ; Kaiser et al. , 2020 ; Hafner et al. , 2020 ) ; using data-augmentations on images to improve sample-efficiency ( Laskin et al. , 2020 ; Kostrikov et al. , 2020 ) . Compute-efficient techniques in machine learning . Most recent progress in deep learning and RL has relied heavily on the increased access to more powerful computational resources . To address this , Mattson et al . ( 2020 ) presented MLPerf , a fair and precise ML benchmark to evaluate model training time on standard datasets , driving scalability alongside performance , following a recent focus on mitigating the computational cost of training ML models . Several techniques , such as pruning and quantization ( Han et al. , 2015 ; Frankle & Carbin , 2019 ; Blalock et al. , 2020 ; Iandola et al. , 2016 ; Tay et al. , 2019 ) have been developed to address compute and memory requirements . Raghu et al . ( 2017 ) proposed freezing earlier layers to remove computationally expensive backward passes in supervised learning tasks , motivated by the bottom-up convergence of neural networks . This intuition was further extended to recurrent neural networks ( Morcos et al. , 2018 ) and continual learning ( Pellegrini et al. , 2019 ) , and Yosinski et al . ( 2014 ) study the transferability of frozen and fine-tuned CNN parameters . Fang et al . ( 2019 ) store low-dimensional embeddings of input observations in scene memory for long-horizon tasks . We focus on the feasibility of freezing neural network layers in deep RL and show that this idea can improve the compute- and memory-efficiency of many offpolicy algorithms using standard RL benchmarks . 3 BACKGROUND . We formulate visual control task as a partially observable Markov decision process ( POMDP ; Sutton & Barto 2018 ; Kaelbling et al . 1998 ) . Formally , at each timestep t , the agent receives a highdimensional observation ot , which is an indirect representation of the state st , and chooses an action at based on its policy π . The environment returns a reward rt and the agent transitions to the next observation ot+1 . The return Rt = ∑∞ k=0 γ krt+k is the total accumulated rewards from timestep t with a discount factor γ ∈ [ 0 , 1 ) . The goal of RL is to learn a policy π that maximizes the expected return over trajectories . By following the common practice in DQN ( Mnih et al. , 2015 ) , we handle the partial observability of environment using stacked input observations , which are processed through the convolutional layers of an encoder fψ . Soft Actor-Critic . SAC ( Haarnoja et al. , 2018 ) is an off-policy actor-critic method based on the maximum entropy RL framework ( Ziebart , 2010 ) , which encourages the robustness to noise and exploration by maximizing a weighted objective of the reward and the policy entropy . To update the parameters , SAC alternates between a soft policy evaluation and a soft policy improvement . At the soft policy evaluation step , a soft Q-function , which is modeled as a neural network with parameters θ , is updated by minimizing the following soft Bellman residual : LSACQ ( θ , ψ ) = Eτt∼B [ ( Qθ ( fψ ( ot ) , at ) − rt − γEat+1∼πφ [ Qθ̄ ( fψ̄ ( ot+1 ) , at+1 ) − α log πφ ( at+1|fψ ( ot+1 ) ) ] ) 2 ] , where τt = ( ot , at , rt , ot+1 ) is a transition , B is a replay buffer , θ̄ , ψ̄ are the delayed parameters , and α is a temperature parameter . At the soft policy improvement step , the policy π with its parameter φ is updated by minimizing the following objective : LSACπ ( φ ) = Eot∼B , at∼πφ [ α log πφ ( at|fψ ( ot ) ) −Qθ ( fψ ( ot ) , at ) ] . ( 1 ) Here , the policy is modeled as a Gaussian with mean and covariance given by neural networks to handle continuous action spaces . Deep Q-learning . DQN algorithm ( Mnih et al. , 2015 ) learns a Q-function , which is modeled as a neural network with parameters θ , by minimizing the following Bellman residual : LDQN ( θ , ψ ) = Eτt∼B [ ( Qθ ( fψ ( ot ) , at ) − rt − γmax a Qθ̄ ( fψ̄ ( ot+1 ) , a ) ) 2 ] , ( 2 ) where τt = ( ot , at , rt , ot+1 ) is a transition , B is a replay buffer , and θ̄ , ψ̄ are the delayed parameters . Rainbow DQN integrates several techniques , such as double Q-learning ( Van Hasselt et al. , 2016 ) and distributional DQN ( Bellemare et al. , 2017 ) . For exposition , we refer the reader to Hessel et al . ( 2018 ) for more detailed explanations of Rainbow DQN . 4 LEVER : LATENT VECTOR EXPERIENCE REPLAY . In this section , we present LeVER : Latent Vector Experience Replay , which can be used in conjunction with most modern off-policy RL algorithms , such as SAC ( Haarnoja et al. , 2018 ) and Rainbow DQN ( Hessel et al. , 2018 ) . Our main idea is to freeze lower layers during training and only update higher layers , which eliminates the computational overhead of computing gradients and updating in lower layers . We additionally improve the memory-efficiency of off-policy RL algorithms by storing low-dimensional latent vectors in the replay buffer instead of high-dimensional pixel observations . See Figure 1 and Appendix A for more details of our method .

Summary: This manuscript proposes to reduce the intensive computation and memory requirement in reinforcement learning trainings by freezing the parameters of lower layers early. Besides, the authors also propose to store the low-dimensional latent vectors rather than the high-dimensional images in the replay buffer for experience replay. The effectiveness of the proposed techniques is evaluated on DeepMind Control environments and Atari. The motivation for this work is strong, and the results are impressive. However, the proposed technique is described in a very general way without clearly defined applicable conditions and specific design methods. Below are detailed comments and questions.