Scaling Algorithm Distillation for Continuous Control with Mamba

TL;DR

This paper introduces Mamba, a scalable sequence modeling approach using S6 models for Algorithm Distillation in continuous control, enabling longer context handling and improved reinforcement learning performance.

Contribution

The paper proposes Mamba, a novel model combining S6 layers with Algorithm Distillation to scale ICRL to longer contexts in continuous environments.

Findings

Mamba outperforms transformer-based AD in complex continuous tasks.

Scaling AD with Mamba improves ICRL performance significantly.

Mamba achieves competitive results with state-of-the-art online meta RL methods.

Abstract

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.