Combining Pre-Trained Models for Enhanced Feature Representation in Reinforcement Learning

TL;DR

This paper introduces Weight Sharing Attention (WSA), a novel method for combining multiple pre-trained model embeddings to improve state representations in reinforcement learning, achieving competitive results with enhanced efficiency.

Contribution

The paper presents WSA, a new architecture that effectively combines multiple pre-trained models' embeddings for reinforcement learning, addressing efficiency and performance tradeoffs.

Findings

WSA achieves comparable performance to end-to-end models on Atari games.

Scaling the number of models influences agent performance during and after training.

WSA demonstrates improved generalization capabilities in RL tasks.

Abstract

The recent focus and release of pre-trained models have been a key components to several advancements in many fields (e.g. Natural Language Processing and Computer Vision), as a matter of fact, pre-trained models learn disparate latent embeddings sharing insightful representations. On the other hand, Reinforcement Learning (RL) focuses on maximizing the cumulative reward obtained via agent's interaction with the environment. RL agents do not have any prior knowledge about the world, and they either learn from scratch an end-to-end mapping between the observation and action spaces or, in more recent works, are paired with monolithic and computationally expensive Foundational Models. How to effectively combine and leverage the hidden information of different pre-trained models simultaneously in RL is still an open and understudied question. In this work, we propose Weight Sharing Attention (WSA), a new architecture to combine embeddings of multiple pre-trained models to shape an enriched state representation, balancing the tradeoff between efficiency and performance. We run an extensive comparison between several combination modes showing that WSA obtains comparable performance on multiple Atari games compared to end-to-end models. Furthermore, we study the generalization capabilities of this approach and analyze how scaling the number of models influences agents' performance during and after training.