Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation
Authors: Emilio Parisotto
Russ Salakhutdinov
Venue: ICLR 2021 Poster (2021)
Keywords: Deep Reinforcement Learning
Memory
Transformers
Distillation

Abstract: Many real-world applications such as robotics provide hard constraints on power and compute that limit the viable model complexity of Reinforcement Learning (RL) agents. Similarly, in many distributed RL settings, acting is done on un-accelerated hardware such as CPUs, which likewise restricts model size to prevent intractable experiment run times. These "actor-latency" constrained settings present a major obstruction to the scaling up of model complexity that has recently been extremely successful in supervised learning. To be able to utilize large model capacity while still operating within the limits imposed by the system during acting, we develop an "Actor-Learner Distillation" (ALD) procedure that leverages a continual form of distillation that transfers learning progress from a large capacity learner model to a small capacity actor model. As a case study, we develop this procedure in the context of partially-observable environments, where transformer models have had large improvements over LSTMs recently, at the cost of significantly higher computational complexity. With transformer models as the learner and LSTMs as the actor, we demonstrate in several challenging memory environments that using Actor-Learner Distillation largely recovers the clear sample-efficiency gains of the transformer learner model while maintaining the fast inference and reduced total training time of the LSTM actor model.

Published as a conference paper at ICLR 2021
EFFICIENT TRANSFORMERS IN REINFORCEMENT LEARNING USING
DISTILLATION
Emilio Parisotto & Ruslan Salakhutdinov
Machine Learning Department Careeiie Mellon University Piutburgh, PA 15213, USA (eparisot, rsala kuu
Becchuuuuueeeeee Bcs cmu edu
ABSTRACT
Many real-world applications such as robotics provide hard constraints on power and compute that limit the viable model complexity of Reinforcement Learning (RL) agents. Similarly, in many distributed RL settings, acting is done on un- accelerated hardware such as CPUs, which likewise restricts model size to pre- vent intractable experiment run times. These "actor-latency" constrained settings present a major obstruction to the scaling up of model complexity that has recently been extremely successful in supervised learning. To be able to utilize large model capacity while still operating within the limits imposed by the system during act- ing, develop "Actor-Learne Distillation" (ALD) procedure that leverages continual form of distillation that transfers learning progress from large capacity learner model to small capacity actor model. As case study, develop this procedure in the context of partially-observable environments, where transformer models have had large improvements over LSTMs recently, the cost of signif- icantly higher computational complexity. With transformer models as the learner and LSTMs the actor, we demonstrate in several challenging memory environ- ments that using Actor-Leamnrr Distillation recovers the clear sample-efficiency gains of the transformer learner model while maintaining the fast inference and reduced total training time of the LSTM actor model.
INTRODUCTION
Compared to standard supervised learning domains, reinforcement learning presents unique chal- lenges in that the agent must act while it is learning. In certain application areas, which we term actor-latency-coostrriiee
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 low cnaputatiooall-omplxxii model to be acting during learning. Here, the major constraint is that the acting policy model must execute a single inference step within fixed budget of time, which we denote by Tactor the amount of compute resources used during learning is in contrast not highly constrained. This setting is ubiquitous within real-world applica- tion 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 provide actions at the control frequency of the robot's motors.
However, for many of these strictly actor-latency constrained settings there are orthogonal chal- lenges involved which prevent ease of experimentation, such as the requirement own and main- tenance 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 on- policy reinforcement learning (Mnih et al. 2016)] (Schulman et al.) 2017] |Espehott et al.) (018). Here a central learner process receives data from 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 dif- ferent machines. With 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 this constraint is increase

---

