---
license: apache-2.0
task_categories:
- time-series-forecasting
- other
tags:
- physics
- pde
- ode
- simulation
- world-models
- neural-surrogate
- trust-signal
- reaction-diffusion
- euler-equations
- rigid-body
pretty_name: "Hybrid Neural World Models"
size_categories:
- 10K<n<100K
viewer: false
---

# 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 (Δt_save = 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/dataset_config.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:

| dim   | quantity              |
|-------|------------------------|
| 0–2   | position `(x, y, z)`   |
| 3–5   | linear velocity        |
| 6–8   | angular velocity       |

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

```python
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

| split    | Oregonator | Euler 2D | Ball 3D |
|----------|-----------:|---------:|--------:|
| train    | 1 200      | 800      | 1 000   |
| val      | 150        | 100      | 200     |
| test     | 150        | 100      | 200     |
| ood_near | 250        | 200      | 200     |
| ood_far  | 250        | 200      | 200     |

(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.

| split    | seed offset | seed range (Oregonator)  |
|----------|------------:|-------------------------:|
| train    | 0           | 0 ... 1 199              |
| val      | 100 000     | 100 000 ... 100 149      |
| test     | 200 000     | 200 000 ... 200 149      |
| ood_near | 300 000     | 300 000 ... 300 249      |
| ood_far  | 400 000     | 400 000 ... 400 249      |

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](https://x.com/lossfunk) — open an issue on the repo or DM on X.