Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Masked Generative Policy for Robotic Control

Lipeng Zhuang, Shiyu Fan, Florent P. Audonnet, Yingdong Ru, Gerardo Aragon Camarasa, Paul Henderson

TL;DR

This work tackles visuomotor imitation under high-dimensional sensing by replacing diffusion and strictly autoregressive policies with Masked Generative Policy (MGP), which treats actions as discrete tokens generated in parallel by a masked transformer. It introduces an action tokenizer via a VQ-VAE and two sampling regimes: MGP-Short for fast, short-horizon control and MGP-Long for long-horizon, non-Markovian tasks with adaptive token refinement guided by posterior-confidence estimation. Empirical results on 150 tasks across Meta-World and LIBERO show that MGP-Long delivers global, dynamic planning with robustness to missing observations and non-Markovian dynamics, while MGP-Short achieves state-of-the-art speed and competitive accuracy. Real-world demonstrations on towel-sorting validate the approach’s practical applicability, highlighting substantial gains in both inference speed and reliability over diffusion and autoregressive baselines.

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.

Masked Generative Policy for Robotic Control

TL;DR

This work tackles visuomotor imitation under high-dimensional sensing by replacing diffusion and strictly autoregressive policies with Masked Generative Policy (MGP), which treats actions as discrete tokens generated in parallel by a masked transformer. It introduces an action tokenizer via a VQ-VAE and two sampling regimes: MGP-Short for fast, short-horizon control and MGP-Long for long-horizon, non-Markovian tasks with adaptive token refinement guided by posterior-confidence estimation. Empirical results on 150 tasks across Meta-World and LIBERO show that MGP-Long delivers global, dynamic planning with robustness to missing observations and non-Markovian dynamics, while MGP-Short achieves state-of-the-art speed and competitive accuracy. Real-world demonstrations on towel-sorting validate the approach’s practical applicability, highlighting substantial gains in both inference speed and reliability over diffusion and autoregressive baselines.

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.

Paper Structure

This paper contains 70 sections, 6 equations, 10 figures, 20 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Diffusion Models for Robotic Manipulation.
  4. Autoregressive Models for Robotic Manipulation
  5. Masked Generative Transformers
  6. Methods
  7. Action Tokenizer
  8. Action Tokenization.
  9. Training.
  10. Masked Generative Transformer
  11. Structure.
  12. Training.
  13. Short horizon Mask-and-Refine Sampling (MGP-Short)
  14. Long horizon Mask-and-Refine Sampling (MGP-Long)
  15. Posterior-confidence estimation.
  16. ...and 55 more sections

Figures (10)

  • Figure 1: Overview of MGP. (1) Four properties of MGP: dynamic adaptation, flexible replanning steps, resilient execution, and global-coherent prediction. (2) Versus prior SOTA, MGP is both faster (lower per-sequence inference time) and better (higher success). (3) MGP also excels on several challenging settings: dynamic environments, observation-missing environments, and non-Markovian, long-horizon tasks.
  • Figure 2: Left: Training Stage 1 - Action Tokenizer and Middle: Training Stage 2 - Masked Generative Transformer and Right: Short-horizon sampling (MGP-Short)
  • Figure 3: Long-horizon sampling (MGP-Long) through Adaptive Token Refinement (ATR).
  • Figure 4: Qualitative results in Dynamic environments
  • Figure 5: Qualitative results of Button Press On/Off Tasks. (a) is the result of MGP-Long; (b) and (c) are results of Short-Horizon baselines.
  • ...and 5 more figures