Efficient Transformers in Reinforcement Learning using Actor-Learner 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 unaccelerated 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 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.

TransDreamer: Reinforcement Learning with Transformer World Models

“…Parisotto et al (2020) address the problem of using transformers in RL and showed that adding gating layers on top of the transformers layers can stabilize training. Subsequent works addressed the increased computational load of using a transformer for an agent's policy (Irie et al, 2021;Parisotto & Salakhutdinov, 2021). Chen et al (2021); Janner et al (2021) take a different approach by modeling the RL problem as a sequence modeling problem and use a transformer to predict actions without additional networks for an actor or critic.…”

Switch Trajectory Transformer with Distributional Value Approximation for Multi-Task Reinforcement Learning

“…In light of the above works, researchers are tempted to investigate the benefit of transformer models in improving reinforcement learning performance. The first line of work applies the transformer model to represent the component in standard RL algorithms, such as policy, models and value functions (Parisotto et al, 2020;Parisotto & Salakhutdinov, 2021). Instead of this, the second line of work (Chen et al, 2021;Janner et al, 2021b) abstracts RL as a sequence modelling problem and efficiently utilize the existing transformer framework widely used in language modelling to solve the RL problem.…”

From eye-blinks to state construction: Diagnostic benchmarks for online representation learning

“…Learned sparse attention mechanisms combined with feed-forward neural networks represent exciting alternatives for training RNNs. The best way to use attention strategies for partially observable reinforcement learning is still evolving (Parisotto et al, 2020b;Parisotto & Salakhutdinov, 2021;Loynd et al, 2020;Chen et al, 2021;Janner et al, 2021). Chen et al (2021) and Janner et al (2021) use transformers in the offline reinforcement learning setting.…”