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

Elia Piccoli, Malio Li, Giacomo Carfi', Vincenzo Lomonaco, Davide Bacciu

Keywords: Reinforcement Learning, Continual RL, State Representation, Pre-Trained Models

TL;DR: Combining multiple pre-trained models as inductive biases to shape state representation in RL

Abstract: Reinforcement Learning (RL) focuses on maximizing the cumulative reward obtained via agents' 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 knowledge 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.

Submission Number: 8