Reprem: Representation Pre-training With Masked Model For Reinforcement Learning

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.

Stats

• citations0
• S2 citations—
• github stars0
• HF likes0
• heat score0.00
• arxiv keycai2023reprem

Related papers

• M\(^3\)PC: Test-time Model Predictive Control For Pretrained Masked Trajectory Model (2024)0.00
• Masked Autoencoding For Scalable And Generalizable Decision Making (2022)0.00
• Data-efficient Reinforcement Learning With Self-predictive Representations (2020)0.00
• CAMEL: Continuous Action Masking Enabled By Large Language Models For Reinforcement Learning (2025)0.00
• The Surprising Ineffectiveness Of Pre-trained Visual Representations For Model-based Reinforcement Learning (2024)0.00
• Lifelong Reinforcement Learning With Modulating Masks (2022)0.00
• Towards Principled Representation Learning From Videos For Reinforcement Learning (2024)0.00
• Representation Matters: Offline Pretraining For Sequential Decision Making (2021)0.00