diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd123dee..2809daaa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add a general probabilistic forecasting interface (`ProbabilisticForecaster`,
+  `ProbabilisticARForecaster`, `ProbabilisticForecasterModule`) and move
+  ownership of the training objective, scoring rule and per-variable std
+  from `ForecasterModule` onto the `Forecaster`.
+  [\#685](https://github.com/mllam/neural-lam/issues/685)
+  @Sir-Sloth-The-Lazy
+
 - Add `PropagationNet` GNN layer that incentivises directional message
   propagation from sender to receiver nodes, and expose it alongside
   `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`,
diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py
index cb87d76d..1bfe9eb5 100644
--- a/neural_lam/models/__init__.py
+++ b/neural_lam/models/__init__.py
@@ -3,7 +3,12 @@
 # Local
 from .forecasters.autoregressive import ARForecaster
 from .forecasters.base import Forecaster
+from .forecasters.probabilistic import (
+    ProbabilisticARForecaster,
+    ProbabilisticForecaster,
+)
 from .module import ForecasterModule
+from .probabilistic_module import ProbabilisticForecasterModule
 from .step_predictors.base import StepPredictor
 from .step_predictors.graph.base import BaseGraphModel
 from .step_predictors.graph.graph_lam import GraphLAM
diff --git a/neural_lam/models/forecasters/__init__.py b/neural_lam/models/forecasters/__init__.py
index 7ea9f6fd..254c4ba0 100644
--- a/neural_lam/models/forecasters/__init__.py
+++ b/neural_lam/models/forecasters/__init__.py
@@ -5,3 +5,4 @@
 # Local
 from .autoregressive import ARForecaster
 from .base import Forecaster
+from .probabilistic import ProbabilisticARForecaster, ProbabilisticForecaster
diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py
index a8135f62..74f2dc77 100644
--- a/neural_lam/models/forecasters/autoregressive.py
+++ b/neural_lam/models/forecasters/autoregressive.py
@@ -1,12 +1,13 @@
 """Forecaster that uses an auto-regressive strategy to unroll a forecast."""
 
-# Standard library
-
 # Third-party
 import torch
 
 # Local
+from ... import metrics
+from ...config import NeuralLAMConfig
 from ...datastore import BaseDatastore
+from ...loss_weighting import get_state_feature_weighting
 from ..step_predictors.base import StepPredictor
 from .base import Forecaster
 
@@ -18,7 +19,11 @@ class ARForecaster(Forecaster):
     """
 
     def __init__(
-        self, predictor: StepPredictor, datastore: BaseDatastore
+        self,
+        predictor: StepPredictor,
+        datastore: BaseDatastore,
+        config: NeuralLAMConfig | None = None,
+        loss: str = "wmse",
     ) -> None:
         """
         Initialize the ARForecaster.
@@ -29,6 +34,14 @@ def __init__(
             The predictor to use for each step.
         datastore : BaseDatastore
             The datastore providing grid metadata and boundary masks.
+        config : NeuralLAMConfig or None
+            Configuration used to compute the constant per-variable std
+            substituted for ``pred_std`` when ``predictor`` does not output
+            its own (see ``per_var_std``). Only required for that case;
+            forecasters used purely for inference can omit it.
+        loss : str, default "wmse"
+            The scoring rule (from ``neural_lam.metrics``) used by
+            ``compute_training_loss`` and stored as ``self.loss``.
         """
         super().__init__()
         self.predictor = predictor
@@ -44,6 +57,31 @@ def __init__(
             "interior_mask", 1.0 - self.boundary_mask, persistent=False
         )
 
+        self.loss = metrics.get_metric(loss)
+
+        # Store per_var_std here if the predictor does not output its own std
+        if not self.predicts_std and config is not None:
+            da_state_stats = datastore.get_standardization_dataarray(
+                category="state"
+            )
+            state_feature_weights = get_state_feature_weighting(
+                config=config, datastore=datastore
+            )
+            diff_std = torch.tensor(
+                da_state_stats.state_diff_std_standardized.values,
+                dtype=torch.float32,
+            )
+            feature_weights_t = torch.tensor(
+                state_feature_weights, dtype=torch.float32
+            )
+            self.register_buffer(
+                "per_var_std",
+                diff_std / torch.sqrt(feature_weights_t),
+                persistent=False,
+            )
+        else:
+            self.per_var_std = None
+
     @property
     def predicts_std(self) -> bool:
         """
@@ -137,10 +175,73 @@ def forward(
 
         prediction = torch.stack(prediction_list, dim=1)
         # If predictor outputs std, stack it; otherwise return None so
-        # ForecasterModule can substitute the constant per_var_std
+        # callers can substitute the constant per_var_std
         if pred_std_list:
             pred_std = torch.stack(pred_std_list, dim=1)
         else:
             pred_std = None
 
         return prediction, pred_std
+
+    def compute_training_loss(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        target_states: torch.Tensor,
+        interior_mask_bool: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        Score the deterministic rollout with ``self.loss``.
+
+        Unrolls a single forecast over the full rollout, scores it against
+        the target states on interior nodes and averages over batch and
+        time.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the rollout length,
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_forcing_vars`` is the forcing feature dimension (already
+            concatenated past/current/future windows).
+        target_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            states at each predicted step, used both as the prediction
+            targets and to overwrite boundary nodes during the rollout.
+            Dims: same as the prediction.
+        interior_mask_bool : torch.Tensor
+            Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior
+            nodes; passed as ``mask`` to ``self.loss`` so that only interior
+            nodes are scored.
+
+        Returns
+        -------
+        batch_loss : torch.Tensor
+            Scalar. The scoring rule applied to the rollout, averaged over
+            batch and time.
+        loss_components : dict of {str: torch.Tensor}
+            Empty; the deterministic objective has no separate components.
+        """
+        prediction, pred_std = self(
+            init_states, forcing_features, target_states
+        )
+        if pred_std is None:
+            pred_std = self.per_var_std
+
+        batch_loss = torch.mean(
+            self.loss(
+                prediction,
+                target_states,
+                pred_std,
+                mask=interior_mask_bool,
+            )
+        )
+        return batch_loss, {}
diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py
index 4d957916..8dfb0002 100644
--- a/neural_lam/models/forecasters/base.py
+++ b/neural_lam/models/forecasters/base.py
@@ -79,3 +79,58 @@ def forward(
             per-variable std is substituted upstream by
             ``ForecasterModule``. Dims: same as ``prediction``.
         """
+
+    @abstractmethod
+    def compute_training_loss(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        target_states: torch.Tensor,
+        interior_mask_bool: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        Compute the training objective for one batch.
+
+        The forecaster owns its complete training objective: which forecasts
+        to produce from the batch, which loss terms to compute from them and
+        how to combine those terms into a single scalar, using its own
+        ``self.loss`` scoring rule and ``self.per_var_std`` fallback std. The
+        wrapping ``ForecasterModule`` only injects the interior mask, logs
+        the returned components and optimizes the returned loss.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the rollout length,
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_forcing_vars`` is the forcing feature dimension (already
+            concatenated past/current/future windows).
+        target_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            states at each predicted step, used both as the prediction
+            targets and to overwrite boundary nodes during forecasting.
+            Dims: same as the prediction.
+        interior_mask_bool : torch.Tensor
+            Shape ``(num_grid_nodes,)``, boolean. ``True`` for interior
+            nodes; passed as ``mask`` to ``self.loss`` so that only interior
+            nodes are scored.
+
+        Returns
+        -------
+        batch_loss : torch.Tensor
+            Scalar. The full training loss for the batch, to take gradients
+            of.
+        loss_components : dict of {str: torch.Tensor}
+            Scalar loss-related quantities to log alongside the loss, keyed
+            by component name. The wrapping module prefixes the names with
+            the training phase. Empty when the objective has no separate
+            components worth logging.
+        """
diff --git a/neural_lam/models/forecasters/probabilistic.py b/neural_lam/models/forecasters/probabilistic.py
new file mode 100644
index 00000000..1c5aeeb4
--- /dev/null
+++ b/neural_lam/models/forecasters/probabilistic.py
@@ -0,0 +1,227 @@
+"""Forecasters producing probabilistic (ensemble) forecasts."""
+
+# Standard library
+from abc import abstractmethod
+
+# Third-party
+import torch
+
+# Local
+from ...config import NeuralLAMConfig
+from ...datastore import BaseDatastore
+from ..step_predictors.base import StepPredictor
+from .autoregressive import ARForecaster
+from .base import Forecaster
+
+
+class ProbabilisticForecaster(Forecaster):
+    """
+    Forecaster whose forecasts are samples from a predictive distribution.
+
+    Adds the capability that probabilistic evaluation and ensemble-based
+    objectives build on: sampling an ensemble of forecasts. How the
+    members are produced (auto-regressive sampling, diffusion, ...) is
+    left to subclasses; consumers only rely on the shape of the returned
+    ensemble.
+
+    When ``sample_ensemble`` returns a ``per_member_std``, it is each
+    member's own predicted std, not a std describing the spread across
+    members. The predictive distribution is then a mixture of ``S``
+    Gaussians, one per member: ``p(x) = mean_s N(x; ensemble[:, s],
+    per_member_std[:, s]**2)``, not a single Gaussian. In particular, the
+    variance of that mixture is not the average of the per-member
+    variances: it also includes the spread between the member means.
+    """
+
+    @abstractmethod
+    def sample_ensemble(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        boundary_states: torch.Tensor,
+        num_members: int,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Sample an ensemble of forecasts.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the rollout length,
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_forcing_vars`` is the forcing feature dimension (already
+            concatenated past/current/future windows).
+        boundary_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            state values used only to overwrite boundary nodes at each
+            predicted step, identically in every member. Dims: same as one
+            member.
+        num_members : int
+            Number of ensemble members ``S`` to sample.
+
+        Returns
+        -------
+        ensemble : torch.Tensor
+            Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``.
+            The sampled forecasts, stacked along the ensemble dimension
+            ``S``.
+        per_member_std : torch.Tensor or None
+            Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``.
+            Each member's own predicted std (see the class docstring for
+            why the ensemble is then a mixture, not this averaged with the
+            others), when the forecaster predicts an std, otherwise
+            ``None``. Dims: same as ``ensemble``.
+        """
+
+
+class ProbabilisticARForecaster(ARForecaster, ProbabilisticForecaster):
+    """
+    Auto-regressive forecaster for step predictors that sample their output.
+
+    Each call to the wrapped predictor draws a fresh sample of the next
+    state, so the inherited ``ARForecaster.forward`` unrolls one sampled
+    trajectory. This class adds ensemble forecasting on top: unrolling
+    several trajectories and stacking them along an ensemble dimension.
+
+    ``compute_training_loss`` is intentionally left abstract here (it does
+    not fall back to ``ARForecaster``'s single-rollout objective, which
+    would silently train on one stochastic sample). There is no default
+    objective that fits every stochastic model: scoring the ensemble mean
+    with a pointwise metric only rewards the mean being right, giving the
+    model no incentive to keep a calibrated spread, and risks training it
+    to collapse the ensemble to a point estimate. Concrete subclasses must
+    define an objective appropriate to how they are meant to be trained
+    (e.g. an ensemble scoring rule such as CRPS, or a variational
+    objective).
+    """
+
+    def __init__(
+        self,
+        predictor: StepPredictor,
+        datastore: BaseDatastore,
+        config: NeuralLAMConfig | None = None,
+        loss: str = "wmse",
+    ) -> None:
+        """
+        Initialize the ProbabilisticARForecaster.
+
+        Parameters
+        ----------
+        predictor : StepPredictor
+            The predictor to use for each step. Each call should draw a
+            fresh sample of the next state.
+        datastore : BaseDatastore
+            The datastore providing grid metadata and boundary masks.
+        config : NeuralLAMConfig or None
+            Configuration used to compute the constant per-variable std
+            substituted for ``pred_std`` when ``predictor`` does not output
+            its own (see ``per_var_std``). Only required for that case;
+            forecasters used purely for inference can omit it.
+        loss : str, default "wmse"
+            The scoring rule (from ``neural_lam.metrics``) used by
+            ``compute_training_loss`` and stored as ``self.loss``.
+        """
+        super().__init__(predictor, datastore, config=config, loss=loss)
+
+    def sample_ensemble(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        boundary_states: torch.Tensor,
+        num_members: int,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """
+        Sample an ensemble of forecasts.
+
+        Unrolls ``num_members`` independent forecasts, each sampling fresh
+        randomness at every step, and stacks them along a new ensemble
+        dimension after the batch dimension.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start each rollout from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the rollout length,
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_forcing_vars`` is the forcing feature dimension (already
+            concatenated past/current/future windows).
+        boundary_states : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            state values used only to overwrite boundary nodes at each AR
+            step, identically in every member. Dims: same as one member.
+        num_members : int
+            Number of ensemble members ``S`` to sample.
+
+        Returns
+        -------
+        ensemble : torch.Tensor
+            Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``.
+            The sampled forecasts, stacked along the ensemble dimension
+            ``S``.
+        per_member_std : torch.Tensor or None
+            Shape ``(B, S, pred_steps, num_grid_nodes, num_state_vars)``.
+            Each member's own predicted std (see the class docstring for
+            why the ensemble is then a mixture, not this averaged with the
+            others), when the wrapped predictor outputs an std, otherwise
+            ``None``. Dims: same as ``ensemble``.
+
+        Raises
+        ------
+        ValueError
+            If ``num_members`` is less than 1.
+        """
+        if num_members < 1:
+            raise ValueError(
+                f"num_members must be at least 1, got {num_members}"
+            )
+
+        member_list = []
+        member_std_list = []
+        for _ in range(num_members):
+            prediction, pred_std = self(
+                init_states, forcing_features, boundary_states
+            )
+            member_list.append(prediction)
+            if pred_std is not None:
+                member_std_list.append(pred_std)
+
+        ensemble = torch.stack(member_list, dim=1)
+        # After stacking, ensemble has shape
+        # (B, S, pred_steps, num_grid_nodes, num_state_vars)
+        per_member_std = (
+            torch.stack(member_std_list, dim=1) if member_std_list else None
+        )
+        return ensemble, per_member_std
+
+    @abstractmethod
+    def compute_training_loss(
+        self,
+        init_states: torch.Tensor,
+        forcing_features: torch.Tensor,
+        target_states: torch.Tensor,
+        interior_mask_bool: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """
+        Compute the training objective for one batch.
+
+        Left abstract; see the class docstring for why there is no default
+        objective. Concrete subclasses typically call ``sample_ensemble``
+        and score the resulting members with an objective appropriate to
+        the model (see ``Forecaster.compute_training_loss`` for the
+        signature and general contract).
+        """
diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py
index 71ce7951..e7420b52 100644
--- a/neural_lam/models/module.py
+++ b/neural_lam/models/module.py
@@ -20,7 +20,6 @@
 from .. import metrics, vis
 from ..config import NeuralLAMConfig
 from ..datastore import BaseDatastore
-from ..loss_weighting import get_state_feature_weighting
 from ..weather_dataset import WeatherDataset
 from .forecasters.base import Forecaster
 
@@ -38,7 +37,6 @@ def __init__(
         forecaster: Forecaster,
         config: NeuralLAMConfig,
         datastore: BaseDatastore,
-        loss: str = "wmse",
         lr: float = 1e-3,
         restore_opt: bool = False,
         n_example_pred: int = 1,
@@ -54,13 +52,14 @@ def __init__(
         Parameters
         ----------
         forecaster : Forecaster
-            The forecaster model to use for predictions.
+            The forecaster model to use for predictions. Owns the scoring
+            rule (``forecaster.loss``) and the constant per-variable std
+            fallback (``forecaster.per_var_std``) used for training and for
+            validation/test loss reporting here.
         config : NeuralLAMConfig
             Configuration object for the neural LAM model.
         datastore : BaseDatastore
             Datastore providing grid metadata and data access.
-        loss : str, default "wmse"
-            The loss function to use.
         lr : float, default 1e-3
             Learning rate for the optimizer.
         restore_opt : bool, default False
@@ -79,7 +78,7 @@ def __init__(
         args : argparse.Namespace, optional
             Pre-refactor ``ARModel`` checkpoint hyperparameters. When
             provided, attributes on ``args`` take precedence over the
-            corresponding explicit kwargs (``loss``, ``lr``, ``restore_opt``,
+            corresponding explicit kwargs (``lr``, ``restore_opt``,
             ``n_example_pred``, ``create_gif``, ``val_steps_to_log``,
             ``metrics_watch``, ``var_leads_metrics_watch``) so legacy
             checkpoints round-trip through ``load_from_checkpoint``
@@ -90,10 +89,9 @@ def __init__(
         # inside an argparse Namespace under the single key 'args'. When
         # Lightning calls __init__ during load_from_checkpoint it would
         # otherwise drop 'args' (not in the new signature) and silently fall
-        # back to defaults for loss/lr/create_gif/etc. Unpack the namespace
-        # here so legacy checkpoints round-trip correctly.
+        # back to defaults for lr/create_gif/etc. Unpack the namespace here
+        # so legacy checkpoints round-trip correctly.
         if args is not None:
-            loss = getattr(args, "loss", loss)
             lr = getattr(args, "lr", lr)
             restore_opt = getattr(args, "restore_opt", restore_opt)
             n_example_pred = getattr(args, "n_example_pred", n_example_pred)
@@ -124,6 +122,14 @@ def __init__(
         self.save_hyperparameters(ignore=["datastore", "forecaster"])
         self.datastore = datastore
         self.forecaster = forecaster
+        if forecaster.per_var_std is None and not forecaster.predicts_std:
+            raise ValueError(
+                "forecaster.per_var_std is None but the forecaster does "
+                "not predict its own std (forecaster.predicts_std is "
+                "False), so training/validation/test scoring has no std "
+                "to use. Pass config to the forecaster's constructor so "
+                "it can compute the constant per-variable std."
+            )
         self.matched_metrics: set = set()
 
         # Compute interior_mask_bool directly from datastore
@@ -139,29 +145,6 @@ def __init__(
             persistent=False,
         )
 
-        # Store per_var_std here if predictor does not output std
-        if not self.forecaster.predicts_std:
-            da_state_stats = datastore.get_standardization_dataarray(
-                category="state"
-            )
-            state_feature_weights = get_state_feature_weighting(
-                config=config, datastore=datastore
-            )
-            diff_std = torch.tensor(
-                da_state_stats.state_diff_std_standardized.values,
-                dtype=torch.float32,
-            )
-            feature_weights_t = torch.tensor(
-                state_feature_weights, dtype=torch.float32
-            )
-            self.register_buffer(
-                "per_var_std",
-                diff_std / torch.sqrt(feature_weights_t),
-                persistent=False,
-            )
-        else:
-            self.per_var_std = None
-
         # Standardization statistics used to normalize each batch on-device in
         # `on_after_batch_transfer`. WeatherDataset returns unstandardized
         # data, so state and forcing are normalized here rather than on CPU.
@@ -207,9 +190,6 @@ def __init__(
             self.forcing_mean = None
             self.forcing_std = None
 
-        # Instantiate loss function
-        self.loss = metrics.get_metric(loss)
-
         self.val_metrics: dict[str, list] = {
             "mse": [],
         }
@@ -338,7 +318,7 @@ def on_after_batch_transfer(self, batch, dataloader_idx):
 
     def common_step(self, batch):
         """
-        Perform a common prediction step for training, validation, and testing.
+        Perform a common prediction step for validation and testing.
 
         Parameters
         ----------
@@ -362,6 +342,11 @@ def training_step(self, batch):
         """
         Perform a single training step.
 
+        The training objective is fully assembled by the wrapped forecaster,
+        which owns its own scoring rule; this method injects the interior
+        mask, then logs the loss and any loss components the forecaster
+        returns.
+
         Parameters
         ----------
         batch : tuple
@@ -372,20 +357,18 @@ def training_step(self, batch):
         torch.Tensor
             The computed loss for the training step.
         """
-        prediction, target_states, pred_std, _ = self.common_step(batch)
-        if pred_std is None:
-            pred_std = self.per_var_std
-
-        batch_loss = torch.mean(
-            self.loss(
-                prediction,
-                target_states,
-                pred_std,
-                mask=self.interior_mask_bool,
-            )
+        init_states, target_states, forcing_features, _ = batch
+        batch_loss, loss_components = self.forecaster.compute_training_loss(
+            init_states,
+            forcing_features,
+            target_states,
+            interior_mask_bool=self.interior_mask_bool,
         )
 
-        log_dict = {"train_loss": batch_loss}
+        log_dict = {
+            f"train_{name}": value for name, value in loss_components.items()
+        }
+        log_dict["train_loss"] = batch_loss
         self.log_dict(
             log_dict,
             prog_bar=True,
@@ -448,10 +431,10 @@ def validation_step(self, batch, batch_idx):
         """
         prediction, target_states, pred_std, _ = self.common_step(batch)
         if pred_std is None:
-            pred_std = self.per_var_std
+            pred_std = self.forecaster.per_var_std
 
         time_step_loss = torch.mean(
-            self.loss(
+            self.forecaster.loss(
                 prediction,
                 target_states,
                 pred_std,
@@ -528,10 +511,10 @@ def test_step(self, batch, batch_idx):
             self.test_metrics["output_std"].append(mean_pred_std)
 
         if pred_std is None:
-            pred_std = self.per_var_std
+            pred_std = self.forecaster.per_var_std
 
         time_step_loss = torch.mean(
-            self.loss(
+            self.forecaster.loss(
                 prediction,
                 target_states,
                 pred_std,
@@ -568,7 +551,7 @@ def test_step(self, batch, batch_idx):
             )
             self.test_metrics[metric_name].append(batch_metric_vals)
 
-        spatial_loss = self.loss(
+        spatial_loss = self.forecaster.loss(
             prediction, target_states, pred_std, average_grid=False
         )
         log_spatial_losses = spatial_loss[
@@ -976,15 +959,20 @@ def on_load_checkpoint(self, checkpoint):
         # 1. Broad namespace remap: for pre-refactor checkpoints
         # The old ``ARModel`` was a flat LightningModule. Everything that
         # belonged to the predictor needs to be moved to
-        # 'forecaster.predictor.'
+        # 'forecaster.predictor.', while 'per_var_std' (now owned by the
+        # forecaster itself) moves to 'forecaster.per_var_std' and
+        # 'interior_mask_bool' (still owned by the module) stays as-is.
         old_keys = list(loaded_state_dict.keys())
         for key in old_keys:
-            if not key.startswith("forecaster.") and key not in (
-                "interior_mask_bool",
-                "per_var_std",
-            ):
-                new_key = f"forecaster.predictor.{key}"
-                loaded_state_dict[new_key] = loaded_state_dict.pop(key)
+            if key.startswith("forecaster.") or key == "interior_mask_bool":
+                continue
+            if key == "per_var_std":
+                loaded_state_dict["forecaster.per_var_std"] = (
+                    loaded_state_dict.pop(key)
+                )
+                continue
+            new_key = f"forecaster.predictor.{key}"
+            loaded_state_dict[new_key] = loaded_state_dict.pop(key)
 
         # 2. Specific rename from g2m_gnn.grid_mlp -> encoding_grid_mlp
         # Will be under forecaster.predictor due to the remap above, or
diff --git a/neural_lam/models/probabilistic_module.py b/neural_lam/models/probabilistic_module.py
new file mode 100644
index 00000000..453b00c6
--- /dev/null
+++ b/neural_lam/models/probabilistic_module.py
@@ -0,0 +1,189 @@
+"""Lightning module evaluating probabilistic forecasters as ensembles."""
+
+# Standard library
+import warnings
+
+# Third-party
+import torch
+
+# Local
+from .. import metrics
+from .forecasters.probabilistic import ProbabilisticForecaster
+from .module import ForecasterModule
+
+
+class ProbabilisticForecasterModule(ForecasterModule):
+    """
+    Lightning module for forecasters that sample ensemble forecasts.
+
+    Training is inherited unchanged from ``ForecasterModule``: the wrapped
+    forecaster assembles its own training loss. Validation and testing are
+    ensemble based instead of deterministic: an ensemble is sampled from
+    the forecaster and scored through its ensemble mean (root-mean-squared
+    error of the ensemble mean). The module only assumes that the
+    forecaster can sample ensemble forecasts of the correct shape; it makes
+    no assumption on how the members are produced.
+    """
+
+    # The wrapped forecaster must be able to sample ensemble forecasts
+    forecaster: ProbabilisticForecaster
+
+    def __init__(self, *args, eval_ensemble_size: int, **kwargs):
+        """
+        Initialize the module and store the evaluation ensemble size.
+
+        Parameters
+        ----------
+        *args
+            Positional arguments forwarded to
+            ``ForecasterModule.__init__`` (``forecaster``, ``config``,
+            ``datastore``, ...).
+        eval_ensemble_size : int
+            Number of ensemble members sampled during validation and
+            testing.
+        **kwargs
+            Keyword arguments forwarded to ``ForecasterModule.__init__``
+            (``lr``, ...).
+        """
+        super().__init__(*args, **kwargs)
+        if eval_ensemble_size < 1:
+            raise ValueError(
+                "eval_ensemble_size must be at least 1, "
+                f"got {eval_ensemble_size}"
+            )
+        self.eval_ensemble_size = eval_ensemble_size
+        self.val_metrics = {"ens_mse": []}
+        self.test_metrics = {"ens_mse": []}
+
+    def _ensemble_step(self, batch, phase: str):
+        """
+        Sample an ensemble and score its mean against the target states.
+
+        Shared by ``validation_step`` and ``test_step``: samples
+        ``self.eval_ensemble_size`` members, scores the ensemble mean with
+        plain (unweighted) MSE on interior nodes, logs the root-mean-squared
+        error per configured rollout step and averaged over the rollout
+        under the given phase's prefix.
+
+        Parameters
+        ----------
+        batch : tuple
+            The batch of data.
+        phase : str
+            Logging phase, either ``"val"`` or ``"test"``.
+
+        Returns
+        -------
+        torch.Tensor
+            Per-variable ensemble-mean MSE, shape
+            ``(B, pred_steps, num_state_vars)``, for epoch-end aggregation.
+        """
+        init_states, target_states, forcing_features, _ = batch
+        ensemble, _ = self.forecaster.sample_ensemble(
+            init_states,
+            forcing_features,
+            target_states,
+            num_members=self.eval_ensemble_size,
+        )
+        ensemble_mean = ensemble.mean(dim=1)
+        # metrics.mse ignores the std argument, but requires one
+        std_placeholder = torch.ones(
+            target_states.shape[-1], device=target_states.device
+        )
+
+        time_step_mse = torch.mean(
+            metrics.mse(
+                ensemble_mean,
+                target_states,
+                std_placeholder,
+                mask=self.interior_mask_bool,
+            ),
+            dim=0,
+        )
+        time_step_rmse = torch.sqrt(time_step_mse)
+        mean_rmse = torch.mean(time_step_rmse)
+        self._warn_skipped_val_steps(len(time_step_rmse), phase)
+
+        log_dict = {
+            f"{phase}_loss_unroll{step}": time_step_rmse[step - 1]
+            for step in self.hparams.val_steps_to_log
+            if step <= len(time_step_rmse)
+        }
+        log_dict[f"{phase}_mean_loss"] = mean_rmse
+        self.log_dict(
+            log_dict,
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+            batch_size=batch[0].shape[0],
+        )
+
+        return metrics.mse(
+            ensemble_mean,
+            target_states,
+            std_placeholder,
+            mask=self.interior_mask_bool,
+            sum_vars=False,
+        )
+
+    def validation_step(self, batch, batch_idx):
+        """
+        Perform a single ensemble validation step.
+
+        Scores the ensemble mean against the target states (see
+        ``_ensemble_step``) and collects per-variable ensemble-mean MSE for
+        epoch-end aggregation.
+
+        Parameters
+        ----------
+        batch : tuple
+            The batch of data.
+        batch_idx : int
+            The index of the batch.
+        """
+        entry_mses = self._ensemble_step(batch, "val")
+        self.val_metrics["ens_mse"].append(entry_mses)
+
+    def test_step(self, batch, batch_idx):
+        """
+        Perform a single ensemble test step.
+
+        Scores the ensemble mean against the target states (see
+        ``_ensemble_step``) and collects per-variable ensemble-mean MSE for
+        epoch-end aggregation.
+
+        Parameters
+        ----------
+        batch : tuple
+            The batch of data.
+        batch_idx : int
+            The index of the batch.
+        """
+        entry_mses = self._ensemble_step(batch, "test")
+        self.test_metrics["ens_mse"].append(entry_mses)
+
+    def on_test_epoch_end(self):
+        """
+        Perform actions at the end of the test epoch.
+
+        Aggregates ensemble test metrics. Overrides
+        ``ForecasterModule.on_test_epoch_end``, which also handles spatial
+        loss maps and example plots that ``test_step`` here does not
+        populate.
+        """
+        self.aggregate_and_plot_metrics(self.test_metrics, prefix="test")
+
+        if self.trainer.is_global_zero and self.hparams.metrics_watch:
+            unmatched = set(self.hparams.metrics_watch) - self.matched_metrics
+            if unmatched:
+                warnings.warn(
+                    "The following metrics in --metrics_watch "
+                    "were not found during test phase: "
+                    f"{sorted(unmatched)}. Ensure the metric prefix "
+                    "matches the evaluation mode (expected 'test_')."
+                )
+
+        self.matched_metrics = set()
+
+        for metric_list in self.test_metrics.values():
+            metric_list.clear()
diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py
index f98065c4..d5e86536 100644
--- a/neural_lam/train_model.py
+++ b/neural_lam/train_model.py
@@ -63,7 +63,9 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore):
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(
+        predictor, datastore, config=config, loss=args.loss
+    )
     return ForecasterModule.load_from_checkpoint(
         ckpt_path,
         forecaster=forecaster,
@@ -457,13 +459,14 @@ def main(input_args=None):
         mesh_up_gnn_type=args.mesh_up_gnn_type,
         mesh_down_gnn_type=args.mesh_down_gnn_type,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(
+        predictor, datastore, config=config, loss=args.loss
+    )
 
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss=args.loss,
         lr=args.lr,
         restore_opt=args.restore_opt,
         n_example_pred=args.n_example_pred,
diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
index 2e5f3148..6f114043 100644
--- a/tests/test_checkpoint.py
+++ b/tests/test_checkpoint.py
@@ -50,12 +50,11 @@ def test_saved_checkpoint_excludes_datastore_and_forecaster(tmp_path):
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(predictor, datastore, config=config, loss="mse")
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss="mse",
         lr=1.0e-3,
         n_example_pred=1,
         val_steps_to_log=[1],
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 4b35840e..1941206b 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -235,13 +235,14 @@ def _create_graph():
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
     )
-    forecaster = ARForecaster(predictor, datastore=datastore)
+    forecaster = ARForecaster(
+        predictor, datastore=datastore, config=config, loss=args.loss
+    )
 
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss=args.loss,
         restore_opt=args.restore_opt,
         n_example_pred=args.n_example_pred,
         val_steps_to_log=args.val_steps_to_log,
diff --git a/tests/test_gnn_layers.py b/tests/test_gnn_layers.py
index 04c99003..789297d0 100644
--- a/tests/test_gnn_layers.py
+++ b/tests/test_gnn_layers.py
@@ -73,7 +73,7 @@ def _build_model_and_data(
         output_clamping_upper=config.training.output_clamping.upper,
         **gnn_kwargs,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(predictor, datastore, config=config)
 
     B = 2
     num_grid_nodes = predictor.num_grid_nodes
diff --git a/tests/test_gpu_normalization.py b/tests/test_gpu_normalization.py
index 8d516bfb..b063d626 100644
--- a/tests/test_gpu_normalization.py
+++ b/tests/test_gpu_normalization.py
@@ -26,7 +26,7 @@ def _build_module(datastore):
         )
     )
     predictor = _MockStepPredictor(datastore=datastore, output_std=False)
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(predictor, datastore, config=config)
     return ForecasterModule(
         forecaster=forecaster, config=config, datastore=datastore
     )
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 616d563d..970590be 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -467,13 +467,14 @@ class ModelArgs:
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
     )
-    forecaster = ARForecaster(predictor, datastore=datastore)
+    forecaster = ARForecaster(
+        predictor, datastore=datastore, config=config, loss=args.loss
+    )
 
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss=args.loss,
         restore_opt=args.restore_opt,
         n_example_pred=args.n_example_pred,
         val_steps_to_log=args.val_steps_to_log,
@@ -679,12 +680,11 @@ def _build_metrics_watch_module(datastore, config):
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(predictor, datastore, config=config, loss="mse")
     return ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss="mse",
         lr=1.0e-3,
         restore_opt=False,
         n_example_pred=1,
diff --git a/tests/test_prediction_model_classes.py b/tests/test_prediction_model_classes.py
index 73e2f905..9bc9d9c0 100644
--- a/tests/test_prediction_model_classes.py
+++ b/tests/test_prediction_model_classes.py
@@ -97,13 +97,12 @@ def test_forecaster_module_checkpoint(tmp_path):
         num_future_forcing_steps=1,
         output_std=False,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(predictor, datastore, config=config, loss="mse")
 
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss="mse",
         lr=1e-3,
         restore_opt=False,
         n_example_pred=1,
@@ -133,7 +132,9 @@ def test_forecaster_module_checkpoint(tmp_path):
         num_future_forcing_steps=1,
         output_std=False,
     )
-    load_forecaster = ARForecaster(load_predictor, datastore)
+    load_forecaster = ARForecaster(
+        load_predictor, datastore, config=config, loss="mse"
+    )
 
     # Load from checkpoint
     loaded_model = ForecasterModule.load_from_checkpoint(
@@ -193,8 +194,6 @@ def test_forecaster_module_old_checkpoint(tmp_path):
         num_future_forcing_steps=1,
         output_std=False,
     )
-    forecaster = ARForecaster(predictor, datastore)
-
     # Use distinctive non-default values so we can detect silent fallback
     # to ForecasterModule's defaults during load.
     saved_loss = "mse"
@@ -203,11 +202,14 @@ def test_forecaster_module_old_checkpoint(tmp_path):
     saved_val_steps = [2]
     saved_n_example_pred = 7
 
+    forecaster = ARForecaster(
+        predictor, datastore, config=config, loss=saved_loss
+    )
+
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss=saved_loss,
         lr=saved_lr,
         restore_opt=False,
         n_example_pred=saved_n_example_pred,
@@ -269,7 +271,9 @@ def test_forecaster_module_old_checkpoint(tmp_path):
         num_future_forcing_steps=1,
         output_std=False,
     )
-    load_forecaster = ARForecaster(load_predictor, datastore)
+    load_forecaster = ARForecaster(
+        load_predictor, datastore, config=config, loss=saved_loss
+    )
 
     # Load from hacked old checkpoint
     loaded_model = ForecasterModule.load_from_checkpoint(
@@ -284,7 +288,6 @@ def test_forecaster_module_old_checkpoint(tmp_path):
 
     # Hyperparameters nested in the legacy 'args' namespace must round-trip
     # rather than silently falling back to ForecasterModule defaults.
-    assert loaded_model.hparams.loss == saved_loss
     assert loaded_model.hparams.lr == saved_lr
     assert loaded_model.hparams.val_steps_to_log == saved_val_steps
     assert loaded_model.create_gif is saved_create_gif
diff --git a/tests/test_probabilistic_forecaster.py b/tests/test_probabilistic_forecaster.py
new file mode 100644
index 00000000..aecdb09f
--- /dev/null
+++ b/tests/test_probabilistic_forecaster.py
@@ -0,0 +1,366 @@
+# Third-party
+import pytest
+import torch
+from torch import nn
+
+# First-party
+from neural_lam import config as nlconfig
+from neural_lam import metrics
+from neural_lam.models import (
+    ARForecaster,
+    ForecasterModule,
+    ProbabilisticARForecaster,
+    ProbabilisticForecasterModule,
+    StepPredictor,
+)
+from tests.conftest import init_datastore_example
+
+
+class ZeroStepPredictor(StepPredictor):
+    """Deterministic predictor always predicting the zero state."""
+
+    def forward(self, prev_state, prev_prev_state, forcing):
+        pred_state = torch.zeros_like(prev_state)
+        pred_std = torch.zeros_like(prev_state) if self.output_std else None
+        return pred_state, pred_std
+
+
+class NoisyStepPredictor(StepPredictor):
+    """Stochastic predictor sampling a fresh state at every call."""
+
+    def __init__(self, datastore, **kwargs):
+        super().__init__(datastore, **kwargs)
+        self.noise_scale = nn.Parameter(torch.tensor(1.0))
+
+    def forward(self, prev_state, prev_prev_state, forcing):
+        pred_state = self.noise_scale * torch.randn_like(prev_state)
+        return pred_state, None
+
+
+class ConcreteProbabilisticARForecaster(ProbabilisticARForecaster):
+    """
+    Test-only concrete ``ProbabilisticARForecaster``.
+
+    ``ProbabilisticARForecaster`` leaves ``compute_training_loss`` abstract
+    (no single default objective fits every stochastic model), so tests
+    that only need a working forecaster to instantiate use this example
+    ensemble-mean objective rather than the base class directly.
+    ``sample_ensemble`` always requires an explicit member count, so this
+    class takes its own ``train_num_members`` for the training objective.
+    """
+
+    def __init__(self, *args, train_num_members: int = 2, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.train_num_members = train_num_members
+
+    def compute_training_loss(
+        self,
+        init_states,
+        forcing_features,
+        target_states,
+        interior_mask_bool,
+    ):
+        ensemble, per_member_std = self.sample_ensemble(
+            init_states,
+            forcing_features,
+            target_states,
+            num_members=self.train_num_members,
+        )
+        ensemble_mean = ensemble.mean(dim=1)
+        pred_std = (
+            per_member_std.mean(dim=1)
+            if per_member_std is not None
+            else self.per_var_std
+        )
+        batch_loss = torch.mean(
+            self.loss(
+                ensemble_mean,
+                target_states,
+                pred_std,
+                mask=interior_mask_bool,
+            )
+        )
+        return batch_loss, {}
+
+
+def _example_batch(datastore, B=2, pred_steps=3):
+    """Create constant example input tensors matching the datastore dims."""
+    num_grid_nodes = datastore.num_grid_points
+    d_state = datastore.get_num_data_vars(category="state")
+    num_past_forcing_steps = 1
+    num_future_forcing_steps = 1
+    d_forcing = datastore.get_num_data_vars(category="forcing") * (
+        num_past_forcing_steps + num_future_forcing_steps + 1
+    )
+    init_states = torch.ones(B, 2, num_grid_nodes, d_state)
+    forcing_features = torch.ones(B, pred_steps, num_grid_nodes, d_forcing)
+    target_states = torch.ones(B, pred_steps, num_grid_nodes, d_state) * 5.0
+    return init_states, forcing_features, target_states
+
+
+def test_ar_forecaster_training_loss_matches_direct_score():
+    datastore = init_datastore_example("mdp")
+    predictor = ZeroStepPredictor(datastore=datastore, output_std=False)
+    forecaster = ARForecaster(predictor, datastore, loss="mse")
+
+    init_states, forcing_features, target_states = _example_batch(datastore)
+    score_metric = metrics.get_metric("mse")
+    interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool)
+    d_state = target_states.shape[-1]
+    # per_var_std is normally computed from config; override directly since
+    # this test only cares about the loss computation, not standardization.
+    forecaster.per_var_std = torch.ones(d_state)
+
+    batch_loss, loss_components = forecaster.compute_training_loss(
+        init_states,
+        forcing_features,
+        target_states,
+        interior_mask_bool=interior_mask_bool,
+    )
+
+    prediction, _ = forecaster(init_states, forcing_features, target_states)
+    expected_loss = torch.mean(
+        score_metric(
+            prediction,
+            target_states,
+            forecaster.per_var_std,
+            mask=interior_mask_bool,
+        )
+    )
+
+    assert batch_loss.shape == ()
+    assert loss_components == {}
+    torch.testing.assert_close(batch_loss, expected_loss)
+
+
+def test_sample_ensemble_shapes_and_member_variability():
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+    forecaster = ConcreteProbabilisticARForecaster(predictor, datastore)
+
+    # Override masks to test boundary masking behaviour
+    forecaster.interior_mask = torch.zeros_like(forecaster.interior_mask)
+    forecaster.interior_mask[0, 0] = 1  # One node is interior
+    forecaster.boundary_mask = 1 - forecaster.interior_mask
+
+    B, pred_steps, num_members = 2, 3, 4
+    init_states, forcing_features, target_states = _example_batch(
+        datastore, B=B, pred_steps=pred_steps
+    )
+    num_grid_nodes = datastore.num_grid_points
+    d_state = target_states.shape[-1]
+
+    torch.manual_seed(42)
+    ensemble, per_member_std = forecaster.sample_ensemble(
+        init_states,
+        forcing_features,
+        target_states,
+        num_members=num_members,
+    )
+
+    assert ensemble.shape == (
+        B,
+        num_members,
+        pred_steps,
+        num_grid_nodes,
+        d_state,
+    )
+    assert per_member_std is None
+
+    # Members carry independent samples on the interior node
+    assert not torch.allclose(ensemble[:, 0, :, 0], ensemble[:, 1, :, 0])
+    # Boundary nodes are overwritten with the true state in every member
+    assert torch.all(ensemble[:, :, :, 1:] == 5.0)
+
+
+def test_probabilistic_training_loss_gradient_flow():
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+    forecaster = ConcreteProbabilisticARForecaster(
+        predictor, datastore, loss="mse", train_num_members=2
+    )
+
+    init_states, forcing_features, target_states = _example_batch(datastore)
+    interior_mask_bool = forecaster.interior_mask[0, :, 0].to(torch.bool)
+    d_state = target_states.shape[-1]
+    # per_var_std is normally computed from config; override directly since
+    # this test only cares about the loss computation, not standardization.
+    forecaster.per_var_std = torch.ones(d_state)
+
+    torch.manual_seed(42)
+    batch_loss, loss_components = forecaster.compute_training_loss(
+        init_states,
+        forcing_features,
+        target_states,
+        interior_mask_bool=interior_mask_bool,
+    )
+
+    assert batch_loss.shape == ()
+    assert loss_components == {}
+    assert torch.isfinite(batch_loss)
+
+    batch_loss.backward()
+    assert predictor.noise_scale.grad is not None
+    assert predictor.noise_scale.grad != 0.0
+
+
+def test_sample_ensemble_rejects_empty_member_count():
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+    forecaster = ConcreteProbabilisticARForecaster(predictor, datastore)
+    init_states, forcing_features, target_states = _example_batch(datastore)
+
+    with pytest.raises(ValueError, match="num_members"):
+        forecaster.sample_ensemble(
+            init_states, forcing_features, target_states, num_members=0
+        )
+
+
+def test_probabilistic_ar_forecaster_is_abstract():
+    """ProbabilisticARForecaster leaves compute_training_loss abstract, so
+    it cannot be instantiated directly; only a subclass that defines an
+    objective can."""
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+
+    with pytest.raises(TypeError, match="abstract"):
+        ProbabilisticARForecaster(predictor, datastore)
+
+
+def test_module_training_step_delegates_to_forecaster():
+    datastore = init_datastore_example("mdp")
+    predictor = ZeroStepPredictor(datastore=datastore, output_std=False)
+
+    config = nlconfig.NeuralLAMConfig(
+        datastore=nlconfig.DatastoreSelection(
+            kind=datastore.SHORT_NAME, config_path=datastore.root_path
+        )
+    )
+    forecaster = ARForecaster(predictor, datastore, config=config, loss="mse")
+    model = ForecasterModule(
+        forecaster=forecaster,
+        config=config,
+        datastore=datastore,
+    )
+
+    init_states, forcing_features, target_states = _example_batch(datastore)
+    batch_times = torch.zeros(init_states.shape[0], target_states.shape[1])
+    batch = (init_states, target_states, forcing_features, batch_times)
+
+    batch_loss = model.training_step(batch)
+
+    expected_loss, _ = forecaster.compute_training_loss(
+        init_states,
+        forcing_features,
+        target_states,
+        interior_mask_bool=model.interior_mask_bool,
+    )
+
+    torch.testing.assert_close(batch_loss, expected_loss)
+
+
+class MemberCountRecordingForecaster(ConcreteProbabilisticARForecaster):
+    """ProbabilisticARForecaster recording the requested member count."""
+
+    def sample_ensemble(self, *args, **kwargs):
+        self.last_num_members = kwargs.get("num_members")
+        return super().sample_ensemble(*args, **kwargs)
+
+
+def test_probabilistic_module_validation_scores_ensemble_mean():
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+
+    config = nlconfig.NeuralLAMConfig(
+        datastore=nlconfig.DatastoreSelection(
+            kind=datastore.SHORT_NAME, config_path=datastore.root_path
+        )
+    )
+    forecaster = MemberCountRecordingForecaster(
+        predictor, datastore, config=config
+    )
+    model = ProbabilisticForecasterModule(
+        forecaster=forecaster,
+        config=config,
+        datastore=datastore,
+        eval_ensemble_size=3,
+    )
+
+    B, pred_steps = 2, 3
+    init_states, forcing_features, target_states = _example_batch(
+        datastore, B=B, pred_steps=pred_steps
+    )
+    batch_times = torch.zeros(B, pred_steps)
+    batch = (init_states, target_states, forcing_features, batch_times)
+
+    torch.manual_seed(42)
+    model.validation_step(batch, 0)
+
+    # Validation samples the configured number of evaluation members
+    assert forecaster.last_num_members == 3
+
+    # Ensemble-mean MSE entries are collected for epoch-end aggregation
+    d_state = target_states.shape[-1]
+    (entry_mses,) = model.val_metrics["ens_mse"]
+    assert entry_mses.shape == (B, pred_steps, d_state)
+    assert torch.all(torch.isfinite(entry_mses))
+
+
+def test_probabilistic_module_rejects_empty_eval_ensemble():
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+    config = nlconfig.NeuralLAMConfig(
+        datastore=nlconfig.DatastoreSelection(
+            kind=datastore.SHORT_NAME, config_path=datastore.root_path
+        )
+    )
+    forecaster = ConcreteProbabilisticARForecaster(
+        predictor, datastore, config=config
+    )
+
+    with pytest.raises(ValueError, match="eval_ensemble_size"):
+        ProbabilisticForecasterModule(
+            forecaster=forecaster,
+            config=config,
+            datastore=datastore,
+            eval_ensemble_size=0,
+        )
+
+
+def test_probabilistic_module_test_step_scores_ensemble_mean():
+    datastore = init_datastore_example("mdp")
+    predictor = NoisyStepPredictor(datastore=datastore, output_std=False)
+
+    config = nlconfig.NeuralLAMConfig(
+        datastore=nlconfig.DatastoreSelection(
+            kind=datastore.SHORT_NAME, config_path=datastore.root_path
+        )
+    )
+    forecaster = MemberCountRecordingForecaster(
+        predictor, datastore, config=config
+    )
+    model = ProbabilisticForecasterModule(
+        forecaster=forecaster,
+        config=config,
+        datastore=datastore,
+        eval_ensemble_size=3,
+    )
+
+    B, pred_steps = 2, 3
+    init_states, forcing_features, target_states = _example_batch(
+        datastore, B=B, pred_steps=pred_steps
+    )
+    batch_times = torch.zeros(B, pred_steps)
+    batch = (init_states, target_states, forcing_features, batch_times)
+
+    torch.manual_seed(42)
+    model.test_step(batch, 0)
+
+    # Test samples the configured number of evaluation members
+    assert forecaster.last_num_members == 3
+
+    # Ensemble-mean MSE entries are collected for epoch-end aggregation
+    d_state = target_states.shape[-1]
+    (entry_mses,) = model.test_metrics["ens_mse"]
+    assert entry_mses.shape == (B, pred_steps, d_state)
+    assert torch.all(torch.isfinite(entry_mses))
diff --git a/tests/test_training.py b/tests/test_training.py
index bf1a5884..589e9d89 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -123,13 +123,12 @@ def run_simple_training(
         output_clamping_lower=config.training.output_clamping.lower,
         output_clamping_upper=config.training.output_clamping.upper,
     )
-    forecaster = ARForecaster(predictor, datastore)
+    forecaster = ARForecaster(predictor, datastore, config=config, loss="mse")
 
     model = ForecasterModule(
         forecaster=forecaster,
         config=config,
         datastore=datastore,
-        loss="mse",
         lr=1.0e-3,
         restore_opt=False,
         n_example_pred=1,