Density Matters: Normalized Successor Measures for Zero-Shot Control

Behavior foundation models (BFMs) promise zero-shot control: pretrain on offline data, then reuse the learned representation for downstream tasks without task-specific training. Successor-measure methods such as Forward–Backward representations (FB) realize this promise for diverse locomotion behaviors, but can struggle when success depends on precise interactions. For example, in pick-and-place tasks, FB policies can reach an object but fail to maintain a grasp or carry it to a specified goal. We attribute these failures to control bottlenecks, transitions that are necessary passages to valuable regions where the choice of action substantially affects return. We show that the low-rank squared approximation used by FB can suppress the differences in action values at these states, even at its global optimum, causing the agent to choose incorrect actions. To preserve these differences with limited capacity, we introduce Density-FB, which treats the bilinear output as a score, normalizes the scores into a distribution over future states with a softmax, and learns this distribution with a categorical Bellman objective. Across OGBench, MetaWorld, and ManiSkill, Density-FB improves manipulation success over FB by a median of $64.4\%$, $30.0\%$, and $8.5\%$ percentage points, respectively. Additionally, Density-FB preserves correct action rankings with smaller embeddings than vanilla FB.