WorldFlux Extensibility Guide¶

This document describes how to extend WorldFlux safely while keeping API contracts stable.

WorldFlux follows a contract-first approach: each model family must define runtime I/O contracts before wider integration.

Table of Contents¶

1. Implementation Verification
2. Extensibility Assessment
3. World Model Classification
4. Adding New Models
5. Future Architecture Roadmap
6. Change Type Analysis

Implementation Verification¶

Unified API Verification¶

WorldFlux provides a unified API through the WorldModel base class so all model families share a consistent interface:

class WorldModel(ABC):
    config: WorldModelConfig

    def encode(self, obs: Tensor | dict[str, Tensor] | WorldModelInput, deterministic: bool = False) -> State: ...
    def transition(self, state: State, action: ActionPayload | Tensor | None, conditions: ConditionPayload | None = None, deterministic: bool = False) -> State: ...
    def update(self, state: State, action: ActionPayload | Tensor | None, obs: Tensor | dict[str, Tensor] | WorldModelInput, conditions: ConditionPayload | None = None) -> State: ...
    def decode(self, state: State, conditions: ConditionPayload | None = None) -> ModelOutput: ...
    def rollout(self, initial_state: State, action_sequence: ActionSequence | Tensor | None, conditions: ConditionPayload | None = None, deterministic: bool = False, mode: str = "autoregressive") -> Trajectory: ...
    def initial_state(self, batch_size: int, device: ...) -> State: ...
    def loss(self, batch: Batch) -> LossOutput: ...

Design Patterns¶

1. Registry Pattern (src/worldflux/core/registry.py)

@WorldModelRegistry.register("dreamer", DreamerV3Config)
class DreamerV3WorldModel(nn.Module): ...

1. Universal State Representation (src/worldflux/core/state.py)

2. State supports model-specific tensor keys via tensors + meta.

3. Polymorphic Latent Spaces (src/worldflux/core/latent_space.py)

4. GaussianLatentSpace
5. CategoricalLatentSpace

6. SimNormLatentSpace

7. Unified Trainer (src/worldflux/training/trainer.py)

8. Common training flow for model families implementing the WorldModel contract.

Extensibility Assessment¶

Architectural Characteristics¶

Strengths¶

• Additive model-family integration via registry.
• Shared runtime contracts and payload types.
• Unified serialization path (save_pretrained / from_pretrained).
• External plugin discovery hooks.

Areas for Enhancement¶

• Additional support patterns for long-horizon video/sequence use cases.
• Broader reusable abstractions for iterative decoders/samplers.
• More examples for advanced external plugin packaging.

World Model Classification¶

Taxonomy of Architecture Families¶

Five-Layer Pluggable Core¶

WorldFlux standardizes model composition around replaceable components:

1. ObservationEncoder
2. DynamicsModel
3. ActionConditioner
4. Decoder (optional)
5. RolloutExecutor (open-loop execution)

Factory-level overrides are supported:

create_world_model(..., component_overrides={"action_conditioner": "my.plugin.component"})

Planner Metadata Contract¶

Planner outputs must include:

• extras["wf.planner.horizon"]

External Plugin Hooks¶

Third-party packages can register through entry-point groups:

• worldflux.models
• worldflux.components

Plugin manifests should provide compatibility metadata (plugin_api_version, worldflux_version_range, capabilities).

Minimal Installable Plugin Example¶

This repository includes a minimal plugin package:

• examples/plugins/minimal_plugin

Install and run smoke validation:

uv pip install -e examples/plugins/minimal_plugin
uv run python examples/plugins/smoke_minimal_plugin.py

The sample uses both entry-point groups:

• worldflux.models
• worldflux.components

It registers:

• model alias: minimalplugin-dreamer
• component id: minimal_plugin.zero_action_conditioner

Required Components by Family¶

Adding New Models¶

Step 1: Create Config Class¶

# src/worldflux/core/config.py
@dataclass
class MyModelConfig(WorldModelConfig):
    model_type: str = "mymodel"
    custom_param: int = 256

Step 2: Implement World Model¶

# src/worldflux/models/mymodel/world_model.py
@WorldModelRegistry.register("mymodel", MyModelConfig)
class MyWorldModel(nn.Module):
    def encode(self, obs: Tensor, deterministic: bool = False) -> State: ...
    def transition(self, state: State, action: Tensor, ...) -> State: ...
    def update(self, state: State, action: Tensor, obs: Tensor) -> State: ...
    def decode(self, state: State) -> ModelOutput: ...
    def rollout(self, initial_state: State, actions: Tensor, ...) -> Trajectory: ...
    def loss(self, batch: Batch) -> LossOutput: ...

Step 3: Export and Register¶

• Export symbols in src/worldflux/models/mymodel/__init__.py
• Ensure registry and factory aliases are configured where needed.

Step 4: Add Tests and Docs¶

• Add family tests under tests/test_models/.
• Document required batch/state keys and contract expectations.

Future Architecture Roadmap¶

The following areas are tracked as technical extension directions:

• Tokenization/VQ pathways for sequence-oriented models.
• Additional iterative sampler/decoder abstractions.
• Enhanced support for video-shaped tensors and temporal objectives.
• Integration patterns for large-scale external pretrained families.

These are implementation directions, not release commitments.

Change Type Analysis¶

Additive Changes (Backward Compatible)¶

Potentially Breaking Changes¶

Compatibility Guidelines¶

• Add new config fields with safe defaults.
• Keep optional extensions additive where possible.
• Provide explicit migration notes for contract-affecting changes.

Extension Points¶