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 work proposes LeVER, a method that modifies general off-policy RL algorithms with a fixed layer freezing policy for early embedding layers (in this particular case, a few early layers of a CNN). As a direct consequence, the method enables to store embeddings in the experience replay buffer rather than observations, with a potential decrease in memory required, as well as providing a boost in clock time due to fewer gradient computations needed for every update. The method is benchmarked with a couple of off-policy RL algorithms against a few different environments.