RePreM: Representation Pre-training with Masked Model for Reinforcement Learning

2023/03/03 by Cai, Yuanying, Zhang, Chuheng, Shen, Wei +3
#Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG)

paper · doi:10.48550/arxiv.2303.01668

Abstract

Inspired by the recent success of sequence modeling in RL and the use of masked language model for pre-training, we propose a masked model for pre-training in RL, RePreM (Representation Pre-training with Masked Model), which trains the encoder combined with transformer blocks to predict the masked states or actions in a trajectory. RePreM is simple but effective compared to existing representation pre-training methods in RL. It avoids algorithmic sophistication (such as data augmentation or estimating multiple models) with sequence modeling and generates a representation that captures long-term dynamics well. Empirically, we demonstrate the effectiveness of RePreM in various tasks, including dynamic prediction, transfer learning, and sample-efficient RL with both value-based and actor-critic methods. Moreover, we show that RePreM scales well with dataset size, dataset quality, and the scale of the encoder, which indicates its potential towards big RL models.

Related

• Masked Autoencoding for Scalable and Generalizable Decision Making
• Masked Trajectory Models for Prediction, Representation, and Control
• Masked Generative Priors Improve World Models Sequence Modelling Capabilities
• Learning to Modulate pre-trained Models in RL
• Mask-based Latent Reconstruction for Reinforcement Learning