Skip to content

World Model — Spike Prediction + Planning

Three components: (1) online-learnable spike predictor for codec integration, (2) stochastic state-transition model, (3) greedy action planner.

SpikePredictor — Online Autoregressive Codec

The core workhorse. Predicts multi-channel spike patterns from recent history using a linear autoregressive model trained online via LMS (Least Mean Squares). No backprop, no batches — updates one sample at a time.

Codec integration: Encoder and decoder both maintain identical SpikePredictor instances. Both see the same history. Prediction error (XOR of actual vs predicted) is what gets transmitted. At the decoder, XOR recovers the original. Deterministic: same history → same prediction → lossless roundtrip.

Parameter Default Meaning
n_channels (required) Number of spike channels
history_len 8 Context window (K past timesteps)
lr 0.01 LMS learning rate
threshold 0.5 Binary prediction threshold

Codec functions:

  • predict_and_xor_world_model(spikes, n_channels, ...) → (errors, correct_count) — Encoder
  • xor_and_recover_world_model(errors, n_channels, ...) → spikes — Decoder

PredictiveWorldModel — Linear Gaussian State-Space

Probabilistic predictive model implemented as a Linear Gaussian State-Space Model (LGSSM) with Kalman filter (forward), RTS smoother (backward), and EM parameter learner. References: Kalman 1960, Rauch-Tung-Striebel 1965, Shumway & Stoffer 1982, Bishop 2006 §13.3.

Model parameters and returned moments are finite float64 arrays with fail-closed shape and covariance validation. The Python inference path uses Cholesky solves and Joseph-form covariance updates without explicit matrix inverses. The forward filter is cross-wired to Mojo, Go, Rust, Julia, and Python backends; backend="auto" follows that stable availability-aware order. RTS smoothing and the EM M-step remain explicit Python/NumPy responsibilities.

Provides predict_next_state() (deterministic mean), predict_next_state_with_cov() (mean + covariance), forecast() / forecast_with_cov() for multi-step rollouts.

In controlled EM fits, B and D are fixed but their B @ u_t and D @ u_t contributions are subtracted from the sufficient statistics. Lag-one smoother covariances have the documented Cov[x_t, x_{t+1} | y] orientation. See the predictive-model detail page for contracts, backend boundaries, source-bound benchmark evidence, and verification.

SCPlanner — Greedy Action Selection

Uses PredictiveWorldModel for random-shooting planning: sample N candidate actions, predict outcomes, pick the one closest to the goal state.

  • propose_action(current, goal, n_candidates) — Best single action
  • plan_sequence(current, goal, horizon) — Greedy multi-step plan

Usage

Python
from sc_neurocore.world_model import SpikePredictor
from sc_neurocore.world_model.spike_predictor import (
    predict_and_xor_world_model,
    xor_and_recover_world_model,
)
import numpy as np

# Lossless codec roundtrip
spikes = (np.random.rand(100, 32) < 0.3).astype(np.int8)
errors, correct = predict_and_xor_world_model(spikes, n_channels=32)
recovered = xor_and_recover_world_model(errors, n_channels=32)
assert np.array_equal(spikes, recovered)  # Always true
print(f"Prediction accuracy: {correct / (100 * 32):.1%}")

# Planning
from sc_neurocore.world_model import PredictiveWorldModel, SCPlanner
model = PredictiveWorldModel(state_dim=4, action_dim=2)
planner = SCPlanner(world_model=model)
plan = planner.plan_sequence(
    current_state=np.array([0.1, 0.2, 0.3, 0.4]),
    goal_state=np.array([0.9, 0.8, 0.7, 0.6]),
    horizon=5,
)

sc_neurocore.world_model

sc_neurocore.world_model -- Tier: research (experimental / research).

SpikePredictor dataclass

Online autoregressive spike pattern predictor.

Learns to predict spike[t] from spike[t-K:t] per channel. Weight matrix W of shape (N, N*K) maps flattened history to per-channel firing probabilities. Binary prediction via threshold.

LMS update after each timestep.

W += lr * outer(error, history)

where error = actual - predicted_prob.

Parameters

