Scaling Algorithm Distillation for Continuous Control with Mamba
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.
Code (0)
등록된 구현이 없습니다.
Tasks
continuous-controlContinuous ControlIn-Context Reinforcement LearningMambaMeta Reinforcement Learningreinforcement-learningReinforcement LearningMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
StableMamba: Distillation-free Scaling of Large SSMs for Images and Videos
State-space models (SSMs), exemplified by S4, have introduced a novel context modeling method by integrating state-space techniques into deep learning. However, they struggle with global context modeling due to their dat…
Action Recognitionimage-classificationImage ClassificationKnowledge Distillation+2Attention to Mamba: A Recipe for Cross-Architecture Distillation
State Space Models (SSMs) such as Mamba have become a popular alternative to Transformer models, due to their reduced memory consumption and higher throughput at generation compared to their Attention-based counterparts.…
Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners
Recent advancements have demonstrated that the performance of large language models (LLMs) can be significantly enhanced by scaling computational resources at test time. A common strategy involves generating multiple Cha…
MambaMathematical ReasoningPathMamba: A Hybrid Mamba-Transformer for Topologically Coherent Road Segmentation in Satellite Imagery
Achieving both high accuracy and topological continuity in road segmentation from satellite imagery is a critical goal for applications ranging from urban planning to disaster response. State-of-the-art methods often rel…
Road SegmentationMamba as Decision Maker: Exploring Multi-scale Sequence Modeling in Offline Reinforcement Learning
Sequential modeling has demonstrated remarkable capabilities in offline reinforcement learning (RL), with Decision Transformer (DT) being one of the most notable representatives, achieving significant success. However, R…
MambaOpenAI GymReinforcement Learning (RL)