Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation
This addresses the actor-latency constraints in RL for applications like robotics, enabling use of large models on limited hardware, though it is incremental as it adapts existing distillation methods to RL.
The paper tackled the problem of deploying large transformer models in reinforcement learning under compute constraints by proposing Actor-Learner Distillation, which transfers knowledge from a large transformer learner to a small LSTM actor, recovering transformer sample-efficiency gains while maintaining fast inference and reduced training time in memory environments.
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 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.