n_channels : int Number of spike channels. history_len : int Number of past timesteps to use as context (K). lr : float LMS learning rate. threshold : float Probability threshold for binary prediction. seed : int RNG seed for weight initialization.

Source code in src/sc_neurocore/world_model/spike_predictor.py
Python
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
@dataclass
class SpikePredictor:
    """Online autoregressive spike pattern predictor.

    Learns to predict spike[t] from spike[t-K:t] per channel.
    Weight matrix W of shape (N, N*K) maps flattened history to
    per-channel firing probabilities. Binary prediction via threshold.

    Training: LMS update after each timestep.
        W += lr * outer(error, history)
    where error = actual - predicted_prob.

    Parameters
    ----------
    n_channels : int
        Number of spike channels.
    history_len : int
        Number of past timesteps to use as context (K).
    lr : float
        LMS learning rate.
    threshold : float
        Probability threshold for binary prediction.
    seed : int
        RNG seed for weight initialization.
    """

    n_channels: int
    history_len: int = 8
    lr: float = 0.01
    threshold: float = 0.5
    seed: int = 42

    def __post_init__(self) -> None:
        """Seed the RNG and initialise the predictor weights."""
        rng = np.random.RandomState(self.seed)
        n_features = self.n_channels * self.history_len
        # Small random weights — predict from history
        self.W = rng.randn(self.n_channels, n_features) * 0.01
        self.bias = np.zeros(self.n_channels)
        # Circular buffer for history
        self._history = np.zeros((self.history_len, self.n_channels), dtype=np.float64)
        self._t = 0

    def _features(self) -> np.ndarray[Any, Any]:
        """Flatten history buffer into feature vector."""
        # Ordered: oldest first
        indices = [(self._t + i) % self.history_len for i in range(self.history_len)]
        return self._history[indices].ravel()

    def predict_probs(self) -> np.ndarray[Any, Any]:
        """Predict per-channel firing probabilities from history."""
        features = self._features()
        logits = self.W @ features + self.bias
        # Sigmoid activation
        probs: np.ndarray[Any, Any] = 1.0 / (1.0 + np.exp(-np.clip(logits, -20, 20)))
        return probs

    def predict(self) -> np.ndarray[Any, Any]:
        """Predict binary spike pattern."""
        return (self.predict_probs() > self.threshold).astype(np.int8)

    def update(self, actual: np.ndarray[Any, Any]) -> None:
        """Update weights with observed spike pattern (LMS rule).

        Parameters
        ----------
        actual : ndarray of shape (n_channels,), binary
        """
        features = self._features()
        probs = self.predict_probs()
        error = actual.astype(np.float64) - probs

        # LMS weight update
        self.W += self.lr * np.outer(error, features)
        self.bias += self.lr * error

        # Push actual into history buffer
        self._history[self._t % self.history_len] = actual.astype(np.float64)
        self._t += 1

    def reset(self) -> None:
        """Reset to initial state (same seed → same weights)."""
        self.__post_init__()

__post_init__()

Seed the RNG and initialise the predictor weights.

Source code in src/sc_neurocore/world_model/spike_predictor.py
Python
63
64
65
66
67
68
69
70
71
72
def __post_init__(self) -> None:
    """Seed the RNG and initialise the predictor weights."""
    rng = np.random.RandomState(self.seed)
    n_features = self.n_channels * self.history_len
    # Small random weights — predict from history
    self.W = rng.randn(self.n_channels, n_features) * 0.01
    self.bias = np.zeros(self.n_channels)
    # Circular buffer for history
    self._history = np.zeros((self.history_len, self.n_channels), dtype=np.float64)
    self._t = 0

predict_probs()

Predict per-channel firing probabilities from history.

Source code in src/sc_neurocore/world_model/spike_predictor.py
Python
80
81
82
83
84
85
86
def predict_probs(self) -> np.ndarray[Any, Any]:
    """Predict per-channel firing probabilities from history."""
    features = self._features()
    logits = self.W @ features + self.bias
    # Sigmoid activation
    probs: np.ndarray[Any, Any] = 1.0 / (1.0 + np.exp(-np.clip(logits, -20, 20)))
    return probs

predict()

