Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learner

TL;DR

This paper extends the Decision Pre-Trained Transformer to multi-domain environments, demonstrating improved generalization and scalability in in-context reinforcement learning for training versatile agents.

Contribution

It introduces a scalable multi-domain DPT trained with Flow Matching, enhancing generalization and performance over prior methods like Algorithm Distillation.

Findings

Achieves better generalization to unseen tasks.

Demonstrates stronger online and offline inference performance.

Trains across hundreds of diverse tasks.

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.