Hybrid Neural World Models

Authors: Pranav Lakshmanan, Paras Chopra (Lossfunk) arXiv: 2605.28317 Submitted: 27 May 2026 Categories: cs.LG, cs.AI, math.NA, physics.comp-ph

Abstract

Neural surrogates promise large speedups over classical solvers for physical dynamics but fail silently at sharp dynamical events such as shocks, fronts, and contact. This paper presents hybrid neural world models for physical dynamics: a recipe for training multi-horizon surrogates in physical state space, where a single network with continuous horizon conditioning predicts any future state at horizon T in one forward pass. The trained surrogate implicitly encodes discontinuity locations, recoverable as a per-trajectory error map that concentrates on shocks, fronts, and contacts without any supervision for their location.

Key Contributions

- Multi-horizon shortcut surrogate architecture: A single network trained with direct supervision against textbook reference solvers, using continuous horizon conditioning to predict any future state at horizon T in a single forward pass.

  • Implicit error map: Although no part of the training data, loss function, or architecture supervises discontinuity location, the trained surrogate encodes it implicitly as a per-trajectory error map recoverable from forward passes alone.
  • Two operating modes: Mode 1 (surrogate alone) achieves 26x–72x CPU speedups against textbook solvers. Mode 2 (trust-aware fallback) uses the error map to gate a reference-solver fallback, roughly halving the surrogate's residual error.

    Method Details

    The architecture consists of:

    1. Multi-horizon shortcut surrogate: A single neural network trained with direct supervision against reference solvers. The network uses continuous horizon conditioning (rather than discrete time steps) to produce predictions at any arbitrary horizon T in one forward pass. Training targets ground-truth trajectories from classical solvers (e.g., ODE/PDE integrators).

    2. Error map extraction: The error map is derived from "step-doubling" — running the surrogate at both T and 2T, then comparing consistency. Regions of high inconsistency indicate shocks, fronts, or contacts. This requires no labeled discontinuities and no calibration set.

    3. Mode 1 — Surrogate alone: Maximum throughput mode using only the neural surrogate. Suitable for smooth dynamical regions.

    4. Mode 2 — Trust-aware fallback: The error map flags uncertain trajectories; uncertain regions are deferred to a classical reference solver. This provides rigorous accuracy guarantees at reduced computational cost.

    The method was evaluated across three physics environments: Oregonator (reaction-diffusion PDE), Euler 2D (compressible flow PDE), and Ball 3D (rigid-body collision ODE).

    Key Results

    | Environment | Mode 1 Speedup vs Classical Solver | Mode 2 RMSE Reduction | |-------------|-----------------------------------|----------------------| | Oregonator (reaction-diffusion) | 26x–72x CPU | ~50% reduction | | Euler 2D (compressible flow) | 26x–72x CPU | ~50% reduction | | Ball 3D (rigid-body collision) | 26x–72x CPU | ~50% reduction |

    The error map is competitive with or better than standard label-free baselines including deep ensembles, learned error heads, gradient-magnitude indicators, and locally-adaptive conformal prediction — while using only a single trained network.

    Limitations and Future Work

    - Solver-vectorisation caveat: The method assumes batched state vectorisation, limiting applicability to scenarios where many trajectories can be evaluated in parallel.

  • Per-environment training: Separate surrogate training is required for each physical environment; no zero-shot generalization across physics domains.
  • Trust-signal failure modes: The step-doubling error map may misfire in highly chaotic regimes where prediction divergence is not attributable to discontinuities.
  • Validation-set calibration: Mode 2's threshold for deferral requires a small calibration set, introducing minor overhead.
  • Future work: Extending to higher-dimensional PDEs, learning cross-domain surrogate composition, and removing the calibration-set requirement.

    Relevance to Patrick's Research

    This work directly contributes to the neural simulator / world model sub-area of world models research. The key insight — that a single-horizon-conditioned network implicitly learns to detect its own failure modes at discontinuities — is architecturally interesting and potentially applicable to other sequential prediction domains (video prediction, robotics simulation). The "Mode 2" hybrid approach (neural + classical fallback) is a practical pattern for deploying neural simulators in safety-critical physical reasoning tasks. The 26x–72x speedup numbers are compelling for real-time simulation use cases.