Predict binary spike pattern.

Source code in src/sc_neurocore/world_model/spike_predictor.py
Python
88
89
90
def predict(self) -> np.ndarray[Any, Any]:
    """Predict binary spike pattern."""
    return (self.predict_probs() > self.threshold).astype(np.int8)

update(actual)

Update weights with observed spike pattern (LMS rule).

Parameters

actual : ndarray of shape (n_channels,), binary

Source code in src/sc_neurocore/world_model/spike_predictor.py
Python
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def update(self, actual: np.ndarray[Any, Any]) -> None:
    """Update weights with observed spike pattern (LMS rule).

    Parameters
    ----------
    actual : ndarray of shape (n_channels,), binary
    """
    features = self._features()
    probs = self.predict_probs()
    error = actual.astype(np.float64) - probs

    # LMS weight update
    self.W += self.lr * np.outer(error, features)
    self.bias += self.lr * error

    # Push actual into history buffer
    self._history[self._t % self.history_len] = actual.astype(np.float64)
    self._t += 1

reset()

Reset to initial state (same seed → same weights).

Source code in src/sc_neurocore/world_model/spike_predictor.py
Python
111
112
113
def reset(self) -> None:
    """Reset to initial state (same seed → same weights)."""
    self.__post_init__()

SCPlanner dataclass

A planner that uses a PredictiveWorldModel to select actions.

Source code in src/sc_neurocore/world_model/planner.py
Python
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
@dataclass
class SCPlanner:
    """
    A planner that uses a PredictiveWorldModel to select actions.
    """

    world_model: PredictiveWorldModel

    def propose_action(
        self,
        current_state: np.ndarray[Any, Any],
        goal_state: np.ndarray[Any, Any],
        n_candidates: int = 10,
    ) -> np.ndarray[Any, Any]:
        """
        Propose the best action among n_candidates based on predicted outcome.
        """
        best_action = None
        min_dist = float("inf")

        for _ in range(n_candidates):
            # Sample a random action
            candidate_action = np.random.uniform(0, 1, self.world_model.action_dim)

            # Predict next state
            predicted_state = self.world_model.predict_next_state(current_state, candidate_action)

            # Evaluate distance to goal
            dist = np.linalg.norm(predicted_state - goal_state)

            if dist < min_dist:
                min_dist = dist  # type: ignore[assignment]
                best_action = candidate_action

        return best_action  # type: ignore[return-value]

    def plan_sequence(
        self,
        current_state: np.ndarray[Any, Any],
        goal_state: np.ndarray[Any, Any],
        horizon: int = 5,
    ) -> List[np.ndarray[Any, Any]]:
        """
        Simple greedy planning for a sequence of actions.
        """
        plan = []
        curr_s = current_state
        for _ in range(horizon):
            action = self.propose_action(curr_s, goal_state)
            plan.append(action)
            curr_s = self.world_model.predict_next_state(curr_s, action)
        return plan

propose_action(current_state, goal_state, n_candidates=10)

Propose the best action among n_candidates based on predicted outcome.

Source code in src/sc_neurocore/world_model/planner.py
Python
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def propose_action(
    self,
    current_state: np.ndarray[Any, Any],
    goal_state: np.ndarray[Any, Any],
    n_candidates: int = 10,
) -> np.ndarray[Any, Any]:
    """
    Propose the best action among n_candidates based on predicted outcome.
    """
    best_action = None
    min_dist = float("inf")

    for _ in range(n_candidates):
        # Sample a random action
        candidate_action = np.random.uniform(0, 1, self.world_model.action_dim)

        # Predict next state
        predicted_state = self.world_model.predict_next_state(current_state, candidate_action)

        # Evaluate distance to goal
        dist = np.linalg.norm(predicted_state - goal_state)

        if dist < min_dist:
            min_dist = dist  # type: ignore[assignment]
            best_action = candidate_action

    return best_action  # type: ignore[return-value]

plan_sequence(current_state, goal_state, horizon=5)

Simple greedy planning for a sequence of actions.

