Calibrated Predictive Safety for Heterogeneous Robots: An Action-Conditioned JEPA Framework with Model-Based Safety Shields

Abstract: Vision-language-action policies generalize broadly but provide no execution-time guarantees; classical model-based planners respect kinematic and geometric constraints but generalize poorly. We study whether an action-conditioned Joint-Embedding Predictive Architecture (JEPA) world model can predict, before execution, both task progress and physical risk for candidate action chunks, and whether coupling these predictions to an embodiment-specific model-based safety shield yields a deployable pipeline for heterogeneous robots. We propose a receding-horizon decision pipeline: (1) a proposer produces K candidate action chunks; (2) an action-conditioned JEPA rolls each candidate forward in a frozen-encoder latent space conditioned on an embodiment embedding; (3) calibrated risk and progress heads score each rollout and report uncertainty; (4) a deterministic per-embodiment safety shield filters inadmissible candidates; (5) a fallback ladder handles empty-admissible-set cases. The learned ranking only reorders admissible candidates; enforcement guarantees come from the deterministic shield and fallback ladder. We evaluate with a pre-registered protocol in simulation (LIBERO-Long). In 600-episode configurations the full framework improved success over a shield-only baseline and reduced collision false negatives at matched recall. Deployment-efficiency measurements on target on-robot and edge accelerators are included. Real-robot experiments and an offline reranking significance test remain future work; see the paper for disclosures.

• The paper introduces a layered framework in which an action-conditioned JEPA ranks candidate robot actions while a deterministic, embodiment-specific shield alone enforces admissibility and safety constraints.
• The framework improves collision-risk prediction to AUROC 0.88 and false-negative rate 0.14, while calibration reduces expected calibration error from 0.14 to 0.04 and closed-loop testing reaches 62% success with 3% collisions.
• The results show a 7-point success gain over shield-only control and sub-second decision latency, but leave open whether learned ranking reliably outperforms reranking alone under distribution shift and real-robot conditions.

This paper presents a framework for composing learned action screening with deterministic runtime safety enforcement for heterogeneous robots (2608.17496). The central architectural claim is a strict separation of responsibilities: an action-conditioned Joint-Embedding Predictive Architecture (JEPA) world model provides a calibrated ranking over candidate actions, while every safety guarantee originates from a deterministic, per-embodiment shield and a fallback ladder that the learned components cannot override.

Problem formulation and the shield-dominance contract

The system operates on a receding horizon: at each control step, a proposer — a VLA policy sampled at temperature, an MPPI/CEM optimiser warm-started from the VLA mean, or a parameterised skill library — emits $K \in \{8,16\}$ candidate action chunks. Each candidate is rolled out in latent space by an action-conditioned predictor, scored by risk and progress heads, ranked by a scalar cost $J_k$, and filtered by a hard admissibility predicate $C_{\mathrm{emb}}$ evaluated against the embodiment specification (kinematics, geometry, actuation limits, dynamics envelope). The paper states its key design contract as Proposition 1 (shield dominance): for any learned parameters, the executed action satisfies $C_{\mathrm{emb}} = 1$, so the enforced constraint set is invariant to the quality, calibration, or failure of the learned components. The authors are explicit about what this proposition does not say: it is a statement about which constraint set is enforced, not a certificate of physical safety, and it says nothing about hazards outside the embodiment model — unmodelled payloads, sensor blind spots, non-rigid obstacles are excluded by Assumption 1.

When the admissible set is empty, an ordered fallback ladder takes over: certified safe stop, reverse recovery to the last state with non-empty admissible set, replanning with enlarged candidate budget, and human escalation. The stopping-distance feasibility check ($c_6$) is the recursive-feasibility condition that makes this ladder well-founded; without it the robot could enter states admitting no safe continuation.

Method

The predictor rolls latents forward autoregressively under each candidate chunk, conditioned on an embodiment embedding via AdaLN/FiLM, with targets from an EMA target encoder under stop-gradient and no pixel decoder at any point. A head network emits task progress, collision, stuck, and failure probabilities plus a scalar uncertainty, trained with focal binary cross-entropy (chosen because false negatives dominate the cost structure), a calibration loss combining temperature scaling, ensemble/MC-dropout variance, and a conformal wrapper providing distribution-free miscoverage guarantees on the calibration split. Failure trajectories in training data are treated as mandatory rather than optional: with only successful demonstrations the risk heads have no positive class.

