Improving Computational Efficiency in Visual Reinforcement Learning via Stored Embeddings

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 [ 30 , 37 , 1 ] , grandmaster level in Starcraft II [ 50 ] and grasping a diverse set of objects with impressive success rates and generalization with robots in the real world [ 21 ] . Modern off-policy RL algorithms [ 30 , 15 , 11 , 12 , 39 , 22 , 24 ] have improved the sample-efficiency of agents that process high-dimensional pixel inputs with convolutional neural networks ( CNNs ; LeCun et al . 25 ) using past experiential data that is typically stored as raw observations in a replay buffer [ 28 ] . 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 [ 13 , 8 ] , quantization [ 13 , 17 ] and freezing [ 53 , 36 ] have been proposed in supervised learning and unsupervised learning for various purposes ( see Section 2 for more details ) . In computer vision , Raghu et al . [ 36 ] and Brock et al . [ 5 ] showed that the computational cost of updating CNNs can be reduced by freezing lower layers earlier in training , and Han et al . [ 13 ] introduced a deep compression , which reduces the memory requirement of neural networks by producing a sparse network . In natural language processing , several approaches [ 46 , 42 ] have studied improving the computational efficiency of Transformers [ 49 ] . 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 Stored Embeddings for Efficient Reinforcement Learning ( SEER ) , a simple technique to reduce computational overhead and memory requirements that is compatible with various 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . off-policy RL algorithms [ 10 , 15 , 39 ] . 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 SEER to adaptively increase replay capacity , resulting in improved sample-efficiency of off-policy RL algorithms in constrained-memory settings . SEER achieves these improvements without sacrificing performance due to early convergence of CNN encoders . The main contributions of this paper are as follows : • We present SEER , a compute- and memory-efficient technique that can be used in conjunction with most modern off-policy RL algorithms [ 10 , 15 ] . • We show that SEER significantly reduces computation while matching the original performance of existing RL algorithms on both continuous control tasks from DeepMind Control Suite [ 45 ] and discrete control tasks from Atari games [ 2 ] . • We show that SEER improves the sample-efficiency of RL agents in constrained-memory settings by enabling an increased replay buffer capacity . 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 [ 44 ] . Off-policy RL algorithms are typically based on Q-Learning [ 51 ] which estimates the optimal value functions for the task at hand , while actor-critic based off-policy methods [ 27 , 38 , 10 ] are also commonly used . In this paper we will consider Deep Q-Networks ( DQN ; Mnih et al . 30 ) , which combine the function approximation capability of deep convolutional neural networks ( CNNs ; LeCun et al . 25 ) with Q-Learning along with the usage of the experience replay buffer [ 28 ] as well as off-policy actor-critic methods [ 27 , 10 ] , 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 [ 21 ] , the sample-inefficiency of off-policy RL algorithms has been a research topic of wide interest and importance [ 23 , 20 ] . To address this , several improvements in pixelbased off-policy RL have been proposed recently : algorithmic improvements such as Rainbow [ 15 ] and its data-efficient versions [ 48 ] ; using ensemble approaches based on bootstrapping [ 34 , 26 ] ; combining RL algorithms with auxiliary predictive , reconstruction and contrastive losses [ 18 , 16 , 33 , 52 , 39 , 40 ] ; using world models for auxiliary losses and/or synthetic rollouts [ 43 , 9 , 20 , 12 ] ; using data-augmentations on images [ 24 , 22 ] . 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 . [ 29 ] 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 [ 13 , 8 , 4 , 17 , 46 ] have been developed to address compute and memory requirements . Raghu et al . [ 36 ] and Brock et al . [ 5 ] 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 [ 31 ] and continual learning [ 35 ] , and Yosinski et al . [ 53 ] study the transferability of frozen and fine-tuned CNN parameters . Fang et al . [ 7 ] 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 off-policy algorithms using standard RL benchmarks . 3 Background . We formulate visual control task as a partially observable Markov decision process ( POMDP ; Sutton & Barto 44 , Kaelbling et al . 19 ) . Formally , at each timestep t , the agent receives a high-dimensional 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 [ 30 ] , 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 [ 10 ] is an off-policy actor-critic method based on the maximum entropy RL framework [ 55 ] , which encourages 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 ) ] . Here , the policy is modeled as a Gaussian with mean and covariance given by neural networks . Deep Q-learning . DQN algorithm [ 30 ] 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 ] , 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 [ 47 ] and distributional DQN [ 3 ] . For exposition , we refer the reader to Hessel et al . [ 15 ] for more detailed explanations of Rainbow DQN . 4 SEER : Stored Embeddings for Efficient Reinforcement Learning . In this section , we present SEER : Stored Embeddings for Efficient Reinforcement Learning , which can be used in conjunction with most modern off-policy RL algorithms , such as SAC [ 10 ] and Rainbow DQN [ 15 ] . 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 ? ? for more details of our method . 4.1 Freezing encoder for saving computation and memory . We process high-dimensional image input with an encoder fψ to obtain zt = fψ ( ot ) , which is used as input for policy πφ and Q-functionQθ as described in Section 3 . In off-policy RL , we store transitions ( ot , at , ot+1 , rt ) in the replay buffer B to improve sample-efficiency by reusing experience from the past . However , processing high-dimensional image input ot is computationally expensive . To handle this issue , after Tf updates , we freeze the parameters of encoder ψ , and only update the policy and Q-function . We remark that this simple technique can save computation without performance degradation because the encoder is modeled as deep convolutional neural networks , while a shallow MLP is used for policy and Q-function . Freezing lower layers of neural networks also has been investigated in supervised learning based on the observation that neural networks converge to their final representations from the bottom-up , i.e. , lower layers converge very early in training [ 36 ] . For the first time , we show the feasibility and effectiveness of this idea in pixel-based reinforcement learning ( see Figure 7a for supporting experimental results ) and present solutions to its RL-specific implementation challenges . Moreover , in order to save memory , we consider storing ( compressed ) latent vectors instead of high-dimensional image inputs . Specifically , each experience in B is replaced by the latent transition ( zt , at , zt+1 , rt ) , and the replay capacity is increased to Ĉ ( see Section 4.2 for more details ) . Thereafter , for each subsequent environment interaction , the latent vectors zt = fψ ( ot ) and zt+1 = fψ ( ot+1 ) are computed prior to storing ( zt , at , zt+1 , rt ) in B . During agent updates , the sampled latent vectors are directly passed into the policy πφ and Q-function Qθ , bypassing the encoder convolutional layers . Since the agent samples and trains with latent vectors after freezing , we only store the latent vectors and avoid the need to maintain large image observations in B .

Summary: - This paper presents a method for compute- and memory-efficient reinforcement learning where the visual encoder is frozen partway into training.  After freezing, latent vectors are stored in the replay buffer instead of images (and any existing images are replaced by them).  This leads to both better compute and memory utilization. - The authors demonstrate their method by comparing to Rainbow on Atari and CURL on DM Control.  On DM control, their method reduces the compute by a considerable margin.  On Atari, the results are less clear cut, but the compute cost is reduced. - When they also impose a memory constraint, the effectiveness of their method is further increased.