Source code in src/sc_neurocore/world_model/planner.py
Python
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def plan_sequence(
    self,
    current_state: np.ndarray[Any, Any],
    goal_state: np.ndarray[Any, Any],
    horizon: int = 5,
) -> List[np.ndarray[Any, Any]]:
    """
    Simple greedy planning for a sequence of actions.
    """
    plan = []
    curr_s = current_state
    for _ in range(horizon):
        action = self.propose_action(curr_s, goal_state)
        plan.append(action)
        curr_s = self.world_model.predict_next_state(curr_s, action)
    return plan

PredictiveWorldModel dataclass

Forecast latent-state means and covariances through an LGSSM.

Parameters

state_dim : int Positive latent-state dimension. action_dim : int Non-negative action dimension. seed : int, default=42 Seed used to initialise the stable random LGSSM.

Notes

This class preserves the historical planning-facing API. Use :class:LinearGaussianSSM, :class:KalmanFilter, and :class:RTSSmoother when observations are available.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
@dataclass
class PredictiveWorldModel:
    """Forecast latent-state means and covariances through an LGSSM.

    Parameters
    ----------
    state_dim : int
        Positive latent-state dimension.
    action_dim : int
        Non-negative action dimension.
    seed : int, default=42
        Seed used to initialise the stable random LGSSM.

    Notes
    -----
    This class preserves the historical planning-facing API. Use
    :class:`LinearGaussianSSM`, :class:`KalmanFilter`, and
    :class:`RTSSmoother` when observations are available.

    """

    state_dim: int
    action_dim: int
    seed: int = 42

    def __post_init__(self) -> None:
        """Initialise a validated stable state-transition model."""
        self.state_dim = _require_dimension(
            self.state_dim,
            name="state_dim",
            allow_zero=False,
        )
        self.action_dim = _require_dimension(
            self.action_dim,
            name="action_dim",
            allow_zero=True,
        )
        self.model = LinearGaussianSSM.random(
            state_dim=self.state_dim,
            obs_dim=self.state_dim,
            control_dim=self.action_dim,
            seed=self.seed,
        )
        self._mu = self.model.mu_0.copy()
        self._Sigma = self.model.Sigma_0.copy()

    def reset(self) -> None:
        """Reset the stored belief moments to the model prior."""
        self._mu = self.model.mu_0.copy()
        self._Sigma = self.model.Sigma_0.copy()

    def predict_next_state(
        self,
        current_state: npt.ArrayLike,
        action: npt.ArrayLike,
    ) -> FloatArray:
        """Predict the next latent-state mean.

        Parameters
        ----------
        current_state : array-like, shape (state_dim,)
            Current state estimate.
        action : array-like, shape (action_dim,)
            Current control input. A scalar is accepted when ``action_dim=1``.

        Returns
        -------
        numpy.ndarray, shape (state_dim,)
            Conditional mean ``A x_t + B u_t``.

        Raises
        ------
        ValueError
            If an input has an incompatible shape or non-finite value.

        """
        state = _normalise_vector(
            current_state,
            name="current_state",
            length=self.state_dim,
        )
        control = _normalise_vector(
            action,
            name="action",
            length=self.action_dim,
            allow_scalar=True,
        )
        return np.asarray(self.model.A @ state + self.model.B @ control)

    def predict_next_state_with_cov(
        self,
        current_state: npt.ArrayLike,
        current_cov: npt.ArrayLike,
        action: npt.ArrayLike,
    ) -> tuple[FloatArray, FloatArray]:
        """Predict the next latent-state mean and covariance.

        Parameters
        ----------
        current_state : array-like, shape (state_dim,)
            Current state estimate.
        current_cov : array-like, shape (state_dim, state_dim)
            Symmetric positive-semidefinite current covariance.
        action : array-like, shape (action_dim,)
            Current control input.

        Returns
        -------
        tuple of numpy.ndarray
            Mean ``A x_t + B u_t`` and covariance ``A P_t A^T + Q``.

        """
        state_covariance = _normalise_state_covariance(
            current_cov,
            state_dim=self.state_dim,
            name="current_cov",
        )
        mean = self.predict_next_state(current_state, action)
        covariance = _symmetrise(self.model.A @ state_covariance @ self.model.A.T + self.model.Q)
        return mean, covariance

    def forecast(
        self,
        initial_state: npt.ArrayLike,
        actions: list[npt.ArrayLike],
    ) -> list[FloatArray]:
        """Forecast a deterministic mean trajectory.

        Parameters
        ----------
        initial_state : array-like, shape (state_dim,)
            State before the first action.
        actions : list of array-like
            Ordered actions, one per returned state.

        Returns
        -------
        list of numpy.ndarray
            Independent state arrays after each action.

        """
        state = _normalise_vector(
            initial_state,
            name="initial_state",
            length=self.state_dim,
        )
        trajectory: list[FloatArray] = []
        for action in actions:
            state = self.predict_next_state(state, action)
            trajectory.append(state.copy())
        return trajectory

    def forecast_with_cov(
        self,
        initial_state: npt.ArrayLike,
        initial_cov: npt.ArrayLike,
        actions: list[npt.ArrayLike],
    ) -> list[tuple[FloatArray, FloatArray]]:
        """Forecast a mean and covariance trajectory.

        Parameters
        ----------
        initial_state : array-like, shape (state_dim,)
            State before the first action.
        initial_cov : array-like, shape (state_dim, state_dim)
            Initial symmetric positive-semidefinite covariance.
        actions : list of array-like
            Ordered actions, one per returned state.

        Returns
        -------
        list of tuple of numpy.ndarray
            Independent ``(mean, covariance)`` pairs after each action.

        """
        state = _normalise_vector(
            initial_state,
            name="initial_state",
            length=self.state_dim,
        )
        covariance = _normalise_state_covariance(
            initial_cov,
            state_dim=self.state_dim,
            name="initial_cov",
        )
        trajectory: list[tuple[FloatArray, FloatArray]] = []
        for action in actions:
            state, covariance = self.predict_next_state_with_cov(
                state,
                covariance,
                action,
            )
            trajectory.append((state.copy(), covariance.copy()))
        return trajectory

