Partner as Task: Belief-conditioned Successor Features for Zero-shot Coordination

Zero-shot coordination (ZSC) requires cooperation with partners unseen during training. We view each partner as a task: the reward remains fixed, but the partner changes the ego-agent's effective dynamics and hence how it should act. Learning these dynamics is difficult when the partner policy is unobserved, and reward is sparse. We introduce belief-conditioned successor features (BSFs), which summarize discounted future features conditioned on a belief over the partner's latent type. The belief is inferred from interaction history and augments the state. We show that the value of a belief-independent policy is linear in the belief, connecting belief inference to task conditioning in standard successor features. We approximate BSFs in MA-JEPA, an on-policy self-play (SP) algorithm that combines a learned history encoder with a temporal difference latent prediction objective. This objective provides reward-free supervision for long-term predictive representations. Across Overcooked and SMAC, MA-JEPA outperforms self-play, population-based, and agent-modeling baselines, with the largest cross-play gains over the strongest baseline under sparse rewards ($77\%$) and RGB observations ($32\%$).