Phase 13.1 — WorldModel: predictive environment model for model-based planning

[web3guru888]

Phase 13.1 — WorldModel: predictive environment model for model-based planning

This issue tracks the specification and implementation of WorldModel, the first sub-phase of Phase 13: World Modeling & Model-Based Planning.

Phase 13 context — Discussion #367 opened the question of which direction to take ASI-Build next. World Modeling was selected as Phase 13 because it is foundational to model-based planning, dream rollouts, and surprise-driven exploration — capabilities that naturally complement the multi-agent coalition infrastructure built in Phase 12.

---

Motivation

A WorldModel enables the agent to simulate "what happens if" before committing to an action. Instead of relying solely on trial-and-error in the environment, the agent rolls out imagined trajectories through its internal model, selects the best plan, and executes with far greater sample efficiency.

Key benefits:

• Dyna-style planning: interleave real and imagined experience
• Zero-shot transfer: updated world model generalises across tasks
• Surprise detection: high prediction error flags distributional shift
• Scalable: model is a lightweight neural net, not a full simulator

---

Enumerations

from enum import Enum, auto

class TransitionBackend(Enum):
    """Backend used for world-model forward pass."""
    MLP = auto()           # simple feed-forward MLP (fast, low-capacity)
    LSTM = auto()          # recurrent — tracks temporal context
    TRANSFORMER = auto()   # attention over observation history window
    ENSEMBLE = auto()      # ensemble of N MLPs for epistemic uncertainty

class PredictionTarget(Enum):
    """What the model predicts on each step."""
    NEXT_OBS = auto()      # predict next observation vector
    REWARD = auto()        # predict scalar reward
    DONE = auto()          # predict episode-termination flag
    ALL = auto()           # joint head — obs + reward + done

---

Data Classes

from dataclasses import dataclass, field
from typing import Sequence

@dataclass(frozen=True)
class ModelInput:
    """Single (obs, action) pair passed to WorldModel."""
    observation: tuple[float, ...]   # current obs vector
    action: int                      # discrete action index

@dataclass(frozen=True)
class ModelOutput:
    """Predicted next state."""
    next_observation: tuple[float, ...]  # predicted next obs
    predicted_reward: float              # predicted reward
    predicted_done: bool                 # predicted termination
    prediction_error: float              # MSE vs last ground-truth (0 if unavailable)

@dataclass(frozen=True)
class DreamRollout:
    """Result of an imagined trajectory through the world model."""
    steps: tuple[ModelOutput, ...]   # ordered sequence of imagined steps
    total_imagined_reward: float     # sum of predicted rewards
    surprise_score: float            # mean prediction_error across steps

@dataclass(frozen=True)
class WorldModelConfig:
    backend: TransitionBackend = TransitionBackend.MLP
    obs_dim: int = 64
    action_dim: int = 8
    hidden_dim: int = 256
    ensemble_size: int = 5           # used when backend == ENSEMBLE
    rollout_horizon: int = 10        # steps per dream rollout
    learning_rate: float = 3e-4
    surprise_threshold: float = 0.05 # error above this triggers re-planning
    max_buffer_size: int = 100_000   # replay buffer for model training

---

Protocol

from typing import Protocol, runtime_checkable

@runtime_checkable
class WorldModel(Protocol):
    """Predictive model of environment dynamics."""

    async def predict(self, inp: ModelInput) -> ModelOutput:
        """One-step forward pass: (obs, action) → next state prediction."""
        ...

    async def dream_rollout(
        self,
        initial_obs: tuple[float, ...],
        action_sequence: Sequence[int],
    ) -> DreamRollout:
        """Simulate a trajectory through the model."""
        ...

    async def update(
        self,
        inp: ModelInput,
        actual_next_obs: tuple[float, ...],
        actual_reward: float,
        actual_done: bool,
    ) -> float:
        """Online update from a real transition; returns scalar loss."""
        ...

    async def surprise(self, inp: ModelInput, actual_next_obs: tuple[float, ...]) -> float:
        """Return prediction error for surprise detection."""
        ...

    async def snapshot(self) -> dict:
        """Return serialisable model weights + config."""
        ...

---

InMemoryWorldModel Implementation

import asyncio
import math
import random
from collections import deque

