Scaling Algorithm Distillation for Continuous Control with Mamba
This addresses the practical limitation of Algorithm Distillation for continuous control in Meta RL, though it appears incremental as it applies an existing model architecture (Mamba) to a known bottleneck.
The paper tackles the problem of scaling Algorithm Distillation for continuous control by replacing transformers with Mamba (S6) models, which scale linearly with sequence length instead of quadratically. They demonstrate Mamba's superiority over transformers in four complex Meta RL environments and show that scaling to long contexts makes it competitive with a state-of-the-art online meta RL baseline.
Algorithm Distillation (AD) was recently proposed as a new approach to perform In-Context Reinforcement Learning (ICRL) by modeling across-episodic training histories autoregressively with a causal transformer model. However, due to practical limitations induced by the attention mechanism, experiments were bottlenecked by the transformer's quadratic complexity and limited to simple discrete environments with short time horizons. In this work, we propose leveraging the recently proposed Selective Structured State Space Sequence (S6) models, which achieved state-of-the-art (SOTA) performance on long-range sequence modeling while scaling linearly in sequence length. Through four complex and continuous Meta Reinforcement Learning environments, we demonstrate the overall superiority of Mamba, a model built with S6 layers, over a transformer model for AD. Additionally, we show that scaling AD to very long contexts can improve ICRL performance and make it competitive even with a SOTA online meta RL baseline.