PraLak/Hybrid_Neural_World_Models

Hybrid Neural World Models

Training, validation, test, and out-of-distribution (OOD) trajectories for the three physical systems used in Hybrid Neural World Models (Pranav Lakshmanan, Paras Chopra). The accompanying code, checkpoints, and paper define a single neural surrogate that predicts states at any horizon plus a step-doubling trust signal that flags when its forecasts can be trusted.

This repo contains the raw trajectory data only. Models / training code live separately.

What's inside

.
├── oregonator/                          reaction-diffusion PDE (Belousov-Zhabotinsky)
│   ├── oregonator_train.h5              1200 traj × 201 steps × 2 ch × 256 × 256
│   ├── oregonator_val.h5                 150 traj
│   ├── oregonator_test.h5                150 traj
│   ├── oregonator_ood_near.h5            250 traj   (mild parameter shift)
│   ├── oregonator_ood_far.h5             250 traj   (stronger parameter shift)
│   └── dataset_config.json
├── euler2d/                             compressible flow PDE
│   ├── euler2d_v2_train.h5
│   ├── euler2d_v2_val.h5
│   ├── euler2d_v2_test.h5
│   ├── euler2d_v2_ood_near.h5
│   └── euler2d_v2_ood_far.h5
└── ball3d/                              rigid-body bouncing ODE (MuJoCo)
    ├── ball3d_train.h5                  1000 traj × 101 steps × 9 dims
    ├── ball3d_val.h5
    ├── ball3d_test.h5
    ├── ball3d_ood_near.h5
    └── ball3d_ood_far.h5

Environments

1. Reaction-diffusion fronts (Oregonator / BZ)

Two-variable Tyson Oregonator in non-dimensional units (D = 1, ε = 0.05, q = 0.002, f = 2.0) integrated on a 256 × 256 periodic grid with explicit-FV spatial stencils and Strang splitting. Each trajectory captures 201 saved frames (Δtsave = 0.05, total time = 10.0) of the two chemical concentrations (`u`, `v`). Initial conditions are a mix of spiral, target-pattern, and random superpositions — see `oregonator/datasetconfig.json` for ratios and seed offsets.

State shape per saved frame: (2, 256, 256) float32.

2. Compressible gas flow (Euler 2D)

2D compressible Euler equations solved with an MUSCL-Hancock + HLLC scheme on a 128 × 128 grid. Initial conditions cover Schulz-Rinne Riemann configurations and shock-interaction setups. Each trajectory: 100 saved frames of (ρ, ρvx, ρvy, E) — i.e. four conservative fields stored flattened to 16 384 per timestep (reshape to (128, 128, 4) if you want per-channel imagery).

3. Rigid-body bouncing (Ball 3D)

MuJoCo simulation of a single ball bouncing inside a 1 × 1 × 0.6 m box with elastic-ish contact. State is a 9-vector per frame:

Trajectories: 101 frames at fixed Δt.

OOD splits

• —*_ood_near.h5 — parameters drawn from a moderately shifted distribution relative to train. Used to test robustness in the paper's "OOD-near" cells.
• —*_ood_far.h5 — stronger shift. Used in "OOD-far" cells. Errors are larger; the trust signal stays informative (AUROC ≥ 0.65) and is the main argument for the trust-gated fallback (Mode 2).

Quick start

import h5py

with h5py.File("oregonator/oregonator_test.h5", "r") as f:
    states  = f["states"]       # (N, T, 2, 256, 256), float32
    params  = f["params"]       # (N, 4)              IC + physics params
    seeds   = f["seeds"]        # (N,)
    ic_types= f["ic_types"]     # (N,)                spiral / target / random
    print(states.shape, states.dtype)
    snapshot_u = states[0, 40, 0]  # u-concentration of traj 0 at t=40

with h5py.File("ball3d/ball3d_test.h5", "r") as f:
    states = f["states"][:]     # (200, 101, 9)

Euler trajectories live under per-trajectory groups inside the file; iterate keys to enumerate.

Split sizes

(Counts are number of independent trajectories. Each trajectory holds 100-201 saved frames depending on the environment.)

Reproducing splits

Splits do not share initial conditions: each split is generated from a disjoint range of integer seeds. The seed offset below is the first seed used for that split; consecutive trajectories use offset, offset+1, offset+2, ... — so train uses seeds 0 through 1 199, val uses 100 000 through 100 149, etc. This guarantees a regenerable, leak-free split assignment.

See oregonator/dataset_config.json for full per-env metadata (grid sizes, integration step, physics parameters, IC mixture).

Intended use

• —Benchmarking neural surrogates that predict any horizon in a single forward pass.
• —Studying label-free uncertainty / trust-signal methods for physics models.
• —Cross-system generalization studies (one set of trajectories per env, same splits across all three).

Citation

If you use these datasets please cite:

@misc{hybrid_neural_world_models_2026,
  title   = {Hybrid Neural World Models for Physical Dynamics},
  author  = {Lakshmanan, Pranav and Chopra, Paras},
  year    = {2026},
}

Contact

lossfunk — open an issue on the repo or DM on X.