Diffusion-MPC in Discrete Domains: Feasibility Constraints, Horizon Effects, and Critic Alignment: Case study with Tetris

TL;DR

This paper investigates the application of diffusion-based model predictive control in discrete domains using Tetris, highlighting the importance of feasibility constraints, the effects of planning horizon, and critic alignment issues.

Contribution

It introduces feasibility masking for discrete sampling, analyzes reranking strategies, and examines compute scaling effects, providing practical insights for diffusion-MPC in discrete environments.

Findings

Feasibility masking improves score and survival rates.

Naive DQN reranking is often misaligned with actual quality.

Shorter horizons outperform longer ones in sparse reward settings.

Abstract

We study diffusion-based model predictive control (Diffusion-MPC) in discrete combinatorial domains using Tetris as a case study. Our planner samples candidate placement sequences with a MaskGIT-style discrete denoiser and selects actions via reranking. We analyze three key factors: (1) feasibility-constrained sampling via logit masking over valid placements, (2) reranking strategies using a heuristic score, a pretrained DQN critic, and a hybrid combination, and (3) compute scaling in candidate count and planning horizon. We find that feasibility masking is necessary in discrete domains, removing invalid action mass (46%) and yielding a 6.8% improvement in score and 5.6% improvement in survival over unconstrained sampling. Naive DQN reranking is systematically misaligned with rollout quality, producing high decision regret (mean 17.6, p90 36.6). Shorter planning horizons outperform longer ones under sparse and delayed rewards, suggesting uncertainty compounding in long imagined rollouts. Overall, compute choices (K, H) determine dominant failure modes: small K limits candidate quality, while larger H amplifies misranking and model mismatch. Our findings highlight structural challenges of diffusion planners in discrete environments and provide practical diagnostics for critic integration.