WorldFlux
Unified Interface for World Models in Reinforcement Learning
One API. Multiple Architectures. Clear Contracts.
WorldFlux provides a unified Python interface for world models used in reinforcement learning.
Features
- Unified API: Common interface across model families
- Simple Usage: One-liner model creation with
create_world_model() - Training Infrastructure: Training loop with callbacks, checkpointing, and logging
- Type Safe: Full type annotations and mypy compatibility
Quick Start
uv sync --extra dev
uv run python examples/quickstart_cpu_success.py --quick
Official CPU-success docs: CPU Success Path
This landing page highlights the public onboarding surface. Start with the CPU smoke path before treating any workflow as a benchmark or proof claim.
Unified Comparison Demo
uv sync --extra dev --extra training
uv run python examples/compare_unified_training.py --quick
This canonical public demo shows DreamerV3 and TD-MPC2 under the same unified
API, the same TrainingConfig contract, and the same quick verification flow.
It writes summary.json, per-family imagination artifacts, and per-family
quick_verify.json outputs.
API Glimpse
from worldflux import create_world_model
import torch
model = create_world_model(
"dreamerv3:size12m",
obs_shape=(3, 64, 64),
action_dim=4,
)
obs = torch.randn(1, 3, 64, 64)
state = model.encode(obs)
actions = torch.randn(15, 1, 4)
trajectory = model.rollout(state, actions)
print(trajectory.rewards.shape)
Available Models
| Family | Presets | Status |
|---|---|---|
| DreamerV3 | size12m, size25m, size50m, size100m, size200m | Reference-family |
| TD-MPC2 | 5m, 19m, 48m, 317m | Reference-family |
This table highlights the supported MVP presets. For the public catalog including advanced proof-oriented presets, run:
worldflux models list --surface public --verbose
Reference-family models map to maintained upstream families and internal proof-mode parity workflows. Public proof claims require published evidence bundles.
Experimental and skeleton families remain available behind explicit opt-in:
worldflux models list --surface all --verbose
Documentation
Getting Started
Minimal onboarding for installation and first execution.
- Installation
- Quick Start
- Train Your First Model
- DreamerV3 vs TD-MPC2
- Reproduce DreamerV3 and TD-MPC2 Locally
- Model Choice in
worldflux init - CPU Success Path
API Reference
Implementation-aligned API contracts and autogenerated symbols.
Reference
Operational documentation and project quality checks.
- Benchmarks
- Observation Shape and Action Dim
- Unified Comparison
- Parity Harness
- Documentation Stack
- Release Checklist
- Publishing
- Tutorial Rollout Policy
- Troubleshooting
- WASR Metrics
Architecture
Installation
uv tool install worldflux
worldflux init my-world-model
Try It Now
The fastest way to get started is our interactive Colab notebook.
Contributing
Contributions are welcome. See our Contributing Guide.