Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation

1 INTRODUCTION . Compared to standard supervised learning domains , reinforcement learning presents unique challenges in that the agent must act while it is learning . In certain application areas , which we term actor-latency-constrained settings , there exists maximum latency constraints on the acting policy which limit its model size . These constraints on latency preclude typical solutions to reducing the computational cost of high capacity models , such as model compression or off-policy reinforcement learning , as it strictly requires a low computational-complexity model to be acting during learning . Here , the major constraint is that the acting policy model must execute a single inference step within a fixed budget of time , which we denote by Tactor – the amount of compute or resources used during learning is in contrast not highly constrained . This setting is ubiquitous within real-world application areas : for example , in the context of learning policies for robotic platforms , because of inherent limitations in compute ability due to power and weight considerations it is unlikely a large model could run fast enough to provide actions at the control frequency of the robot ’ s motors . However , for many of these strictly actor-latency constrained settings there are orthogonal challenges involved which prevent ease of experimentation , such as the requirement to own and maintenance real robot hardware . In order to develop a solution to actor-latency constrained settings without needing to deal with substantial externalities , we focus on the related area of distributed onpolicy reinforcement learning ( Mnih et al. , 2016 ; Schulman et al. , 2017 ; Espeholt et al. , 2018 ) . Here a central learner process receives data from a series of parallel actor processes interacting with the environment . The actor processes run step-wise policy inference to collect trajectories of interaction to provide to the learner , and they can be situated adjacent to the accelerator or distributed on different machines . With a large model capacity , the bottleneck in experiment run-times for distributed learning quickly becomes actor inference , as actors are commonly run on CPUs or devices without significant hardware acceleration available . The simplest solution to this constraint is to increase the number of parallel actors , often resulting in excessive CPU resource usage and limiting the total number of experiments that can be run on a compute cluster . Therefore , while not a hard constraint , experiment run-time in this setting is largely dominated by actor inference speed . The distributed RL setting therefore presents itself as a accessible test-bed for solutions to actor-latency constrained reinforcement learning . Within the domain of distributed RL , an area where reduced actor-latency during learning could make a significant impact is in the use of large Transformers ( Vaswani et al. , 2017 ) to solve partiallyobservable environments ( Parisotto et al. , 2019 ) . Transformers ( Vaswani et al. , 2017 ) have rapidly emerged as the state-of-the-art architecture across a wide variety of sequence modeling tasks ( Brown et al. , 2020 ; Radford et al. , 2019 ; Devlin et al. , 2019 ) owing to their ability to arbitrarily and instantly access information across time as well as their superior scaling properties compared to recurrent architectures . Recently , their application to reinforcement learning domains has shown results surpassing previous state-of-the-art architectures while matching standard LSTMs in robustness to hyperparameter settings ( Parisotto et al. , 2019 ) . However , a downside to the Transformer compared to LSTM models is its significant computational cost . In this paper , we present a solution to actor-latency constrained settings , “ Actor-Learner Distillation ” ( ALD ) , which leverages a continual form of Policy Distillation ( Rusu et al. , 2015 ; Parisotto et al. , 2015 ) to compress , online , a larger “ learner model ” towards a tractable “ actor model ” . In particular , we focus on the distributed RL setting applied to partially-observable environments , where we aim to be able to exploit the transformer model ’ s superior sample-efficiency while still having parity with the LSTM model ’ s computational-efficiency during acting . On challenging memory environments where the transformer has a clear advantage over the LSTM , we demonstrate our Actor-Learner Distillation procedure provides substantially improved sample efficiency while still having experiment run-time comparable to the smaller LSTM . 2 BACKGROUND . A Markov Decision Process ( MDP ) ( Sutton & Barto , 1998 ) is a tuple of ( S , A , T , γ , R ) where S is a finite set of states , A is a finite action space , T ( s′|s , a ) is the transition model , γ ∈ [ 0 , 1 ] is a discount factor and R is the reward function . A stochastic policy π ∈ Π is a mapping from states to a probability distribution over actions . The value function V π ( s ) of a policy π for a particular state s is defined as the expected future discounted return of starting at state s and executing π : V π ( s ) = ∑∞ t=0 γ trt . The optimal policy π∗ is defined as the policy with the maximum value at every state , i.e . ∀s ∈ S , π ∈ Π we have V π∗ ( s ) > V π ( s ) . The optimal policy is guaranteed to exist ( Sutton & Barto , 1998 ) . In this work , we focus on Partially-Observable MDPs ( POMDPs ) which define environments which can not observe the state directly and instead must reason over observations . POMDPs require the agent to reason over histories of observations in order to make an informed decision over which action to choose . This motivates the use of memory models such as LSTMs ( Hochreiter & Schmidhuber , 1997 ) or transformers ( Vaswani et al. , 2017 ) . In this work , we use V-MPO ( Song et al. , 2020 ) as the main reinforcement learning algorithm . V-MPO , an on-policy value-based extension of the Maximum a Posteriori Policy Optimisation algorithm ( Abdolmaleki et al. , 2018 ) , used an EM-style policy optimization update along with regularizing KL constraints to obtain state-of-the-art results across a wide variety of environments . Similar to Song et al . ( 2020 ) , we use the IMPALA ( Espeholt et al. , 2018 ) distributed RL framework to parallelize acting and learning . Actor processes step through the environment and collect trajectories of data to send to the learner process , which batches actor trajectories and optimizes reinforcement learning objectives to update the policy parameters . The updated parameters are then communicated back to actor processes in order to maintain on-policyness . In the following , we first refer to a single observation and associated statistics ( reward , etc . ) as a step . We refer to an environment step as the acquisition of a step from the environment after an action is taken . The total number of environment steps is a measure of the sample complexity of an RL algorithm . We refer to a step processed by an RL algorithm as an agent step . The total number of agent steps is a measure of the computational complexity of the RL algorithm . We use steps per second ( SPS ) to measure the speed of an algorithm . We refer to Learner SPS as the total number of agent steps processed by the learning algorithm per second . We refer to Actor SPS as the total number of environment steps acquired by a single actor per second . 3 ACTOR-LEARNER DISTILLATION . Within the setting of distributed RL , a major challenge in applying transformers to reinforcement learning is their significant computational cost owing to their high actor latency . As an example , in Fig . 1 we present some results comparing a transformer to an LSTM on the I-Maze memory environment ( see Sec . 5.1 ) . In Fig . 1 , Left , a 4-layer Gated Transformer-XL ( GTrXL ) ( Parisotto et al. , 2019 ) agent significantly surpasses a 32-dim . LSTM in terms of data efficiency , with the xaxis as environment steps . However , in the center diagram , when the x-axis is switched to wall-clock time the LSTM becomes the clearly more efficient model . In order to work towards an ideal model with both high sample efficiency and low experiment run time , we analyzed which parts of the distributed actor-critic system led to the transformer ’ s main computational bottlenecks . As shown in the table on the right-hand-side of Fig . 1 , we found the percentage of time that the asynchronous learner process was spending waiting for new trajectory data was substantially higher for the transformer than for the LSTM . Further looking at actor SPS between models revealed that the cause of this idling is mainly due to the transformer ’ s actor SPS being over 15 times slower than the LSTM . This substantial decrease in SPS could even render the transformer not viable in actor-latency-constrained settings , e.g . robotic control environments , which entails a hard constraint on actor inference speed . These findings suggest that considerable improvements in training speed could be attainable if the actors produced data faster . While de-coupling actor and learner speeds using an off-policy RL algorithm could be a suitable solution for many applications , we desired a solution that respected strict constraints on actor latency , which precludes ever running a large model during inference . As an alternative solution which respected this constraint , we designed a model compression procedure termed “ Actor-Learner Distillation ” ( ALD ) which continually performs distillation between a largecapacity “ Learner ” model , which is trained using RL but never run directly in the environment , and a fast “ Actor ” model , which is trained using distillation from the learner and is used to collect data . As we use an actor-critic training paradigm in this work , each of these models effectively comprises two functions , the policy and value function . We denote the actor model byMA = ( πA , V πA ) and learner model byML = ( πL , V πL ) . The actor model is chosen with the appropriate model capacity to compute an inference step on the target hardware as fast as possible , i.e . within time Tactor . In contrast , there are far fewer constraints on learner modelML , as it is typically running on a central process with accelerated hardware . While there are no restrictions besides latency concerns on the model class of actor and learner in the context of ALD , in this work we choose to focus on the case of the actor model being an LSTM and the learner model being a transformer . As there is a clear data efficiency advantage to the transformer and a clear computational advantage to the LSTM , the success of ALD at extracting the benefits of both models can be more clearly recognized . Actor-Learner Distillation proceeds with the learner modelML being trained using a standard reinforcement learning algorithm , with the exception of the environment data being generated by the actor model . The actor model is trained using a policy distillation loss : LπALD = Es∼πA [ DKL ( πA ( ·|s ) ||πL ( ·|s ) ) ] = Es∼πA [ ∑ a∈A πA ( a|s ) log πL ( a|s ) πA ( a|s ) ] ( 1 ) Similar to previous distillation work in the multitask setting ( Teh et al. , 2017 ) , we employ this loss bidirectionally : the actor is trained towards the learner policy and the learner policy is regularized towards the actor policy . This regularization loss on the learner was seen to enable smoother optimization . As actor and learner policy will naturally be different at some points during training , an off-policy RL algorithm for the learner could be thought to be required if the divergence is large enough . We found leveraging off-policy data to be effective for the actor , and sampled trajectory data from a replay buffer when computing Eq . 1 during actor model optimization . Beyond a distillation loss , we experimented with a value distillation loss : LVALD = Es∼πA [ 1 2 ( V πL ( s ) − V πA ( s ) ) 2 ] ( 2 ) This loss functions to encourage actor model representations to model task reward structure . The use of value distillation in order to improve representations in a policy network has also been explored in concurrent work ( Cobbe et al. , 2020 ) . Unlike the policy distillation loss , this loss is only used to train the actor model . The final Actor-Learner Distillation loss is : LALD = απL π ALD + αV L V ALD ( 3 ) where απ and αV are mixing coefficients to control the contribution of each loss .

Summary: The paper proposes an original idea to use distillation to speed up modern distributed RL settings, when data collection is done on CPUs with the learning happening on accelerated hardware, e.g. GPU. More specifically, the authors propose to use a transformer for the learner and distil the policy into LSTM actors. With this, they achieve much faster wallclock time compared to the Transformer for Actors setup, however, losing in sample-efficiency.