[Paper Review] Masked Autoencoding for Scalable and Generalizable Decision Making

Fangchen Liu, Hao Liu+2 more|arXiv (Cornell University)|Nov 23, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes Masked Decision Prediction (MaskDP), a self-supervised pretraining method for reinforcement learning that uses masked autoencoding on state-action trajectories to learn generalizable, scalable agents. By randomly masking states and actions and reconstructing them, MaskDP enables zero-shot transfer to goal-reaching and offline RL tasks, outperforming prior methods in zero-shot settings and showing strong scalability with model size.

ABSTRACT

We are interested in learning scalable agents for reinforcement learning that can learn from large-scale, diverse sequential data similar to current large vision and language models. To this end, this paper presents masked decision prediction (MaskDP), a simple and scalable self-supervised pretraining method for reinforcement learning (RL) and behavioral cloning (BC). In our MaskDP approach, we employ a masked autoencoder (MAE) to state-action trajectories, wherein we randomly mask state and action tokens and reconstruct the missing data. By doing so, the model is required to infer masked-out states and actions and extract information about dynamics. We find that masking different proportions of the input sequence significantly helps with learning a better model that generalizes well to multiple downstream tasks. In our empirical study, we find that a MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching, and it can zero-shot infer skills from a few example transitions. In addition, MaskDP transfers well to offline RL and shows promising scaling behavior w.r.t. to model size. It is amenable to data-efficient finetuning, achieving competitive results with prior methods based on autoregressive pretraining.

Motivation & Objective

• To develop a scalable, self-supervised pretraining method for reinforcement learning that generalizes across diverse downstream tasks.
• To enable zero-shot transfer to goal-reaching and offline RL without task-specific fine-tuning.
• To explore the effectiveness of non-autoregressive, masked token prediction in sequential decision-making, inspired by vision and language models.
• To investigate how varying mask ratios during pretraining improves generalization and downstream performance.
• To demonstrate that MaskDP can learn from diverse, unlabeled sequential data and achieve competitive performance with autoregressive pretraining methods.

Proposed method

• MaskDP applies masked autoencoding to state-action trajectories by randomly masking a portion of state and action tokens.
• The model reconstructs the masked tokens using a transformer-based architecture, learning both forward and inverse dynamics from the same training signal.
• Multiple mask ratios (15%, 35%, 75%, 95%) are used during pretraining, with each sequence sampled using a random ratio to improve generalization.
• The model is pre-trained on diverse, unlabeled sequential data, then fine-tuned with minimal data for downstream tasks.
• A total loss is used during training, combining masked and unmasked token reconstruction, which empirically converges faster than using only masked loss.
• The method is evaluated on zero-shot goal-reaching, sequential goal tasks, and offline RL benchmarks, with performance compared against autoregressive baselines like Goal-GPT.

Experimental results

Research questions

• RQ1 Can masked autoencoding on state-action trajectories enable zero-shot transfer to diverse downstream RL tasks?
• RQ2 How does varying the mask ratio during pretraining affect generalization and downstream performance?
• RQ3 Does a non-autoregressive, masked prediction approach outperform autoregressive pretraining in offline RL and goal-reaching?
• RQ4 Can MaskDP scale effectively with model size, and does it maintain performance gains across larger architectures?
• RQ5 How does MaskDP compare to autoregressive pretraining methods like Goal-GPT in terms of data efficiency and zero-shot transfer?

Key findings

• MaskDP achieves performance that exceeds or matches both training-from-scratch and pretraining-based methods in single-goal reaching tasks, even in zero-shot settings.
• In sequential multiple-goal reaching, MaskDP outperforms iterative baselines significantly, successfully planning multi-step trajectories without closed-loop execution.
• MaskDP shows strong transfer to offline RL, matching the performance of specialized offline RL methods and achieving state-of-the-art results on the Walker domain.
• The model demonstrates strong scalability: larger models fine-tuned from MaskDP pretraining consistently outperform smaller models, unlike in Goal-GPT where gains are minimal.
• Using a mixed mask ratio strategy (15%, 35%, 75%, 95%) significantly outperforms fixed-ratio baselines, with 95% masking proving most effective for meaningful reconstruction.
• Training with total loss (masked + unmasked) leads to faster convergence than using only masked loss, even on clean expert data.