Two design decisions deserve emphasis. First, cross-embodiment learning is applied asymmetrically: encoder, predictor, and risk-head trunk are shared across platforms and conditioned on the embodiment embedding, but the shield is never shared — stopping distance, footprint, joint limits, and traversable slope are per-platform specifications, so onboarding a new robot is a specification-and-verification task, not a retraining task. Second, execution is partitioned edge–cloud asynchronously: the proposer, shield, and safe-stop controller run on-robot in hard real time, world-model rollouts run off-loop, and on link loss the system degrades to proposer-plus-shield — weaker in task performance but with an unchanged enforced constraint set. Compute per control step is fixed at $1 + K$ network evaluations, unlike iterative planning whose cost grows with optimisation rounds.

Evaluation protocol and results

The paper pre-registers a four-level protocol before execution and reports honestly which levels have been run. Level 2 (frozen-encoder risk prediction) and Level 4 (closed-loop) were executed on LIBERO/LIBERO-Long in simulation across 3 seeds; Level 3 (action-conditioned prediction and offline reranking) was not executed; no real-robot deployment was attempted. The primary safety statistic throughout is collision-prediction false-negative rate (FNR) at fixed recall of 0.90, chosen so all methods are compared at equal caution.

Predictive performance (Level 2). The full model achieves AUROC $0.88 \pm 0.01$ and FNR $0.14 \pm 0.02$, roughly halving the FNR of the V-JEPA-style feature baseline (0.21 → 0.14). The most instructive ablation concerns calibration: removing the calibration loss leaves AUROC (0.87 vs. 0.88) and Brier score (0.09 vs. 0.08) nearly unchanged but degrades ECE by roughly $3.5\times$ (0.14 vs. 0.04). This demonstrates that discriminative accuracy and calibration are dissociable — a model can rank collisions correctly while its probabilities are unusable as thresholds, precisely the failure mode a success-rate metric alone cannot detect.

Closed-loop performance (Level 4). On LIBERO-Long (600 episodes per configuration), the full framework reaches 62% success versus 55% for the model-based-shield-only baseline (+7 points, $p \approx 0.014$) and reduces collision rate to 3%. However, the margin over the reranking-only baseline (+3 points, 62% vs. 59%) is not statistically significant ($J_k$0), a limitation the authors report explicitly rather than round away: the central thesis requires the composition to beat both halves individually, and at this sample size it beats one conclusively and the other only directionally. The table also includes rejection and fallback columns specifically to prevent the degenerate reading where a frozen controller appears optimal — random candidate selection achieves a lower collision rate (7%) than reranking alone (9%) purely by passing through the shield, while succeeding only 31% of the time. Notably, the rule-based shield without ranking slightly hurts success (53% vs. 54% base), showing that veto-only caution costs task progress.

Deployment efficiency. Measured on a Jetson AGX Orin (on-robot) and RTX 4090 edge server: end-to-end decision latency p50/p99 of 590/870 ms on-robot and 165/310 ms edge-assisted, with the shield evaluating in 14 ms on-robot regardless of configuration — the only latency in the safety-critical path. The on-robot p50 corresponds to roughly 1.7 Hz, well below typical control rates; the asynchronous architecture accommodates this but does not remove the gap, and the edge path trades lower median latency for a network-driven tail (0.4% timeout frequency).

Limitations

The paper concedes several points plainly. The guarantee is conditional and narrow: it excludes unmodelled hazards entirely, and the JEPA contributes no guarantee whatsoever. System capability is bounded by the proposer's recall — the oracle row (74% success) quantifies the ceiling gap. Conservatism has a measurable cost in stuck and intervention rates, with no setting optimising all metrics simultaneously. Label quality on real corpora is heuristic (contact inferred from proprioception), bounding real-data risk-head accuracy. Most consequentially, calibration is split-dependent: deployment in new environments voids calibratedness, since conformal bounds hold only under exchangeability, which deployment routinely violates. The statistical weakness of the B5-vs-B4 comparison stands until either episode budgets increase or a paired-episode test replaces the unpaired comparison used here.

Conclusion

The paper's durable contribution is a formulation rather than a headline number: proposal, ranking, and admissibility are separated into three distinct stages, only the third carries a guarantee, and it carries one precisely because it contains no learned parameters. The empirical evidence supports the composition partially — significant gains over shield-only baselines, a demonstrated dissociation between discrimination and calibration, and measured deployment feasibility — while leaving open whether the learned ranking's contribution over reranking alone can be established at deployable sample sizes, and whether calibrated risk estimates survive the distribution shift inherent in real deployment.