class InMemoryWorldModel:
    """NumPy/pure-Python reference implementation (no torch dependency)."""

    def __init__(self, config: WorldModelConfig) -> None:
        self._cfg = config
        self._weights: dict = {}   # placeholder for weight tensors
        self._buffer: deque = deque(maxlen=config.max_buffer_size)
        self._step_count = 0
        self._lock = asyncio.Lock()

    # ------------------------------------------------------------------
    # Protocol implementation
    # ------------------------------------------------------------------

    async def predict(self, inp: ModelInput) -> ModelOutput:
        async with self._lock:
            return await self._forward(inp)

    async def dream_rollout(
        self,
        initial_obs: tuple[float, ...],
        action_sequence: Sequence[int],
    ) -> DreamRollout:
        obs = initial_obs
        steps: list[ModelOutput] = []
        async with self._lock:
            for action in action_sequence:
                out = await self._forward(ModelInput(observation=obs, action=action))
                steps.append(out)
                if out.predicted_done:
                    break
                obs = out.next_observation
        total_reward = sum(s.predicted_reward for s in steps)
        surprise = sum(s.prediction_error for s in steps) / max(len(steps), 1)
        return DreamRollout(
            steps=tuple(steps),
            total_imagined_reward=total_reward,
            surprise_score=surprise,
        )

    async def update(
        self,
        inp: ModelInput,
        actual_next_obs: tuple[float, ...],
        actual_reward: float,
        actual_done: bool,
    ) -> float:
        async with self._lock:
            self._buffer.append((inp, actual_next_obs, actual_reward, actual_done))
            loss = await self._sgd_step(inp, actual_next_obs, actual_reward, actual_done)
            self._step_count += 1
            _wm_update_total.inc()
            _wm_buffer_size.set(len(self._buffer))
            return loss

    async def surprise(self, inp: ModelInput, actual_next_obs: tuple[float, ...]) -> float:
        out = await self.predict(inp)
        err = _mse(out.next_observation, actual_next_obs)
        _wm_surprise_score.set(err)
        return err

    async def snapshot(self) -> dict:
        return {
            "config": vars(self._cfg),
            "step_count": self._step_count,
            "buffer_size": len(self._buffer),
        }

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    async def _forward(self, inp: ModelInput) -> ModelOutput:
        """Stub forward pass — implementer replaces with real net."""
        noise = [random.gauss(0, 0.01) for _ in range(self._cfg.obs_dim)]
        next_obs = tuple(o + n for o, n in zip(inp.observation, noise))
        return ModelOutput(
            next_observation=next_obs,
            predicted_reward=0.0,
            predicted_done=False,
            prediction_error=0.0,
        )

    async def _sgd_step(self, inp, actual_next_obs, actual_reward, actual_done) -> float:
        """Stub SGD — implementer replaces with autograd update."""
        pred = await self._forward(inp)
        loss = _mse(pred.next_observation, actual_next_obs)
        return loss

def _mse(a: tuple[float, ...], b: tuple[float, ...]) -> float:
    return sum((x - y) ** 2 for x, y in zip(a, b)) / max(len(a), 1)

def build_world_model(config: WorldModelConfig | None = None) -> WorldModel:
    return InMemoryWorldModel(config or WorldModelConfig())

---

CognitiveCycle Integration

class CognitiveCycle:
    def __init__(self, ..., world_model: WorldModel) -> None:
        self._world_model = world_model

    async def _model_based_step(
        self,
        obs: tuple[float, ...],
        candidate_actions: list[int],
    ) -> int:
        """Select best action via dream rollout."""
        best_action, best_reward = candidate_actions[0], float("-inf")
        for action in candidate_actions:
            rollout = await self._world_model.dream_rollout(
                initial_obs=obs,
                action_sequence=[action] * self._cfg.rollout_horizon,
            )
            if rollout.total_imagined_reward > best_reward:
                best_reward = rollout.total_imagined_reward
                best_action = action
            if rollout.surprise_score > self._world_model_cfg.surprise_threshold:
                await self._trigger_replan()
        return best_action

    async def _update_world_model(
        self,
        inp: ModelInput,
        actual_next_obs: tuple[float, ...],
        actual_reward: float,
        actual_done: bool,
    ) -> None:
        loss = await self._world_model.update(inp, actual_next_obs, actual_reward, actual_done)
        _wm_train_loss.set(loss)

---

Prometheus Metrics

Metric	Type	Description
wm_update_total	Counter	Total update() calls (real transitions consumed)
wm_train_loss	Gauge	Most recent SGD loss
wm_surprise_score	Gauge	Most recent surprise score
wm_buffer_size	Gauge	Current replay-buffer fill level
wm_rollout_steps_total	Counter	Cumulative dream-rollout steps executed

from prometheus_client import Counter, Gauge

_wm_update_total      = Counter("wm_update_total", "Real transitions consumed")
_wm_train_loss        = Gauge("wm_train_loss", "Most recent SGD loss")
_wm_surprise_score    = Gauge("wm_surprise_score", "Most recent surprise score")
_wm_buffer_size       = Gauge("wm_buffer_size", "Replay buffer fill level")
_wm_rollout_steps_total = Counter("wm_rollout_steps_total", "Dream rollout steps executed")

PromQL examples:

# Surprise spikes (potential distribution shift)
wm_surprise_score > 0.05

