WorldFlux¶

Unified Interface for World Models in Reinforcement Learning

One API. Multiple Architectures. Infinite Imagination.

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

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
DreamerV3	`size12m`, `size25m`, `size50m`, `size100m`, `size200m`
TD-MPC2	`5m`, `19m`, `48m`, `317m`
JEPA	`base`
Token	`base`
Diffusion	`base`

Documentation¶

- **Getting Started** --- Minimal onboarding for installation and first execution. [:octicons-arrow-right-24: Installation](getting-started/installation.md) [:octicons-arrow-right-24: Quick Start](getting-started/quickstart.md) [:octicons-arrow-right-24: CPU Success Path](getting-started/cpu-success.md) - **API Reference** --- Implementation-aligned API contracts and autogenerated symbols. [:octicons-arrow-right-24: Factory Functions](api/factory.md) [:octicons-arrow-right-24: WorldModel Base Class](api/protocol.md) [:octicons-arrow-right-24: Training](api/training.md) - **Reference** --- Operational documentation and project quality checks. [:octicons-arrow-right-24: Benchmarks](reference/benchmarks.md) [:octicons-arrow-right-24: Unified Comparison](reference/unified-comparison.md) [:octicons-arrow-right-24: Parity Harness](reference/parity.md) [:octicons-arrow-right-24: Documentation Stack](reference/docs-stack.md) [:octicons-arrow-right-24: Release Checklist](reference/release-checklist.md) [:octicons-arrow-right-24: Publishing](reference/publishing.md) [:octicons-arrow-right-24: Tutorial Placeholder Policy](reference/tutorial-policy.md) [:octicons-arrow-right-24: Troubleshooting](reference/troubleshooting.md) [:octicons-arrow-right-24: WASR Metrics](reference/wasr.md)

Architecture¶

graph LR
    subgraph Input
        A[Observation]
    end

    subgraph WorldModel["World Model"]
        B[Encoder]
        C[State]
        D[Dynamics]
        E[Decoder]
    end

    subgraph Output
        F[Predictions]
    end

    A --> B
    B --> C
    C --> D
    D --> C
    C --> E
    E --> F

    style C fill:#e1f5fe
    style D fill:#fff3e0

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.

License¶

Apache License 2.0 - see LICENSE and NOTICE for details.