__post_init__()

Initialise a validated stable state-transition model.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def __post_init__(self) -> None:
    """Initialise a validated stable state-transition model."""
    self.state_dim = _require_dimension(
        self.state_dim,
        name="state_dim",
        allow_zero=False,
    )
    self.action_dim = _require_dimension(
        self.action_dim,
        name="action_dim",
        allow_zero=True,
    )
    self.model = LinearGaussianSSM.random(
        state_dim=self.state_dim,
        obs_dim=self.state_dim,
        control_dim=self.action_dim,
        seed=self.seed,
    )
    self._mu = self.model.mu_0.copy()
    self._Sigma = self.model.Sigma_0.copy()

reset()

Reset the stored belief moments to the model prior.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
74
75
76
77
def reset(self) -> None:
    """Reset the stored belief moments to the model prior."""
    self._mu = self.model.mu_0.copy()
    self._Sigma = self.model.Sigma_0.copy()

predict_next_state(current_state, action)

Predict the next latent-state mean.

Parameters

current_state : array-like, shape (state_dim,) Current state estimate. action : array-like, shape (action_dim,) Current control input. A scalar is accepted when action_dim=1.

Returns

numpy.ndarray, shape (state_dim,) Conditional mean A x_t + B u_t.

Raises

ValueError If an input has an incompatible shape or non-finite value.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
def predict_next_state(
    self,
    current_state: npt.ArrayLike,
    action: npt.ArrayLike,
) -> FloatArray:
    """Predict the next latent-state mean.

    Parameters
    ----------
    current_state : array-like, shape (state_dim,)
        Current state estimate.
    action : array-like, shape (action_dim,)
        Current control input. A scalar is accepted when ``action_dim=1``.

    Returns
    -------
    numpy.ndarray, shape (state_dim,)
        Conditional mean ``A x_t + B u_t``.

    Raises
    ------
    ValueError
        If an input has an incompatible shape or non-finite value.

    """
    state = _normalise_vector(
        current_state,
        name="current_state",
        length=self.state_dim,
    )
    control = _normalise_vector(
        action,
        name="action",
        length=self.action_dim,
        allow_scalar=True,
    )
    return np.asarray(self.model.A @ state + self.model.B @ control)

