Deriving a Data-Efficient, Generalist Agent Through World Models

Abstract

Machine learning researchers long to derive a generalist agent whose problem-solving abilities span over multiple tasks. At the same time, we are also interested in data efficiency, i.e., training agents with fewer, limited training samples. To these two distinctive ends, recent lines of research on world models (WMs) for image-based reinforcement learning (RL) tasks have also shown a stark divergence. The first line is to scale a singular world model to operate on many environments at the cost of requiring billions of training experiences. The second focuses on training one specialized world model per environment, but within a limited budget of one hundred thousand (100k) interactions. In this research, we propose mae, the transformer WM agent that thrives in both ends, i.e., operates in multiple environments using a single set of backbone world model parameters with the limited hundred thousand (100k) training samples per environment. On the standard Atari 100k benchmark, we achieve superhuman performance in 9 of the 26 games, obtaining human-normalized mean, interquartile mean (IQM), and median scores of 1.15, 0.56, and 0.46, respectively, using a single WM backbone across all games and game-specific RL networks. While competitive, these scores fall modestly short of specialized, game-specific transformer WM agents. Remarkably, mae learns a rewarding policy without observing any positive reward in the target environment, and ablations provide evidence of cross-environment transfer. Lastly, we validate mae’s applicability in a 3D, first-person view environment with a long-horizon task.