From 259911da439dd6328c4b040387ac281daa006b26 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 4 Aug 2026 15:02:08 -0700 Subject: [PATCH 1/5] Add world model framework and ORT runner Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 41fde59e-190b-4448-b949-6a86bf9f57a6 --- README.md | 18 + docs/index.md | 1 + docs/world_models.md | 101 ++++++ pyproject.toml | 1 + src/mobius/__init__.py | 7 +- src/mobius/_configs/__init__.py | 2 + src/mobius/_configs/_world_model.py | 64 ++++ .../integrations/onnxruntime/__init__.py | 16 + .../integrations/onnxruntime/world_model.py | 337 ++++++++++++++++++ .../onnxruntime/world_model_test.py | 138 +++++++ src/mobius/models/__init__.py | 2 + src/mobius/models/world_model.py | 87 +++++ src/mobius/models/world_model_test.py | 196 ++++++++++ src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_world_model.py | 82 +++++ 15 files changed, 1054 insertions(+), 1 deletion(-) create mode 100644 docs/world_models.md create mode 100644 src/mobius/_configs/_world_model.py create mode 100644 src/mobius/integrations/onnxruntime/__init__.py create mode 100644 src/mobius/integrations/onnxruntime/world_model.py create mode 100644 src/mobius/integrations/onnxruntime/world_model_test.py create mode 100644 src/mobius/models/world_model.py create mode 100644 src/mobius/models/world_model_test.py create mode 100644 src/mobius/tasks/_world_model.py diff --git a/README.md b/README.md index f9f2320b6..c7217021e 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ multi-component export for pipelines. | **Audio** | Wav2Vec2, HuBERT, WavLM, SpeechT5 | | **Vision** | ViT, BEiT, DeiT, DINOv2, Swin, CLIP, SigLIP | | **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage, HunyuanDiT, CogVideoX | +| **World Models** | Directly declared state-transition graphs with stateful ONNX Runtime rollout | | **Adapters** | T2I-Adapter, IP-Adapter | Supports **290+ Transformers model types** and **10 Diffusers component types** @@ -97,6 +98,23 @@ pkg = build("meta-llama/Llama-3.2-1B", See the [EP quickstart](docs/ep_quickstart.md) and [full EP reference](docs/execution_providers.md) for all supported EPs and options. +**World models** use a directly declared single-step state-transition graph +and a stateful ONNX Runtime runner: + +```python +from mobius import MLPWorldModel, WorldModelConfig, build_from_module + +config = WorldModelConfig( + observation_shape=(64,), + action_shape=(6,), + state_shape=(128,), +) +pkg = build_from_module(MLPWorldModel(config), config, task="world-model") +``` + +See the [world-model guide](docs/world_models.md) for weight loading and +ONNX Runtime rollout. + ### CLI ```sh diff --git a/docs/index.md b/docs/index.md index 103db3353..2181f58c6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -11,6 +11,7 @@ Build ONNX models directly from HuggingFace model IDs with automatic weight down :caption: User Guide getting-started +world_models cli_reference module-architecture model-catalog diff --git a/docs/world_models.md b/docs/world_models.md new file mode 100644 index 000000000..394b84e27 --- /dev/null +++ b/docs/world_models.md @@ -0,0 +1,101 @@ +# World models + +Mobius can build world models directly as ONNX graphs without tracing a +PyTorch `forward()` method. The initial framework defines a deterministic, +single-step state-transition contract: + +```text +(observation, action, state) + -> (next_state, observation_prediction, reward, continuation) +``` + +All tensors have a dynamic leading batch dimension. `reward` and +`continuation` have shape `[batch, 1]`; `continuation` is a probability. + +## Build the reference model + +`MLPWorldModel` is a minimal directly declared implementation. Its input shapes +exclude the batch dimension: + +```python +from safetensors.torch import load_file + +from mobius import MLPWorldModel, WorldModelConfig, build_from_module + +config = WorldModelConfig( + observation_shape=(64,), + action_shape=(6,), + state_shape=(128,), + hidden_size=512, + num_hidden_layers=3, +) +module = MLPWorldModel(config) +package = build_from_module(module, config, task="world-model") + +weights = load_file("world_model.safetensors") +package.apply_weights(module.preprocess_weights(weights)) +package.save("world-model-onnx") +``` + +The ONNX graph is constructed through `onnxscript.nn.Module` and +`onnx_ir.GraphBuilder`; PyTorch is used only as a possible source of weight +tensors. + +## Run with ONNX Runtime + +Install the optional runtime dependency: + +```bash +pip install -e ".[runtime]" +``` + +`WorldModelRunner` validates the graph contract and retains `next_state` +between calls: + +```python +import numpy as np + +from mobius.integrations.onnxruntime import WorldModelRunner + +runner = WorldModelRunner.from_path( + "world-model-onnx/model.onnx", + providers=["CPUExecutionProvider"], +) + +observation = np.zeros((1, 64), dtype=np.float32) +action = np.zeros((1, 6), dtype=np.float32) + +# The first call creates a zero state. Later calls reuse next_state. +result = runner.step(observation, action) +result = runner.step(observation, action) + +trajectory = runner.rollout( + [observation, observation], + [action, action], +) +``` + +Pass `state=` to `step()` or `initial_state=` to `rollout()` when the model +uses a learned or externally sampled initial state. Each `rollout()` starts +from zero state unless `initial_state` is provided; repeated `step()` calls +continue from the runner's retained state. + +## Implement a custom world model + +A custom module must use the task's forward and return contracts: + +```python +from onnxscript import nn + + +class MyWorldModel(nn.Module): + def forward(self, op, observation, action, state): + # Directly declare ONNX operations here. + ... + return next_state, observation_prediction, reward, continuation +``` + +Use a custom `ModelTask` when an architecture needs multiple recurrent-state +tensors, stochastic latent outputs, separate observe/imagine graphs, or a +different prediction contract. Keep sampling outside the graph unless it must +be part of the deployed model. diff --git a/pyproject.toml b/pyproject.toml index cb4bbfea7..5d588ed48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ testing = [ "jsonschema>=4.0", ] ort-genai = ["onnxruntime-genai", "Pillow"] +runtime = ["onnxruntime"] docs = ["furo", "myst-parser", "sphinx", "sphinx_copybutton"] [tool.setuptools.packages.find] diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 3799ea335..0962f505a 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -22,6 +22,7 @@ "ModelRegistration", "ModelRegistry", "ModelTask", + "MLPWorldModel", "MMSConfig", "OPSET_VERSION", "Sam2Config", @@ -29,6 +30,8 @@ "VisionConfig", "VisionLanguageConfig", "WhisperConfig", + "WorldModelConfig", + "WorldModelTask", "YolosConfig", "apply_weights", "build", @@ -76,6 +79,7 @@ VisionConfig, VisionLanguageConfig, WhisperConfig, + WorldModelConfig, YolosConfig, ) from mobius._constants import OPSET_VERSION @@ -91,4 +95,5 @@ from mobius._weight_loading import apply_weights from mobius.integrations.gguf import build_from_gguf from mobius.integrations.nemo import build_from_nemo -from mobius.tasks import CausalLMTask, ModelTask +from mobius.models import MLPWorldModel +from mobius.tasks import CausalLMTask, ModelTask, WorldModelTask diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 68ea2b426..afb241a16 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -79,6 +79,7 @@ TTSConfig, VisionConfig, ) +from mobius._configs._world_model import WorldModelConfig __all__ = [ "DEFAULT_INT", @@ -117,6 +118,7 @@ "VisionConfig", "VisionLanguageConfig", "WhisperConfig", + "WorldModelConfig", "YolosConfig", "Zamba2Config", "_as_int", diff --git a/src/mobius/_configs/_world_model.py b/src/mobius/_configs/_world_model.py new file mode 100644 index 000000000..17eea7d2e --- /dev/null +++ b/src/mobius/_configs/_world_model.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Configuration for directly declared world models.""" + +from __future__ import annotations + +import dataclasses +import math + +from mobius._configs._base import BaseModelConfig + + +@dataclasses.dataclass +class WorldModelConfig(BaseModelConfig): + """Configuration shared by single-step world-model graphs. + + The three shapes exclude the leading batch dimension. The default + :class:`~mobius.models.MLPWorldModel` flattens each value internally, while + custom modules may preserve their original ranks. + """ + + observation_shape: tuple[int, ...] = (1,) + action_shape: tuple[int, ...] = (1,) + state_shape: tuple[int, ...] = (1,) + hidden_size: int = 128 + num_hidden_layers: int = 2 + hidden_act: str | None = "silu" + residual_state: bool = True + + @property + def observation_size(self) -> int: + """Flattened observation size.""" + return math.prod(self.observation_shape) + + @property + def action_size(self) -> int: + """Flattened action size.""" + return math.prod(self.action_shape) + + @property + def state_size(self) -> int: + """Flattened recurrent-state size.""" + return math.prod(self.state_shape) + + def validate(self) -> None: + """Validate dimensions required by the world-model task and reference model.""" + for name, shape in ( + ("observation_shape", self.observation_shape), + ("action_shape", self.action_shape), + ("state_shape", self.state_shape), + ): + if not shape: + raise ValueError(f"{name} must contain at least one dimension") + if any( + not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 for dim in shape + ): + raise ValueError(f"{name} must contain only positive integer dimensions") + if self.hidden_size <= 0: + raise ValueError("hidden_size must be positive") + if self.num_hidden_layers <= 0: + raise ValueError("num_hidden_layers must be positive") + if self.hidden_act is None: + raise ValueError("hidden_act must be set") diff --git a/src/mobius/integrations/onnxruntime/__init__.py b/src/mobius/integrations/onnxruntime/__init__.py new file mode 100644 index 000000000..f1bc37672 --- /dev/null +++ b/src/mobius/integrations/onnxruntime/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""ONNX Runtime helpers for Mobius models.""" + +from mobius.integrations.onnxruntime.world_model import ( + WorldModelRunner, + WorldModelSession, + WorldModelStepOutput, +) + +__all__ = [ + "WorldModelRunner", + "WorldModelSession", + "WorldModelStepOutput", +] diff --git a/src/mobius/integrations/onnxruntime/world_model.py b/src/mobius/integrations/onnxruntime/world_model.py new file mode 100644 index 000000000..c9e2dc746 --- /dev/null +++ b/src/mobius/integrations/onnxruntime/world_model.py @@ -0,0 +1,337 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Stateful ONNX Runtime execution for the world-model task contract.""" + +from __future__ import annotations + +import os +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, Protocol + +import ml_dtypes +import numpy as np +from numpy.typing import ArrayLike + +from mobius.tasks._world_model import WorldModelTask + + +class _NodeMetadata(Protocol): + @property + def name(self) -> str: ... + + @property + def shape(self) -> Sequence[int | str | None]: ... + + @property + def type(self) -> str: ... + + +class WorldModelSession(Protocol): + """Minimal inference-session interface required by :class:`WorldModelRunner`.""" + + def get_inputs(self) -> Sequence[_NodeMetadata]: ... + + def get_outputs(self) -> Sequence[_NodeMetadata]: ... + + def run( + self, + output_names: Sequence[str], + input_feed: dict[str, np.ndarray], + ) -> Sequence[Any]: ... + + +@dataclass(frozen=True) +class WorldModelStepOutput: + """Outputs from one world-model transition.""" + + next_state: np.ndarray + observation_prediction: np.ndarray + reward: np.ndarray + continuation: np.ndarray + + +_ORT_TYPE_TO_NUMPY: dict[str, np.dtype] = { + "tensor(float)": np.dtype(np.float32), + "tensor(float16)": np.dtype(np.float16), + "tensor(double)": np.dtype(np.float64), + "tensor(bfloat16)": np.dtype(ml_dtypes.bfloat16), + "tensor(int64)": np.dtype(np.int64), + "tensor(int32)": np.dtype(np.int32), + "tensor(int16)": np.dtype(np.int16), + "tensor(int8)": np.dtype(np.int8), + "tensor(uint64)": np.dtype(np.uint64), + "tensor(uint32)": np.dtype(np.uint32), + "tensor(uint16)": np.dtype(np.uint16), + "tensor(uint8)": np.dtype(np.uint8), + "tensor(bool)": np.dtype(np.bool_), +} + + +class WorldModelRunner: + """Run a Mobius world-model graph while preserving recurrent state.""" + + def __init__(self, session: WorldModelSession): + self._session = session + self._inputs = {node.name: node for node in session.get_inputs()} + self._outputs = {node.name: node for node in session.get_outputs()} + self._validate_contract() + self._state: np.ndarray | None = None + + @classmethod + def from_path( + cls, + model_path: str | os.PathLike[str], + *, + providers: Sequence[str] | None = None, + session_options: object | None = None, + ) -> WorldModelRunner: + """Create a runner from an ONNX file without making ORT a core dependency.""" + try: + import onnxruntime as ort + except ModuleNotFoundError as exc: + if exc.name != "onnxruntime": + raise + raise ImportError( + "WorldModelRunner.from_path() requires onnxruntime; " + "install mobius-onnx[runtime]" + ) from exc + + kwargs: dict[str, object] = {} + if providers is not None: + kwargs["providers"] = list(providers) + if session_options is not None: + kwargs["sess_options"] = session_options + return cls(ort.InferenceSession(os.fspath(model_path), **kwargs)) + + @property + def session(self) -> WorldModelSession: + """Underlying inference session.""" + return self._session + + @property + def state(self) -> np.ndarray | None: + """Recurrent state used by the next step.""" + return self._state + + def reset( + self, + state: ArrayLike | None = None, + *, + batch_size: int | None = None, + ) -> np.ndarray: + """Reset recurrent state to an explicit value or a zero tensor.""" + if state is not None: + prepared = self._prepare_input("state", state) + if batch_size is not None and prepared.shape[0] != batch_size: + raise ValueError( + f"state batch dimension is {prepared.shape[0]}, expected {batch_size}" + ) + else: + prepared = self._make_zero_state(batch_size or 1) + self._state = prepared + return prepared + + def step( + self, + observation: ArrayLike, + action: ArrayLike, + *, + state: ArrayLike | None = None, + ) -> WorldModelStepOutput: + """Execute one transition and retain its ``next_state``.""" + observation_array = self._prepare_input("observation", observation) + action_array = self._prepare_input("action", action) + batch_size = observation_array.shape[0] + if action_array.shape[0] != batch_size: + raise ValueError( + "observation and action batch dimensions differ: " + f"{batch_size} != {action_array.shape[0]}" + ) + + if state is not None: + current_state = self._prepare_input("state", state) + elif self._state is not None: + current_state = self._state + else: + current_state = self._make_zero_state(batch_size) + if current_state.shape[0] != batch_size: + raise ValueError( + "observation and state batch dimensions differ: " + f"{batch_size} != {current_state.shape[0]}" + ) + + values = self._session.run( + list(WorldModelTask.output_names), + { + "observation": observation_array, + "action": action_array, + "state": current_state, + }, + ) + if len(values) != len(WorldModelTask.output_names): + raise RuntimeError( + f"session returned {len(values)} outputs, " + f"expected {len(WorldModelTask.output_names)}" + ) + + arrays = { + name: self._validate_output(name, value) + for name, value in zip(WorldModelTask.output_names, values, strict=True) + } + if arrays["next_state"].shape != current_state.shape: + raise ValueError( + "next_state shape is incompatible with recurrent state: " + f"{arrays['next_state'].shape} != {current_state.shape}" + ) + if arrays["observation_prediction"].shape != observation_array.shape: + raise ValueError( + "observation_prediction shape is incompatible with observation: " + f"{arrays['observation_prediction'].shape} != {observation_array.shape}" + ) + expected_scalar_shape = (batch_size, 1) + for name in ("reward", "continuation"): + if arrays[name].shape != expected_scalar_shape: + raise ValueError( + f"{name} has shape {arrays[name].shape}, expected {expected_scalar_shape}" + ) + output = WorldModelStepOutput( + next_state=arrays["next_state"], + observation_prediction=arrays["observation_prediction"], + reward=arrays["reward"], + continuation=arrays["continuation"], + ) + self._state = output.next_state + return output + + def rollout( + self, + observations: Sequence[ArrayLike], + actions: Sequence[ArrayLike], + *, + initial_state: ArrayLike | None = None, + ) -> tuple[WorldModelStepOutput, ...]: + """Run a sequence of transitions along the leading time dimension.""" + if len(observations) != len(actions): + raise ValueError("observations and actions must contain the same number of steps") + if len(observations) == 0: + raise ValueError("rollout requires at least one step") + if initial_state is not None: + self.reset(initial_state) + else: + first_observation = self._prepare_input("observation", observations[0]) + self.reset(batch_size=first_observation.shape[0]) + + return tuple( + self.step(observation, action) + for observation, action in zip(observations, actions, strict=True) + ) + + def _validate_contract(self) -> None: + expected_inputs = set(WorldModelTask.input_names) + expected_outputs = set(WorldModelTask.output_names) + actual_inputs = set(self._inputs) + actual_outputs = set(self._outputs) + if actual_inputs != expected_inputs: + raise ValueError( + "world-model input contract mismatch: " + f"expected {sorted(expected_inputs)}, got {sorted(actual_inputs)}" + ) + if actual_outputs != expected_outputs: + raise ValueError( + "world-model output contract mismatch: " + f"expected {sorted(expected_outputs)}, got {sorted(actual_outputs)}" + ) + state_input = self._inputs["state"] + next_state_output = self._outputs["next_state"] + if state_input.type != next_state_output.type: + raise ValueError( + "next_state dtype must match state dtype: " + f"{next_state_output.type!r} != {state_input.type!r}" + ) + if not self._shapes_are_recurrently_compatible( + state_input.shape, + next_state_output.shape, + ): + raise ValueError( + "next_state shape must be compatible with state shape: " + f"{tuple(next_state_output.shape)} != {tuple(state_input.shape)}" + ) + + def _prepare_input(self, name: str, value: ArrayLike) -> np.ndarray: + node = self._inputs[name] + array = np.ascontiguousarray(value, dtype=self._numpy_dtype(node)) + self._validate_shape(name, array, node) + return array + + def _validate_output(self, name: str, value: Any) -> np.ndarray: + node = self._outputs[name] + array = np.asarray(value) + expected_dtype = self._numpy_dtype(node) + if array.dtype != expected_dtype: + raise TypeError(f"{name} has dtype {array.dtype}, expected {expected_dtype}") + self._validate_shape(name, array, node) + return array + + def _make_zero_state(self, batch_size: int) -> np.ndarray: + if batch_size <= 0: + raise ValueError("batch_size must be positive") + node = self._inputs["state"] + shape: list[int] = [] + for axis, dim in enumerate(node.shape): + if axis == 0: + if isinstance(dim, int) and dim != batch_size: + raise ValueError( + f"state model batch dimension is fixed at {dim}, got {batch_size}" + ) + shape.append(batch_size) + elif isinstance(dim, int): + shape.append(dim) + else: + raise ValueError( + "cannot create zero state because a non-batch state dimension " + f"is dynamic ({dim!r}); pass an explicit state to reset() or step()" + ) + return np.zeros(shape, dtype=self._numpy_dtype(node)) + + @staticmethod + def _numpy_dtype(node: _NodeMetadata) -> np.dtype: + try: + return _ORT_TYPE_TO_NUMPY[node.type] + except KeyError as exc: + raise TypeError(f"unsupported ONNX Runtime tensor type {node.type!r}") from exc + + @staticmethod + def _validate_shape(name: str, array: np.ndarray, node: _NodeMetadata) -> None: + expected = tuple(node.shape) + if array.ndim != len(expected): + raise ValueError( + f"{name} has rank {array.ndim}, expected rank {len(expected)} " + f"with shape {expected}" + ) + for axis, (actual_dim, expected_dim) in enumerate( + zip(array.shape, expected, strict=True) + ): + if isinstance(expected_dim, int) and actual_dim != expected_dim: + raise ValueError( + f"{name} dimension {axis} is {actual_dim}, expected {expected_dim}" + ) + + @staticmethod + def _shapes_are_recurrently_compatible( + state_shape: Sequence[int | str | None], + next_state_shape: Sequence[int | str | None], + ) -> bool: + if len(state_shape) != len(next_state_shape): + return False + return all( + state_dim == next_state_dim + if isinstance(state_dim, int) and isinstance(next_state_dim, int) + else True + for state_dim, next_state_dim in zip( + state_shape, + next_state_shape, + strict=True, + ) + ) diff --git a/src/mobius/integrations/onnxruntime/world_model_test.py b/src/mobius/integrations/onnxruntime/world_model_test.py new file mode 100644 index 000000000..ad2e7df95 --- /dev/null +++ b/src/mobius/integrations/onnxruntime/world_model_test.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import dataclasses + +import numpy as np +import pytest + +from mobius.integrations.onnxruntime import WorldModelRunner + + +@dataclasses.dataclass +class _Node: + name: str + shape: list[int | str | None] + type: str = "tensor(float)" + + +class _FakeSession: + def __init__(self): + self.inputs = [ + _Node("observation", ["batch", 4]), + _Node("action", ["batch", 2]), + _Node("state", ["batch", 3]), + ] + self.outputs = [ + _Node("next_state", ["batch", 3]), + _Node("observation_prediction", ["batch", 4]), + _Node("reward", ["batch", 1]), + _Node("continuation", ["batch", 1]), + ] + self.last_feed = None + + def get_inputs(self): + return self.inputs + + def get_outputs(self): + return self.outputs + + def run(self, output_names, input_feed): + self.last_feed = input_feed + batch = input_feed["observation"].shape[0] + return [ + input_feed["state"] + 1.0, + input_feed["observation"] * 2.0, + np.zeros((batch, 1), dtype=np.float32), + np.ones((batch, 1), dtype=np.float32), + ] + + +def test_step_initializes_and_preserves_state(): + session = _FakeSession() + runner = WorldModelRunner(session) + observation = np.ones((2, 4), dtype=np.float64) + action = np.ones((2, 2), dtype=np.float64) + + first = runner.step(observation, action) + assert session.last_feed["observation"].dtype == np.float32 + np.testing.assert_array_equal( + session.last_feed["state"], + np.zeros((2, 3), dtype=np.float32), + ) + np.testing.assert_array_equal(first.next_state, np.ones((2, 3), dtype=np.float32)) + + second = runner.step(observation, action) + np.testing.assert_array_equal( + session.last_feed["state"], + first.next_state, + ) + np.testing.assert_array_equal(second.next_state, np.full((2, 3), 2.0, np.float32)) + + +def test_rollout_uses_initial_state(): + runner = WorldModelRunner(_FakeSession()) + observations = np.ones((3, 1, 4), dtype=np.float32) + actions = np.ones((3, 1, 2), dtype=np.float32) + initial_state = np.full((1, 3), 5.0, dtype=np.float32) + + outputs = runner.rollout( + observations, + actions, + initial_state=initial_state, + ) + + assert len(outputs) == 3 + np.testing.assert_array_equal(outputs[-1].next_state, np.full((1, 3), 8.0)) + assert runner.state is outputs[-1].next_state + + +def test_rollout_defaults_to_fresh_zero_state(): + runner = WorldModelRunner(_FakeSession()) + runner.reset(np.full((1, 3), 10.0, dtype=np.float32)) + + outputs = runner.rollout( + np.ones((2, 1, 4), dtype=np.float32), + np.ones((2, 1, 2), dtype=np.float32), + ) + + np.testing.assert_array_equal(outputs[-1].next_state, np.full((1, 3), 2.0)) + + +def test_rejects_contract_mismatch(): + session = _FakeSession() + session.inputs.pop() + + with pytest.raises(ValueError, match="input contract mismatch"): + WorldModelRunner(session) + + +def test_rejects_incompatible_recurrent_state_contract(): + session = _FakeSession() + session.outputs[0].type = "tensor(double)" + + with pytest.raises(ValueError, match="next_state dtype must match"): + WorldModelRunner(session) + + +def test_rejects_mismatched_batches_before_inference(): + session = _FakeSession() + runner = WorldModelRunner(session) + + with pytest.raises(ValueError, match="batch dimensions differ"): + runner.step( + np.zeros((2, 4), dtype=np.float32), + np.zeros((1, 2), dtype=np.float32), + ) + assert session.last_feed is None + + +def test_reset_requires_explicit_dynamic_state_tail(): + session = _FakeSession() + session.inputs[-1].shape = ["batch", "state_size"] + runner = WorldModelRunner(session) + + with pytest.raises(ValueError, match="non-batch state dimension is dynamic"): + runner.reset(batch_size=2) diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 7d7377f91..9ca9e2ce3 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -149,6 +149,7 @@ "Wav2Vec2ForCTCModel", "Wav2Vec2Model", "WhisperForConditionalGeneration", + "MLPWorldModel", "XLMCausalLMModel", "Zamba2CausalLMModel", "mimi_default_config", @@ -309,5 +310,6 @@ from mobius.models.wav2vec2 import Wav2Vec2Model from mobius.models.wav2vec2_ctc import Wav2Vec2ForCTCModel from mobius.models.whisper import WhisperForConditionalGeneration +from mobius.models.world_model import MLPWorldModel from mobius.models.xlm import XLMCausalLMModel from mobius.models.zamba2 import Zamba2CausalLMModel diff --git a/src/mobius/models/world_model.py b/src/mobius/models/world_model.py new file mode 100644 index 000000000..3b7dd80e3 --- /dev/null +++ b/src/mobius/models/world_model.py @@ -0,0 +1,87 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Minimal directly declared world-model implementation.""" + +from __future__ import annotations + +import torch +from onnxscript import OpBuilder, nn + +from mobius._configs import WorldModelConfig +from mobius.components import Linear, get_activation + + +class MLPWorldModel(nn.Module): + """Deterministic MLP reference model for the world-model task contract.""" + + default_task = "world-model" + config_class = WorldModelConfig + category = "World Model" + + def __init__(self, config: WorldModelConfig): + super().__init__() + config.validate() + self.config = config + input_size = config.observation_size + config.action_size + config.state_size + self.input_layer = Linear(input_size, config.hidden_size) + self.hidden_layers = nn.ModuleList( + [ + Linear(config.hidden_size, config.hidden_size) + for _ in range(config.num_hidden_layers - 1) + ] + ) + self.state_head = Linear(config.hidden_size, config.state_size) + self.observation_head = Linear(config.hidden_size, config.observation_size) + self.reward_head = Linear(config.hidden_size, 1) + self.continuation_head = Linear(config.hidden_size, 1) + self._activation = get_activation(config.hidden_act) + + def forward(self, op: OpBuilder, observation, action, state): + observation_flat = op.Flatten(observation, axis=1) + action_flat = op.Flatten(action, axis=1) + state_flat = op.Flatten(state, axis=1) + + hidden = self._activation( + op, + self.input_layer( + op, + op.Concat(observation_flat, action_flat, state_flat, axis=1), + ), + ) + for layer in self.hidden_layers: + hidden = self._activation(op, layer(op, hidden)) + + next_state_flat = self.state_head(op, hidden) + if self.config.residual_state: + next_state_flat = op.Add(state_flat, next_state_flat) + + observation_prediction_flat = self.observation_head(op, hidden) + reward = self.reward_head(op, hidden) + continuation = op.Sigmoid(self.continuation_head(op, hidden)) + + next_state = self._reshape_batch( + op, + next_state_flat, + state, + self.config.state_shape, + ) + observation_prediction = self._reshape_batch( + op, + observation_prediction_flat, + observation, + self.config.observation_shape, + ) + return next_state, observation_prediction, reward, continuation + + @staticmethod + def _reshape_batch(op: OpBuilder, value, batch_source, shape: tuple[int, ...]): + batch = op.Shape(batch_source, start=0, end=1) + return op.Reshape(value, op.Concat(batch, list(shape), axis=0)) + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Return weights unchanged; provided for parity with other Mobius models.""" + return state_dict diff --git a/src/mobius/models/world_model_test.py b/src/mobius/models/world_model_test.py new file mode 100644 index 000000000..a42325608 --- /dev/null +++ b/src/mobius/models/world_model_test.py @@ -0,0 +1,196 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import numpy as np +import pytest +import torch +import torch.nn.functional as functional +from onnxscript import nn + +from mobius import ( + MLPWorldModel, + WorldModelConfig, + WorldModelTask, + build_from_module, +) +from mobius.integrations.onnxruntime import WorldModelRunner +from mobius.tasks import TASK_REGISTRY, get_task + + +class _TorchMLPWorldModel(torch.nn.Module): + def __init__(self, config: WorldModelConfig): + super().__init__() + input_size = config.observation_size + config.action_size + config.state_size + self.config = config + self.input_layer = torch.nn.Linear(input_size, config.hidden_size) + self.hidden_layers = torch.nn.ModuleList( + [ + torch.nn.Linear(config.hidden_size, config.hidden_size) + for _ in range(config.num_hidden_layers - 1) + ] + ) + self.state_head = torch.nn.Linear(config.hidden_size, config.state_size) + self.observation_head = torch.nn.Linear(config.hidden_size, config.observation_size) + self.reward_head = torch.nn.Linear(config.hidden_size, 1) + self.continuation_head = torch.nn.Linear(config.hidden_size, 1) + + def forward(self, observation, action, state): + hidden = functional.silu( + self.input_layer( + torch.cat( + ( + observation.flatten(start_dim=1), + action.flatten(start_dim=1), + state.flatten(start_dim=1), + ), + dim=1, + ) + ) + ) + for layer in self.hidden_layers: + hidden = functional.silu(layer(hidden)) + + next_state = self.state_head(hidden) + if self.config.residual_state: + next_state = state.flatten(start_dim=1) + next_state + return ( + next_state.reshape_as(state), + self.observation_head(hidden).reshape_as(observation), + self.reward_head(hidden), + torch.sigmoid(self.continuation_head(hidden)), + ) + + +def _config() -> WorldModelConfig: + return WorldModelConfig( + observation_shape=(2, 2), + action_shape=(2,), + state_shape=(3,), + hidden_size=8, + num_hidden_layers=2, + ) + + +class TestWorldModelConfig: + def test_flattened_sizes(self): + config = _config() + assert config.observation_size == 4 + assert config.action_size == 2 + assert config.state_size == 3 + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("observation_shape", ()), + ("action_shape", (0,)), + ("state_shape", (True,)), + ("hidden_size", 0), + ("num_hidden_layers", 0), + ("hidden_act", None), + ], + ) + def test_invalid_config_raises(self, field, value): + config = _config() + setattr(config, field, value) + with pytest.raises(ValueError): + config.validate() + + +class TestWorldModelTask: + def test_registered(self): + assert TASK_REGISTRY["world-model"] is WorldModelTask + assert isinstance(get_task("world-model"), WorldModelTask) + + def test_graph_contract(self): + config = _config() + package = build_from_module( + MLPWorldModel(config), + config, + task="world-model", + ) + model = package["model"] + + assert [value.name for value in model.graph.inputs] == list(WorldModelTask.input_names) + assert [value.name for value in model.graph.outputs] == list( + WorldModelTask.output_names + ) + assert list(model.graph.inputs[0].shape)[1:] == [2, 2] + assert list(model.graph.inputs[1].shape)[1:] == [2] + assert list(model.graph.inputs[2].shape)[1:] == [3] + assert model.graph.name == "world_model_step" + + def test_rejects_wrong_module_output_contract(self): + class InvalidWorldModel(nn.Module): + def forward(self, op, observation, action, state): + return op.Identity(state) + + with pytest.raises(TypeError, match="must return"): + WorldModelTask().build(InvalidWorldModel(), _config()) + + +def test_mlp_world_model_matches_pytorch_and_preserves_state(tmp_path): + torch.manual_seed(7) + config = _config() + reference = _TorchMLPWorldModel(config).eval() + package = build_from_module( + MLPWorldModel(config), + config, + task="world-model", + execution_provider="cpu", + ) + package.apply_weights(dict(reference.state_dict())) + package.save(str(tmp_path), progress_bar=False) + + runner = WorldModelRunner.from_path( + tmp_path / "model.onnx", + providers=["CPUExecutionProvider"], + ) + rng = np.random.default_rng(11) + observation = rng.standard_normal((2, 2, 2)).astype(np.float32) + action = rng.standard_normal((2, 2)).astype(np.float32) + state = rng.standard_normal((2, 3)).astype(np.float32) + + with torch.no_grad(): + expected = reference( + torch.from_numpy(observation), + torch.from_numpy(action), + torch.from_numpy(state), + ) + actual = runner.step(observation, action, state=state) + + for actual_value, expected_value in zip( + ( + actual.next_state, + actual.observation_prediction, + actual.reward, + actual.continuation, + ), + expected, + strict=True, + ): + np.testing.assert_allclose( + actual_value, + expected_value.numpy(), + rtol=1e-5, + atol=1e-6, + ) + + next_observation = rng.standard_normal((2, 2, 2)).astype(np.float32) + next_action = rng.standard_normal((2, 2)).astype(np.float32) + with torch.no_grad(): + expected_next = reference( + torch.from_numpy(next_observation), + torch.from_numpy(next_action), + expected[0], + ) + actual_next = runner.step(next_observation, next_action) + + np.testing.assert_allclose( + actual_next.next_state, + expected_next[0].numpy(), + rtol=1e-5, + atol=1e-6, + ) + assert runner.state is actual_next.next_state diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index e7d13d982..d7f947ffc 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -67,6 +67,7 @@ "VAETask", "VideoDenoisingTask", "VisionLanguageTask", + "WorldModelTask", "build_decoder_from_embeds", "build_embedding_from_features", "get_task", @@ -128,6 +129,7 @@ QwenVLTask, VisionLanguageTask, ) +from mobius.tasks._world_model import WorldModelTask # --------------------------------------------------------------------------- # Task registry @@ -179,6 +181,7 @@ "ssm2-text-generation": SSM2CausalLMTask, "tts": TTSTask, "video-denoising": VideoDenoisingTask, + "world-model": WorldModelTask, } diff --git a/src/mobius/tasks/_world_model.py b/src/mobius/tasks/_world_model.py new file mode 100644 index 000000000..f70aa6970 --- /dev/null +++ b/src/mobius/tasks/_world_model.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Single-step world-model task.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import WorldModelConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class WorldModelTask(ModelTask): + """Build a stateful one-step world-model graph. + + Inputs: + - observation: ``[batch, *observation_shape]`` + - action: ``[batch, *action_shape]`` + - state: ``[batch, *state_shape]`` + + Outputs: + - next_state: recurrent state for the next invocation + - observation_prediction: prediction using ``observation_shape`` + - reward: scalar reward per batch item, shaped ``[batch, 1]`` + - continuation: continuation probability, shaped ``[batch, 1]`` + """ + + input_names: ClassVar[tuple[str, ...]] = ("observation", "action", "state") + output_names: ClassVar[tuple[str, ...]] = ( + "next_state", + "observation_prediction", + "reward", + "continuation", + ) + model_roles: ClassVar[dict[str, str]] = {"model": "encoder"} + + def build( + self, + module: nn.Module, + config: WorldModelConfig, + ) -> ModelPackage: + config.validate() + batch = ir.SymbolicDim("batch") + graph, builder = _make_graph(name="world_model_step") + + observation = builder.input( + "observation", + dtype=config.dtype, + shape=[batch, *config.observation_shape], + ) + action = builder.input( + "action", + dtype=config.dtype, + shape=[batch, *config.action_shape], + ) + state = builder.input( + "state", + dtype=config.dtype, + shape=[batch, *config.state_shape], + ) + + outputs = module( + builder.op, + observation=observation, + action=action, + state=state, + ) + if not isinstance(outputs, (tuple, list)) or len(outputs) != len(self.output_names): + raise TypeError( + f"{type(module).__name__} must return " + "(next_state, observation_prediction, reward, continuation)" + ) + + for value, name in zip(outputs, self.output_names, strict=True): + builder.add_output(value, name) + + return ModelPackage({"model": _make_model(graph)}, config=config) From 0787d907079fedafc16e94cbde80b3cf8a89583e Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 4 Aug 2026 15:04:49 -0700 Subject: [PATCH 2/5] Remove ORT-specific world model runtime Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 41fde59e-190b-4448-b949-6a86bf9f57a6 --- README.md | 7 +- docs/world_models.md | 41 +-- pyproject.toml | 1 - .../integrations/onnxruntime/__init__.py | 16 - .../integrations/onnxruntime/world_model.py | 337 ------------------ .../onnxruntime/world_model_test.py | 138 ------- src/mobius/models/world_model_test.py | 36 +- 7 files changed, 28 insertions(+), 548 deletions(-) delete mode 100644 src/mobius/integrations/onnxruntime/__init__.py delete mode 100644 src/mobius/integrations/onnxruntime/world_model.py delete mode 100644 src/mobius/integrations/onnxruntime/world_model_test.py diff --git a/README.md b/README.md index c7217021e..b8c6e974a 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ multi-component export for pipelines. | **Audio** | Wav2Vec2, HuBERT, WavLM, SpeechT5 | | **Vision** | ViT, BEiT, DeiT, DINOv2, Swin, CLIP, SigLIP | | **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage, HunyuanDiT, CogVideoX | -| **World Models** | Directly declared state-transition graphs with stateful ONNX Runtime rollout | +| **World Models** | Directly declared state-transition graphs | | **Adapters** | T2I-Adapter, IP-Adapter | Supports **290+ Transformers model types** and **10 Diffusers component types** @@ -98,8 +98,7 @@ pkg = build("meta-llama/Llama-3.2-1B", See the [EP quickstart](docs/ep_quickstart.md) and [full EP reference](docs/execution_providers.md) for all supported EPs and options. -**World models** use a directly declared single-step state-transition graph -and a stateful ONNX Runtime runner: +**World models** use a directly declared single-step state-transition graph: ```python from mobius import MLPWorldModel, WorldModelConfig, build_from_module @@ -113,7 +112,7 @@ pkg = build_from_module(MLPWorldModel(config), config, task="world-model") ``` See the [world-model guide](docs/world_models.md) for weight loading and -ONNX Runtime rollout. +custom model contracts. ### CLI diff --git a/docs/world_models.md b/docs/world_models.md index 394b84e27..2a11907f8 100644 --- a/docs/world_models.md +++ b/docs/world_models.md @@ -41,44 +41,9 @@ The ONNX graph is constructed through `onnxscript.nn.Module` and `onnx_ir.GraphBuilder`; PyTorch is used only as a possible source of weight tensors. -## Run with ONNX Runtime - -Install the optional runtime dependency: - -```bash -pip install -e ".[runtime]" -``` - -`WorldModelRunner` validates the graph contract and retains `next_state` -between calls: - -```python -import numpy as np - -from mobius.integrations.onnxruntime import WorldModelRunner - -runner = WorldModelRunner.from_path( - "world-model-onnx/model.onnx", - providers=["CPUExecutionProvider"], -) - -observation = np.zeros((1, 64), dtype=np.float32) -action = np.zeros((1, 6), dtype=np.float32) - -# The first call creates a zero state. Later calls reuse next_state. -result = runner.step(observation, action) -result = runner.step(observation, action) - -trajectory = runner.rollout( - [observation, observation], - [action, action], -) -``` - -Pass `state=` to `step()` or `initial_state=` to `rollout()` when the model -uses a learned or externally sampled initial state. Each `rollout()` starts -from zero state unless `initial_state` is provided; repeated `step()` calls -continue from the runner's retained state. +Runtime execution is intentionally outside Mobius. A runtime only needs to +implement this tensor contract and feed each `next_state` back as the following +step's `state`; it may use any ONNX-compatible execution backend. ## Implement a custom world model diff --git a/pyproject.toml b/pyproject.toml index 5d588ed48..cb4bbfea7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,6 @@ testing = [ "jsonschema>=4.0", ] ort-genai = ["onnxruntime-genai", "Pillow"] -runtime = ["onnxruntime"] docs = ["furo", "myst-parser", "sphinx", "sphinx_copybutton"] [tool.setuptools.packages.find] diff --git a/src/mobius/integrations/onnxruntime/__init__.py b/src/mobius/integrations/onnxruntime/__init__.py deleted file mode 100644 index f1bc37672..000000000 --- a/src/mobius/integrations/onnxruntime/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""ONNX Runtime helpers for Mobius models.""" - -from mobius.integrations.onnxruntime.world_model import ( - WorldModelRunner, - WorldModelSession, - WorldModelStepOutput, -) - -__all__ = [ - "WorldModelRunner", - "WorldModelSession", - "WorldModelStepOutput", -] diff --git a/src/mobius/integrations/onnxruntime/world_model.py b/src/mobius/integrations/onnxruntime/world_model.py deleted file mode 100644 index c9e2dc746..000000000 --- a/src/mobius/integrations/onnxruntime/world_model.py +++ /dev/null @@ -1,337 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Stateful ONNX Runtime execution for the world-model task contract.""" - -from __future__ import annotations - -import os -from collections.abc import Sequence -from dataclasses import dataclass -from typing import Any, Protocol - -import ml_dtypes -import numpy as np -from numpy.typing import ArrayLike - -from mobius.tasks._world_model import WorldModelTask - - -class _NodeMetadata(Protocol): - @property - def name(self) -> str: ... - - @property - def shape(self) -> Sequence[int | str | None]: ... - - @property - def type(self) -> str: ... - - -class WorldModelSession(Protocol): - """Minimal inference-session interface required by :class:`WorldModelRunner`.""" - - def get_inputs(self) -> Sequence[_NodeMetadata]: ... - - def get_outputs(self) -> Sequence[_NodeMetadata]: ... - - def run( - self, - output_names: Sequence[str], - input_feed: dict[str, np.ndarray], - ) -> Sequence[Any]: ... - - -@dataclass(frozen=True) -class WorldModelStepOutput: - """Outputs from one world-model transition.""" - - next_state: np.ndarray - observation_prediction: np.ndarray - reward: np.ndarray - continuation: np.ndarray - - -_ORT_TYPE_TO_NUMPY: dict[str, np.dtype] = { - "tensor(float)": np.dtype(np.float32), - "tensor(float16)": np.dtype(np.float16), - "tensor(double)": np.dtype(np.float64), - "tensor(bfloat16)": np.dtype(ml_dtypes.bfloat16), - "tensor(int64)": np.dtype(np.int64), - "tensor(int32)": np.dtype(np.int32), - "tensor(int16)": np.dtype(np.int16), - "tensor(int8)": np.dtype(np.int8), - "tensor(uint64)": np.dtype(np.uint64), - "tensor(uint32)": np.dtype(np.uint32), - "tensor(uint16)": np.dtype(np.uint16), - "tensor(uint8)": np.dtype(np.uint8), - "tensor(bool)": np.dtype(np.bool_), -} - - -class WorldModelRunner: - """Run a Mobius world-model graph while preserving recurrent state.""" - - def __init__(self, session: WorldModelSession): - self._session = session - self._inputs = {node.name: node for node in session.get_inputs()} - self._outputs = {node.name: node for node in session.get_outputs()} - self._validate_contract() - self._state: np.ndarray | None = None - - @classmethod - def from_path( - cls, - model_path: str | os.PathLike[str], - *, - providers: Sequence[str] | None = None, - session_options: object | None = None, - ) -> WorldModelRunner: - """Create a runner from an ONNX file without making ORT a core dependency.""" - try: - import onnxruntime as ort - except ModuleNotFoundError as exc: - if exc.name != "onnxruntime": - raise - raise ImportError( - "WorldModelRunner.from_path() requires onnxruntime; " - "install mobius-onnx[runtime]" - ) from exc - - kwargs: dict[str, object] = {} - if providers is not None: - kwargs["providers"] = list(providers) - if session_options is not None: - kwargs["sess_options"] = session_options - return cls(ort.InferenceSession(os.fspath(model_path), **kwargs)) - - @property - def session(self) -> WorldModelSession: - """Underlying inference session.""" - return self._session - - @property - def state(self) -> np.ndarray | None: - """Recurrent state used by the next step.""" - return self._state - - def reset( - self, - state: ArrayLike | None = None, - *, - batch_size: int | None = None, - ) -> np.ndarray: - """Reset recurrent state to an explicit value or a zero tensor.""" - if state is not None: - prepared = self._prepare_input("state", state) - if batch_size is not None and prepared.shape[0] != batch_size: - raise ValueError( - f"state batch dimension is {prepared.shape[0]}, expected {batch_size}" - ) - else: - prepared = self._make_zero_state(batch_size or 1) - self._state = prepared - return prepared - - def step( - self, - observation: ArrayLike, - action: ArrayLike, - *, - state: ArrayLike | None = None, - ) -> WorldModelStepOutput: - """Execute one transition and retain its ``next_state``.""" - observation_array = self._prepare_input("observation", observation) - action_array = self._prepare_input("action", action) - batch_size = observation_array.shape[0] - if action_array.shape[0] != batch_size: - raise ValueError( - "observation and action batch dimensions differ: " - f"{batch_size} != {action_array.shape[0]}" - ) - - if state is not None: - current_state = self._prepare_input("state", state) - elif self._state is not None: - current_state = self._state - else: - current_state = self._make_zero_state(batch_size) - if current_state.shape[0] != batch_size: - raise ValueError( - "observation and state batch dimensions differ: " - f"{batch_size} != {current_state.shape[0]}" - ) - - values = self._session.run( - list(WorldModelTask.output_names), - { - "observation": observation_array, - "action": action_array, - "state": current_state, - }, - ) - if len(values) != len(WorldModelTask.output_names): - raise RuntimeError( - f"session returned {len(values)} outputs, " - f"expected {len(WorldModelTask.output_names)}" - ) - - arrays = { - name: self._validate_output(name, value) - for name, value in zip(WorldModelTask.output_names, values, strict=True) - } - if arrays["next_state"].shape != current_state.shape: - raise ValueError( - "next_state shape is incompatible with recurrent state: " - f"{arrays['next_state'].shape} != {current_state.shape}" - ) - if arrays["observation_prediction"].shape != observation_array.shape: - raise ValueError( - "observation_prediction shape is incompatible with observation: " - f"{arrays['observation_prediction'].shape} != {observation_array.shape}" - ) - expected_scalar_shape = (batch_size, 1) - for name in ("reward", "continuation"): - if arrays[name].shape != expected_scalar_shape: - raise ValueError( - f"{name} has shape {arrays[name].shape}, expected {expected_scalar_shape}" - ) - output = WorldModelStepOutput( - next_state=arrays["next_state"], - observation_prediction=arrays["observation_prediction"], - reward=arrays["reward"], - continuation=arrays["continuation"], - ) - self._state = output.next_state - return output - - def rollout( - self, - observations: Sequence[ArrayLike], - actions: Sequence[ArrayLike], - *, - initial_state: ArrayLike | None = None, - ) -> tuple[WorldModelStepOutput, ...]: - """Run a sequence of transitions along the leading time dimension.""" - if len(observations) != len(actions): - raise ValueError("observations and actions must contain the same number of steps") - if len(observations) == 0: - raise ValueError("rollout requires at least one step") - if initial_state is not None: - self.reset(initial_state) - else: - first_observation = self._prepare_input("observation", observations[0]) - self.reset(batch_size=first_observation.shape[0]) - - return tuple( - self.step(observation, action) - for observation, action in zip(observations, actions, strict=True) - ) - - def _validate_contract(self) -> None: - expected_inputs = set(WorldModelTask.input_names) - expected_outputs = set(WorldModelTask.output_names) - actual_inputs = set(self._inputs) - actual_outputs = set(self._outputs) - if actual_inputs != expected_inputs: - raise ValueError( - "world-model input contract mismatch: " - f"expected {sorted(expected_inputs)}, got {sorted(actual_inputs)}" - ) - if actual_outputs != expected_outputs: - raise ValueError( - "world-model output contract mismatch: " - f"expected {sorted(expected_outputs)}, got {sorted(actual_outputs)}" - ) - state_input = self._inputs["state"] - next_state_output = self._outputs["next_state"] - if state_input.type != next_state_output.type: - raise ValueError( - "next_state dtype must match state dtype: " - f"{next_state_output.type!r} != {state_input.type!r}" - ) - if not self._shapes_are_recurrently_compatible( - state_input.shape, - next_state_output.shape, - ): - raise ValueError( - "next_state shape must be compatible with state shape: " - f"{tuple(next_state_output.shape)} != {tuple(state_input.shape)}" - ) - - def _prepare_input(self, name: str, value: ArrayLike) -> np.ndarray: - node = self._inputs[name] - array = np.ascontiguousarray(value, dtype=self._numpy_dtype(node)) - self._validate_shape(name, array, node) - return array - - def _validate_output(self, name: str, value: Any) -> np.ndarray: - node = self._outputs[name] - array = np.asarray(value) - expected_dtype = self._numpy_dtype(node) - if array.dtype != expected_dtype: - raise TypeError(f"{name} has dtype {array.dtype}, expected {expected_dtype}") - self._validate_shape(name, array, node) - return array - - def _make_zero_state(self, batch_size: int) -> np.ndarray: - if batch_size <= 0: - raise ValueError("batch_size must be positive") - node = self._inputs["state"] - shape: list[int] = [] - for axis, dim in enumerate(node.shape): - if axis == 0: - if isinstance(dim, int) and dim != batch_size: - raise ValueError( - f"state model batch dimension is fixed at {dim}, got {batch_size}" - ) - shape.append(batch_size) - elif isinstance(dim, int): - shape.append(dim) - else: - raise ValueError( - "cannot create zero state because a non-batch state dimension " - f"is dynamic ({dim!r}); pass an explicit state to reset() or step()" - ) - return np.zeros(shape, dtype=self._numpy_dtype(node)) - - @staticmethod - def _numpy_dtype(node: _NodeMetadata) -> np.dtype: - try: - return _ORT_TYPE_TO_NUMPY[node.type] - except KeyError as exc: - raise TypeError(f"unsupported ONNX Runtime tensor type {node.type!r}") from exc - - @staticmethod - def _validate_shape(name: str, array: np.ndarray, node: _NodeMetadata) -> None: - expected = tuple(node.shape) - if array.ndim != len(expected): - raise ValueError( - f"{name} has rank {array.ndim}, expected rank {len(expected)} " - f"with shape {expected}" - ) - for axis, (actual_dim, expected_dim) in enumerate( - zip(array.shape, expected, strict=True) - ): - if isinstance(expected_dim, int) and actual_dim != expected_dim: - raise ValueError( - f"{name} dimension {axis} is {actual_dim}, expected {expected_dim}" - ) - - @staticmethod - def _shapes_are_recurrently_compatible( - state_shape: Sequence[int | str | None], - next_state_shape: Sequence[int | str | None], - ) -> bool: - if len(state_shape) != len(next_state_shape): - return False - return all( - state_dim == next_state_dim - if isinstance(state_dim, int) and isinstance(next_state_dim, int) - else True - for state_dim, next_state_dim in zip( - state_shape, - next_state_shape, - strict=True, - ) - ) diff --git a/src/mobius/integrations/onnxruntime/world_model_test.py b/src/mobius/integrations/onnxruntime/world_model_test.py deleted file mode 100644 index ad2e7df95..000000000 --- a/src/mobius/integrations/onnxruntime/world_model_test.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -from __future__ import annotations - -import dataclasses - -import numpy as np -import pytest - -from mobius.integrations.onnxruntime import WorldModelRunner - - -@dataclasses.dataclass -class _Node: - name: str - shape: list[int | str | None] - type: str = "tensor(float)" - - -class _FakeSession: - def __init__(self): - self.inputs = [ - _Node("observation", ["batch", 4]), - _Node("action", ["batch", 2]), - _Node("state", ["batch", 3]), - ] - self.outputs = [ - _Node("next_state", ["batch", 3]), - _Node("observation_prediction", ["batch", 4]), - _Node("reward", ["batch", 1]), - _Node("continuation", ["batch", 1]), - ] - self.last_feed = None - - def get_inputs(self): - return self.inputs - - def get_outputs(self): - return self.outputs - - def run(self, output_names, input_feed): - self.last_feed = input_feed - batch = input_feed["observation"].shape[0] - return [ - input_feed["state"] + 1.0, - input_feed["observation"] * 2.0, - np.zeros((batch, 1), dtype=np.float32), - np.ones((batch, 1), dtype=np.float32), - ] - - -def test_step_initializes_and_preserves_state(): - session = _FakeSession() - runner = WorldModelRunner(session) - observation = np.ones((2, 4), dtype=np.float64) - action = np.ones((2, 2), dtype=np.float64) - - first = runner.step(observation, action) - assert session.last_feed["observation"].dtype == np.float32 - np.testing.assert_array_equal( - session.last_feed["state"], - np.zeros((2, 3), dtype=np.float32), - ) - np.testing.assert_array_equal(first.next_state, np.ones((2, 3), dtype=np.float32)) - - second = runner.step(observation, action) - np.testing.assert_array_equal( - session.last_feed["state"], - first.next_state, - ) - np.testing.assert_array_equal(second.next_state, np.full((2, 3), 2.0, np.float32)) - - -def test_rollout_uses_initial_state(): - runner = WorldModelRunner(_FakeSession()) - observations = np.ones((3, 1, 4), dtype=np.float32) - actions = np.ones((3, 1, 2), dtype=np.float32) - initial_state = np.full((1, 3), 5.0, dtype=np.float32) - - outputs = runner.rollout( - observations, - actions, - initial_state=initial_state, - ) - - assert len(outputs) == 3 - np.testing.assert_array_equal(outputs[-1].next_state, np.full((1, 3), 8.0)) - assert runner.state is outputs[-1].next_state - - -def test_rollout_defaults_to_fresh_zero_state(): - runner = WorldModelRunner(_FakeSession()) - runner.reset(np.full((1, 3), 10.0, dtype=np.float32)) - - outputs = runner.rollout( - np.ones((2, 1, 4), dtype=np.float32), - np.ones((2, 1, 2), dtype=np.float32), - ) - - np.testing.assert_array_equal(outputs[-1].next_state, np.full((1, 3), 2.0)) - - -def test_rejects_contract_mismatch(): - session = _FakeSession() - session.inputs.pop() - - with pytest.raises(ValueError, match="input contract mismatch"): - WorldModelRunner(session) - - -def test_rejects_incompatible_recurrent_state_contract(): - session = _FakeSession() - session.outputs[0].type = "tensor(double)" - - with pytest.raises(ValueError, match="next_state dtype must match"): - WorldModelRunner(session) - - -def test_rejects_mismatched_batches_before_inference(): - session = _FakeSession() - runner = WorldModelRunner(session) - - with pytest.raises(ValueError, match="batch dimensions differ"): - runner.step( - np.zeros((2, 4), dtype=np.float32), - np.zeros((1, 2), dtype=np.float32), - ) - assert session.last_feed is None - - -def test_reset_requires_explicit_dynamic_state_tail(): - session = _FakeSession() - session.inputs[-1].shape = ["batch", "state_size"] - runner = WorldModelRunner(session) - - with pytest.raises(ValueError, match="non-batch state dimension is dynamic"): - runner.reset(batch_size=2) diff --git a/src/mobius/models/world_model_test.py b/src/mobius/models/world_model_test.py index a42325608..02b7e8308 100644 --- a/src/mobius/models/world_model_test.py +++ b/src/mobius/models/world_model_test.py @@ -4,6 +4,7 @@ from __future__ import annotations import numpy as np +import onnxruntime as ort import pytest import torch import torch.nn.functional as functional @@ -15,7 +16,6 @@ WorldModelTask, build_from_module, ) -from mobius.integrations.onnxruntime import WorldModelRunner from mobius.tasks import TASK_REGISTRY, get_task @@ -130,7 +130,7 @@ def forward(self, op, observation, action, state): WorldModelTask().build(InvalidWorldModel(), _config()) -def test_mlp_world_model_matches_pytorch_and_preserves_state(tmp_path): +def test_mlp_world_model_matches_pytorch(tmp_path): torch.manual_seed(7) config = _config() reference = _TorchMLPWorldModel(config).eval() @@ -143,8 +143,8 @@ def test_mlp_world_model_matches_pytorch_and_preserves_state(tmp_path): package.apply_weights(dict(reference.state_dict())) package.save(str(tmp_path), progress_bar=False) - runner = WorldModelRunner.from_path( - tmp_path / "model.onnx", + session = ort.InferenceSession( + str(tmp_path / "model.onnx"), providers=["CPUExecutionProvider"], ) rng = np.random.default_rng(11) @@ -158,15 +158,17 @@ def test_mlp_world_model_matches_pytorch_and_preserves_state(tmp_path): torch.from_numpy(action), torch.from_numpy(state), ) - actual = runner.step(observation, action, state=state) + actual = session.run( + list(WorldModelTask.output_names), + { + "observation": observation, + "action": action, + "state": state, + }, + ) for actual_value, expected_value in zip( - ( - actual.next_state, - actual.observation_prediction, - actual.reward, - actual.continuation, - ), + actual, expected, strict=True, ): @@ -185,12 +187,18 @@ def test_mlp_world_model_matches_pytorch_and_preserves_state(tmp_path): torch.from_numpy(next_action), expected[0], ) - actual_next = runner.step(next_observation, next_action) + actual_next = session.run( + list(WorldModelTask.output_names), + { + "observation": next_observation, + "action": next_action, + "state": actual[0], + }, + ) np.testing.assert_allclose( - actual_next.next_state, + actual_next[0], expected_next[0].numpy(), rtol=1e-5, atol=1e-6, ) - assert runner.state is actual_next.next_state From 078d57652bbe045f4910b7f9fb3f21129134cfd6 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 4 Aug 2026 15:10:16 -0700 Subject: [PATCH 3/5] Remove incomplete world model README exposure Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 41fde59e-190b-4448-b949-6a86bf9f57a6 --- README.md | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/README.md b/README.md index b8c6e974a..f9f2320b6 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,6 @@ multi-component export for pipelines. | **Audio** | Wav2Vec2, HuBERT, WavLM, SpeechT5 | | **Vision** | ViT, BEiT, DeiT, DINOv2, Swin, CLIP, SigLIP | | **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage, HunyuanDiT, CogVideoX | -| **World Models** | Directly declared state-transition graphs | | **Adapters** | T2I-Adapter, IP-Adapter | Supports **290+ Transformers model types** and **10 Diffusers component types** @@ -98,22 +97,6 @@ pkg = build("meta-llama/Llama-3.2-1B", See the [EP quickstart](docs/ep_quickstart.md) and [full EP reference](docs/execution_providers.md) for all supported EPs and options. -**World models** use a directly declared single-step state-transition graph: - -```python -from mobius import MLPWorldModel, WorldModelConfig, build_from_module - -config = WorldModelConfig( - observation_shape=(64,), - action_shape=(6,), - state_shape=(128,), -) -pkg = build_from_module(MLPWorldModel(config), config, task="world-model") -``` - -See the [world-model guide](docs/world_models.md) for weight loading and -custom model contracts. - ### CLI ```sh From 90ec84cccbb343727876e959be4d9f2cc439cf3d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang Date: Tue, 4 Aug 2026 15:10:50 -0700 Subject: [PATCH 4/5] Remove standalone world model guide Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 41fde59e-190b-4448-b949-6a86bf9f57a6 --- docs/index.md | 1 - docs/world_models.md | 66 -------------------------------------------- 2 files changed, 67 deletions(-) delete mode 100644 docs/world_models.md diff --git a/docs/index.md b/docs/index.md index 2181f58c6..103db3353 100644 --- a/docs/index.md +++ b/docs/index.md @@ -11,7 +11,6 @@ Build ONNX models directly from HuggingFace model IDs with automatic weight down :caption: User Guide getting-started -world_models cli_reference module-architecture model-catalog diff --git a/docs/world_models.md b/docs/world_models.md deleted file mode 100644 index 2a11907f8..000000000 --- a/docs/world_models.md +++ /dev/null @@ -1,66 +0,0 @@ -# World models - -Mobius can build world models directly as ONNX graphs without tracing a -PyTorch `forward()` method. The initial framework defines a deterministic, -single-step state-transition contract: - -```text -(observation, action, state) - -> (next_state, observation_prediction, reward, continuation) -``` - -All tensors have a dynamic leading batch dimension. `reward` and -`continuation` have shape `[batch, 1]`; `continuation` is a probability. - -## Build the reference model - -`MLPWorldModel` is a minimal directly declared implementation. Its input shapes -exclude the batch dimension: - -```python -from safetensors.torch import load_file - -from mobius import MLPWorldModel, WorldModelConfig, build_from_module - -config = WorldModelConfig( - observation_shape=(64,), - action_shape=(6,), - state_shape=(128,), - hidden_size=512, - num_hidden_layers=3, -) -module = MLPWorldModel(config) -package = build_from_module(module, config, task="world-model") - -weights = load_file("world_model.safetensors") -package.apply_weights(module.preprocess_weights(weights)) -package.save("world-model-onnx") -``` - -The ONNX graph is constructed through `onnxscript.nn.Module` and -`onnx_ir.GraphBuilder`; PyTorch is used only as a possible source of weight -tensors. - -Runtime execution is intentionally outside Mobius. A runtime only needs to -implement this tensor contract and feed each `next_state` back as the following -step's `state`; it may use any ONNX-compatible execution backend. - -## Implement a custom world model - -A custom module must use the task's forward and return contracts: - -```python -from onnxscript import nn - - -class MyWorldModel(nn.Module): - def forward(self, op, observation, action, state): - # Directly declare ONNX operations here. - ... - return next_state, observation_prediction, reward, continuation -``` - -Use a custom `ModelTask` when an architecture needs multiple recurrent-state -tensors, stochastic latent outputs, separate observe/imagine graphs, or a -different prediction contract. Keep sampling outside the graph unless it must -be part of the deployed model. From 83a8b61106eb0f335044bc87fe77ef951c5a9d26 Mon Sep 17 00:00:00 2001 From: Xiaoyu <85524621+xiaoyu-work@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:01:28 -0700 Subject: [PATCH 5/5] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Xiaoyu <85524621+xiaoyu-work@users.noreply.github.com> --- src/mobius/tasks/_world_model.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/mobius/tasks/_world_model.py b/src/mobius/tasks/_world_model.py index f70aa6970..d6b6a3863 100644 --- a/src/mobius/tasks/_world_model.py +++ b/src/mobius/tasks/_world_model.py @@ -48,18 +48,20 @@ def build( batch = ir.SymbolicDim("batch") graph, builder = _make_graph(name="world_model_step") + observation_name, action_name, state_name = self.input_names + observation = builder.input( - "observation", + observation_name, dtype=config.dtype, shape=[batch, *config.observation_shape], ) action = builder.input( - "action", + action_name, dtype=config.dtype, shape=[batch, *config.action_shape], ) state = builder.input( - "state", + state_name, dtype=config.dtype, shape=[batch, *config.state_shape], )