predict_next_state_with_cov(current_state, current_cov, action)

Predict the next latent-state mean and covariance.

Parameters

current_state : array-like, shape (state_dim,) Current state estimate. current_cov : array-like, shape (state_dim, state_dim) Symmetric positive-semidefinite current covariance. action : array-like, shape (action_dim,) Current control input.

Returns

tuple of numpy.ndarray Mean A x_t + B u_t and covariance A P_t A^T + Q.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def predict_next_state_with_cov(
    self,
    current_state: npt.ArrayLike,
    current_cov: npt.ArrayLike,
    action: npt.ArrayLike,
) -> tuple[FloatArray, FloatArray]:
    """Predict the next latent-state mean and covariance.

    Parameters
    ----------
    current_state : array-like, shape (state_dim,)
        Current state estimate.
    current_cov : array-like, shape (state_dim, state_dim)
        Symmetric positive-semidefinite current covariance.
    action : array-like, shape (action_dim,)
        Current control input.

    Returns
    -------
    tuple of numpy.ndarray
        Mean ``A x_t + B u_t`` and covariance ``A P_t A^T + Q``.

    """
    state_covariance = _normalise_state_covariance(
        current_cov,
        state_dim=self.state_dim,
        name="current_cov",
    )
    mean = self.predict_next_state(current_state, action)
    covariance = _symmetrise(self.model.A @ state_covariance @ self.model.A.T + self.model.Q)
    return mean, covariance

forecast(initial_state, actions)

Forecast a deterministic mean trajectory.

Parameters

initial_state : array-like, shape (state_dim,) State before the first action. actions : list of array-like Ordered actions, one per returned state.

Returns

list of numpy.ndarray Independent state arrays after each action.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def forecast(
    self,
    initial_state: npt.ArrayLike,
    actions: list[npt.ArrayLike],
) -> list[FloatArray]:
    """Forecast a deterministic mean trajectory.

    Parameters
    ----------
    initial_state : array-like, shape (state_dim,)
        State before the first action.
    actions : list of array-like
        Ordered actions, one per returned state.

    Returns
    -------
    list of numpy.ndarray
        Independent state arrays after each action.

    """
    state = _normalise_vector(
        initial_state,
        name="initial_state",
        length=self.state_dim,
    )
    trajectory: list[FloatArray] = []
    for action in actions:
        state = self.predict_next_state(state, action)
        trajectory.append(state.copy())
    return trajectory

forecast_with_cov(initial_state, initial_cov, actions)

Forecast a mean and covariance trajectory.

Parameters

initial_state : array-like, shape (state_dim,) State before the first action. initial_cov : array-like, shape (state_dim, state_dim) Initial symmetric positive-semidefinite covariance. actions : list of array-like Ordered actions, one per returned state.

Returns

list of tuple of numpy.ndarray Independent (mean, covariance) pairs after each action.

Source code in src/sc_neurocore/world_model/_predictive_world_model.py
Python
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def forecast_with_cov(
    self,
    initial_state: npt.ArrayLike,
    initial_cov: npt.ArrayLike,
    actions: list[npt.ArrayLike],
) -> list[tuple[FloatArray, FloatArray]]:
    """Forecast a mean and covariance trajectory.

    Parameters
    ----------
    initial_state : array-like, shape (state_dim,)
        State before the first action.
    initial_cov : array-like, shape (state_dim, state_dim)
        Initial symmetric positive-semidefinite covariance.
    actions : list of array-like
        Ordered actions, one per returned state.

    Returns
    -------
    list of tuple of numpy.ndarray
        Independent ``(mean, covariance)`` pairs after each action.

    """
    state = _normalise_vector(
        initial_state,
        name="initial_state",
        length=self.state_dim,
    )
    covariance = _normalise_state_covariance(
        initial_cov,
        state_dim=self.state_dim,
        name="initial_cov",
    )
    trajectory: list[tuple[FloatArray, FloatArray]] = []
    for action in actions:
        state, covariance = self.predict_next_state_with_cov(
            state,
            covariance,
            action,
        )
        trajectory.append((state.copy(), covariance.copy()))
    return trajectory