Paper deep dive
MA-JEPA: Joint-Embedding World Models for Multi-Agent Reinforcement Learning
Brandon Gary Kaplowitz, Osaze James Obahor, Christian Schroeder de Witt
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:World models improve sample efficiency by training policies on imagined trajectories, but their usefulness depends on learning representations that capture the information needed for future control. We study whether self-supervised joint-embedding prediction (JEPA) can provide this learning signal for multi-agent reinforcement learning. We introduce MA-JEPA, a stochastic world model that replaces observation reconstruction with prediction of target representations, enabling model-based multi-agent reinforcement learning with centralized training and decentralized execution. A categorical latent state and a causal Transformer are trained with posterior and action-conditioned dynamics prediction objectives and are then used for actor-critic learning from latent imagination. A training-only joint predictor conditions on all agents' local states and actions to predict each agent's next local observation embedding. These predictions are passed through the same local posterior used during real interaction with a centralized critic that is used only for value learning, with execution remaining decentralized. Our experiments show that this architecture performs strongly on SMAC, matching or exceeding the strongest reported comparator mean win rate on four of eight evaluated maps.
Tags
Links
- Source: https://arxiv.org/abs/2609.33563v1
- Canonical: https://arxiv.org/abs/2609.33563v1
Trouble viewing inline? Open PDF directly →
Full Text
81,350 characters extracted from source content.
Expand or collapse full text
MA-JEPA: Joint-Embedding World Models for Multi-Agent Reinforcement Learning Brandon Gary Kaplowitz* Osaze James Obahor* Christian Schroeder de Witt Affiliation: Department of Engineering Science, University of Oxford Affiliation: *Equal contribution. Correspondence to brandon.kaplowitz@eng.ox.ac.uk Abstract World models improve sample efficiency by training policies on imagined trajectories, but their usefulness depends on learning representations that capture the information needed for future control. We study whether self-supervised joint-embedding prediction (JEPA) can provide this learning signal for multi-agent reinforcement learning. We introduce MA-JEPA, a stochastic world model that replaces observation reconstruction with prediction of target representations, enabling model-based multi-agent reinforcement learning with centralized training and decentralized execution. A categorical latent state and a causal Transformer are trained with posterior and action-conditioned dynamics prediction objectives and are then used for actor-critic learning from latent imagination. A training-only joint predictor conditions on all agents’ local states and actions to predict each agent’s next local observation embedding. These predictions are passed through the same local posterior used during real interaction with a centralized critic that is used only for value learning, with execution remaining decentralized. Our experiments show that this architecture performs strongly on SMAC, matching or exceeding the strongest reported comparator mean win rate on four of eight evaluated maps. 1 Introduction Model-based reinforcement learning (MBRL) learns an environment model and uses it to generate trajectories for policy improvement. By training an actor and critic on these imagined trajectories, methods can reuse each real transition for many learning updates (Hafner et al., 2019; Hafner et al., 2020; Hafner et al., 2025). The benefit, however, depends on the quality of the world model. Errors can accumulate as the model is rolled forward, and a policy trained inside the model may exploit predictions that do not match the real environment (Janner et al., 2019; Moerland et al., 2023). Multi-agent MBRL adds a second problem, where an agent’s transition probabilities generally depend on the actions of the other agents. For example, in the cooperative combat environments of the StarCraft Multi-Agent Challenge (SMAC), whether an opponent remains alive after one step depends on which target the whole team attacks. A local world model conditioned only on one agent’s action cannot fully predict this outcome. A centralized model can condition on the joint action, but directly using a centralized latent state for policy learning leads to a different problem, where the actor is then trained on information that is unavailable during decentralized execution. Existing multi-agent model-based reinforcement learning does not fully address the problem (Egorov and Shpilman, 2022; Venugopal et al., 2024; Zhang et al., 2025; Deihim et al., 2025). In this paper, we directly tackle this problem by introducing MA-JEPA, the first JEPA model-based multi-agent reinforcement learning (MARL) method. It achieves centralized training with decentralized execution (CTDE) by training on the same latent representation of observations in centralized training that it does in decentralized execution. (Oliehoek et al., 2008; Kraemer and Banerjee, 2016) Our agent explicitly maintains an independent recurrent state in a parameter-shared local world model. This enables us both to train agents using local latent states and to maintain parallel, independent copies of the world model during execution, each with information specific to an individual agent. We train this model with a joint-embedding predictive architecture (JEPA) objective that predicts target observation embeddings from latent states and action histories, without reconstructing observations, unlike leading model-based MARL methods (Assran et al., 2023; Bardes et al., 2024; Assran et al., 2025). During multi-agent training, a joint predictor conditions on all agents’ local states and actions to predict an embedding of each agent’s next local observation. The joint predictor does not construct the actor state directly. Instead, its prediction is passed through the same local posterior that processes an encoded real observation. During real interaction, this posterior combines the local dynamics state with the observed input’s embedding. During centralized imagination, it combines the local dynamics state with the embedding predicted from the joint action. The actor, therefore, receives the same type of locally constructed state in both cases. A centralized critic uses all local states for value learning. Importantly, neither it nor the joint predictor runs during execution, meaning only agent-specific information is used during execution. We evaluate the full method on SMAC at matched interaction budgets and compare the effects of outcome and embedding gradients on the local representations. Our contributions include: 1. A model-based architecture for CTDE that predicts each agent’s next local observation embedding from all agents’ local states and actions, then uses the same local posterior as during execution to construct the actor state. 2. To our knowledge, the first JEPA-based MARL approach, which trains local and joint dynamics by predicting target embeddings instead of reconstructing observations, enabling the same representation-learning approach in single-agent and multi-agent imagination. 2 Preliminaries 2.1 Model-Based Multi-Agent Reinforcement Learning We consider a cooperative decentralized partially observable Markov decision process (Dec-POMDP) with N agents. At timestep t, the environment is in state xt∈x_t , agent i receives a local observation ti∈io_t^i ^i, and selects an action ti∈ia_t^i ^i. The individual actions form the joint action t1:N=(t1,…,tN),a_t^1:N=(a_t^1,…,a_t^N), (1) after which the environment transitions according to p(xt+1∣xt,t1:N)p(x_t+1 x_t,a_t^1:N) and produces a reward rt+1r_t+1. Because the environment state is not available during decentralized execution, agent i acts from its local interaction history τti=(1i,1i,…,ti),ti∼πi(⋅∣τti). _t^i=(o_1^i,a_1^i,…,o_t^i), _t^i π^i(· _t^i). (2) Under CTDE, information from all agents may be used to learn models and value functions, while the executed policy remains conditioned only on information available in τti _t^i. Throughout the paper, superscript i indexes agents and t1:N=(t1,…,tN)x_t^1:N=(x_t^1,…,x_t^N) denotes a synchronized collection of agent-indexed quantities. We use hats for quantities generated by a learned model and sg(⋅)sg(·) for stop-gradient. Model-based reinforcement learning learns environment dynamics from collected experience and uses the learned model to generate trajectories for policy improvement (Moerland et al., 2023). In the multi-agent setting, imagined rollouts are generated from replay-inferred states while actions are sampled from the decentralized policies, and the model predicts the quantities required to advance the rollout, including rewards and continuation factors r^t+1i r_t+1^i and c^t+1i c_t+1^i. We define c^t+1i c_t+1^i as the discounted continuation factor, i.e., it incorporates both discounting and episode termination. A critic may use centralized information during training. Let Vi(^t1:N)V^i( s_t^1:N) denote the value estimate for agent i from the synchronized imagined states. Imagined λ-returns are Gtλ,i=r^t+1i+c^t+1i[(1−λ)Vi(^t+11:N)+λGt+1λ,i],G_t^λ,i= r_t+1^i+ c_t+1^i [(1-λ)V^i( s_t+1^1:N)+λ G_t+1^λ,i ], (3) where λ∈[0,1]λ∈[0,1] controls the trade-off between value bootstrapping and longer imagined returns (Sutton, 1988; Hafner et al., 2025). The actor is optimized from imagined returns while centralized information remains restricted to training. 2.2 Joint-Embedding Predictive Architectures Joint-embedding predictive architectures (JEPAs) learn representations by predicting the embedding of a target from a related context instead of reconstructing the target itself (LeCun, 2022; Assran et al., 2023). Let x denote a context and y a prediction target. An online encoder EϕE_φ, a target encoder Eϕ¯E_ φ, and a predictor PθP_θ produce x=Eϕ(x),¯y=sg(Eϕ¯(y)),^y=Pθ(x),e_x=E_φ(x), e_y=sg\! (E_ φ(y) ), e_y=P_θ(e_x), (4) with objective ℒJEPA=D(^y,¯y),L_JEPA=D\! ( e_y, e_y ), (5) where D measures discrepancy in representation space. The target encoder is typically updated as an exponential moving average (EMA) of the online encoder (Grill et al., 2020; Assran et al., 2023). Since predictive agreement can admit collapsed representations, practical JEPA methods additionally employ architectural asymmetry or representation regularization; we use SIGReg (Balestriero and LeCun, 2025). For sequential decision making, prediction is additionally conditioned on the actions connecting the current context to a future observation. At horizon h, this takes the general form ^t+h=Pθ(t,t:t+h−1),¯t+h=sg(Eϕ¯(t+h)), e_t+h=P_θ\! (c_t,a_t:t+h-1 ), e_t+h=sg\! (E_ φ(o_t+h) ), (6) where tc_t denotes the available predictive context. Action-conditioned joint-embedding models have been used for latent prediction, planning, and control (Assran et al., 2025; Maes et al., 2026). For the remainder of the paper, ti=Eϕ(ti)e_t^i=E_φ(o_t^i) denotes the online observation embedding of agent i, ¯ti e_t^i its stopped EMA target, and ^ti e_t^i a predicted observation embedding. We reserve tiz_t^i for the stochastic latent state inferred by the local world model. In MA-JEPA the joint model predicts observation embeddings, which are subsequently processed by the local world-model posterior to construct the latent state used by the policy. 3 Related Work 3.1 Joint-Embedding Predictive Architectures Joint-embedding predictive architectures learn representations by predicting target embeddings rather than reconstructing observations (LeCun, 2022). I-JEPA applies this principle to spatially masked image regions, demonstrating that predictive objectives can learn semantic visual representations without pixel-level reconstruction (Assran et al., 2023). V-JEPA extends the approach to video by predicting masked spatiotemporal representations, while V-JEPA 2 scales video pretraining and introduces an action-conditioned predictor for physical reasoning and planning (Bardes et al., 2024; Assran et al., 2025). These works establish joint-embedding prediction as an effective alternative to reconstructive generative modeling for learning structured representations of visual dynamics. Recent work has brought this principle closer to reinforcement learning and world modeling. LeWorldModel (LeWM) trains an action-conditioned JEPA end-to-end from pixels and uses SIGReg to stabilize the learned latent space, showing that future representation prediction can support efficient model-based control without an observation decoder (Maes et al., 2026). TD-JEPA instead develops policy-conditioned long-horizon latent prediction for zero-shot reinforcement learning, using temporal-difference structure to learn representations of future dynamics across policies (Bagatella et al., 2025). Existing applications have primarily focused on single-agent representation learning, planning, and control. MA-JEPA extends this architectural family to MARL, using the learned embedding space both as the predictive target of each local world model and as the interface through which a centralized joint model captures multi-agent interactions. 3.2 Model-Based Reinforcement Learning Modern model-based reinforcement learning learns compact world models and uses their predictions for planning or policy optimization. PlaNet introduced stochastic latent dynamics for planning from high-dimensional observations, and the Dreamer family instead trains actor-critic policies directly on latent imagined trajectories (Hafner et al., 2019; Hafner et al., 2020; Hafner et al., 2025). More recent world models increasingly use sequence-modeling architectures. IRIS represents visual observations with discrete tokens and models their dynamics autoregressively with a Transformer, while STORM combines stochastic latent states with Transformer dynamics for efficient imagination (Micheli et al., 2023; Zhang et al., 2023). Extending world models to MARL additionally requires modeling interactions between agents while preserving decentralized policies. Existing methods have largely adapted established single-agent world-model architectures to this setting. MAMBA extends Dreamer-style latent imagination to cooperative MARL, using parameter-shared local world models together with inter-agent communication (Egorov and Shpilman, 2022). DMAWM similarly builds on DreamerV2, separating local agent modules from a shared environment model and aligning imagined latent-state distributions with those inferred from local observations (Xue et al., 2026). In contrast, MATWM develops a multi-agent extension of the STORM architecture, combining stochastic Transformer dynamics with teammate prediction, decentralized imagination, and a semi-centralized critic (Deihim et al., 2025). MARIE also adopts Transformer-based world modeling, coupling decentralized agent dynamics with a centralized Perceiver that aggregates information across agents (Zhang et al., 2025). MABL instead introduces global and local latent variables so that centralized information can improve world-model learning while policies remain decentralized (Venugopal et al., 2024). MA-JEPA follows a different architectural lineage. Rather than adapting a Dreamer- or STORM-style world model to the multi-agent setting, we are the first paper, to our knowledge, to leverage joint-embedding predictive world modeling in a MARL setting. MA-JEPA combines joint-embedding predictive learning with centralized imagination and decentralized latent-state based policies for execution. Unlike MARIE, which trains policies on reconstructed observations, our agents act on latent states maintained by their local world models. These states summarize each agent’s observation-action history, allowing its policy to use the temporal representations learned through dynamics prediction. Both local and joint models learn by predicting observation representations without reconstructing observations, removing the requirement that these representations also preserve the information needed for reconstruction. During training, a separate joint model captures interactions using all agents’ local states and joint actions. While DMAWM generates imagined latent states through a shared environment model and trains them to be aligned with locally inferred states, our joint model predicts each agent’s next observation embedding. These embeddings pass through the same local posterior used for real observations. The actor’s latent state is therefore constructed by the same local mechanism during imagination and execution, requiring no additional context from the centralized training for local execution and removing a potential architectural source of training–execution mismatch. 4 Theoretical Background We study identification and transition estimation in an action-conditioned Gaussian model under assumptions similar to Klindt et al. (2026). Appendix E gives consistency results for empirical successor distributions and population optima of unrestricted autoregressive models under fixed representations and collection policies. Action-conditioned Gaussian identification. Following the notation of Klindt et al. (2026), let z,z′∈ℝnz,z ^n be consecutive physical states. Let x=g(z)x=g(z) be the joint observation and y=f(x)=h(z)∈ℝny=f(x)=h(z) ^n the learned representation, where f and g are measurable. Assume z∼(0,In)z (0,I_n) and let A be a joint action in a finite set A, independent of z. For each action a∈a , let Ba∈ℝn×nB_a ^n× n and assume z′|z,A=a∼(Baz,In−BaBa⊤),‖Ba‖op≤1.z z,A=a (B_az,I_n-B_aB_a ), \|B_a\|_op≤ 1. Define κB=λmin(A[BABA⊤])−A‖BA‖op4 _B= _ (E_A[B_AB_A ])-E_A\|B_A\|_op^4 and assume κB>0 _B>0. Let Jω(y,a)J_ω(y,a) predict the next representation. Consider minimizing the population loss ℒ=‖h(z′)−Jω(h(z),A)‖2L=E\|h(z )-J_ω(h(z),A)\|^2 over encoder and predictor classes, subject to h(z)=0Eh(z)=0 and Cov(h(z))=InCov(h(z))=I_n. Assume these classes contain the encoder h0(z)=Q0zh_0(z)=Q_0z and predictor Jω0(y,a)=Q0BaQ0⊤yJ_ _0(y,a)=Q_0B_aQ_0 y for some orthogonal matrix Q0Q_0. For each observed action a, define C^m,a C_m,a as the empirical mean of y′y⊤y y over samples with action a. Define V^m,a V_m,a as the empirical mean of (y′−C^m,ay)(y′−C^m,ay)⊤(y - C_m,ay)(y - C_m,ay) over those same samples. For each observed action, set P^m(⋅∣y,a)=(C^m,ay,V^m,a) P_m(· y,a)=N( C_m,ay, V_m,a). For an action not yet observed, define P^m(⋅∣y,a) P_m(· y,a) to be any fixed Gaussian distribution. preprint-pratenddefaultcategory.tex Theorem 4.1 (Identification and transition consistency). The minimum of the population loss ℒL over the admissible encoder–predictor pairs is n−trA[BABA⊤]n-trE_A[B_AB_A ]. For every globally minimizing pair, there is an orthogonal matrix Q such that its encoder satisfies h(z)=Qzh(z)=Qz almost surely. For any such encoder, let y′=h(z′)y =h(z ). The conditional distribution of y′y given y=h(z)y=h(z) and A=aA=a admits the Gaussian version P(⋅∣y,a)=(QBaQ⊤y,In−QBaBa⊤Q⊤).P(· y,a)=N(QB_aQ y,I_n-QB_aB_a Q ). Hold this globally minimizing encoder fixed and collect m fresh independent and identically distributed triples (y,A,y′)(y,A,y ) from the joint distribution induced by the assumed model for (z,A,z′)(z,A,z ) and this fixed encoder. As m→∞m→∞, P^m(⋅∣y,a) P_m(· y,a) converges weakly to P(⋅∣y,a)P(· y,a) almost surely for each fixed y∈ℝny ^n and every action a∈a with Pr(A=a)>0 (A=a)>0. For each such action satisfying ‖Ba‖op<1\|B_a\|_op<1, this convergence also holds in total variation. preprint-pratenddefaultcategory.tex preprint-pratenddefaultcategory.tex See proof in .preprint-pratenddefaultcategory.tex Under the stated assumptions, every global minimizer of the whitened population squared-error objective recovers the physical state up to an orthogonal transformation. Holding this encoder fixed, empirical transition-moment estimates consistently recover the action-conditioned Gaussian successor law. This joint-action extension of Klindt et al. (2026) motivates predictive representations for centralized world modeling. 5 Method We introduce MA-JEPA, a joint-embedding predictive model-based MARL algorithm for centralized training with decentralized execution. Each agent maintains a local latent world model and acts from its own inferred state. During training, a joint predictor models the effect of synchronized agent states and actions by predicting each agent’s next observation embedding. These predictions are passed through the same local posterior used for real observations, enabling joint imagination without constructing a centralized actor state. A centralized critic is used for policy learning. 5.1 Local Predictive World Model Decentralized execution. Agent i encodes its local observation and updates tis_t^i from its own observation–action history. The shared policy then selects ti∼πψ(⋅∣ti)a_t^i _ψ(· _t^i) without access to other agents. We denote the observation encoder by EϕE_φ and the local latent dynamics by parameters θ, comprising a causal Transformer TθT_θ, a history-conditioned prior pθp_θ, and an observation-conditioned posterior qθq_θ. Each agent’s action-observation history is summarized by a deterministic history state tih_t^i and a categorical stochastic latent state tiz_t^i, with ti=cat(ti,vec(ti))s_t^i=cat\! (h_t^i,vec(z_t^i) ) (7) forming the local state used by the policy. Local state inference. The encoder represents the current observation, while the causal Transformer summarizes the preceding stochastic states and actions. The posterior combines the observation embedding and history state to infer the current stochastic latent state. ti _t^i =Eϕ(ti), =E_φ(o_t^i), (8) ti _t^i =Tθ(<ti,<ti), =T_θ(z_<t^i,a_<t^i), ti _t^i ∼qθ(⋅∣ti,ti). q_θ(· _t^i,e_t^i). The prior pθ(ti∣ti)p_θ(z_t^i _t^i) predicts the stochastic state from local history alone and is trained jointly with the posterior. The deterministic state tih_t^i summarizes information from the preceding latent–action history, whereas tiz_t^i incorporates information revealed by the current observation. This separation allows the local dynamics to advance without observing the next environment input while retaining an observation-conditioned correction when one becomes available. Parameters are shared across agents, while latent histories are maintained independently. Learning joint-embeddings. We train two local predictors toward the EMA target of the embedding ¯ti e_t^i defined in Sec. 2.2. The posterior predictor PpostP_post maps the complete local state tis_t^i into the target embedding space, while the dynamics predictor PdynP_dyn maps only the deterministic history state tih_t^i into the same space: ℒpost=dcos(Ppost(ti),¯ti),ℒdyn=dcos(Pdyn(ti),¯ti).L_post=d_ \! (P_post(s_t^i), e_t^i ), _dyn=d_ \! (P_dyn(h_t^i), e_t^i ). (9) The two predictors impose complementary constraints on the local representation. PpostP_post is evaluated after the current observation has been incorporated and therefore encourages tis_t^i to retain the information represented by the target encoder. In contrast, PdynP_dyn receives only tih_t^i and must predict the same target from the preceding latent–action history. The resulting representation is therefore both descriptive of the current observation and predictable from local dynamics. Crucially, this embedding space later forms the interface through which the joint model supplies predicted observations to the local world model. The complete local objective additionally contains prior–posterior latent regularization, local action-availability prediction, and SIGReg. Reward and continuation are predicted by the joint model. The appendix gives the full objective for the vector-observation SMAC configuration. Decentralized policy. Each agent acts only from its local latent state, ti∼πψ(⋅∣ti).a_t^i _ψ(· _t^i). (10) Policy parameters are shared across agents, and each agent maintains its own local history. 5.2 Joint Prediction of Multi-Agent Dynamics Centralized training. The joint model JωJ_ω combines synchronized local states t1:Ns_t^1:N with the joint action t1:Na_t^1:N, forming agent-specific interaction representations tiu_t^i from which it predicts the next local embeddings ^t+1i e_t+1^i. Local transition prediction does not account for the simultaneous actions of other agents. During centralized training, we therefore introduce a joint predictor JωJ_ω conditioned on the synchronized local states t1:Ns_t^1:N and joint action t1:Na_t^1:N. It produces an agent-specific interaction representation, t1:N=Jω(t1:N,t1:N),u_t^1:N=J_ω\! (s_t^1:N,a_t^1:N ), (11) from which ^t+1i=Pe(ti) e_t+1^i=P_e(u_t^i) (12) predicts the next local observation embedding of agent i. Conditioning on the joint action allows the predictor to model interaction effects that cannot be inferred from the focal agent’s local history alone. The output remains agent-specific: the model predicts the representation of each agent’s next local observation rather than constructing a shared actor state. Joint predictor training. The predicted embedding is supervised by both the EMA target and the stopped online encoding ℒemb _emb =dcos(^t+1i,¯t+1i), =d_ \! ( e_t+1^i, e_t+1^i ), (13) ℒint _int =SmoothL1(^t+1i,sg(Eϕ(t+1i))). =SmoothL1\! ( e_t+1^i,sg(E_φ(o_t+1^i)) ). The two targets serve different roles. The EMA target aligns joint prediction with the representation learned by the local JEPA objectives. The stopped online encoding matches the real encoding that the local posterior receives during environment interaction. This objective reduces the mismatch between encoded observations and joint-model predictions before both are consumed by the same posterior during imagination. Together with auxiliary transition predictions, the joint objective is ℒjoint=βeℒemb+βIℒint+ℒauxJ.L_joint= _eL_emb+ _IL_int+L_aux^J. (14) Here, ℒauxJL_aux^J includes reward, continuation, agent-alive, and legal-action prediction. In the full method, factual reward, continuation, and agent-alive losses also update the local input representations at scale 1.0. The factual one-step embedding loss supplies an additional local-input gradient at scale 0.1. All other joint losses stop at those inputs. The joint predictor and its heads retain their ordinary loss weights. The baseline sets both additional gradient scales to zero. Multi-step prediction. We supervise future embeddings directly from replay of realized past states. Let ℋH denote the set of prediction horizons and t,hi χ_t,h^i the intervening action context used to predict the future of agent i. For each h∈ℋh , a head predicts ^t+h|ti=Ph(ti,t,hi),ℒMS=∑h∈ℋαhdcos(^t+h|ti,¯t+hi). e_t+h t^i=P_h\! (u_t^i, χ_t,h^i ), _MS= _h _h\,d_ \! ( e_t+h t^i, e_t+h^i ). (15) We train each multi-step head to predict future embeddings from the recorded focal-agent action sequence. For horizons greater than one, an action-discrimination loss substitutes each distinct legal alternative for the final focal-agent action while keeping the joint root and earlier actions fixed. It encourages the factual prediction to exceed each alternative’s cosine similarity to the observed future embedding by a margin of 0.1. The loss has coefficient 0.1 and favors distinct predictions even for actions with equivalent outcomes. The appendix gives the construction of t,hi χ_t,h^i and the objective. These heads provide auxiliary training supervision. Imagination uses the recurrent one-step joint predictor. 5.3 Joint Imagination and Decentralized Policy Learning Shared posterior interface. During imagination, agents sample actions from their local policies and advance their local deterministic states. The joint predictor then produces the next observation embedding from the synchronized imagined states and joint action ^ti a_t^i ∼πψ(⋅∣^ti), _ψ(· s_t^i), (16) ^t+1i h_t+1^i =Tθ(^≤ti,^≤ti), =T_θ( z_≤ t^i, a_≤ t^i), ^t+1i e_t+1^i =Pe(Jω(sg(^t1:N),^t1:N)i). =P_e\! (J_ω(sg( s_t^1:N), a_t^1:N)_i ). The predicted embedding replaces the encoder output in the same posterior used for real observations t+1i _t+1^i ∼qθ(⋅∣t+1i,Eϕ(t+1i)), q_θ(· _t+1^i,E_φ(o_t+1^i)), (17) ^t+1i z_t+1^i ∼qθ(⋅∣^t+1i,^t+1i). q_θ(· h_t+1^i, e_t+1^i). The joint predictor therefore does not generate the next actor state directly. It predicts only the observation-side input to the local posterior. The posterior combines this prediction with the independently advanced local dynamics state in the same manner as it combines a real encoded observation with t+1ih_t+1^i. Consequently, centralized information changes the predicted consequence of the joint action without changing the mechanism that constructs the decentralized policy state. Centralized value learning. The critic uses a central value network VξV_ξ that takes as inputs the stopped local states of all agents during training, Vξi=Vξ(sg(^t1:N))i.V_ξ^i=V_ξ\! (sg( s_t^1:N) )_i. (18) Using rewards and discounted continuations predicted by the joint model, we compute imagined λ-returns from equation 3. The centralized critic attends to the agents’ local latent states together with roster and liveness masks. Each actor conditions on its own tis_t^i. The actor is optimized from the imagined returns and the critic from the corresponding value targets, with gradients stopped at the world-model states. At execution time, each agent uses only the shared local encoder, latent dynamics, posterior, and policy. The joint predictor, direct multi-step objectives, and centralized critic are restricted to training. 6 Experiments 6.1 Experimental Setup We evaluate whether joint-embedding world models support effective policy learning from latent imagination in the StarCraft Multi-Agent Challenge (SMAC, Samvelyan et al., 2019). The main comparison covers eight SMAC v1 tasks under StarCraft I 4.10.0, including homogeneous and heterogeneous teams, asymmetric engagements, and different coordination demands. The gradient ablations also include 3s_vs_4z. A detailed description of SMAC can be found in Appendix B.1. We compare MA-JEPA against DMAWM, MAMBA, MAPPO, QMIX, and MAT at matched interaction budgets. MA-JEPA hyperparameters are given in Appendix D.9.1, and baseline configurations in Appendix D.9. MAPPO is additionally trained to 1M steps as a higher-budget reference. All algorithms were benchmarked under NVIDIA A100 PCIe GPUs. 6.2 Results Table 1 reports win rates, and Figure 1 shows learning curves for the eight tasks. These results demonstrate effective control through joint-embedding-based imagination. Performance remains task-dependent. We match or achieve the highest reported mean win rate on four of the eight tasks and reach 95.0% on so_many_baneling. On 3s_vs_4z, our mean win rate is 74% with a sample standard deviation of 7.5 percentage points, compared with 95.7% and 2.3 percentage points for DMAWM. We hypothesize that errors in imagined kiting trajectories contribute to this variability across seeds. In follow-up work, we plan to explore the geometry of the latent space embedding to understand exactly what is and is not learned by the world model in these “kiting” settings, what the optimal solution looks like, and identify optimal and learned policy features like recurrent action sequences. Figure 1: Win rate versus environment steps on six SMAC v1 tasks at the indicated budgets. Solid curves show means across available training seeds, with ±1± 1 sample-standard-deviation bands. Recorded evaluations are connected by linear interpolation without smoothing; The dotted horizontal line shows MAPPO’s mean win rate at 1M steps. Table 1: SMAC win rate (%) at matched nominal environment-step budgets, reported as mean (sample standard deviation) across available training seeds from the set 1302, 2771, and 7636. Results use 100-episode evaluations at the budget checkpoint. Bold denotes the highest available mean win rate for each map. Rows are grouped into Easy, Hard, and Super Hard maps. Map Steps MA-JEPA (ours) DMAWM MAPPO QMIX MAT Easy 2s_vs_1sc 100k 97.3 (2.5) 95.7 (5.9) 100.0 (0.0) 13.7 (7.4) 51.3 (27.1) 2s3z 100k 78.7 (14.5) 88.0 (7.0) 8.0 (2.6) 31.3 (24.5) 16.5 (9.2) 8m 100k 92.0 (1.0) 92.0 (4.6) 65.7 (33.1) 83.0 (10.1) 26.0 (22.6) M 100k 91.0 (3.0) 90.0 (4.6) 20.7 (17.0) 3.0 (2.6) 0.0 (0.0) so_many_baneling 100k 95.0 (6.1) 99.5 (0.7) 77.0 (7.2) 21.7 (18.0) 32.0 (24.3) Hard 3s_vs_4z 100k 74.0 (7.5) 95.7 (2.3) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 2c_vs_64zg 200k 78.3 (9.5) 75.3 (8.9) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) Super [-2pt]Hard corridor 400k 66.0 (3.5) 50.0 (44.9) 0.0 (0.0) 0.0 (0.0) 0.0 (0.0) 6.3 Ablations Figure 2 shows that the effects of the two gradient routes depend on the task. Removing factual JEPA gradients reduces performance on 2s3z but improves it on 3s_vs_4z. Conversely, removing outcome gradients improves performance on 2s3z, but yields lower mean performance and substantially greater seed variability on 3s_vs_4z. Enabling both routes provides a more balanced configuration across these tasks, although it does not achieve the highest mean on either task individually. We isolate how centralized prediction losses shape the representations used by decentralized policies. All four variants retain the local world model, joint predictor, outcome heads, and latent-imagination policy training. The ablation changes which factual joint losses additionally backpropagate into the local features supplied to the joint model. We vary two gradient routes in a 2×22× 2 factorial design. The outcome route propagates gradients from factual reward, continuation, and agent-alive prediction into the local input features. The JEPA route propagates the factual joint embedding-prediction error into those features at scale 0.1. This error compares the predicted embedding with a stopped exponential-moving-average target using cosine loss. The scale applies to the additional local-input gradient, not to the ordinary training loss of the joint predictor. All four variants retain the action-discrimination loss with margin 0.1 and loss coefficient 0.1. 7 Conclusion and Future Work We introduced MA-JEPA, a joint-embedding predictive world model for multi-agent reinforcement learning under centralized training and decentralized execution. Rather than reconstructing observations, the local world model learns predictive representations that support latent imagination and decentralized control. A training-only joint model captures interactions between agents by predicting each agent’s next local observation embedding, which is processed by the same local posterior used during real interaction. Our experiments show that this architecture performs strongly on SMAC at budgets of 100k and 200k environment steps, matching or exceeding the strongest reported comparator mean win rate on four of eight evaluated maps. These results suggest that joint-embedding prediction provides a viable alternative to reconstructive world models for multi-agent control. An important direction for future work is to understand which properties of the learned representations make them effective for world-model prediction and control. In particular, we are interested in whether predictive objectives recover representations that separate persistent physical dynamics, agent identity, and interaction-dependent information, and how such structure affects generalization across environments and teams. Extending MA-JEPA to broader continuous-control and heterogeneous multi-agent settings is another natural direction. References Assran et al. (2025) M. Assran, A. Bardes, D. Fan, Q. Garrido, R. Howes, M. Komeili, M. Muckley, A. Rizvi, C. Roberts, K. Sinha, A. Zholus, S. Arnaud, A. Gejji, A. Martin, F. Robert Hogan, D. Dugas, P. Bojanowski, V. Khalidov, P. Labatut, F. Massa, M. Szafraniec, K. Krishnakumar, Y. Li, X. Ma, S. Chandar, F. Meier, Y. LeCun, M. Rabbat, and N. Ballas V-JEPA 2: self-supervised video models enable understanding, prediction and planning. External Links: 2506.09985 Cited by: §1, §2.2, §3.1. Assran et al. (2023) M. Assran, Q. Duval, I. Misra, P. Bojanowski, P. Vincent, M. Rabbat, Y. LeCun, and N. Ballas Self-supervised learning from images with a joint-embedding predictive architecture. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 15619–15629. Cited by: §1, §2.2, §2.2, §3.1. Bagatella et al. (2025) M. Bagatella, M. Pirotta, A. Touati, A. Lazaric, and A. Tirinzoni TD-JEPA: latent-predictive representations for zero-shot reinforcement learning. External Links: 2510.00739 Cited by: §3.1. Balestriero and LeCun (2025) R. Balestriero and Y. LeCun LeJEPA: provable and scalable self-supervised learning without the heuristics. External Links: 2511.08544 Cited by: §D.1, §2.2. Bardes et al. (2024) A. Bardes, Q. Garrido, J. Ponce, X. Chen, M. Rabbat, Y. LeCun, M. Assran, and N. Ballas Revisiting feature prediction for learning visual representations from video. Transactions on Machine Learning Research. Cited by: §1, §3.1. Deihim et al. (2025) A. Deihim, E. Alonso, and D. Apostolopoulou Transformer world model for sample efficient multi-agent reinforcement learning. External Links: 2506.18537 Cited by: §1, §3.2. Egorov and Shpilman (2022) V. Egorov and A. Shpilman Scalable multi-agent model-based reinforcement learning. In Proceedings of the 21st International Conference on Autonomous Agents and Multiagent Systems, p. 381–390. Cited by: §D.9.5, §1, §3.2. Grill et al. (2020) J. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. D. Guo, M. G. Azar, B. Piot, K. Kavukcuoglu, R. Munos, and M. Valko Bootstrap your own latent: a new approach to self-supervised learning. In Advances in Neural Information Processing Systems, Vol. 33, p. 21271–21284. Cited by: §2.2. Hafner et al. (2020) D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi Dream to control: learning behaviors by latent imagination. In International Conference on Learning Representations, Cited by: §1, §3.2. Hafner et al. (2019) D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson Learning latent dynamics for planning from pixels. In Proceedings of the 36th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 97, p. 2555–2565. Cited by: §1, §3.2. Hafner et al. (2025) D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap Mastering diverse control tasks through world models. Nature 640 (8059), p. 647–653. External Links: Document Cited by: §1, §2.1, §3.2. Janner et al. (2019) M. Janner, J. Fu, M. Zhang, and S. Levine When to trust your model: model-based policy optimization. In Advances in Neural Information Processing Systems, Vol. 32. Cited by: §1. Klindt et al. (2026) D. Klindt, Y. LeCun, and R. Balestriero When does LeJEPA learn a world model?. External Links: 2605.26379, Document, Link Cited by: §4, §4, §4. Kraemer and Banerjee (2016) L. Kraemer and B. Banerjee Multi-agent reinforcement learning as a rehearsal for decentralized planning. Neurocomputing 190, p. 82–94. External Links: ISSN 0925-2312, Document, Link Cited by: §1. LeCun (2022) Y. LeCun A path towards autonomous machine intelligence. Note: Version 0.9.2 External Links: Link Cited by: §2.2, §3.1. Maes et al. (2026) L. Maes, Q. Le Lidec, D. Scieur, Y. LeCun, and R. Balestriero LeWorldModel: stable end-to-end joint-embedding predictive architecture from pixels. External Links: 2603.19312 Cited by: §2.2, §3.1. Micheli et al. (2023) V. Micheli, E. Alonso, and F. Fleuret Transformers are sample-efficient world models. In International Conference on Learning Representations, Cited by: §3.2. Moerland et al. (2023) T. M. Moerland, J. Broekens, A. Plaat, and C. M. Jonker Model-based reinforcement learning: a survey. Foundations and Trends in Machine Learning 16 (1), p. 1–118. External Links: Document Cited by: §1, §2.1. Oliehoek et al. (2008) F. A. Oliehoek, M. T. J. Spaan, and N. Vlassis Optimal and approximate q-value functions for decentralized pomdps. J. Artif. Int. Res. 32 (1), p. 289–353. External Links: ISSN 1076-9757 Cited by: §1. Rashid et al. (2018) T. Rashid, M. Samvelyan, C. Schroeder de Witt, G. Farquhar, J. Foerster, and S. Whiteson QMIX: monotonic value function factorisation for deep multi-agent reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 80, p. 4295–4304. External Links: Link Cited by: §D.9.4. Samvelyan et al. (2019) M. Samvelyan, T. Rashid, C. Schroeder de Witt, G. Farquhar, N. Nardelli, T. G. J. Rudner, C. Hung, P. H. S. Torr, J. Foerster, and S. Whiteson The StarCraft multi-agent challenge. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, p. 2186–2188. Cited by: Appendix B, §6.1. Sutton (1988) R. S. Sutton Learning to predict by the methods of temporal differences. Machine Learning 3, p. 9–44. Cited by: §2.1. Venugopal et al. (2024) A. Venugopal, S. Milani, F. Fang, and B. Ravindran MABL: bi-level latent-variable world model for sample-efficient multi-agent reinforcement learning. In Proceedings of the 23rd International Conference on Autonomous Agents and Multiagent Systems, p. 1865–1873. Cited by: §1, §3.2. Wen et al. (2022) M. Wen, J. G. Kuba, R. Lin, W. Zhang, Y. Wen, J. Wang, and Y. Yang Multi-agent reinforcement learning is a sequence modeling problem. In Advances in Neural Information Processing Systems, Vol. 35, p. 16509–16521. External Links: Document, Link Cited by: §D.9.3. Xue et al. (2026) D. Xue, J. Jiang, S. Zhang, W. Guo, L. Yuan, Z. Zhang, and Y. Yu Learning disentangled multi-agent world model for decentralized control. In Forty-Third International Conference on Machine Learning, External Links: Link Cited by: §D.9.2, §3.2. Yu et al. (2022) C. Yu, A. Velu, E. Vinitsky, J. Gao, Y. Wang, A. Bayen, and Y. Wu The surprising effectiveness of PPO in cooperative multi-agent games. In Advances in Neural Information Processing Systems, Vol. 35, p. 24611–24624. External Links: Document, Link Cited by: §D.9.3. Zhang et al. (2023) W. Zhang, G. Wang, J. Sun, Y. Yuan, and G. Huang STORM: efficient stochastic transformer based world models for reinforcement learning. In Advances in Neural Information Processing Systems, Vol. 36. Cited by: §3.2. Zhang et al. (2025) Y. Zhang, C. Bai, B. Zhao, J. Yan, X. Li, and X. Li Decentralized transformers with centralized aggregation are sample-efficient multi-agent world models. Transactions on Machine Learning Research. Cited by: §1, §3.2. Appendix A Comparing Multi-Agent World Model Architectures Table 2: Comparison of MA-JEPA with multi-agent world model architectures. Aspect MARIE MAMBA MBVD MATWM DMAWM MA-JEPA Multi-Agent Design Yes Yes Yes Yes Yes Yes Backbone Architecture Transformer GRU GRU Transformer GRU + Transformer Transformer Latent Representation VQ-VAE Categorical VAE VAE Categorical VAE Categorical RSSM Categorical latent Critic Type Centralized Centralized Centralized Semi-centralized Centralized Centralized Agent Training PPO-style PPO-style Deep Q-learning DreamerV3-style PPO-style PPO-style Appendix B Environment and Task Details We evaluate MA-JEPA on SMAC [Samvelyan et al., 2019] Our task selection follows the evaluation settings used by closely related multi-agent world-model methods, enabling direct comparison under established benchmark configurations. B.1 SMACv1 The StarCraft Multi-Agent Challenge (SMAC) consists of cooperative micromanagement scenarios in StarCraft I. Each allied unit is controlled by an individual agent and receives only a local observation of the battlefield. Observations contain information about visible allied and enemy units, including their relative positions, distances, health, shield values, and unit types. Agents select from discrete movement, stop, and unit-targeted attack actions. Invalid actions are excluded using the environment-provided legal-action mask. SMAC is particularly suitable for evaluating multi-agent world models because an agent’s local transition can depend strongly on the simultaneous actions of its teammates. For example, whether an enemy unit remains alive after an attack cannot in general be determined from the focal agent’s action alone. MA-JEPA therefore uses the joint action during centralized world-model training while retaining a decentralized policy conditioned only on each agent’s local interaction history. The main comparison covers six SMAC scenarios, namely 2s3z, 8m, M, so_many_baneling, 3s_vs_5z, and 2c_vs_64zg. These scenarios cover homogeneous and heterogeneous teams and range from relatively simple coordination problems to difficult scenarios requiring precise joint behavior. Table 3: SMACv1 scenarios used in our experiments. Scenario Difficulty Training budget 2m_vs_1z Easy 100K 2s_vs_1sc Easy 100K 2s3z Easy 100K 3m Easy 100K 3s_vs_3z Easy 100K 8m Easy 100K M Easy 100K so_many_baneling Easy 100K 3s_vs_4z Hard 100K 3s_vs_5z Hard 200K 2c_vs_64zg Hard 200K corridor Super Hard 400K Appendix C Additional Ablations Figure 2: Ablation of gradient routes from joint prediction losses into local representations. MA-JEPA enables outcome gradients at scale 1.0 and factual JEPA gradients at scale 0.1. “w/o factual JEPA,” “w/o outcomes,” and “w/o both” disable the corresponding gradients into local representations; the joint prediction objectives remain active. Connected markers show matched training seeds. Squares and error bars indicate the mean and sample standard deviation across three seeds, also reported numerically above each configuration. Results use 100-episode endpoint evaluations. Across the three ablation maps in Figure 2, MA-JEPA achieves the highest average win rate: 77.0%, compared with 71.1% without factual JEPA gradients, 69.8% without outcome gradients, and 73.9% without either route. Although individual ablations perform better on particular maps, their benefits do not transfer consistently across tasks. Combining both routes therefore provides the strongest overall performance across this evaluation set. One possible explanation is that the two objectives provide complementary training signals for local representations. Outcome gradients emphasize features relevant to reward and termination, while factual JEPA gradients encourage features that support prediction of future latent states. Outcome supervision alone may leave useful information underconstrained, whereas predictive supervision alone need not prioritize features relevant to control. Combining task-relevant outcome gradients with a weaker predictive gradient may balance these demands. The map-dependent results are consistent with this interpretation. Appendix D Implementation Details D.1 Replay Sampling MA-JEPA maintains one replay buffer with two independently sampled views. The world-model view supplies model supervision, while the behavior view supplies the roots for imagined policy learning. Both views use the same mixture of uniform and exponential-recency sampling: P(v)=η1||+(1−η)dage(v)∑w∈dage(w),η=0.5,d=0.9998,P(v)=η 1|V|+(1-η) d^age(v) _w d^age(w), η=0.5, d=0.9998, (19) where V is the set of eligible sequence starts and age is measured by insertion order. The two views use independent random streams but may contain overlapping sequences. Independent sampling therefore separates the two batches without removing recency bias from the policy-root distribution. Each replay sequence contains 192 context records followed by 64 learning records. The context records reconstruct the model’s temporal history; no loss is computed on them. For world-model training, we reconstruct local temporal states from stored latent–action inputs and stochastic states using the current temporal network. For policy and value training, performed after the world-model update, we re-encode the context observations and infer their posterior states using the updated model. In both cases, the resulting local features and recorded joint actions reconstruct the joint Transformer’s temporal history. Histories reset at episode boundaries. States supplied to policy and value learning are treated as fixed inputs, preventing these updates from changing the world model. Sketched Isotropic Gaussian Regularization. To prevent representational collapse, we regularize the online observation embeddings with SIGReg [Balestriero and LeCun, 2025]. For each agent slot, SIGReg samples random unit directions ama_m and compares the empirical one-dimensional projected embedding distribution un,m=am⊤enu_n,m=a_m e_n with a standard normal distribution using their characteristic functions. For projection m, the discrepancy is m=N∫|1N∑n=1Neitun,m−e−t2/2|2e−t2/2dt,T_m=N | 1N _n=1^Ne^itu_n,m-e^-t^2/2 |^2e^-t^2/2\,dt, (20) and ℒSIGRegL_SIGReg is obtained by averaging this statistic over random projections and nonempty agent slots. In practice, the integral is approximated on a finite frequency grid. By encouraging many random one-dimensional projections to resemble (0,1)N(0,1), SIGReg promotes approximately isotropic, non-collapsed embeddings without requiring an observation reconstruction objective. D.2 Local World-Model Objective Let eti=Eϕ(oti)e_t^i=E_φ(o_t^i) denote the online observation embedding and e¯ti=sg(Eϕ¯(oti)) e_t^i=sg(E_ φ(o_t^i)) its stopped target, where ϕ¯ φ is an exponential moving average of ϕφ. The local state is sti=[hti;vec(zti)]s_t^i=[h_t^i;vec(z_t^i)]. Unless stated otherwise, per-transition losses are averaged over their valid entries; invalid entries contribute zero, and an empty set of valid entries contributes zero. For the local objectives, validity is given by the replay presence and learning-activity masks, not necessarily physical controllability. The two local JEPA objectives are ℒpost _post =valid[dcos(Ppost(sti),e¯ti)], =E_valid\! [d_ (P_post(s_t^i), e_t^i) ], (21) ℒdyn _dyn =valid[dcos(Pdyn(hti),e¯ti)], =E_valid\! [d_ (P_dyn(h_t^i), e_t^i) ], (22) where dcos(x,y)=1−simcos(x,y)d_ (x,y)=1-sim_ (x,y). The posterior and prior are qti=qθ(⋅∣hti,eti)q_t^i=q_θ(· h_t^i,e_t^i) and pti=pθ(⋅∣hti)p_t^i=p_θ(· h_t^i), respectively. Both use categorical probabilities with uniform mixing coefficient 0.010.01. The two directed regularizers are ℒdyn−reg _dyn-reg =valid[maxτ,DKL(sg(qti)∥pti)], =E_valid\! [ \! \τ,D_KL (sg(q_t^i)\|p_t^i ) \ ], (23) ℒrep−reg _rep-reg =valid[maxτ,DKL(qti∥sg(pti))], =E_valid\! [ \! \τ,D_KL (q_t^i\|sg(p_t^i) ) \ ], (24) with free-nat threshold τ=1τ=1. KL divergences sum over the categorical variables before applying this threshold. The complete local objective is ℒlocal= _local= 2ℒpost+2ℒdyn+0.05ℒSIGReg 2L_post+2L_dyn+0.05L_SIGReg (25) +ℒmasklocal+ℒdyn−reg+0.1ℒrep−reg. +L_mask^local+L_dyn-reg+0.1L_rep-reg. There are no local reward or continuation losses: these outcomes are predicted by the joint model. The vector-observation configuration contains neither an observation-reconstruction objective nor a spatial prediction objective. Representation regularization. SIGReg is applied to online embeddings, grouping samples by agent slot across batch and time. For each nonempty group, it compares the empirical characteristic functions of 256 random unit-vector projections with the standard Gaussian characteristic function. The comparison uses 17 equally spaced points in [0,3][0,3] and Gaussian weighted quadrature. The projected discrepancy is multiplied by the group’s valid sample count before averaging over nonempty groups. This sample-count scaling is part of the objective. Action-availability supervision. The local availability head predicts a binary vector from stis_t^i. For each vector, binary cross-entropy is averaged separately over legal and illegal actions, then averaged equally over the nonempty classes. The joint availability head uses the same balanced reduction. D.3 One-Step Joint Prediction and Posterior Alignment The joint model receives synchronized local states, the joint action, roster presence, and controllability. Agent self-attention first mixes information within a transition; a causal Transformer then maintains temporal context. Writing its cache explicitly as ktk_t, we obtain (kt+1,ut1:N) (k_t+1,u_t^1:N) =Jω(kt,sg(st1:N),at1:N,mpres,t1:N,mctrl,t1:N), =J_ω\! (k_t,sg(s_t^1:N),a_t^1:N,m_pres,t^1:N,m_ctrl,t^1:N ), (26) e^t+1i e_t+1^i =Pe(uti). =P_e(u_t^i). (27) No fixed agent-identity or agent-position embeddings are added. The detached local-state input specifies the baseline gradient boundary. The embedding and interface objectives are ℒemb _emb =valid[dcos(e^t+1i,e¯t+1i)], =E_valid\! [d_ ( e_t+1^i, e_t+1^i) ], (28) ℒint _int =valid[SmoothL1(e^t+1i,sg(et+1i))]. =E_valid\! [SmoothL1( e_t+1^i,sg(e_t+1^i)) ]. (29) Smooth-L1L_1 is averaged over embedding coordinates and uses a quadratic-to-linear transition at absolute error one. It penalizes pointwise discrepancy with the online encoder output; it does not establish equality of the real and imagined embedding distributions. We additionally align the posterior distributions induced by real and predicted embeddings while holding the factual history fixed: ℒalign=valid[DKL( _align=E_valid\! [D_KL\! ( sg(qθ(⋅∣ht+1i,et+1i))∥ (q_θ(· h_t+1^i,e_t+1^i))\,\| (30) qθ−(⋅∣sg(ht+1i),e^t+1i))], q_θ^-(· (h_t+1^i), e_t+1^i) ) ], where θ−θ^- denotes a frozen copy of the current local parameters, not a separate learned posterior. The KL uses the same uniform-mixed categorical distributions as execution. Gradients pass through the predicted embedding to the joint model, but not to the factual history, target posterior, or local posterior parameters along this path. The one-step joint objective is ℒjoint= _joint= 2ℒemb+ℒint+0.05ℒalign 2L_emb+L_int+0.05L_align (31) +ℒrjoint+ℒcjoint+ℒmaskjoint+ℒalive. +L_r^joint+L_c^joint+L_mask^joint+L_alive. Reward prediction uses a symlog two-hot categorical output. Continuation and controllability use binary outputs. With discounted continuation, the continuation target is γ(1−dt+1i)γ(1-d_t+1^i), where d denotes the replay terminal flag and γ=1−1/333γ=1-1/333. Embedding, interface, alignment, and next-mask losses require a controllable source, a present target, and no reset crossing. We still predict when death occurs after the transition. Reward and continuation losses also supervise dead but present roster slots. Learning when agents are alive is supervised only on unmasked slots. These masks are normalized separately. D.4 Direct Multi-Step JEPA and Action Discrimination Direct prediction heads are supervised at ℋdirect=1,2,4,8H_direct=\1,2,4,8\. The factual interaction feature utiu_t^i already contains the current joint action at1:Na_t^1:N. For h>1h>1, the corresponding head also receives the focal agent’s recorded action tail: e^t+h|ti=Ph(uti,at+1:t+h−1i). e_t+h t^i=P_h(u_t^i,a_t+1:t+h-1^i). (32) The tail is empty for h=1h=1. Later teammate actions are not supplied by a teammate-plan model. These predictions are evaluated directly from factual roots; they are neither chained between horizons nor used to construct PPO states. For ordered horizons (h0,h1,h2,h3)=(1,2,4,8)(h_0,h_1,h_2,h_3)=(1,2,4,8), define αhk=ρk∑j=03ρj,ρ=0.75. _h_k= ρ^k _j=0^3ρ^j, ρ=0.75. (33) The direct loss is ℒMS=∑h∈ℋdirectαhvalid,h[dcos(e^t+h|ti,e¯t+hi)].L_MS= _h _direct _hE_valid,h\! [d_ ( e_t+h t^i, e_t+h^i) ]. (34) The available root indices leave room for the maximum horizon in the sampled suffix. Each horizon additionally requires a legal factual prefix, present and controllable action sources, and no reset crossing. Death exactly at the target is retained. For each h>1h>1, action discrimination replaces the final focal-agent action at+h−1ia_t+h-1^i with every distinct alternative legal at that factual decision. The joint root and earlier focal actions remain unchanged. Each alternative incurs ℓAD,hi=[ _AD,h^i= [ m−simcos(e^t+h|ti,e¯t+hi) m-sim_ ( e_t+h t^i, e_t+h^i) (35) +simcos(e^t+h|ti,alt,e¯t+hi)]+,m=0.1. +sim_ ( e_t+h t^i,alt, e_t+h^i) ]_+, m=0.1. We first average over legal alternatives at each root and then over valid roots. Horizon weights are renormalized over 2,4,8\2,4,8\ for the action loss. The resulting auxiliary is 2ℒMS+0.1ℒAD2L_MS+0.1L_AD. The alternatives are discriminative negatives, not observed counterfactual targets; different legal actions need not always yield different observable outcomes. D.5 Self-Fed Recurrent Supervision In addition to direct prediction, we train recursive joint predictions under recorded joint actions. Starting from a factual root, the joint model predicts the next observation embedding, the local posterior constructs the next latent state, and that predicted state becomes the input to the following transition. Real future observations are used as targets, not injected into the predicted trajectory. We sample up to eight eligible synchronized roots without replacement from the world-model batch. A root requires two valid subsequent transitions; later endpoint validity is assessed separately. The rollout extends for five transitions, with supervision at ℋSF=2,4,5H_SF=\2,4,5\. Terminal arrivals are retained, while reset crossings, disappearing slots, and out-of-range endpoints are masked. Individual death does not itself define an episode boundary. The joint model and reconstructed joint history use inference mode. Local transition parameters are frozen. Gradients through the joint temporal cache and the permitted local input paths are truncated by detaching recurrent state at the start of every two-step chunk. The joint predictor still detaches its local-state input at each step; this is not unrestricted end-to-end backpropagation through the full five-step simulator. The trajectory-alignment term retains the input Jacobian through the frozen local transition within each chunk. Unlike one-step alignment, trajectory alignment compares factual and recursively generated histories: ℒtraj,h=valid,h[DKL( _traj,h=E_valid,h\! [D_KL\! ( sg(qθ(⋅∣ht+hi,et+hi))∥ (q_θ(· h_t+h^i,e_t+h^i))\,\| (36) qθ−(⋅∣h^t+hi,e^t+hi))]. q_θ^-(· h_t+h^i, e_t+h^i) ) ]. The total recurrent auxiliary is ℒSF=0.1|ℋSF|∑h∈ℋSF(CLOSE _SF= 0.1|H_SF| _h _SF ( 2ℒemb,h+ℒint,h+ℒr,hjoint 2L_emb,h+L_int,h+L_r,h^joint (37) OPEN+ℒc,hjoint+ℒmask,hjoint+ℒalive,h+0.1ℒtraj,h). +L_c,h^joint+L_mask,h^joint+L_alive,h+0.1L_traj,h ). Each endpoint loss is normalized over its valid roots and agents; endpoints receive equal weight. The trajectory-KL coefficient is therefore 0.010.01 relative to the endpoint-averaged trajectory loss. Shared reward and continuation losses include dead but present slots; the other endpoint losses require a controllable root. This supervised rollout uses the joint availability head with thresholded masks. Recorded actions continue to drive the rollout even when excluded by its predicted mask. This differs from the local-head, Bernoulli-sampled masks used for PPO imagination below. D.6 Joint Imagination and Decentralized Execution After each world-model update, we reconstruct the independent behavior batch with the updated model. Its 64 suffix positions supply candidate imagination roots; episode-final roots are excluded by validity masks. Each valid root produces H=5H=5 imagined transitions. At every transition, each agent samples an action from the shared actor using only its own local state and current action mask. The local history advances using that agent’s latent and action. In parallel, the joint predictor consumes the synchronized states, joint action, and joint temporal cache to predict the next local embedding. The next stochastic state is then sampled as z^t+1i∼qθ(⋅∣h^t+1i,e^t+1i). z_t+1^i q_θ(· h_t+1^i, e_t+1^i). (38) The local prior remains a training regularizer, but it is not the sampling distribution used for these joint imagined transitions. Imagined action availability. Root masks come from replay. At subsequent states, the local availability head predicts probabilities from the completed local state. Each action’s availability is sampled as a Bernoulli variable. An empty predicted mask falls back to the no-op action, and non-controllable agents receive no-op-only support. The effective masks actually used to sample actions are stored with the imagined batch and are not resampled during PPO updates. SMAC roster and team outcomes. The fixed roster is retained after individual deaths. Imagined controllability is monotonic: an agent remains controllable only if it was controllable previously and the predicted next controllability probability is at least 0.50.5. Joint reward and continuation predictions are averaged over present slots and broadcast back to those slots. The all-dead team is absorbing: the transition into this condition retains its final reward, while continuation and later rewards are zero. These are the shared-reward, fixed-roster rules of the SMAC adapter, not assumptions about all benchmark environments. Execution. Real interaction uses the observation encoder, local history model, local posterior, and actor, together with the environment-provided legal-action mask. Parameters are shared, but each agent maintains its own history. Neither the joint predictor nor the centralized critic is used during execution. There is no teammate-belief network or residual belief adapter in this configuration. D.7 PPO and Replay-Value Learning One imagined batch is generated before actor–critic optimization. Its local features, centralized critic context, actions, effective masks, old policy logits, advantages, returns, and trajectory weights remain fixed throughout the update. The world model receives no PPO gradients. Let V¯ti V_t^i be the target critic prediction and let btib_t^i indicate value-valid states. Define the effective discounted continuation c~t+1i=c^t+1ibt+1i c_t+1^i= c_t+1^ib_t+1^i. Generalized advantages and return targets are δti _t^i =r^t+1i+c~t+1iV¯t+1i−V¯ti, = r_t+1^i+ c_t+1^i V_t+1^i- V_t^i, (39) Ati A_t^i =bti(δti+λc~t+1iAt+1i),AHi=0, =b_t^i ( _t^i+λ c_t+1^iA_t+1^i ), A_H^i=0, (40) Rti R_t^i =Ati+V¯ti,λ=0.95. =A_t^i+ V_t^i, λ=0.95. (41) The continuation already includes discounting; it is not multiplied by γ a second time. Critic-valid states include dead but present SMAC agents, whereas actor-valid decisions additionally require controllability. An individual death therefore removes subsequent actor decisions without discarding the earlier shared-team return. Losses use cumulative continuation weights, wti=bti∏j=1tc~jiw_t^i=b_t^i _j=1^t c_j^i, with an empty product of one. Actor advantages are normalized once using the weighted mean and variance over valid decisions in the frozen batch. Writing the normalized advantage as A~ti A_t^i and the stored mask as MtiM_t^i, the probability ratio is rti(ψ)=πψ(ati∣sti,Mti)πold(ati∣sti,Mti).r_t^i(ψ)= _ψ(a_t^i s_t^i,M_t^i) _old(a_t^i s_t^i,M_t^i). (42) For numerical stability, its log is clipped to [−20,20][-20,20] before exponentiation. The actor minimizes ℒactor=−actor,w[min( _actor=-E_actor,w [ ( OPENrtiA~ti,clip(rti,1−ϵ,1+ϵ)A~ti) r_t^i A_t^i,clip(r_t^i,1-ε,1+ε) A_t^i ) (43) +βentℋ(πψ(⋅∣sti,Mti))], + _entH ( _ψ(· s_t^i,M_t^i) ) ], where actor,wE_actor,w is the normalized weighted mean over actor-valid decisions, ϵ=0.2ε=0.2, and βent=0.003 _ent=0.003. Entropy is not divided by the size of the action support. The centralized critic predicts a symlog two-hot value distribution and is trained with that distribution’s target loss, rather than a clipped mean-squared-error value objective: ℒcritic=value,w[ℓtwohot(Vξ,sg(R))]+0.3ℒreplay−value.L_critic=E_value,w [ _twohot(V_ξ,sg(R)) ]+0.3L_replay-value. (44) The replay-value term uses real rewards and current imagined root-return bootstraps to construct λ-returns with λreplay=0.95 _replay=0.95. It trains only the critic; recorded replay actions are not PPO actor samples. Terminal arrivals retain their rewards and stop bootstrapping, and traces never cross resets. When the adapter distinguishes nonterminal truncation, the final real observation supplies its bootstrap. The supplied SMAC adapter uses the same upstream termination flag for episode-final and terminal markers, so that distinct truncation case is not exercised by this adapter. Five actor steps and five critic steps reuse the same full imagined batch, with actor and critic updates interleaved. No new trajectories or action masks are sampled between these steps. The target critic is copied from the updated critic after each active learner call; the additional slow-critic regularizer is disabled. D.8 Optimizer Ownership and Update Order Parameters are partitioned into four disjoint groups: the local world model, joint world model, actor, and centralized critic. The model objective is ℒWM=ℒlocal+ℒjoint+2ℒMS+0.1ℒAD+ℒSF.L_WM=L_local+L_joint+2L_MS+0.1L_AD+L_SF. (45) The local group contains the encoder, local dynamics and posterior, local JEPA predictors, and local availability head. The joint group contains the interaction and temporal predictor, joint output heads, and direct multi-step heads. In the baseline, joint losses do not update the local-world parameters. Frozen local modules can still transmit gradients with respect to their inputs on the alignment and self-fed paths described above. A learner call first differentiates ℒWML_WM and updates the two model groups, leaving actor and critic optimizer states untouched. It then reconstructs the behavior view with the updated model, generates and freezes an imagined batch, and performs the separate PPO actor and critic steps. This is not a single combined world-model–actor–critic backward pass. All groups use adaptive gradient clipping followed by RMS preconditioning and momentum. The local and joint world-model learning rates are 10−410^-4; actor and critic learning rates are 3×10−53× 10^-5. The clipping coefficient is 0.30.3, momentum coefficients are (0.9,0.999)(0.9,0.999), optimizer ϵε is 10−2010^-20, and weight decay and learning-rate warm-up are disabled in this configuration. The target encoder is updated after a finite model update using ϕ¯←0.99ϕ¯+0.01ϕ. φ← 0.99 φ+0.01φ. (46) D.9 Baseline Hyperparameters We benchmark DMAWM, MARIE, MAMBA, MAPPO, QMIX, and MAT using implementations released by their respective authors. The tables below report the effective settings used by our benchmark launchers, verified against the released implementations and saved run configurations. Unless stated otherwise, all baselines use StarCraft I 4.10.0, AI difficulty 7, and step_mul=8. Training budgets count real environment transitions across collectors; evaluation interactions, imagined transitions, and individual-agent actions do not contribute to this count. D.9.1 MA-JEPA Hyperparameters We use a shared configuration across all SMAC tasks, changing only the task, number of agents, training seed, and interaction budget. Table 4: MA-JEPA hyperparameters used in our SMAC experiments. Hyperparameter Value used World model Observation encoder layers / width 3 / 1024 Deterministic state dimension 4096 Categorical latent 32×6432× 64 Latent uniform mixture 0.01 Local Transformer layers / width / heads 2 / 512 / 8 Local / joint context length 64 / 16 Joint agent-interaction / temporal layers 2 / 12 Joint Transformer width / heads 256 / 4 Joint dropout 0.1 World-model learning rate 10−410^-4 Adaptive gradient clipping 0.3 Target-encoder EMA rate 0.01 SIGReg coefficient 0.05 Direct JEPA prediction horizons 1,2,4,8\1,2,4,8\ Action-margin / loss coefficient 0.1 / 0.1 Self-fed horizons / BPTT length 2,4,5\2,4,5\ / 2 Outcome / factual JEPA input-gradient scales 1.0 / 0.1 Policy learning Actor hidden layers / width 3 / 512 Central critic layers / width 2 / 256 Actor / critic learning rate 3×10−53× 10^-5 / 3×10−53× 10^-5 Actor / critic PPO epochs 5 / 5 PPO clipping 0.2 Entropy coefficient 0.003 Discount / return-trace λ 1−1/3331-1/333 / 0.95 Imagination horizon 5 Auxiliary replay-value loss coefficient 0.3 Imagined action mask Local; Bernoulli sampling D.9.2 DMAWM DMAWM follows Xue et al. [2026] and the authors’ released implementation.11 1 https://github.com/DiXue98/DMAWM/tree/2acaaeb82805b55d275e3ce08ac8f713ec9afbb2 Table 5: DMAWM hyperparameters used in our benchmark. Hyperparameter Value used Reinforcement learning Optimizer Adam Actor / critic learning rate 3×10−53× 10^-5 / 3×10−53× 10^-5 Actor / critic gradient clipping 100 / 100 Entropy coefficient 0.01 PPO epochs / clipping 5 / 0.2 Discount γ / GAE λ 0.99 / 0.95 Replay-value loss coefficient 1.0 Importance-weight cap 4 Actor / critic sharing True / True World model World-model learning rate 10−410^-4 World-model gradient clipping 1000 Batch size / sequence length 16 / 64 Replay burn-in / imagination horizon 4 / 4 Replay capacity / prefill 250,000250,000 / 5000 Training ratio 128 Dynamics / representation KL scales 1.0 / 0.1 KL free bits 1.0 Deterministic latent width 512 Categorical latent 32×3232× 32 Latent uniform mixture 0.01 Transition / global Transformer layers 3 / 3 Transformer heads 8 Observation decoder hidden width / layers 1024 / 2 Collection environments 16 Actor / critic MLP width / layers 256 / 2 Critic Transformer layers / heads 1 / 8 D.9.3 MAPPO and MAT MAPPO uses the official implementation of Yu et al. [2022].22 2 https://github.com/marlbenchmark/on-policy/tree/de66d7a4b23fac2513f56f96f73b3f5cb96695ac MAT uses the official implementation of Wen et al. [2022].33 3 https://github.com/PKU-MARL/Multi-Agent-Transformer/tree/be3f49c8264d454c1fe2c41582a2bfc98498c8 Table 6: MAPPO and MAT hyperparameters used in our benchmark. Hyperparameter MAPPO MAT Policy Recurrent or F Transformer Optimizer Adam Adam Actor / critic learning rate 5×10−45× 10^-4 each Joint 5×10−45× 10^-4 Rollout length per collector 400 100 Parallel collection environments 8 32 Transitions per rollout batch 3200 3200 Entropy coefficient 0.01 0.01 Discount γ 0.99 0.99 GAE λ 0.95 0.95 PPO epochs 15 or 5 15 PPO clipping 0.2 or 0.05 0.05 Minibatches per PPO epoch 1 1 Maximum gradient norm 10 10 MLP / recurrent hidden width 64 — Recurrent layers 1 — Transformer blocks / heads — 1 / 1 Transformer embedding width — 64 Value normalization / PopArt True / False True / False Policy / value active masks True / False True / False D.9.4 QMIX We evaluate QMIX [Rashid et al., 2018] using the optimized PyMARL2 implementation.44 4 https://github.com/hijkzzz/pymarl2/tree/8ccac7c5a134422a2e3009be735d23cdf8ce2f8 The settings combine PyMARL2’s QMIX, default, and SMAC configurations. Table 7: QMIX/PyMARL2 hyperparameters used in our benchmark. Hyperparameter Value used Agent architecture / hidden width Recurrent / 64 Optimizer / learning rate Adam / 10−310^-3 Discount γ / TD λ 0.99 / 0.6 Target-network update interval 200 episodes Maximum gradient norm 10 Training batch size 128 episodes Replay capacity 5000 episodes Parallel collection environments 8 Exploration ϵε 1.0→0.051.0→ 0.05 Linear epsilon annealing duration 100K environment steps Mixer embedding / hypernetwork width 32 / 64 Agent-ID / previous-action inputs True / True Layer normalization False Orthogonal initialization False Prioritized replay False D.9.5 MAMBA MAMBA follows Egorov and Shpilman [2022] and the authors’ released StarCraft. 55 5 https://github.com/jbr-ai-labs/mamba/tree/2c97258f71bf1c421c40ce14fd2f7c3fe7fe19f Table 8: MAMBA hyperparameters used in our benchmark. Hyperparameter Value used World-model / actor / critic optimizer Adam / Adam / Adam World-model learning rate 2×10−42× 10^-4 Actor / critic learning rate 5×10−45× 10^-4 / 5×10−45× 10^-4 Actor weight decay 10−510^-5 World-model / policy gradient clipping 100 / 100 Replay capacity / minimum replay 250,000250,000 / 500 World-model / policy batch size 40 / 40 Replay sequence length 20 Imagination horizon 15 World-model optimizer steps per cycle 60 Policy updates per cycle 4 PPO epochs / clipping 5 / 0.2 Discount γ / return λ 0.99 / 0.95 Initial entropy coefficient 0.001 Entropy multiplier per actor step 0.99998 Deterministic latent width 256 Categorical latent 32×3232× 32 Actor / critic hidden width / layers 256 / 2 Collection environments 1 Collected-sample gate 1; updates after complete episodes Appendix E Consistency of Learned Joint Successor Law with Fixed Representations Consistency of the joint successor law. Fix MA-JEPA’s parameters and collection policy. Consider independent and identically distributed episodes of T<∞T<∞ steps. Let tz_t denote the joint latent state, t∈a_t the joint action, and tm_t the joint legal-action mask. For a fixed history-window length w≥1w≥ 1, define the padded context t=(t,t−w+1:t,t−w:t−1,t−w+1:t)∈c_t=(t,z_t-w+1:t,a_t-w:t-1,m_t-w+1:t) and the complete successor t+1=(t+1,r~t+1,et+1,t+1)∈y_t+1=(z_t+1, r_t+1,e_t+1,m_t+1) . Here r~ r is a quantized reward and e∈0,1e∈\0,1\ indicates termination. All alphabets are finite. Let P(⋅∣,)P(· ,a) denote the conditional distribution of the complete successor under the fixed parameters and collection policy. We suppress the conditioning arguments in the bounds below. Theorem E.1 (Joint-law consistency). From n sampled episodes, let n(,)n(c,a) count occurrences of context–action pair (,)(c,a). At each observed pair, let P^n(⋅∣,) P_n(· ,a) be the empirical distribution of its complete successors. At an unobserved pair, define P^n P_n to be any fixed probability distribution on Y. As n→∞n→∞, P^n P_n converges to P almost surely in total variation, uniformly over all context–action pairs with positive probability of occurring under the fixed collection policy. For fixed n, n0≥1n_0≥ 1, and 0<δ<10<δ<1, the following bound holds with probability at least 1−δ1-δ, simultaneously for all pairs with n(,)≥n0n(c,a)≥ n_0. TV(P^n,P)≤min1,||log2+log(2||||/δ)2n0.TV( P_n,P)≤ \1, |Y| 2+ (2|C||A|/δ)2n_0 \. Conditional on observing a given pair exactly s≥1s≥ 1 times, its expected total-variation error satisfies [TV(P^n,P)∣n(,)=s]≤12(||−1)/s.E [TV( P_n,P) n(c,a)=s ]≤ 12 (|Y|-1)/s. At each context–action pair with positive occurrence probability, consider an unrestricted autoregressive model of the complete successor y. Its population teacher-forced cross-entropy is minimized uniquely by the joint distribution P(⋅∣,)P(· ,a). preprint-pratenddefaultcategory.tex See proof in .preprint-pratenddefaultcategory.tex Appendix F Proofs of Theoretical Results