Robotics Research Hub
One Future, Every Robot: Label-Efficient Collective-State Prediction with Decentralized JEPA

One Future, Every Robot: Label-Efficient Collective-State Prediction with Decentralized JEPA

Alan-Barsag Gazzaev, Alexey Garvilov, Sergey Muravyov · N/A · 2026

reinforcement-learning planning
PDF arXiv

Framework

N/A

License

N/A

Stars

N/A

Summary

Can every robot in a swarm predict the same future collective state from only local observations and bandwidth-limited messages?

Abstract Summary

Can every robot in a swarm predict the same future collective state from only local observations and bandwidth-limited messages? We formulate this as decentralized shared-state prediction and introduce Collective-State JEPA (CS-JEPA), a recurrent joint-embedding predictive architecture whose output at every robot represents one common future token field. At deployment, each robot uses a 16-frame local history and one 64-float recurrent message per directed edge; there is no global pooling, target encoder, episode clock, or recorded future action. After pretraining without downstream collective labels, frozen representations are evaluated with ridge probes fitted on 6, 12, or 24 globally labeled episodes. Against raw-future reconstruction with the same receiver anchor and deployment capacity but 9,607 additional training-only parameters, a prospectively registered five-seed follow-up improves prediction-error and inter-robot-agreement label-budget AUC on in-distribution, ring, mutual-kNN, and unseen-size families up to 108 robots. Every effect favors CS-JEPA in 5/5 outer seeds. In a separate sealed eight-seed follow-up, matched action-conditioned predictors receive each candidate four-step plan before producing receiver-local predictive representations. CS-JEPA reduces branch-value MSE by 45.5% and improves within-context candidate-score Pearson correlation by 0.1291, with both effects favorable in 8/8 seeds, including at unseen N=32. These results support common-future JEPA targets as a label-efficient primitive for decentralized swarm prediction under topology and size shift, with additional evidence of planning-relevant value estimation.

Key Points

  • Can every robot in a swarm predict the same future collective state from only local observations and bandwidth-limited messages?
  • We formulate this as decentralized shared-state prediction and introduce Collective-State JEPA (CS-JEPA), a recurrent joint-embedding predictive architecture whose output at every robot represents one common future token field.
  • At deployment, each robot uses a 16-frame local history and one 64-float recurrent message per directed edge; there is no global pooling, target encoder, episode clock, or recorded future action.
  • After pretraining without downstream collective labels, frozen representations are evaluated with ridge probes fitted on 6, 12, or 24 globally labeled episodes.
  • Against raw-future reconstruction with the same receiver anchor and deployment capacity but 9,607 additional training-only parameters, a prospectively registered five-seed follow-up improves prediction-error and inter-robot-agreement label-budget AUC on in-distribution, ring, mutual-kNN, and unseen-size families up to 108 robots.

Problem Formulation

At time $t$, the active swarm is a dynamic graph $\mathcal{G}_{t}=(\mathcal{V}_{t},\mathcal{E}_{t})$. Robot $i$ observes a receiver-indexed local view $o_{t}^{i}$ and receives messages only from $\mathcal{N}_{t}(i)$. Its allowed information is Here $o_{t}^{i}$ contains only robot $i$’s normalized world position, velocity, two-dimensional task vector, and active bit, all in a shared normalized world frame.

Problem Formulation — One Future, Every Robot: Label-Efficient Collective-State Prediction with Decentralized JEPA
Fig. 2: Absolute label-efficiency curves on the anchor-matched follow-up. Points average five outer-seed estimates after averaging the five fixed label-subset draws and task-balanced evaluation episodes within each seed; bars are 95% outer-seed bootstrap intervals. Lower is better. The same frozen checkpoints, probes, budgets, and episodes produce the registered paired AUC effects in Table III . Each panel uses its own vertical scale; comparisons are within panel.

Collective-State JEPA

A frozen local encoder maps robot $i$’s self observation to $e_{t}^{i}\in\mathbb{R}^{64}$. At every history step, each active robot broadcasts only its previous recurrent state. Receiver $i$ computes We set $\bar{h}_{t-1}^{,i}=0$ when $\mathcal{N}_{t}(i)=\varnothing$, matching the masked-mean implementation.

Updates are synchronous, use one message round per environment step, and send 64 float32 values (256 bytes) per directed edge. With 16 frames, the measured warm-up is 15 message steps.

Experimental Protocol

For the target-isolation follow-up, we froze the protocol and generated a new data family only afterward. It contains three tasks (flocking, formation, coverage), 200-step episodes, process noise 0.02, and no observation noise. Training and ID use bounded-small-world graphs with $N\in\{10,18\}$ and maximum degree four. Table II shows the splits.

Ring and mutual-$k$NN are topology shifts; the size split uses $N\in\{36,72,108\}$, unseen and up to $6\times$ the largest training swarm.

Results

All four registered comparisons satisfy the prespecified joint criterion (Table III ). Every accuracy and agreement effect is positive in 5/5 outer seeds, and every lower confidence bound is above zero. The smallest effect is ID agreement (0.001342); the largest effects are 0.057156 accuracy and 0.098817 agreement, both on ring. Figure 3 exposes every seed rather than treating episodes as independent replications.

Sources and Demonstrations

Figures are reproduced from Gazzaev et al. See the full paper for experimental details and the project page, when the paper links one, for demonstrations and videos.

Share

Related Papers

AcousticDiffusion: Semantically Conditioned Audio-Guided Diffusion Policy for Search-and-Rescue Assistance
arXiv preprint

AcousticDiffusion: Semantically Conditioned Audio-Guided Diffusion Policy for Search-and-Rescue Assistance

Iana Zhura, Didar Seyidov, Dmitrii Plotnikov et al.

Navigating toward human callers is an important capability for rescue robots operating where visual contact is degraded or occluded. We present AcousticDiffusion, a semantically conditioned, audio-guided diffusion policy for human-directed navigation.

locomotion vision reinforcement-learning planning
PDF Advanced
Actively Resolving Contextual Uncertainty for Underspecified Tasks in Natural Language
arXiv preprint

Actively Resolving Contextual Uncertainty for Underspecified Tasks in Natural Language

Zachary Ravichandran, Jonathan Diller, Fernando Cladera et al.

Foundation models provide robots with the ability to interpret natural language and reason about environmental context, yet most language-conditioned policies assume that goals are well-specified and that task-relevant information is provided upfront via a prior map.

vla reinforcement-learning planning
Code PDF Advanced
GitHub ★ —