# Model learning rate (transitions/min)
rate(wm_update_total[1m])

# Training convergence
wm_train_loss

---

mypy Compliance

Class/Function	Issues	Resolution
ModelInput.observation	tuple[float, ...] vs list	always wrap in tuple() at boundary
DreamRollout.steps	mutable list built internally	cast to tuple(steps) before freeze
InMemoryWorldModel._forward	stub returns fixed obs_dim	assert len(next_obs) == cfg.obs_dim
build_world_model	return type WorldModel (Protocol)	# type: ignore[return-value] if needed

---

Test Targets (12 minimum)

1. test_predict_returns_model_output — output type + field types
2. test_predict_obs_dim_preserved — len(next_obs) == obs_dim
3. test_dream_rollout_length — steps == min(horizon, done_step)
4. test_dream_rollout_terminates_on_done — early-exit when predicted_done=True
5. test_update_returns_scalar_loss — isinstance(loss, float)
6. test_update_increments_buffer — buffer grows after each real transition
7. test_surprise_zero_on_perfect_prediction — surprise(inp, predicted_obs) ≈ 0
8. test_surprise_high_on_bad_prediction — large error vector → high surprise
9. test_snapshot_keys — snapshot contains config, step_count, buffer_size
10. test_build_world_model_returns_protocol — isinstance(wm, WorldModel)
11. test_concurrent_predict_safe — 50 concurrent predict() calls, no race
12. test_dream_rollout_accumulates_reward — summed predicted_reward correct

---

Implementation Order (14 steps)

1. Define TransitionBackend + PredictionTarget enums
2. Define ModelInput, ModelOutput, DreamRollout, WorldModelConfig dataclasses
3. Define WorldModel Protocol
4. Implement _mse() helper + build_world_model() factory
5. Implement InMemoryWorldModel.__init__() with asyncio lock + deque buffer
6. Implement _forward() stub (noise-perturbed copy, obs_dim-safe)
7. Implement predict() — lock + _forward()
8. Implement update() — buffer append + _sgd_step() + metrics
9. Implement _sgd_step() stub
10. Implement dream_rollout() — iterative _forward() with early-exit on done
11. Implement surprise() — predict + MSE
12. Implement snapshot()
13. Add 5 Prometheus metrics + PromQL examples
14. Write 12 test targets

---

Phase 13 Roadmap

Sub-phase	Component	Status
13.1	WorldModel	🟡 This issue
13.2	DreamRolloutPlanner	⏳ Planned
13.3	ModelBasedPolicyOptimizer	⏳ Planned
13.4	SurpriseDetector	⏳ Planned
13.5	WorldModelDashboard	⏳ Planned

Phase 12 recap — Coalition infrastructure complete:

• 12.1 AgentRegistry (Phase 12.1 — AgentRegistry: distributed agent identity and capability registry #352) ✅
• 12.2 NegotiationEngine (Phase 12.2 — NegotiationEngine: task-offer negotiation between distributed agents #355) ✅
• 12.3 CollaborationChannel (Phase 12.3 — CollaborationChannel: real-time pub/sub workspace for cooperating agents #358) ✅
• 12.4 ConsensusVoting (Phase 12.4 — ConsensusVoting: threshold-quorum decision ratification for agent coalitions #361) ✅
• 12.5 CoalitionFormation (Phase 12.5 — CoalitionFormation: dynamic capability-matched agent coalitions 🎉 Phase 12 capstone #364) ✅

---

Related: Discussion #367 (Phase 13 direction) | Discussion #369 (Show & Tell) | Discussion #370 (Q&A)
Wiki: Phase-13-World-Model (coming soon)

[web3guru888]

Implementation notes

_forward() stub — obs_dim-safe

async def _forward(self, inp: ModelInput) -> ModelOutput:
    import random
    cfg = self._cfg
    # Perturb current obs by small Gaussian noise
    noise = [random.gauss(0, 0.01) for _ in range(cfg.obs_dim)]
    # Safely handle obs shorter than obs_dim (pad with zeros)
    obs = list(inp.observation) + [0.0] * max(0, cfg.obs_dim - len(inp.observation))
    next_obs = tuple(o + n for o, n in zip(obs[:cfg.obs_dim], noise))
    return ModelOutput(
        next_observation=next_obs,
        predicted_reward=0.0,
        predicted_done=False,
        prediction_error=0.0,
    )

Key: always pad/truncate to obs_dim before returning — avoids mypy tuple-length errors downstream.

---

dream_rollout() — early-exit on predicted_done

async def dream_rollout(self, initial_obs, action_sequence) -> DreamRollout:
    obs = initial_obs
    steps: list[ModelOutput] = []
    async with self._lock:
        for action in action_sequence:
            out = await self._forward(ModelInput(observation=obs, action=action))
            steps.append(out)
            _wm_rollout_steps_total.inc()
            if out.predicted_done:
                break  # episode ends in imagination — stop early
            obs = out.next_observation
    total_reward = sum(s.predicted_reward for s in steps)
    surprise = sum(s.prediction_error for s in steps) / max(len(steps), 1)
    return DreamRollout(
        steps=tuple(steps),
        total_imagined_reward=total_reward,
        surprise_score=surprise,
    )

Key: hold _lock for the whole rollout to prevent mid-rollout weight updates corrupting the trajectory.

---

_sgd_step() — pure-Python MSE gradient

async def _sgd_step(
    self,
    inp: ModelInput,
    actual_next_obs: tuple[float, ...],
    actual_reward: float,
    actual_done: bool,
) -> float:
    pred = await self._forward(inp)
    loss = _mse(pred.next_observation, actual_next_obs)
    # Real implementation: compute gradients via autograd and call optimizer.step()
    # Stub: just return loss for metrics tracking
    _wm_train_loss.set(loss)
    return loss

Note: replace stub with torch.nn.Module + torch.optim.Adam when integrating real backend.

---

ENSEMBLE backend — epistemic uncertainty

When backend == ENSEMBLE, maintain N independent _forward functions. Predict with each, compute variance across predictions as uncertainty estimate:

predictions = [await member._forward(inp) for member in self._ensemble]
mean_obs = tuple(
    sum(p.next_observation[i] for p in predictions) / len(predictions)
    for i in range(self._cfg.obs_dim)
)
variance = sum(
    _mse(p.next_observation, mean_obs) for p in predictions
) / len(predictions)
# High variance → high epistemic uncertainty → explore this region

---

mypy table

Issue	Fix
tuple[float, ...] vs list[float]	always tuple(x for x in ...) at boundary
Sequence[int] parameter	from collections.abc import Sequence
Protocol runtime check	@runtime_checkable + isinstance(wm, WorldModel)
deque[Any]	annotate as deque[tuple[ModelInput, ...]]

---

Test isolation skeleton

import asyncio
import pytest
from asi_build.world_model import (
    WorldModelConfig, ModelInput, build_world_model, WorldModel
)

@pytest.fixture
def wm():
    cfg = WorldModelConfig(obs_dim=4, action_dim=2, rollout_horizon=3)
    return build_world_model(cfg)

@pytest.mark.asyncio
async def test_predict_obs_dim_preserved(wm):
    inp = ModelInput(observation=(0.1, 0.2, 0.3, 0.4), action=0)
    out = await wm.predict(inp)
    assert len(out.next_observation) == 4

@pytest.mark.asyncio
async def test_dream_rollout_length(wm):
    obs = (0.0, 0.0, 0.0, 0.0)
    rollout = await wm.dream_rollout(obs, [0, 1, 0])
    assert 1 <= len(rollout.steps) <= 3

@pytest.mark.asyncio
async def test_concurrent_predict_safe(wm):
    inp = ModelInput(observation=(0.5, 0.5, 0.5, 0.5), action=1)
    results = await asyncio.gather(*[wm.predict(inp) for _ in range(50)])
    assert all(len(r.next_observation) == 4 for r in results)

---

14-step impl order: enums → dataclasses → Protocol → helpers → __init__ → _forward → predict → update/_sgd_step → dream_rollout → surprise → snapshot → metrics → tests → wiki

[web3guru888]

Phase 13.2 filed: DreamPlanner (#371) 🎉

Issue #371 — DreamPlanner: model-predictive planning in imagined rollout space — is now open.

Phase 13 sub-phase tracker

Sub-phase	Component	Issue	Status
13.1	WorldModel	#368	🟡 in progress
13.2	DreamPlanner	#371	🟡 this run
13.3	CuriosityEngine	⏳	planned
13.4	LatentStateEncoder	⏳	planned
13.5	ModelEnsembleCoordinator	⏳	planned

DreamPlanner → WorldModel dependency

DreamPlanner is a direct consumer of WorldModel:

• DreamPlanner.plan() calls WorldModel.dream_rollout() to evaluate candidate action sequences
• DreamPlanner._beam_search() calls WorldModel.predict() per beam node
• DreamPlanner checks WorldModel.surprise() via the terminal_surprise field to abort plans in regions of high model uncertainty (surprise_abort_threshold)

This makes #371 dependent on #368 — the WorldModel Protocol must be stable before the DreamPlanner can be fully wired in.

Planning strategies available

Strategy	Description	Best for
RANDOM_SHOOTING	Sample N random sequences, pick best	Fast, noisy environments
CEM	Cross-Entropy Method iterative refinement	Smooth reward landscapes
BEAM_SEARCH	Deterministic top-K tree expansion	Discrete action spaces
GREEDY	Single rollout argmax	Baseline / ablation

See #371 for full spec and test targets.