Masked Generative Policy for Robotic Control

TL;DR

The paper introduces Masked Generative Policy (MGP), a transformer-based framework for robotic control that improves success rates and inference speed in complex, non-Markovian tasks through novel masked token generation and refinement strategies.

Contribution

It proposes a new masked transformer approach for visuomotor imitation learning, enabling rapid, coherent, and adaptive control in complex robotic tasks, outperforming prior diffusion and autoregressive methods.

Findings

Achieves 9% higher success rate across 150 tasks.

Reduces inference time by up to 35x.

Solves non-Markovian scenarios where others fail.

Abstract

We present Masked Generative Policy (MGP), a novel framework for visuomotor imitation learning. We represent actions as discrete tokens, and train a conditional masked transformer that generates tokens in parallel and then rapidly refines only low-confidence tokens. We further propose two new sampling paradigms: MGP-Short, which performs parallel masked generation with score-based refinement for Markovian tasks, and MGP-Long, which predicts full trajectories in a single pass and dynamically refines low-confidence action tokens based on new observations. With globally coherent prediction and robust adaptive execution capabilities, MGP-Long enables reliable control on complex and non-Markovian tasks that prior methods struggle with. Extensive evaluations on 150 robotic manipulation tasks spanning the Meta-World and LIBERO benchmarks show that MGP achieves both rapid inference and superior success rates compared to state-of-the-art diffusion and autoregressive policies. Specifically, MGP increases the average success rate by 9% across 150 tasks while cutting per-sequence inference time by up to 35x. It further improves the average success rate by 60% in dynamic and missing-observation environments, and solves two non-Markovian scenarios where other state-of-the-art methods fail.