Problem
Research questions and friction points this paper is trying to address.
in-context reinforcement learning
generalization
scalability
multi-domain
unseen tasks
Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learner
🤖 AI Summary
This work addresses the limited generalization of existing context-based reinforcement learning methods in multi-task settings, particularly their inability to adapt to unseen tasks. The authors extend Decision Transformers to large-scale multi-task contextual reinforcement learning and incorporate flow matching into the training procedure, enabling agents to directly acquire new tasks during inference. This approach provides the first demonstration of the scalability of Decision Transformers in such settings and offers a theoretical interpretation of flow matching as Bayesian posterior sampling. Experimental results show that agents trained on hundreds of tasks significantly outperform prior methods—such as Algorithm Distillation—on held-out test sets, exhibiting superior generalization performance in both online and offline inference scenarios.
📝 Abstract
Recent progress in in-context reinforcement learning (ICRL) has demonstrated its potential for training generalist agents that can acquire new tasks directly at inference. Algorithm Distillation (AD) pioneered this paradigm and was subsequently scaled to multi-domain settings, although its ability to generalize to unseen tasks remained limited. The Decision Pre-Trained Transformer (DPT) was introduced as an alternative, showing stronger in-context reinforcement learning abilities in simplified domains, but its scalability had not been established. In this work, we extend DPT to diverse multi-domain environments, applying Flow Matching as a natural training choice that preserves its interpretation as Bayesian posterior sampling. As a result, we obtain an agent trained across hundreds of diverse tasks that achieves clear gains in generalization to the held-out test set. This agent improves upon prior AD scaling and demonstrates stronger performance in both online and offline inference, reinforcing ICRL as a viable alternative to expert distillation for training generalist agents.
Innovation
Methods, ideas, or system contributions that make the work stand out.
Decision Pre-Trained Transformer
In-Context Reinforcement Learning
Flow Matching
Multi-domain Generalization
Bayesian Posterior Sampling
🔎 Similar Papers
No similar papers found.
