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 a method for "actor-latency constrained" settings: Recently, transformers have been shown to be powerful models in RL which, in particular, exhibited better sample complexity in settings in which long-term credit assignment in partial observability was required (e.g. the T-maze). However, they are computationally expensive. Consequently, the authors propose to train transformers on the learner, supported by hardware acceleration, but also train a smaller LSTM agent which can be efficiently executed on the actors.