AIM: Action-Supervised Intent as a Search Prior for World Model Planning

Latent world models let an agent plan goal-directed behavior from visual observations by evaluating candidate action sequences inside learned dynamics. Under a finite compute budget, however, effective sampling-based planning depends not only on how accurately candidates are evaluated but also on whether the search starts in the region of action space relevant to the current task. Planners such as CEM and MPPI typically start from a generic proposal distribution, so even when offline trajectories already record the actions that connect the current observation to a future goal, the planner must rediscover the coarse structure of goal-directed behavior. A substantial share of the sampling budget can therefore be spent locating a useful action region before the world model begins to refine candidate plans. To address this, we propose AIM (Action-Supervised Intent Module), which uses offline goal-conditioned action supervision to build an informed search prior for a pretrained, frozen world model. Given frozen current and goal latents, AIM predicts a compact temporal intent for the connecting action block—the leading coefficients of its orthonormal discrete cosine transform (DCT)—and, through the inverse transform, uses it only to initialize the planner's action-sequence mean, leaving the representation, dynamics, rollout, sampling rule, and planning objective unchanged. This frequency-domain parameterization gives a compact, coarse description of goal-directed behavior and splits the initialization error exactly into a prediction term and a truncation term, which ranks intent dimensions before any planning run. Across three frozen world-model hosts, AIM consistently improves finite-budget planning, raising the five-task average from $47.5$ to $92.9$ with a single 300-candidate CEM iteration in the largest-gain setting, while retaining its advantage across larger search budgets and under MPPI.