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: This work presents a simple technique (Flare) to incorporate explicit temporal information to enable effective RL policy learning in challenging continuous control environments using pixel-based state representations. The approach is inspired from recent advances in the video recognition approaches which employ optical flow information and late fusion to incorporate temporal information. Typically, RL algorithms employ a frame-stacking heuristic to incorporate temporal information (early fusion). Though computing optical flow is slow and can been prohibitive for real-time applications, the authors present a simple alternative to using optical flow, i.e. difference of latent state vectors as a proxy for explicitly encoding motion information with latent vectors representing the observed state (late fusion). Experimental results on challenging continuous control tasks in the DMControl Suite show that Flare can achieve up to 1.9x higher scores than a baseline algorithm (RAD) which uses the frame-stacking heuristic to incorporate temporal information.