Published as a conference paper at ICLR 2021
the number of parallel actors, often resulting in excessive CPU resource usage and limiting the total number of experiments that can be run on 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 al.] 2017) to solve partially- observable environments (Parisotto et al., 2019). Transformers (Vaswani et al. (017) have rapidly emerged as the state-of- the-art 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 in- stantly access information across time as well as their superior scaling properties compared to recur- rent 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 al. 2019). However, downside to the Transformer compared to LSTM models is its significant computational cost.
In this paper, present solution to actor-latency constrained settings. "Actor-Learne Distillation" (ALD), which leverages a continual form of Policy Distillation (Rusu et al. 2013 Parisotto et al.) 2013) to compress, online, a larger "learne model" towards 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-effciencyy while still having parity with cerritaationalfffiiieeee during acting. On challenging memory environments where the transformer has a clear advantage over the LSTM, we demonstrate our Actor-Learner Dis- tillation procedure provides substantially improved sample efficiency while still having experiment run run-time comparable the smaller LSTM.
the LSTM model's
2 BACKGROUND
y, A Markov Decision Process (MDP) (Sutton & Barto, 1998) is a tuple of A, T.7, R) where S
is finite set of states, A is a finite action space, T(s s) is the transition model, ? E [0, 1] is discount factor and R is the reward function. A stochastic policy π E IT is a mapping from states to probability distribution over actions. The value function (*(s) of a policy π for particular state is defined as the expected future discounted return of starting at state 8 and executing क: (5(s) Ei y'r.. The optimal policy n² is defined as the policy with the maximum value at every state, i.e. Vs E S, n E II we have /" (s) > V (s). The optimal policy guaranteed to exist (Sutton& & Barto, 1998). In this work, we focus on -brseally-Obseraablee MDPs (POMDPs) which define environments which cannot 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, on-policy value- based extension of the Maximum Posteriori Policy Optimisation algo- rithm (Abdolmaleki et al. 2018), used an EM-style policy optimization update along with regular- izing KL constraints 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 par- allelize 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 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 RL algorithm 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) measure the speed of 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 single actor per second.
2

---

Published as a conference paper at ICLR 2021
(Data-Efficiecyy)
I-aaze (Wall-Clock Time)
1.0
1.0
32-0im
Mcan Egicode Succes
4-Layer GT/XL
0.8
0.8
% Idle
Actor
Model
0.6
0.6
for Traj.
SPS
0.4
0.4
LSTM GTrXL
34%
1053
84%
70
0.2
0.2
2
10 12
Env. Steps (Million)
Time (minutes)
Figure 1: Left: On the I-Maze environment (see Sec. 5.1), there is a clear sample efficiency advan-
tage of the Gated Transformer-X XL (GTxXL) (Parisotto et al.) 2019), which rapidly reaches 100% success rate compared to the LSTM. Center: When the x-axIS is changed to wall-clock time, the LSTM becomes the more efficient model. Right: Table showing, for both 32-dim. LSTM and 4- layer GTrXL, (1) percentage of time the learner spends waiting for new on-policy trajectory data from distributed actor processes, (2) how many environment steps are processed per second (SPS) on a single actor process. We the learner is spending a majority of its time idling and this is due to the order-of-magnn -addr--f-magniuudess slower actor inference. Plots averaged over seeds, all run using reference machine A (Appendix B)
that for the much more computationally expensive transformer,
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. I we present some results comparing a transformer to an LSTM on the I-Maze memory environment (see Sec. 5.I). In Fig. [ Left, a 4-layer Gated Transformer-Xi (GTXXL) (Parisotto eet al. 2019) agent significantly surpasses a 32-dim. LSTM in terms of data efficiency, with the x- axis 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 sidd-ee of Fig. I 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 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 laae--ltteec-nsttriiee which entails a hard constraint on actor inference speed.
settings, c.g. robotic control environments,
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, strict constraints on actor latency, which precludes ever running a large model during inference. As an alternative solution which respected this constraint, we designed model compression procedure termed "Actor-Learne" Distillation" (ALD) which continually performs distillation between large- capacity "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.
desired solution that respected
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 by MM (TA, VX) and learner model by M₁ (Æ VΣ). The actor model chosen with the appropriate model capacity to compute an inference step the target hardware fast as possible, i.e. within time Tactor- In contrast, there are far fewer constraints on learner model ML, as it typically running on a central process with accelerated hardware. While there 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 transformer. As there is 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 be more elearly recognized.
restrictions besides latency concerns on the
Actor-Learner Distillation proceeds with the learner model Mr being trained using a standard re- inforcement learning algorithm, with the exception of the environment data being generated by the
3

---

Published as a conference paper at ICLR 2021
Updated Parsiaates
Ti-lenghh Fajectoris
Acdrrs
Policy
Buxched Trajectoriss
Actor
Suxched Tajectorics
Runner
Quman
Poplay
Baxchod Trjectorios
Policy
Oiti Processes
Nyp
Paromiters
Distil
Figure 2: Top: An overview of distributed Actor-Leamer Distillation, showing the processes as boxes and communication as arrows (data flow shown as blue arrows, parameter flow as orange).
actor model. The actor model trained using a policy distillation loss:
Ta(als) log TL (a Aa(ass)
LALD = ESNTA [PKu(Fa((lo]]TT
E
Σ LaE.A
(I)
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 opti- mization. As actor and learner policy will naturally be different at some points during training, an off-policy RI. algorithm for the learner could be thought to be required if the divergence is large enough. We found leveraging off- off-policy data to be effective for the actor, and sampled trajectory data from replay buffer when computing Eq. | during actor model optimization.
Beyond a distillation loss, we experimented with a value distillation loss:
LXLD = EENTA 1
(VZ
This loss functions to encourage actor model representations 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 only used to train the actor model. The final Actor-Leames Distillation loss is:
LALD L&LADD + av LXLD
where On and av are mixing coefficients to control the contribution of each loss.
DISTRIBUTED ACTOR-LEARNEL DISTILLATION
In order to make more efficient use of computational resources, we develop a distributed system for Actor-Learner Distillation based on IMPALA (Espeholt et al. 2018). An overview of this system is shown in Fig. and we detail the function of each distinct process below, the order of data flow:
Actor: There are NA parallel actors, each executing on single-threaded processes with only CPU resources available. Each Actor process steps through environment steps sequentially until trajec- tory of To time steps is gathered. Actions are sampled from a local copy of the actor policy TA- Once a completed trajectory acquired, it is communicated to Queue process.
Queue: The Queue process receives trajectories asynchronously across actors and accumulates them into batches. The batched trajectories are then passed to the Learner Runner.
The Learner Runner process runs Mr, in inference mode on the incoming batches of data using a hardware accelerator. This is done to provide (1) learning targets to the Distill pro- cesses downstream and (2) to make sure learner model initial memory states Learner. Specifically when Mr a transformer model, we additionally batch over the time dimen- sion to process data even received along with computed learner model outputs to two separate parallel processing streams, the Learnr process and the Replay process.
Learner Runner:
updated for the
rapidly. The Learner Runner then passes the batched trajectories it
Learner: The Learner process receives batched trajectories from the Learner Runner and computes, using hardware accelerator, all relevant RI. objectives (V-MPO in case (Song et al. 2020
4

---

Published as a conference paper at ICLR 2021
Meta-FFtchh (Data- Efficieni)
9x9 I-Miaze (Data-Efficiency)
25
1.0
ALD
Mean Episode Return
Mean Episode Success
A/ n+HHOWWIL
20
0.8
1.5
Ratio
0.4
Ratio 5
Ratio
0.2
Ratio
10
20
30
40
50
knv. Steps (Mill ionn)
Env. (Million)
Figure 3: Left: On the Meta-Fetthh environment with objects, we evaluated the importance of the "Distillation steps per RL step" ratio. Shown on the left, as we increased this ratio we clearly observed an improved sample efficiency in the Actor-Learner Distillation procedure, until saturation ratio of 10. However, higher ratios came at the cost of substantial increase in wall-clock time, as the high number of distillation steps quickly becomes the bottleneck of the system's speed. Right: Parallelizing the distill processes using HOGWILD! improved the DpRL ratio without causing a slowdown in system speed. enabling an increase in data efficiency without a decrease in system speed, as shown in this result on 9x9 I-Maze.
along with the loss in Eq. to regularize the learner model towards the actor model. In the Learner process, only the learner model parameters updated. Similarly to recent work on deep actor-critic algorithms (Song et al. 2020; Luo et al.) 2020), we use a "target network" where we only communi- cate updated learner parameters to the Learner Runner every K2 optimization steps. Although not used here, the Learner process can additionally receive data from the Replay process.
The Replay process manages replay buffer containing a large store of previously collected batched trajectories. Incoming batches of trajectories from the Learner Runner are archived in a large firt-in-ffrst-out queue. The Learner and Distill processes can then request batches of trajectories which are uniformly sampled from this queue.
Replay:
Distill: Distill processes request data from the Replay buffer and use the retrieved trajectories to compute the distillation loss in Eq. 3 which is then used to update the actor model parameters. Similar to the Learner process, we utilize a target network scheme which updates model parameters on the Actor processes every KAA optimization steps.
4.1 IMPROVING DISTILL/ RL STEP RATIO
An observation found early in the development of ALD was the significance the "distillation steps per RL steps" (DpRL.) ratio. The DpRL ratio measures how many agent steps are taken per second on the actor model (which takes "distillation steps") comparatively with the learner model's agent steps per second (which takes "reinforcement learning steps"). This exemplified in the left side of Fig. where in the Meta-Fetch environment with 3 objects (see Sec. 5.2). we ran ALD but set fixed number of Actor agent steps for every Learner agent step. The graph shows a greatly enhanced sample efficiency when the number of actor agent steps is increased in relation to learner agent steps, until saturating at around 10 actor agent steps per learner agent step. Out of all hyperparameters we tested, we found that high DpRL ratio was consistently the most critical parameter in increasing the sample efficiency of ALD. However, increasing the DpRL. ratio naturally increased the total run time of the system, now the distillation process became the main system bottleneck, and this encouraged investigation into parallelized ways of improving actor agent steps.
To avoid a complete system slowdown while maintaining high DpRL ratio, we leveraged the parallelized training procedure HOGWILD! (Recht et al.) 2011) as a replacement to the single- Distill-process system originally tested. In this setting, there all of which asynchronously sample batched trajectories from the replay buffer to update actor model parameters. As there typically more than available accelerators reference machines in App. B). the parallel Distill processes can be evenly distributed between available accelerators to make the best use of the resources at hand. In the right side of Figure 3 the HOGWILD! Learner Distillation variant consistently achieved better sample complexity with settings.
ND parallel Distill processes,
single machine (see e.g-
equivalent time complexity, across wide sweep of
EXPERIMENTS
For experiments, we use the V-MPO algorithm (Sng et al., 2020) as the RL algorithm underlying each of the procedures we test. For ALD, use V-MPO with V-trace correc-
Algorithm Details:
5

---

Published as conference paper at ICLR 2021
9x9 I-Maze (Data-Efficienc))
9x9 I-Maze (Wall-Clock Time)
1.0
1.0
LSTM
Mean Ep sode Success
GTrXL
0.8
0.8
ALD (Ours)
Asymm.
0.6
0.6
0.4
0.4
0.2
0.2
0
3
5
0
2
4
Env. Steps (Million)
Time (minutes)
Figure 4: Results on the I-Maze environment for all models. Left: x-axis as number of environment steps. Right: x-axis as wallclock time. All curves have seeds. Obtained on reference machine A (App. B
Metz--ootc (Data-Efficeenc))
Meta Fetch (Wall-Clock Time)
Meta Fetch (Data Efficiency)
20
Nean Ep Raturn
Nown Eptpode Success
10
0.6
Axymon
30
0.4
2.
2
10
20
30
31
20
30
40
50
Staps (Millon)
Time 11000 seconds)
Erti, Staps (Mi (oo))))
Time (1000 seconds)
Figure 5: Results on the Meta-Fetch environment with 4 objects. Left: y-axis as environment returns (objects fetched correctly per episode). Right: y-axis as success rate (all 4 objects fetched correctly at least once in an episode). All curves have 3 seeds. Obtained on reference machine B (App. B
tions (Espeholt et al. 2018) which worked slightly better in preliminary experiments. For bascline models V-MPO without V-trace worked better, in [Song al. (202)). For all experiments, use single-layer LSTM of varying dimension depending on the environment as the actor model. the Gated Transformer-XL. (Parisotto al., 2019), which initially ob- served had better sample efficiency and optimization stability than standard transformers. We vary the number of transformer layers dependent the difficulty of the environment. More details on the experimental procedure can be found in App. A
For transformers,
For each environment, we individually run the actor and learner model architectures used in ALD as baselines. As an additional baseline besides the learner and actor models in isolation, we introduce a model which maintains distinct policy and value network. Here the policy network has the same architecture as the actor model in ALD (i.e. small LSTM) and the value network has the same architecture the learner model (i.e. transformer). As the value function does not need to be run actor processes, only the policy, this baseline achieves similar run time ALD without requiring any extra insights. This baseline is an instance of Asymmetric Actor-Crite (Asymm. AC) (Pinto et al., 2017), but where instead of the value function having privileged access to state information has privileged access to computational resources. Comparison of ALD to Asymm. AC will reveal whether the introduction of distillation losses has any positive effect over the simpler method of creating independent policy/value functions.
Baselines:
5.1 I-MAZE
The I-Maze environment has the agent start at the top-left corner of grid-world with the shape of an "I" (Fig. 7 Left). The agent must travel from the top-left corner to one of the bottom corners of the I, where the episode ended and the agent receives reward. The particular corner the agent must travel to is revealed in the top-right corner of the I, which contains an "indicator" tile that is either or depending on whether the left or right bottom corner contains the reward. If the agent does not enter terminating corner state in H time steps, the episode ends without reward. The agent entering the correct bottom corner goal based on the indicator receives a reward 1. Entering the incorrect corner results in the episode terminating with reward. As memory is the main concern of these results instead exploration, for larger maze results (i.e. a maze of width 15), we add a shaping reward of 0.01 whenever the agent visits state along the main corridor of the "P". Once the

---

Published as a conference paper at ICLR 2021
9x9 I-Maze (Data-Effiieecy)
15x15 i-aaee (Date- Efficiency)
% of Steps with Mean Success between [0:25 25.0.75
1.0
2.0
Mean Egtzode Success
Mcon Betsad Success
0.0
0.8
[Ours)
0.6
0.0
Adyiters AC
2..
0.4
LSTH O1EXL
0.3
0.2
Vatue Dietlletion No Valie Dilllttion
0.2
ALD (Ours)
Asymm. AC
H
10
Env. Steps (Million)
Env. Steps [Million]
Value distillation consistently provides a slight improvement to results. Results here were taken over hyperparameter sweep with either value distillation enabled or disabled, and the best performing curves were chosen from each setting. Center, Right: Plotting per-seed curves (Center) of results in 15 x 15 I-Maze reveals that the Asymm. AC and LSTM baseline spend a significant amount of time achieving a mean success rate around 0.5 (quantified in the right-hand bar plot). A success of 0.5 suggests both models have learnt to enter a goal but are not yet able to use their memory to make the inference between goal entered and indicator identity. In contrast ALD and GTrXL rapidly learn reach a perfect return of and do not spend much time near 0.5.
Figure 6: Left:
agent has traversed particular state, it cannot gain another shaping reward for returning there until the next episode. The agent has an orientation and observes every pixel directly in front of it until reaches an occluding wall.
The environment provides a clear test on an agent's ability to remember events over long horizons, as the only way to reliably terminate an episode successfully is by remembering the indicator observed near the start of the maze. In Figure we provide complete reward curves for maze of dimension 9 x 9. We plot curves with two different x-axes: environment steps, measuring sample complexity, and wall-clock time. All curves were obtained on Reference Machine A (see Appendix B under identical operating conditions, meaning wall-clock time is fair comparison. We can first observe the clear sample efficiency gains of the transformer over the LSTM as well it's poorer performance when considering wall-clock time. Furthermore, achieves significantly improved sample complexity over the stand-alone LSTM. This confirms the ability of ALD to recover the sample-efficiency gains of the learner model. In terms of wall-clock time, the Actor-Learner Distillation procedure is unmatched by any other procedure, demonstrating its clear practical advantage.
observe that Actor-Learerr Distillation
In contrast to Actor-Leamer Distillation, the Asymm. AC bascline does not seem to perform as well despite equivalent model complexity and equivalent time complexity, it contains both identical actor and learner models. To gain insight into this result, we plotted the per-seed curves of Actor- Learner Distillation and Asymm. AC on the x 15 I-Maze task in Figure We can clearly see an underlying pattern to the seed curves, all models have local optimum at around 0.5 success where each model class spends a varying amount time. A return of 0.5 suggests that the model has learnt to enter goal, but not yet learnt to use the indicator to correctly determine which goal contains the positive reward. We measured how many environment steps each model class spent in this local minima out of the 10 million total, and found that the LSTM and Asymm. AC spend majority of their time there. While it can be expected that the transformer can casily exit this optima due to its ability to directly look back in time, interestingly ALD seems recover the same efficiency escaping this minima as the stand-alone transformer. This suggests that ALD can successfully impart the learner model's inductive biases the actor model during learning in a way that is substantially more effective than just using the transformer as a value function (Asymm. AC). Finally, in the left- hand-side of Fig. 6 we ran an ablation on 9 x I-Maze to determine whether the value distillation loss provided benefit to learning. We performed a hyperparameter sweep for each setting of value distillation enabled or disabled. The results demonstrated that using value distillation provides a slight but significant improvement.
5.2 META-FETCH
The "Meta-Fetch" environment, shown in Figure 7 requires an agent to fetch number of objects distributed randomly on an empty gridworld, with each object required to be fetched in a particular hidden sequence. Crucially, in Meta-Fetch (1) the agent can not sense the unique identity of each object, meaning each object looks the same to the agent, (2) the agent does not know beforehand the sequence of objects it must obtain, and (3) observations consist of local views of pixels in front of the agent's 's current orientation (see Fig. [7). When the agent collects an object, it can either receive positive reward if this object is the next correct in the sequence, receive reward and
7

---

Published as a conference paper at ICLR 2021
the sequence of objects is reset. To prevent cheating (i.e. the agent obtaining reward and then immediately resetting the sequence, and repeating), the agent receives a positive reward only on the first instance of it collecting the next correct object in the sequence. Once all objects in the sequence have been fetched, the objects and rewards are reset and the environment proceeds as it did at the start of the episode (except the agent is now in the position of the last object fetched). Meta-Fetch presents a highly challenging problem for memory architectures because the agent must simultancously learn the identity of objects by mapping their spatial relations using only local egocentric views, as well as learn discrete search strategy to memorize the correct sequence of objects to hit. Meta-Fetch was inspired by the meta-learning environment "Number Pad" used in previous work (Humplik et al.) 2019 Parisotto et al. 2019).
Results are reported in Figure S for a Meta-Fetch environment with randomly located objects. We present plots of average return using both environment steps and wall-clock time as x-axes (left plots), along with an additional plot which measures "success" (right plots). Success Meta- Fetch is defined as the agent completing at least I full object fetch sequence (agents can accomplish multiple full fetch sequences for additional reward). As in I-Maze, we re-confirm the observation that transformers achieve clearly superior sample efficiency over LSTMs, while at the cost of a much slower wall-clock performance. Additionally, unlike I-Maze where Asymm. AC had a more positive effect on sample efficiency, there is less substantial difference between LSTM and Asymm. AC.
However, we can see major difference between the baselines and Actor-Learner Distillation, which early on achieves close to the transformer's s sample efficiency with a much smaller model. Addi- tionally, it seems that in this much more challenging environment, which requires both memory and compositional reasoning, the LSTM actor model in ALD suffers in performance later on when the GTrXL achieves higher rate of increase in reward. Examining the success rate reveals an interest- ing trend where the ALD-trained LSTM achieves better generalization than the LSTM and Asymm. AC models. In particular, we can see that although the LSTT and Asymm. AC have average return that is increasing, their success rate is relatively stagnant (Asymm. AC) or even decreasing (LSTTM). This suggests the increase in return these baselines are observing mainly stem from them learning how to search through a fraction of possible object layouts more efficiently. In contrast, the ALD- trained LSTM's success rate is correlated with its return, suggesting it is learning to succeed in a larger variety of object layouts.
RELATED WORK
Whether a deep neural network could be compressed into simpler form was examined shortly after deep networks were shown to have widespread success (Ba Caruana 2014 Hinton et al. 2015). Distillation objectives revealed that not only could high-capacity teacher networks be transferred to low-capacity student networks (Ba & Caruana 2014), but that training using these imitation ob- jectives produced a superior low-capacity model compared to training from scratch (Hinton et al.) 2013). with this effect generalizing to the case where student and teacher share equivalent capac- ity (Furlanello et al.) 2018). Similar to this work's 's desired goal to leverage the rapid memorization capabilities of transformers while training an LSTM, other work showed that distillation between different architecture classes could enable the transfer of inductive biases (Kuncoro et al.) 2019 Abnar et al. 2020). Within RL, distillation has most often been used as method for stabilizing multitask learning (Rusu et al. 2015
Parisotto al., 2015 Teh et 2017 Berseth et al. 2018).
The method most closely related to ALD is Mix&Match (Czarnecki et al. 2018), where an ensemble of policies is defined such that each successive model in the ensemble has increasing capacity. All models in the ensemble are distilled between each other, and a mixture of the ensemble is used for sampling trajectories, with the mixture coefficients learnt through population- based training (Jader- berg et al. 2017). There are however significant difference from ALD:: the distillation procedure was meant to transfer simple, quick-learning skills from low-capacity model to high-capacity model (the opposite intended direction of the benefit of distillation in ALD) and mixture policy was used to sample trajectories which did not consider actor-latency constraints (and preferred sampling from the larger model if it performed better).
Similar to ALD, various works have looked at exploiting asymmetries between acting and learning in reinforcement learning (Pinto et al. 2017: Humplik et al., 2019). In (Pinto et al.) 2017), since the value function is ultimately not needed during deployment, it was given privileged access to state information during training (whereas the policy must operate directly from pixels). Our baseline
8

