Reinforcement Learning with Latent Flow

1 Introduction . Reinforcement learning ( RL ) [ 41 ] holds the promise of enabling artificial agents to solve a diverse set of tasks in uncertain and unstructured environments . Recent developments in RL with deep neural networks have led to tremendous advances in autonomous decision making . Notable examples include classical board games [ 36 , 37 ] , video games [ 29 , 6 , 45 ] , and continuous control [ 34 , 28 ] . There has been a large body of research on extracting high quality features during the RL process , such as with auxiliary losses [ 20 , 27 , 35 ] or data augmentation [ 25 , 26 ] . However , another important component in RL representation learning has been largely overlooked : a more effective architecture to incorporate temporal features . This becomes especially crucial in an unstructured real-world setup like the home when compact state representations such as calibrated sensory inputs are unavailable . Motivated by this understanding , we explore architectural improvements to better utilize temporal features for the problem of efficient and effective deep RL from pixels . ∗Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Current approaches in deep RL for learning temporal features are largely heuristic in nature . A commonly employed approach is to stack the most recent frames [ 29 ] as inputs to a convolutional neural network ( CNN ) . This can be interpreted as a form of early fusion [ 24 ] , where information from the recent time window is combined at the pixel level for input to the CNN . In contrast , modern video recognition systems use alternate architectures that employ optical flow and late fusion [ 38 ] , where frames are processed individually with CNN layers before fusion and downstream processing . Late fusion is typically beneficial due to better performance , fewer parameters , and the ability to use multi-modal data [ 8 , 21 ] . However , it is not straightforward how to directly extend such architectures to RL . Real-time computation of optical flow for action selection can be computationally infeasible for many applications with fast control loops like robotics . Furthermore , optical flow computation at training time can also be prohibitively expensive . In our experiments , we also observe that a naive late fusion architecture minus the optical flow yields poor results in RL settings ( see Section 6.3 ) . This observation is consistent with recent findings in related domains like visual navigation [ 46 ] . To overcome the above challenges , we develop Flow of Latents for Reinforcement Learning ( Flare ) , a new architecture for deep RL from pixels ( Figure 4 ) . Flare can be interpreted as a structured late fusion architecture . It processes each frame individually to compute latent vectors , similar to a standard late fusion approach . Subsequently , temporal differences between the latent feature vectors are computed and fused along with the latent vectors by concatenation for downstream processing . By incorporating this structure of temporal difference in latent feature space , we provide the learning agent with appropriate inductive bias . We highlight the main empirical contributions from Flare in the following2 : 1 . Flare recovers optimal performance in state-based RL without explicit access to the state velocity , solely with positional state information . 2 . Flare achieves state-of-the-art performance compared to model-free methods on several challenging pixel-based continuous control tasks within the DeepMind control benchmark suite [ 42 ] , while being the most sample efficient model-free pixel-based RL algorithm across these tasks , outperforming the prior model-free state-of-the-art RAD on the 500k and 1M environment step benchmarks respectively ( Figure 1 ) . A video demonstration of Flare achieving SOTA performance on Quadruped Walk is in the supplementary materials . 3 . When augmented over Rainbow DQN , Flare outperforms the baseline on 5 out of 8 challenging Atari games at 100M step benchmark . Notebly , Flare scores 1668 on Montezuma ’ s Revenge , a signifcant gain over the baseline Rainbow DQN ’ s 900 . 2 Related Work . Pixel-Based RL The ability of an agent to autonomously learn control policies from visual inputs can greatly expand the applicability of deep RL [ 10 , 32 ] . Prior works have used CNNs to extend RL algorithms like PPO [ 34 ] , SAC [ 14 ] , and Rainbow [ 19 ] to pixel-based tasks . Such direct extensions have typically required substantially larger number of environment interactions when compared to the state-based environments . In order to improve sample efficiency , recent efforts have studied the use of auxiliary tasks and loss functions [ 50 , 27 , 35 ] , data augmentation [ 26 , 25 ] , and latent space 2Code : https : //github.com/WendyShang/flare dynamics modeling [ 16 , 15 ] . Despite these advances , there is still a large gap between the learning efficiency in state-based and pixel-based environments in a number of challenging benchmark tasks . Our goal in this work is to identify where and how to improve pixel-based performance on this set of challenging control environments . Neural Network Architectures in RL Mnih et al . [ 29 ] combined Q-learning with CNNs to achieve human level performance in Atari games , wherein Mnih et al . [ 29 ] concatenate the most recent 4 frames and use a convolutional neural network to output the Q values . In 2016 , Mnih et al . [ 30 ] proposed to use a shared CNN among frames to extract visual features and aggregate the temporal information with LSTM . The same architectures have been adopted by most works to date [ 27 , 35 , 25 , 26 ] . Recently , new architectures for RL have been explored that explore dense connections [ 39 ] as well as residual connections and instance norm [ 49 ] . However , the development of new architectures to better capture temporal information in a stream of images has received little attention in deep RL , and our work fills this void . Perhaps the closest to our motivation is the work of Amiranashvili et al . [ 1 ] who explicitly use optical flow as an extra input to the RL policy . However , this approach requires additional information and supervision signal to train the flow estimator , which could be unavailable or inaccurate in practice . In contrast , our approach is a simple modification to existing deep RL architectures and does not require any additional auxiliary tasks or supervision signals . Two-Stream Video Classification In video classification tasks , such as activity recognition [ 40 ] , there are a large body of works on how to utilize temporal information [ 9 , 22 , 44 , 7 , 47 , 12 ] . Of particular relevance is the two-stream architecture of Simonyan and Zisserman [ 38 ] , where one CNN stream takes the usual RGB frames , while the other the optical flow computed from the RGB values . The features from both streams are then late-fused to predict the activity class . That the two-stream architecture yields a significant performance gain compared to the single RGB stream counterpart , indicating the explicit temporal information carried by the flow plays an essential role in video understanding . Instead of directly computing the optical flow , we propose to capture the motion information in latent space to avoid computational overheads and potential flow approximation errors . Our approach also could focus on domain-specific motions that might be overlooked in a generic optical flow representation . 3 Background . Soft Actor Critic ( SAC ) [ 14 ] is an off-policy actor-critic RL algorithm for continuous control with an entropy maximization term augmented to its score function to encourage exploration . SAC learns a policy network πψ ( at|ot ) and critic networks Qφ1 ( ot , at ) and Qφ2 ( ot , at ) to estimate state-action values . The critic Qφi ( ot , at ) is optimized to minimize the ( soft ) Bellman residual error : LQ ( φi ) =Eτ∼B [ ( Qφi ( ot , at ) − ( rt + γV ( ot+1 ) ) ) 2 ] , where r is the reward , γ the discount factor , τ = ( ot , at , ot+1 , rt ) is a transition sampled from replay buffer B , and V ( ot+1 ) is the ( soft ) target value estimated by : V ( ot+1 ) =min i Qφ̄i ( ot+1 , at+1 ) − α log πψ ( at+1|ot+1 ) ] , where α is the entropy maximization coefficient . For stability , Qφ̄i is the exponential moving average of Qφi ’ s over training iterations . The policy πψ is trained to maximize the expected return estimated by Q together with the entropy term Lπ ( ψ ) = −Eat∼π [ min i Qφi ( ot , at ) − α log πψ ( at|ot ) ] , where α is also a learnable parameter . Reinforcement Learning with Augmented Data , or RAD [ 26 ] , is a recently proposed training technique . In short , RAD pre-processes raw pixel observations by applying random data augmentations , such as random translation or cropping , for RL training . As simple as it is , RAD has taken many existing RL algorithms , including SAC , to the next level . For example , on many DMControl [ 42 ] benchmarks , while vanilla pixel-based SAC performs poorly , RAD-SAC—i.e . applying data augmentation to pixel-based SAC—achieves state-of-the-art results both in sample efficiency and final performance . In this work , we refer RAD to RAD-SAC and use random translation as data augmentation . Rainbow DQN is an extension of the Deep Q Network ( DQN ) [ 29 ] , which combines multiple follow-up improvements of DQN to a single algorithm [ 19 ] . In summary , DQN [ 29 ] is an off-policy RL algorithm that leverages deep neural networks ( DNN ) to estimate the Q value directly from the pixel space . The follow-up works Rainbow DQN bring together to enhance the original DQN include double Q learning [ 17 ] , prioritized experience replay [ 33 ] , dueling network [ 48 ] , noisy network [ 13 ] , distributional RL [ 5 ] and multi-step returns [ 41 ] . Rainbow DQN is one of the state-ofthe-art RL algorithms on the Atari 2600 benchmark [ 3 ] . We thus adopt an official implementation of Rainbow [ 31 ] as our baseline to directly augment Flare on top . 4 Motivation . We motivate Flare by investigating the importance of temporal information in state-based RL . Our investigation utilizes 5 diverse DMControl [ 42 ] tasks . The full state for these environments includes both the agent ’ s pose information , such as the joints ’ positions and angles , as well as temporal information , such as the joints ’ translational and angular velocities . First , we train two variants with SAC—one where the agent receives the full state as input ( full-state SAC ) , and the other with the temporal information masked out , i.e . the agent only receives the pose information as its input ( position-only SAC ) . The resulting learning curves are in Figure 2 . While the full-state SAC learns the optimal policy quickly , the position-only SAC learns much sub-optimal policies , which often fail entirely . Therefore , we conclude that effective policies can not be learned from positions alone , and that temporal information is crucial for efficient learning . While full-state SAC can receive velocity information from internal sensors in simulation , in the more general case such as learning from pixels , such information is often not readily available . For this reason , we attempt to approximate temporal information as the difference between two consecutive states ’ positions . Concretely , we compute the positional offset δt= ( s p t−spt−1 , spt−1−spt−2 , spt−2−spt−3 ) , and provide the fused vector ( spt , δt ) to the SAC agent . This procedure describes the state-based version of Flare . Results shown in Figure 2 demonstrate that state-based Flare significantly outperforms the position-only SAC . Furthermore , it achieves optimal asymptotic performance and a learning efficiency comparable to full-state SAC in most environments . Given that the position-only SAC utilizes spt compared to Flare that utilizes s p t and δt , we also investigate a variant ( stack SAC ) where the SAC agent takes consecutive positions ( spt , s p t−1 , s p t−2 , s p t−3 ) . Stack SAC reflects the frame-stack heuristic used in pixel-based RL . Results in Figure 3 show that Flare still significantly outperforms stack SAC . It suggests that the well-structured inductive bias in the form of temporal-position fusion is essential for efficient learning . Lastly , since a recurrent structure is an alternative approach to process temporal information , we implement an SAC variant with recurrent modules ( Recurrent SAC ) to compare with Flare . Specifically , we pass a sequence of poses spt , s p t−1 , s p t−2 , s p t−3 through an LSTM cell . The number of the LSTM hidden units h is set to be the same as the dimension of δt in Flare . The trainable parameters of the LSTM cell are updated to minimize the critic loss . Recurrent SAC is more complex to implement and requires longer wall-clock training time , but performs worse than Flare as shown in Figure 3 . Our findings from the state experiments in Figure 2 and Figure 3 suggest that ( i ) temporal information is crucial to learning effective policies in RL , ( ii ) using Flare to approximate temporal information in the absence of sensors that provide explicit measurements is sufficient in most cases , and ( iii ) to incorporate temporal information via naively staking position states or a recurrent module are less effective than Flare . In the next section , we carry over these insights to pixel-space RL .

Summary: In this paper, the authors modify two state-of-the-art reinforcement learning algorithms (RAD-SAC and Rainbow-DQN) by changing the way information from multiple timesteps is aggregated. Instead of stacking frames at a pixel level, the authors propose to encode each frame individually and then pass the encodings and their pairwise differences to the actor/critic networks. It is shown that when learning policies from pixels this way of fusing information substantially outperforms the baseline variants in 3/5 environments from the DeepMind control suite and 5/8 environments from Atari2600. Furthermore, experiments with low-level observations on 5 DeepMind Control Suite tasks support the particular choice of information fusion over using recurrent networks or stacking states directly without including state differences.