---

Published as a conference paper at ICLR 2021
Asymm. AC was derived from this insight that value networks are solely needed during learning, but our variant had a value function that utilized privileged compute resources instead of state access during training. Another closely related method is the training paradigm of Humplii et al. (2019), where a belief network was trained to predict hidden state information in partially-observalee meta- learning environments. The hidden representations from the belief network were then used as an auxiliary input to an agent network trained to solve the task, enabling representation learning to use privileged information during training.
As an alternative to model compression, architecture development can reduce the computational cost of certain model classes. Transformers in particular have had large amount of attention towards designing more efficient alternatives (Tay et al. 22200), as gains from increased model scale have yet to saturate (Brown et al., 2020p. This architectural development has taken many forms: sparsified attention (Child et al., 2019), compressed attention (Rae et al., 2020 Dai et al. 2020), use of KNN instead of soft-attention (Lample al., 2019), etc. (See [Ta et al. (2020) for a comprehensive review). However, architectural development is largely orthogonal to the ALD procedure, as it means we can further scale actor and learner models by a corresponding degree.
CONCLUSION
In conclusion, towards the development of efficient learning methods actor- actor-latency latency-constrrineed settings, we developed the Actor-Learner Distillation (ALD) procedure, which leverages a continual form of model compression between separate actor and learner models. In the context of learn- ing in partially-obeerabbee environments using a distributed actor-critic system, ALD successfully demonstrated the ability to largely capture the sample efficiency gains of a larger transformer learner model while still maintaining the reduced computational cost and lower experiment run time of the LSTM actor model. As supervised learning demonstrates ever increasing gains in performance from increasing model capacity, we believe the development of effective methods of model compression for RL will become more prominent area of study in the future. In future work, we wish to inves- tigate whether integrating the ALD procedure with batched inference for actors would still maintain the same performance increases demonstrated in results, while the same time enabling larger actor models to be used and correspondingly larger learners.
ACKNOWLEDGMENTS
This work was supported in part by the NSF IIS1763562 and ONR Grant N00014181286.. We would also like to acknowledge NVIDIA's GPU support.
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Mar- tin Riedmiller. Maximum a posteriori policy optimisation. In International Conference on Learn- ing Representations, 2018. URL htt.hhttss///perreeeeet
Samira Abnar, Mostafa Dehghani, and Willem Zuidema. Transferring inductive biases through knowledge distillation. arXiv preprint arXiv:2006.0055: 2020.
Jimmy Ba and Rich Caruana. Do deep nets really need to be deep? In Advances in neural informa- 2014.
tion processing systems, pp.
Glen Berseth, Cheng Xie, Paul Cernek, and Michiel Van de Panne. Progressive reinforcement leamn ing with distillation for multi-skilled motion control. In International Conference on Learning Representations, 2018 URL / http:::///ppeeeeeeeeeee
net. fttp://rrmmmiideeee fon
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1904.00994,, 2019.
Karl Cobbe, Jacob Hilton, Oleg Klimov, and John Schulman. Phasic policy gradient. arXiv preprint S5t::2009.044/66 2020.

---

Published as a conference paper at ICLR 2021
Wojciech Czarnecki, Siddhant Jayakumar, Max Jaderberg, Leonard Hasenclever, Yee Whye Teh, Nicolas Heess, Simon Osindero, and Razvan Pascanu. Mix & match agent curricula for rein- forcement learning. In International Conference Machine Learning, pp. 1087-1095, 2018.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoe V Le. Funnel-transformerr Filtering out sequen- tial redundancy for efficient language processing. arXiv preprint arXiv:2006. 032366 2020.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186, 2019.
Lasse Espeholt, Hubert Soyer, Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning. Shane Legg, and Koray Kavukcuoglu. Im- pala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In ICML, pp. 1406-141 1415, 2018. URL htttp::/proceeediit ngs. mlr wtw.s//80/esseeeee/.cooooo
html
Tommaso Furlanello, Zachary Lipton, Michael Tschannen, Laurent Itti, and Anima Anandkumar. Born again neural networks. In International Conference on Machine Learning, pp. 1607-1616, 2018.
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings of the AAAI Confer- ence on Artificial Intelligence, volume 33, pp. 3796-3803, 2019.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.0253/ 2015.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Jan Humplik, Alexandre Galashov, Leonard Hasenclever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint vrXiv:1905.066 2019.
Max Jaderberg, Valentin Dalibard, Simon Osindero, Wojeiech M Czameeki, Jeff Donahue, Ali Razavi, Oriol Vinyals, Tim Green, lain Dunning, Karen Simonyan, et al. Population based train- ing of neural networks. arXiv preprint arXiv:1711.09888 2017
Adhiguna Kuncoro, Chris Dyer, Laura Rimell, Stephen Clark, and Phil Blunsom. Scalable syntax- aware language models using knowledge distillation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 3472-3484, 2019.
Guillaume Lample, Alexandre Sablayrolles, Marc' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. In Advances in Neural Information Processing Systems, pp. 8548-8559, 2019.
Michael Luo, Jiahao Yao, Richard Liaw, Eric Liang. and Ion Stoica. Impact: Importance weighted asynchronous architectures with clipped target networks. In International Conference on Learn- ing Representations, 2020. URL httptttttt//ppeeeeeeteeetee
net. (tte//ooommmmiddBB....5.3
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928-1937, 2016.
Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. In Proceedings of the International Conference on Learning Representations, ICIR", 2015.
Emilio Parisotto, H Francis Song, Jack W Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. arXiv preprint arXiv: arXiv:1900.0m 1910.06644 2019.
10

---

Published as a conference paper at ICLR 2021
Lerrel Pinto, Marcin Andrychowicz, Peter Welinder, Wojeiech Zaremba, and Pieter Abbeel. Asym- metric actor critic for image-based robot learning. arXiv preprint arxiv:1710.06542,2200 2017.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lilli- crap. Compressive transformers for long-range sequence modelling. In International Confer- ence on Learning Representations, 2020. URL https::/ net/ forum?idd BylKikSYDH
Benjamin Recht, Christopher Re, Stephen Wright, and Feng Niu. Hogwild: A lock-free approach to parallelizing stochastic gradient descent. In Advances in neural information processing systems, pp. 693-6993 2011.
A. A. Rusu, G. Colmenarejo, C. Gulcehre, Desjardins, J. Kirkpatrick, R. Pascanu, V. Mnih, K. Kavukcuoglu, and R. Hadsell. Policy distillation. In Proceedings of the International Confer- ence on Learning Representations, ICLR", 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1007.06347, 2017.
H. Francis Song, Abbas Abdolmaleki, Jost Tobias Springenberg, Aidan Clark, Hubert Soyer, Jack W. Rae, Seb Noury, Arun Ahuja, Siqi Liu, Dhruva Tirumala, Nicolas Heess, Dan Belov, Martin Ried- miller, and Matthew M. Botvinick. V-mpo: On-policy maximum posteriori policy optimization for discrete and continuous control. In International Conference on Learning Representations, 2020.
R. Sutton and Barto. Reinforcement Learning: an Introduction. MIT Press, 1998.
Yi Tay, Mostafa Dehghani. Dara Bahri, and Donald Metzler. Efficient transformers: A survey, arXiv preprint ar arXiv:2009.067:2. 2020.
Yee Teh, Victor Bapst, Wojeiech M Czamecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, pp. 4496-4506, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- tion Processing Systems, pp. 5998-6008, 2017.

---

Published as a conference paper at ICLR 2021
Figure 7: Left: Example of the 15 15-size I-Maze environment, with the indicator in this case being red. Center: An example of the agent observation, which a 3-pixel-wide egocentric beam starting in the direction of the agent's orientation. The agent cannot see behind itself or through walls. Right: Example of the Meta-Fetch environment with 4 randomly located objects shown in brown. The agent is shown in orange with its orientation indicated by the direction of the arrow.
APPENDIX
A EXPERIMENT DETAILS
For models, we sweep over the V-MPO target network update frequency K2 10, 100). In initial experiments, we also sweeped the "Initial a" setting over values 0.51, 1.0, 5.0). All experiment runs have 3 unique seeds. For each model, we choose the hyperparameter setting that achieved highest mean return over all seeds. Additionally we use PopArt (Hessel et al. 2019) for the value output.
Hyperparameter
Value
Optimizer
Adam
Learning Rate NA
0.0001
30
ND
8
Batch Size
64
Tv
20
Discount Factor (7)
0.99
Grad. Norm. Clipping
Disabled
Initial
1.0
Initial
(0.1, 0.5, 1.0, 5.0) 0.1
En
Ea
0.004
KL PopArt Decay
{1. 10. 100)
0.0003
Table 1: Common hyperparameters across experiments.
A.1 I-MAZE EXPERIMENTS
Observations, Actions and Metrics: Observations in I-Maze are the single row of pixels starting the agent's s current position and extending in the direction the agent facing. The row of pixels extends as far as the maze dimension. If the agent has a wall in its field of view (represented as black pixels Fig. (7). pixels further away from that point are occluded. For the 9 maze, a reward is only given when the agent enters the correct goal as decided by the indicator. For the 15 x 15 maze, it is the same but there is an additional shaping reward where whenever the agent enters one of the states along the central column of the "I" for the first time within its current episode, receives a small reward of 0.01. This reward is meant to encourage exploration and prevent the very large number of environment steps otherwise necessary for the agent to end up goal when
12

---

Published as a conference paper at ICLR 2021
taking actions randomly. The agent has 4 actions, move forward, turn left, turn right and do nothing. Moving forward into a wall causes no change in state. 9 x 9 I-Maze episodes last for a horizon of 150 time steps while 15 15 I-Maze episodes last for 350 steps. Observations have channel dimensions, one for whether the pixel is free space or a wall, one for the green goal indicator and one for the red goal indicator.
LSTM: We use a single-layer LSTM with a hidden dimension of 32 for all experiments in this domain.
GTrXL: We use 2-layer GTrXL for the experiments in 9 x 9 and a 4-layer GTrXL for the experi- ments in 9 9. We use an embedding size of 256, attention heads, head dimension of 32, a gated initialization bias of 2 and memory length of 64. Other details were followed from Parisotto et al.
(2019).
ALD: We use the corresponding LSTMs and GTrXL described above depending on the environ- ment. During hyperparameter sweeps, we tested KA E (0, 100). We set on 1 and sweep E 1).
Asymm. AC: We use the corresponding LSTMs and GTrXL described above depending on the environment.
A.2 META-FETCH
Observations, Actions and Metries: In Meta-Fetch, an agent acts in an empty 2D gridworld of size 7 x 7 given only local observations. At each step the agent can choose to either move forward or turn left or right. Observations are the 15 pixels in front of the agent pointing in it's orientation, along with the 2 pixels either side of this ray, for a complete observation size of C x 15 x 3 pixels (see Fig. 7 Each observation has C' channels which determine whether that pixel represents free space, wall, an object, or collected object. Every episode has the agent located in new configuration of object locations. We set a maximum episode length of 300 steps for Meta-Fetch.
The goal of the agent in Meta-Fetch is collect, by colliding with, each of the randomly located objects in a specific sequence. This sequence is not known beforehand and is randomized each episode. The only clue to discovering this sequence is that when the agent hits the next right object in the sequence, it gains a reward of and the object is sensed as "collected". Otherwise, if it hit the wrong object in the sequence, then all previously collected objects are reset and the agent must restart the fetch sequence from the beginning once again. Once all objects are collected, they are all reset and the agent can restart the same sequence to once again gain reward for each collected reward. Objects which have been reset before the complete fetch sequence is achieved do not give further reward upon re-collection, prevent the agent from falling into an local optimum of repeatedly collecting a correct object and then immediately collecting an incorrect object to reset the correct object.
LSTM: We use single-layer LSTM with hidden dimension of 128.
GTrXL: We use a 4-layr GTrXL, an embedding size of 256, attention heads, head dimension of 32, gated initialization bias of 2 and a memory length of 64. Other details were followed from Parisotto et al. (2019).
ALD: We use the corresponding LSTMs and GTrXL described above depending on the environ-
ment. During hyperparameter sweeps, we tested KA E {10, 100}. We set ag = and sweep av E {0, 13.
Asymm. AC: We use the corresponding LSTMs and GTrXL described above depending on the environment.
B COMPUTE DETAILS
Reference Machine A: Reference Machine A has a 36-thread Intel(R) Core(TM) 7980XE CPU @ 2.60GHz, 64GB of RAM, and 2 GPUs: GeForce GTX 1080 Ti and TITAN V.
Reference Machine B: Reference Machine has a 40-thread Intel(R) Xeon(R) CPU E5-2630 v4 @ 2.20GHz, 256GB of RAM and GPUs: Tesla P40 and Tesla V100-PCIE-16GBB
13

---

Published as a conference paper at ICLR 2021
9x9 I-Maze (Wall-Clock Time)
1.0
LSTM
GTrXL
ALD (Ours)
Asymm.
0
10
15
20
25
Time (minutes)
Meta-Fetch (Wall-Clock Time)
Meta-Fetch (Wall-Clock Time)
0.8
15
0.6
10
0.4
5
0.2
0.0
0
10
20
30
0
10
20
30
Time (1000 seconds)
Timc (1000 scconds)
Figure 8: Full wall-clock time curves. Top: 9x9 I-Maze success rate.
Left: Meta-Fetch reward.
Right: Meta-Fetch success.
C DESCRIPTION OF HOGWILD!
In order speed-up certain learning aspects our proposed method, we the HOGWILD! (Recht et al. 2011) algorithm. HOGWILD! enables lock-free stochastic gradient optimization to be per- formed by several parallel learning processes. Each process has shared access to the parameter vector of the learning model, and can update it without requiring mutually-exclusive lock. This allows parameter updates to occur in parallel as soon as they are available, but potentially leads to race conditions where the parameter vector is read during it being updated by different process. This means a fraction of the parameter vector can be stale when a learning process calculates the parameter gradients. However, while this parallel overwriting of the parameter vector introduces noise via this staleness, it usually does not significantly inhibit performance and is typically faster than alternatives which use explicit locking prevent race conditions.
14

---

Published as a conference paper at ICLR 2021
9x9 I-Maze (Data-Efile E
1.0
LSTM
Mean Episode Success
Gl GLL
0.8
ALD (Ours)
Asymm, AC
0.6
0.4
0.2
1
2
Env. Steps (Million)
Figure 9: Per-seed curves for 9x9 I-Maze.
15


Reviews:
[{"review_id": "tu531o1GA0L", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The main aim of this paper is to increase the efficiency (in terms of wall clock time) of Transformer based models for reinforcement learning.  \n\nPrevious works have shown Transformer-like models to be highly performant across a range of domains, including recent results on reinforcement learning (with gated Transformer-XLs). However a drawback of these models  (over say, LSTMs) is their relatively slower inference speeds. This is especially a problem in RL settings where Actors are typically run on CPUs (not GPUs) and send trajectories to a central learner. The problem here is two-fold: 1) slower overall training time due to high latency with the learner waiting on actors 2) slower inference post training in latency sensitive deployment settings like robotics/other control based settings. \n\nThe solution proposed here (the method is referred to as \"Actor-learner distillation\") is to instead use LSTMs for acting and Transformers for learning. Typically the learner would send parameters to the actors every update -- as this is not possible in this hybrid approach, a distillation loss is instead suggested as a means of updating the actors (with a replay buffer). The authors successfully show that this transfers both \"good policies\" and the relevant inductive biases from the Transformer as well allowing fast inference as expected. Further, the \"off policy\"-ness of the model does not turn out to be an issue in the domains considered. The results are promising and the the writing very clear. A few points below:\n\n* One question I have is whether the authors ran an experiment replacing the LSTMs in the actors simply with smaller transformer models? For example, 2 layers instead of 4 or smaller embedding sizes? This would also allow faster inference while inherently retaining some of the inductive biases. \n\n* As a second baseline: while this wouldn't solve the issue of slower training, one could achieve faster inference of the trained model simply by distilling the final Transformer into an LSTM (and then optionally fine-tuning). Are the authors able to report numbers for these?\n\n* Further, in between the two extremes of only LSTMs acting or only Transformers one could also consider more hybrid approaches that start with a Transformer to learn a very good starting policy in a few steps and then distill this into an LSTM online similar to work here (https://arxiv.org/pdf/1806.01780.pdf). The worked linked here performs \"online distillation\" combined with a convex combination of model outputs which would then negate the need for a replay buffer (though would require running both models on the actor until fully switching over to the LSTM). \n\n* While the results are conclusive, the work could be improved by running on larger (or 3-dimensional) environments as considered in the original Gated Transformer paper this builds on. \n\n* Lastly, are the authors able to provide possible future directions to be considered that could boost this work further? \n\nOverall this is well written paper exploring a conceptually simple but well-implemented and important idea. As models grow in size, techniques like this will be very important to allow scaling up of RL capabilities. I would recommend accepting this paper. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "RtTDPG1m26P", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe 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.\n\n### Pros\n\n - The paper tackles a practical problem arising when designing an RL learning pipeline.\n - The paper proposes an original idea of applying distillation in the RL setting to solve a problem above.\n - Implicitly (because there is no discussion of this), the paper raises an important question of a trade-off between sample efficiency and computational efficiency in RL.\n\n### Cons\n\n- Confusing positioning of the paper;\n- Implicit assumptions not discussed/challenged in the paper;\n- Confusing results presentation (plotting, result interpretation);\n\n### Reasoning behind the score\n\nI believe this paper has great potential and can be improved in two possible ways. First, which will make it more general and increase its impact, is to turn it into a discussion of the trade-off between sample efficiency and computation complexity. Second, to narrow the scope and improve the positioning/clarity to get an improved version of this submission. I would really like the authors to choose the first route, but it is their decision.  At this moment, due to the pros and cons I described above, I give it a 6, Marginally Above Acceptance Threshold.\n\n### Questions to the authors\n\n- The paper assumes that it's better to lose in sample efficiency while computing faster. Why do you think it should be like this? What is the motivation behind such an assumption? Can you discuss the sample/computation efficiency trade-off in the paper?\n- The paper assumes that LSTMs are less computationally efficient than Transformers. What exactly do you mean by computational efficiency? Yes, your plots show that the particular model instances make LSTM to achieve better performance with smaller wall-clock time. But is it general? One of the reasons behind transformers' success is their ability to parallelise. Can you explicitly state the assumptions for which you are solving the posed problem?\n- The storyline in the abstract and in the intro/conclusion do not quite fit. For example, the abstract goes from distributed RL setting → constrained actors → cannot train more complex models. But in the intro, you start with the 'Transformers are great' story. In addition, in the second paragraph of the intro, you mention 'actor-latency-constrained' settings when 'there exists a max latency constraint on the acting policy'. I don't think you consider this setting later in your experiments or before (in the abstract). What is the exact problem you are solving?\n- \"An agent designer has to make a hard decision between ... better sample efficiency ... or ... lower wall-clock time and reduced computational cost\". Can you elaborate a bit more on the setting? Ideally, describe this trade-off in more detail. Sometimes, sample-efficiency and time are intertwined, if your data collection is slow, then sample inefficiency will bump the wall-clock time.\n- I'm really confused by the presentation of the results. Can you, please, plot the curves till convergence of all the algorithms? For example, the left subplot in Figure 1 shows the transformer converged to 1.0, but not the LSTM. At the same time, the next subplot shows the LSTM converged, but not the transformer. I think this is of utmost importance for better interpretation of the results. Same holds for other figures.\n- For the table on top of the page 3, how do you get the numbers? At which point in training did you do the measurements? Did you average?\n- On page 6 you say \"... we can observe that ALD achieves near sample complexity parity to the transformer reaching a success rate of 1 in less than twice as much steps.\" Why is it a near parity (2X!)?\n- Intro says 'ALD provides sample efficiency on par with transformers'. This is not what the plots show (e.g. 2x difference on Fig 4, 12.5 vs 20 score on Fig 5 and 0.7 vs >0.8 in episode success on the 3rd subplot of Fig 5. What is your definition of 'to be on par'?\n\n### Other comments\n\n- Background\n    - You cite Puterman for MDPs, however, Puterman does not include the discounting coefficient into the MDP tuple.\n    - Can you give a page of Sutton&Barto for your reference for the existence of an optimal policy?\n    - You say you focus on POMDPs, but later use $s$ as inputs to the policy. This feels weird.\n    - End of the first paragraph of the background: 'recurrent LSTMs'. Are there non-recurrent LSTMs?\n    - Great second paragraph describing the glossary, I really like it.\n- Distributed Actor-Learner Distillation\n    - I think, V-MPO comes too late here. I've had the question about the learning algorithm several times, before I found it here. The same hold for IMPALA. I think this should be mentioned in the introduction.\n    - I like the description of the algorithms in the paragraph, however, I would love to have pseudocode in the appendix.\n    - \"which are uniformly sampled from the FIFO queue\": How does FIFO affect sampling?\n- Experiments\n    - \"For each environment, we run the actor and learner model architectures used in ALD as baselines, where each of these models are independently trained using standard RL\". This sentence is a bit confusing. Also, can you explain what does 'standard RL mean?\n    - \"Comparison of ALD to Asymm.\" should be on the next line, I believe.\n    - Why did you decide to do you present the per-seed curves on 15x15 maze while presenting the aggregated results for 9x9? How do these look on 9x9 maze?\n- Related Work\n    - It would be great to have a paragraph giving pointers to the existing distributed RL frameworks and some questions the papers introducing them investigate.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Cgp0XpUAS8j", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe 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. \n\nPositive points:\n(+) I think this is a relevant problem setting\n(+) Clear description of the algorithm which seems well thought-out \n\nPossible weakness:\n(-) Experimental evaluation. \n\nI'm currently recommending a weak rejection. I believe the idea and execution of the paper (in particular the method part) is good. However, a stronger experimental section which also evaluates on more complex visual domains would greatly strengthen the paper. Doing so on tasks with an equal complexity on the required long-term memory might be computationally challenging, however, I personally would already be satisfied with showing that the proposed algorithm doesn't perform worse than the baselines, even on tasks which don't require (as much) memory. \n\nIn other words: The current experiments clearly show why and when ALD can provide an advantage. What is missing for me is a visually more complex experiment that shows me that the actor/learner split and associated off-policy-ness doesn't create additional problems on (visually) more complex environments. If those environments don't require as much memory, there is no reason to expect that ALD can outperform the baseline, which would be ok for me, as long as it doesn't underperform it.\n\nOne additional question (unrelated to evaluation of the paper) I was wondering: The authors mention Teh et al. What I am wondering is why, when updating pi_A, you only use the distillation loss and not also the RL loss? Did you try?\n\nOther minor point (no impact on evaluation): I would have found a brief description of HOGWILD, e.g. in the appendix, to be helpful. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "CBy-U42P3Jh", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "##########################################################################\n\nSummary:\n\nThe paper proposes a solution to actor-latency constrained settings in RL by using policy distillation to compress a large “learner model” towards a more tractable “actor model”. In particular, it proposes to exploit the superior sample efficiency of transformer models while utilising an LSTM-based actor during execution. The proposed procedure, called Actor-Learner Distillation (ALD), provides comparable performance to transformers in terms of sample efficiency, yet produces a wall-clock run-time that's on par with LSTM agents.\n\n##########################################################################\n\nReasons for score:\n\nI'm voting for accepting the paper. I like the idea of policy distillation between transformer-based learner and LSTM-based actor in the distributed RL framework. The paper is written with great clarity and is easy to follow. I don't have major concerns regarding this work but hope that the authors can address my minor concerns in the rebuttal period.\n\n##########################################################################\n\nPros:\n\n1. The work addresses an important problem for distributed RL: whether it is possible to utilise large model capacity while still acting within the computational limits imposed by the problem setting and do all of this during the training. The problem itself is real and practical, e.g., in robotic operations with limited computation resources.\n2. The process of policy distillation between transformer-based learner and LSTM-based actors is described in great detail and provides additional insight into concrete decisions made for the method, such as the value distillation loss for the actor model and Distil/RL step ratio.\n3. This manuscript includes thorough experimental of the proposed method that includes both qualitative analysis and quantitative results. The choices regarding the two environment tasks and baselines are reasonable.\n\n##########################################################################\n\nCons:\n\n1. Regarding the DpRL hyperparameter (actor SPS / learner SPS), it isn't very clear to me how this ratio has been manipulated for instance in Figure 3. Is this achieved using modifications of actor/learner size? I'd appreciate it if the authors could add clarification on this.\n2. On the result figures where the X-axis corresponds to Wall-Clock Time, it isn't clear to me what was used as a basis for choosing the maximum range of the X-axis. For example, Figure 4 (Right) was cut after 5 minutes, where ALD is much higher than the baselines. However, I'd be interested to know when do ALDs and GTrXL (which achieves the same result in eventually) meet on the diagram.\n\n##########################################################################\n\nQuestions during rebuttal period:\n\nPlease address and clarify the cons above", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: I thank the authors for their submission and participation in the author response period. The reviewers unanimously agree that the papers proposes an interesting and original approach to using a costly model on a learner node, while distilling to a cheaper model run on actor nodes to gather experiences in a distributed RL framework. During discussion, R1 and myself emphasized the concern that the experiments in this paper leave open the question whether the approach will work beyond toy environments. However, I side with R2 and R3 in that the paper presents a valuable contribution to the community as it stands, and that the experiments proof the concept to the point that the paper should be accepted. I therefore recommend acceptance.
Decision: Accept (Poster)