Reference

Training

The trainer that warms up reservoir states, then fits every readout in the model algebraically — ridge regression, not gradient descent — in a single forward pass.

training

ESN Training Utilities

This module provides trainers for ESN models that fit readout layers algebraically using ridge regression, rather than stochastic gradient descent.

CLASS	DESCRIPTION
ESNTrainer	Trainer for fitting readout layers in ESN models.

Examples:

>>> from resdag.training import ESNTrainer
>>> trainer = ESNTrainer(model)
>>> trainer.fit(
...     warmup_inputs=(warmup,),
...     train_inputs=(train,),
...     targets={"output": target},
... )

See Also

resdag.layers.readouts.CGReadoutLayer : Readout with CG solver. resdag.core.ESNModel : ESN model class.

ESNTrainer

ESNTrainer(model: ESNModel)

Trainer for ESN models with algebraic readout fitting.

This trainer fits all :class:ReadoutLayer instances in an ESN model using algebraic methods (e.g., ridge regression) rather than gradient descent. The fitting is performed efficiently in a single forward pass using pre-hooks that intercept inputs to each readout.

The training process:

1. Reset reservoir states
2. Warmup phase: synchronize reservoir states with input dynamics
3. Single forward pass with pre-hooks that fit each readout in topological order before it processes its input

Each readout handles its own fitting hyperparameters (e.g., alpha for ridge regression is set during layer construction).

PARAMETER	DESCRIPTION
model	ESN model to train. Must contain at least one :class:ReadoutLayer. TYPE: ESNModel

ATTRIBUTE	DESCRIPTION
model	The ESN model being trained. TYPE: ESNModel

Examples:

Basic training workflow:

>>> from resdag.training import ESNTrainer
>>> from resdag.core import ESNModel
>>>
>>> trainer = ESNTrainer(model)
>>> trainer.fit(
...     warmup_inputs=(warmup_data,),
...     train_inputs=(train_data,),
...     targets={"output": train_targets},
... )

Training with driving inputs:

>>> trainer.fit(
...     warmup_inputs=(warmup_feedback, warmup_driver),
...     train_inputs=(train_feedback, train_driver),
...     targets={"output": targets},
... )

Multi-readout training:

>>> trainer.fit(
...     warmup_inputs=(warmup_data,),
...     train_inputs=(train_data,),
...     targets={
...         "position": position_targets,
...         "velocity": velocity_targets,
...     },
... )

See Also

ESNModel : ESN model class. CGReadoutLayer : Conjugate gradient readout layer. ESNModel.forecast : Forecasting after training.

Notes

• Warmup and training data must have the same number of input tensors.
• Training data and targets must have the same sequence length.
• Target keys must match readout names (user-defined or auto-generated).

Source code in src/resdag/training/trainer.py

def __init__(self, model: ESNModel) -> None:
    self.model = model

fit

fit(warmup_inputs: tuple[Tensor, ...], train_inputs: tuple[Tensor, ...], targets: dict[str, Tensor]) -> None

Fit all readout layers in a single forward pass.

Uses pre-hooks to fit each readout layer just before its forward method executes. This ensures downstream layers receive outputs from already-fitted readouts.

PARAMETER	DESCRIPTION
warmup_inputs	Warmup sequences for state synchronization. Format: (feedback, driver1, driver2, ...). Each tensor shape: (batch, warmup_steps, features). TYPE: tuple of torch.Tensor
train_inputs	Training sequences for fitting. Format: (feedback, driver1, driver2, ...). Each tensor shape: (batch, train_steps, features). Must have same sequence length as targets. TYPE: tuple of torch.Tensor
targets	Dictionary mapping readout name to target tensor. Each target shape: (batch, train_steps, out_features). Names are either user-defined (via name="output" in readout constructor) or auto-generated module names (e.g., "CGReadoutLayer_1"). TYPE: dict of str to torch.Tensor

RAISES	DESCRIPTION
ValueError	If no warmup or training inputs provided. If number of warmup and training inputs don't match. If any readout is missing from targets. If target sequence length doesn't match training inputs.

Notes

After calling fit(), all readouts will have is_fitted=True and the model is ready for inference or forecasting.

Emits a UserWarning if targets dict contains names not matching any readout.

Examples:

>>> trainer = ESNTrainer(model)
>>> trainer.fit(
...     warmup_inputs=(warmup_data,),
...     train_inputs=(train_data,),
...     targets={"output": targets},
... )
>>> print(model.CGReadoutLayer_1.is_fitted)
True

Source code in src/resdag/training/trainer.py

def fit(
    self,
    warmup_inputs: tuple[torch.Tensor, ...],
    train_inputs: tuple[torch.Tensor, ...],
    targets: dict[str, torch.Tensor],
) -> None:
    """
    Fit all readout layers in a single forward pass.

    Uses pre-hooks to fit each readout layer just before its forward
    method executes. This ensures downstream layers receive outputs
    from already-fitted readouts.

    Parameters
    ----------
    warmup_inputs : tuple of torch.Tensor
        Warmup sequences for state synchronization.
        Format: ``(feedback, driver1, driver2, ...)``.
        Each tensor shape: ``(batch, warmup_steps, features)``.
    train_inputs : tuple of torch.Tensor
        Training sequences for fitting.
        Format: ``(feedback, driver1, driver2, ...)``.
        Each tensor shape: ``(batch, train_steps, features)``.
        Must have same sequence length as targets.
    targets : dict of str to torch.Tensor
        Dictionary mapping readout name to target tensor.
        Each target shape: ``(batch, train_steps, out_features)``.
        Names are either user-defined (via ``name="output"`` in readout
        constructor) or auto-generated module names (e.g., ``"CGReadoutLayer_1"``).

    Raises
    ------
    ValueError
        If no warmup or training inputs provided.
        If number of warmup and training inputs don't match.
        If any readout is missing from targets.
        If target sequence length doesn't match training inputs.

    Notes
    -----
    After calling ``fit()``, all readouts will have ``is_fitted=True`` and the model is ready for inference or forecasting.

    Emits a ``UserWarning`` if targets dict contains names not matching any readout.

    Examples
    --------
    >>> trainer = ESNTrainer(model)
    >>> trainer.fit(
    ...     warmup_inputs=(warmup_data,),
    ...     train_inputs=(train_data,),
    ...     targets={"output": targets},
    ... )
    >>> print(model.CGReadoutLayer_1.is_fitted)
    True
    """
    if len(warmup_inputs) == 0:
        raise ValueError("At least one warmup input is required")
    if len(train_inputs) == 0:
        raise ValueError("At least one training input is required")
    if len(warmup_inputs) != len(train_inputs):
        raise ValueError(
            f"warmup_inputs has {len(warmup_inputs)} tensors, "
            f"but train_inputs has {len(train_inputs)} tensors. Must match."
        )

    # Resolve readouts once; reuse for target validation and hook setup.
    readouts = self._discover_readouts()
    self._validate_targets(targets, readouts)

    train_steps = train_inputs[0].shape[1]

    # Validate target shapes
    readout_names = [name for name, _ in readouts]
    for name in readout_names:
        target = targets[name]
        if target.shape[1] != train_steps:
            raise ValueError(
                f"Target for '{name}' has {target.shape[1]} timesteps, "
                f"but train_inputs has {train_steps} timesteps. Must match."
            )

    # Single warmup to sync reservoir states (warmup resets by default).
    # Readouts are unfitted during warmup, so let any guarded readout (e.g.
    # IncrementalRidgeReadoutLayer) flow a value while the reservoir synchronises.
    with self._allow_unfitted_warmup(readouts):
        self.model.warmup(*warmup_inputs)

    # Register pre-hooks that fit each readout when its forward fires.
    # We don't need explicit topological order: ``self.model(*train_inputs)``
    # walks the graph in execution order, so each hook runs exactly when
    # its readout would naturally execute.
    hooks = []
    for name, readout in readouts:
        target = targets[name]

        def make_fit_hook(
            layer: ReadoutLayer, tgt: torch.Tensor
        ) -> Callable[[torch.nn.Module, tuple[torch.Tensor, ...]], None]:
            def hook(module: torch.nn.Module, args: tuple[torch.Tensor, ...]) -> None:
                layer.fit(args[0], tgt)

            return hook

        handle = readout.register_forward_pre_hook(make_fit_hook(readout, target))
        hooks.append(handle)

    try:
        # Single forward pass - hooks fit each readout as it executes
        with torch.no_grad():
            self.model(*train_inputs)
    finally:
        # Always remove hooks
        for h in hooks:
            h.remove()

fit_stream

fit_stream(warmup_inputs: tuple[Tensor, ...], chunks: Iterable[StreamChunk]) -> None

Fit :class:IncrementalRidgeReadoutLayer readouts over a stream of chunks.

The streaming counterpart of :meth:fit. Instead of fitting from one in-memory (B, T, F) block, it warms the reservoir once and then consumes chunks one at a time — e.g. windows yielded by a :class:torch.utils.data.DataLoader — accumulating each readout's ridge sufficient statistics with :meth:~resdag.layers.IncrementalRidgeReadoutLayer.partial_fit and solving once at the end with :meth:~resdag.layers.IncrementalRidgeReadoutLayer.finalize. No more than a single chunk's states are ever held in memory, which is what lets a sequence too long to materialise — or one streamed off disk — be fitted.

Because the reservoir is stateful, the chunks must be contiguous in time and in order: state flows from the end of one chunk into the start of the next, exactly as it would in a single long forward pass. Shuffling the chunks would desynchronise the reservoir and is not supported.

Every readout in the model must be an :class:~resdag.layers.IncrementalRidgeReadoutLayer (the only readout with a partial_fit / finalize interface). The accumulators are reset at the start, so calling fit_stream again re-fits from scratch.

PARAMETER	DESCRIPTION
warmup_inputs	Warmup sequences for state synchronization, format (feedback, driver1, ...), each of shape (batch, warmup_steps, features). The reservoir is reset before this pass (as in :meth:fit). TYPE: tuple of torch.Tensor
chunks	An iterable yielding (inputs, targets) per chunk. inputs is an input tuple (feedback, driver1, ...) of the same arity as warmup_inputs; targets maps each readout name to its target tensor of shape (batch, chunk_steps, out_features). Consumed lazily, so a generator or DataLoader never materialises the whole sequence. TYPE: iterable of (tuple of torch.Tensor, dict of str to torch.Tensor)

RAISES	DESCRIPTION
ValueError	If no warmup inputs are provided, if a chunk's input arity does not match the warmup arity, if any readout is missing from a chunk's targets, or if a target's sequence length does not match its chunk's input length.
TypeError	If any readout in the model is not an IncrementalRidgeReadoutLayer.
RuntimeError	If chunks is empty (no statistics were accumulated, so there is nothing to finalize).

Notes

After fit_stream returns, every readout has is_fitted=True and the model is ready for inference or forecasting.

Examples:

>>> from resdag.layers import IncrementalRidgeReadoutLayer
>>> trainer = ESNTrainer(model)  # model uses IncrementalRidgeReadoutLayer
>>> def chunk_stream():
...     for x, y in dataloader:  # contiguous windows
...         yield (x,), {"output": y}
>>> trainer.fit_stream(
...     warmup_inputs=(warmup_data,),
...     chunks=chunk_stream(),
... )

See Also

fit : Single-pass in-memory fitting. resdag.layers.IncrementalRidgeReadoutLayer : The streaming readout.

Source code in src/resdag/training/trainer.py

def fit_stream(
    self,
    warmup_inputs: tuple[torch.Tensor, ...],
    chunks: Iterable[StreamChunk],
) -> None:
    """Fit :class:`IncrementalRidgeReadoutLayer` readouts over a stream of chunks.

    The streaming counterpart of :meth:`fit`. Instead of fitting from one
    in-memory ``(B, T, F)`` block, it warms the reservoir once and then
    consumes ``chunks`` one at a time — e.g. windows yielded by a
    :class:`torch.utils.data.DataLoader` — accumulating each readout's ridge
    sufficient statistics with
    :meth:`~resdag.layers.IncrementalRidgeReadoutLayer.partial_fit` and solving
    once at the end with
    :meth:`~resdag.layers.IncrementalRidgeReadoutLayer.finalize`. No more than a
    single chunk's states are ever held in memory, which is what lets a
    sequence too long to materialise — or one streamed off disk — be fitted.

    Because the reservoir is stateful, the chunks must be **contiguous in
    time and in order**: state flows from the end of one chunk into the
    start of the next, exactly as it would in a single long forward pass.
    Shuffling the chunks would desynchronise the reservoir and is not
    supported.

    Every readout in the model must be an
    :class:`~resdag.layers.IncrementalRidgeReadoutLayer` (the only readout with a
    ``partial_fit`` / ``finalize`` interface). The accumulators are reset at
    the start, so calling ``fit_stream`` again re-fits from scratch.

    Parameters
    ----------
    warmup_inputs : tuple of torch.Tensor
        Warmup sequences for state synchronization, format
        ``(feedback, driver1, ...)``, each of shape
        ``(batch, warmup_steps, features)``. The reservoir is reset before
        this pass (as in :meth:`fit`).
    chunks : iterable of (tuple of torch.Tensor, dict of str to torch.Tensor)
        An iterable yielding ``(inputs, targets)`` per chunk. ``inputs`` is
        an input tuple ``(feedback, driver1, ...)`` of the same arity as
        ``warmup_inputs``; ``targets`` maps each readout name to its target
        tensor of shape ``(batch, chunk_steps, out_features)``. Consumed
        lazily, so a generator or ``DataLoader`` never materialises the whole
        sequence.

    Raises
    ------
    ValueError
        If no warmup inputs are provided, if a chunk's input arity does not
        match the warmup arity, if any readout is missing from a chunk's
        targets, or if a target's sequence length does not match its chunk's
        input length.
    TypeError
        If any readout in the model is not an ``IncrementalRidgeReadoutLayer``.
    RuntimeError
        If ``chunks`` is empty (no statistics were accumulated, so there is
        nothing to finalize).

    Notes
    -----
    After ``fit_stream`` returns, every readout has ``is_fitted=True`` and
    the model is ready for inference or forecasting.

    Examples
    --------
    >>> from resdag.layers import IncrementalRidgeReadoutLayer
    >>> trainer = ESNTrainer(model)  # model uses IncrementalRidgeReadoutLayer
    >>> def chunk_stream():
    ...     for x, y in dataloader:  # contiguous windows  # doctest: +SKIP
    ...         yield (x,), {"output": y}
    >>> trainer.fit_stream(  # doctest: +SKIP
    ...     warmup_inputs=(warmup_data,),
    ...     chunks=chunk_stream(),
    ... )

    See Also
    --------
    fit : Single-pass in-memory fitting.
    resdag.layers.IncrementalRidgeReadoutLayer : The streaming readout.
    """
    if len(warmup_inputs) == 0:
        raise ValueError("At least one warmup input is required")

    readouts = self._discover_readouts()
    non_incremental = [
        name for name, ro in readouts if not isinstance(ro, IncrementalRidgeReadoutLayer)
    ]
    if non_incremental:
        raise TypeError(
            f"fit_stream() requires every readout to be an IncrementalRidgeReadoutLayer, "
            f"but these are not: {non_incremental}. Use fit() for single-pass readouts, "
            f"or rebuild the model with IncrementalRidgeReadoutLayer."
        )

    # Reset accumulators so a re-fit starts clean, then warm up once. The
    # readouts are unfitted during warmup too, so let their forward flow a
    # value while the reservoir synchronises (shared with ``fit``).
    for _, readout in readouts:
        assert isinstance(readout, IncrementalRidgeReadoutLayer)
        readout.reset_accumulators()
    with self._allow_unfitted_warmup(readouts):
        self.model.warmup(*warmup_inputs)

    # Per-chunk: register pre-hooks that accumulate sufficient statistics
    # for each readout, then run a forward pass that carries the reservoir
    # state forward into the next chunk.
    n_chunks = 0
    for inputs, targets in chunks:
        if len(inputs) != len(warmup_inputs):
            raise ValueError(
                f"Chunk {n_chunks} has {len(inputs)} input tensors, but warmup_inputs "
                f"has {len(warmup_inputs)}. Must match."
            )
        self._validate_targets(targets, readouts)
        chunk_steps = inputs[0].shape[1]

        hooks: list[torch.utils.hooks.RemovableHandle] = []
        for name, readout in readouts:
            target = targets[name]
            if target.shape[1] != chunk_steps:
                for h in hooks:
                    h.remove()
                raise ValueError(
                    f"Target for '{name}' in chunk {n_chunks} has {target.shape[1]} "
                    f"timesteps, but the chunk inputs have {chunk_steps}. Must match."
                )

            def make_partial_fit_hook(
                layer: IncrementalRidgeReadoutLayer, tgt: torch.Tensor
            ) -> Callable[[torch.nn.Module, tuple[torch.Tensor, ...]], None]:
                def hook(module: torch.nn.Module, args: tuple[torch.Tensor, ...]) -> None:
                    layer.partial_fit(args[0], tgt)

                return hook

            assert isinstance(readout, IncrementalRidgeReadoutLayer)
            handle = readout.register_forward_pre_hook(make_partial_fit_hook(readout, target))
            hooks.append(handle)
            # The readout is unfitted during accumulation; its forward only
            # needs to flow a value downstream, so bypass the fitted guard.
            readout._allow_unfitted_forward = True

        try:
            with torch.no_grad():
                self.model(*inputs)
        finally:
            for h in hooks:
                h.remove()
            for _, readout in readouts:
                assert isinstance(readout, IncrementalRidgeReadoutLayer)
                readout._allow_unfitted_forward = False
        n_chunks += 1

    if n_chunks == 0:
        raise RuntimeError(
            "fit_stream() received no chunks; nothing was accumulated to finalize."
        )

    # Solve every readout once from its accumulated statistics.
    for _, readout in readouts:
        assert isinstance(readout, IncrementalRidgeReadoutLayer)
        readout.finalize()

fit_force

fit_force(warmup_inputs: tuple[Tensor, ...], targets: Tensor, *, forgetting: float | None = None, delta: float | None = None, initial_feedback: Tensor | None = None, reset: bool = True) -> Tensor

FORCE-train the readout to autonomously generate a target pattern.

Implements FORCE learning (Sussillo & Abbott 2009, "Generating Coherent Patterns of Activity from Chaotic Neural Networks"): the reservoir is run closed-loop — driven at every step by the network's own readout output fed back as the next input — while the single :class:~resdag.layers.RLSReadout is updated online (one rank-1 RLS step per timestep) toward the target signal d_t at that step. Because the readout is trained inside the feedback loop it will actually run in, the fed-back output progressively locks onto the target; after training the network free-runs and reproduces the pattern autonomously.

When to use FORCE (honest scope)

FORCE's value is training the readout inside the closed loop, so its weights are fit against the states the loop actually visits rather than the teacher-forced states an open-loop fit sees. Its classic advantage is stabilising autonomous generation in regimes where an open-loop teacher-forced ridge fit would diverge in closed loop — high-dimensional or strongly chaotic reservoirs where teacher-forced and free-running state trajectories pull apart and open-loop errors compound.

It is not a free win over open-loop ridge in general. For a clean, low-dimensional periodic target trained over a long window, an open-loop teacher-forced ridge readout (:meth:fit) is already an excellent closed-loop generator — the smooth target leaves little error to compound — and FORCE will match it rather than beat it (empirically, across seeds, FORCE reproduces such a target to a small fraction of its scale and stays competitive with ridge, but any seed-specific win is not a guaranteed property). Reach for FORCE when the reservoir is chaotic enough that the open-loop free run is unstable, or when the target is non-stationary and you want the online / forgetting machinery of :class:RLSReadout; for a stationary, easily-generated target, the simpler :meth:fit is often sufficient.

Per step t (after an optional teacher-forced warmup that settles the reservoir on the target):

1. Feed the current feedback z_t (the network's own last output, or the warmup's for the first step) through the model, reading the readout's input feature vector r_t (reservoir state, possibly concatenated with the input, exactly as the graph presents it).
2. RLS-update the readout toward d_t using r_t — the a-priori error d_t - W_t^\top r_t drives a rank-1 update of W and the inverse-correlation matrix P (see :class:RLSReadout).
3. Recompute the output z_{t+1} = W_{t+1}^\top r_t with the updated weights and feed it back as the next step's input.

The readout's RLS state (P and the weights) is reset at the start of the FORCE pass (via :meth:RLSReadout.reset_state) so a re-run trains from scratch; pass forgetting / delta to override the readout's recursion hyperparameters for the pass.

PARAMETER	DESCRIPTION
warmup_inputs	Teacher-forced warmup driving the reservoir open-loop with the true signal to settle its state before the closed-loop FORCE phase. Format (feedback, driver1, ...); the feedback is typically the leading slice of the target series. Each tensor has shape (batch, warmup_steps, features). May be empty-length in time (warmup_steps == 0) to start FORCE from the reset state, but at least the feedback tensor must be present (its dtype/device/batch and the driver arity are read from it). TYPE: tuple of torch.Tensor
targets	The pattern to generate, shape (batch, T, feedback_dim). Its feature dimension must equal the model's feedback input dimension (and the readout's out_features) — the autoregression contract, since the readout output is fed back as the next feedback. T is the number of closed-loop FORCE steps. TYPE: Tensor
forgetting	Overrides the readout's forgetting factor :math:\lambda \in (0, 1] for this pass. None (default) keeps the readout's configured value. FORCE is classically run with 1.0 (growing window). TYPE: float DEFAULT: None
delta	Overrides the readout's :math:P_0 = I / \delta prior for this pass. None (default) keeps the readout's configured value. A moderate delta (around 1.0) is the standard FORCE choice; a very small delta (huge initial P) makes the early closed-loop updates explode. TYPE: float DEFAULT: None
initial_feedback	Feedback of shape (batch, 1, feedback_dim) seeding the first closed-loop step. If None (default) the last warmup target slice is used (or a zero seed when the warmup is empty). TYPE: Tensor DEFAULT: None
reset	If True (default), reset the reservoir states before the teacher-forced warmup. Set False to continue from the current (already warmed) reservoir state. TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
Tensor	The closed-loop training outputs of shape (batch, T, feedback_dim) — the network's own generated signal at each FORCE step, after that step's RLS update. By the end of training these track targets closely; free-run generation with :meth:ESNModel.forecast continues the pattern from the final reservoir state.

RAISES	DESCRIPTION
TypeError	If the model's readout is not an :class:RLSReadout (FORCE needs the online RLS partial_fit machinery).
ValueError	If the model has more than one readout (FORCE trains a single generator readout), if warmup_inputs is empty, if targets is not 3-D, if the readout's out_features (and hence the fed-back output dimension) does not match targets's feature dimension (the autoregression contract), if the driver arity of warmup_inputs does not match the model, or if initial_feedback has the wrong shape.

Notes

• Single-readout, self-generating only. FORCE here trains exactly one readout whose output is fed straight back as the model's feedback, so the readout must be an RLSReadout and its out_features must equal the feedback dimension. Multi-readout or driven-generation setups are out of scope and rejected with a clear error.
• Driven models. If the model has driver inputs, pass the driver warmup slices in warmup_inputs; during the closed-loop phase the drivers are held at their last warmup value (FORCE classically targets autonomous generation, not driven mapping). A feedback-only model is the common case.
• Reservoir regime. FORCE works best with a chaotic reservoir (spectral radius > 1, e.g. 1.5); the online readout tames the chaos into the target pattern. See the module example.

Examples:

>>> from resdag.models import classic_esn
>>> from resdag import RLSReadout
>>> from resdag.training import ESNTrainer
>>> model = classic_esn(
...     600, feedback_size=1, output_size=1, spectral_radius=1.5,
...     readout=RLSReadout, readout_name="output", readout_bias=False,
... )
>>> trainer = ESNTrainer(model)
>>> trainer.fit_force(
...     warmup_inputs=(target[:, :200, :],),
...     targets=target[:, 200:, :],
...     delta=1.0,
... )
>>> generated = model.forecast(target, horizon=1000)

See Also

fit : Open-loop teacher-forced algebraic readout fitting. resdag.layers.RLSReadout : The online RLS readout FORCE drives. resdag.core.ESNModel.forecast : Free-run generation after FORCE training.

Source code in src/resdag/training/trainer.py

def fit_force(
    self,
    warmup_inputs: tuple[torch.Tensor, ...],
    targets: torch.Tensor,
    *,
    forgetting: float | None = None,
    delta: float | None = None,
    initial_feedback: torch.Tensor | None = None,
    reset: bool = True,
) -> torch.Tensor:
    r"""FORCE-train the readout to autonomously *generate* a target pattern.

    Implements FORCE learning (Sussillo & Abbott 2009, "Generating Coherent
    Patterns of Activity from Chaotic Neural Networks"): the reservoir is run
    **closed-loop** — driven at every step by the network's *own* readout
    output fed back as the next input — while the single
    :class:`~resdag.layers.RLSReadout` is updated online (one rank-1 RLS step
    per timestep) toward the target signal ``d_t`` at that step.  Because the
    readout is trained *inside* the feedback loop it will actually run in, the
    fed-back output progressively locks onto the target; after training the
    network **free-runs** and reproduces the pattern autonomously.

    When to use FORCE (honest scope)
    --------------------------------
    FORCE's value is training the readout *inside the closed loop*, so its
    weights are fit against the states the loop actually visits rather than
    the teacher-forced states an open-loop fit sees. Its classic advantage is
    **stabilising autonomous generation in regimes where an open-loop
    teacher-forced ridge fit would diverge in closed loop** — high-dimensional
    or strongly chaotic reservoirs where teacher-forced and free-running state
    trajectories pull apart and open-loop errors compound.

    It is **not** a free win over open-loop ridge in general. For a *clean,
    low-dimensional periodic* target trained over a long window, an open-loop
    teacher-forced ridge readout (:meth:`fit`) is already an excellent
    closed-loop generator — the smooth target leaves little error to compound
    — and FORCE will *match* it rather than beat it (empirically, across
    seeds, FORCE reproduces such a target to a small fraction of its scale and
    stays competitive with ridge, but any seed-specific win is not a
    guaranteed property). Reach for FORCE when the reservoir is chaotic enough
    that the open-loop free run is unstable, or when the target is
    non-stationary and you want the online / forgetting machinery of
    :class:`RLSReadout`; for a stationary, easily-generated target, the
    simpler :meth:`fit` is often sufficient.

    Per step ``t`` (after an optional teacher-forced warmup that settles the
    reservoir on the target):

    1. Feed the current feedback ``z_t`` (the network's own last output, or
       the warmup's for the first step) through the model, reading the
       readout's input feature vector ``r_t`` (reservoir state, possibly
       concatenated with the input, exactly as the graph presents it).
    2. RLS-update the readout toward ``d_t`` using ``r_t`` — the a-priori
       error ``d_t - W_t^\top r_t`` drives a rank-1 update of ``W`` and the
       inverse-correlation matrix ``P`` (see :class:`RLSReadout`).
    3. Recompute the output ``z_{t+1} = W_{t+1}^\top r_t`` with the *updated*
       weights and feed it back as the next step's input.

    The readout's RLS state (``P`` and the weights) is reset at the start of
    the FORCE pass (via :meth:`RLSReadout.reset_state`) so a re-run trains
    from scratch; pass ``forgetting`` / ``delta`` to override the readout's
    recursion hyperparameters for the pass.

    Parameters
    ----------
    warmup_inputs : tuple of torch.Tensor
        Teacher-forced warmup driving the reservoir *open-loop* with the true
        signal to settle its state before the closed-loop FORCE phase.
        Format ``(feedback, driver1, ...)``; the feedback is typically the
        leading slice of the target series.  Each tensor has shape
        ``(batch, warmup_steps, features)``.  May be empty-length in time
        (``warmup_steps == 0``) to start FORCE from the reset state, but at
        least the feedback tensor must be present (its dtype/device/batch and
        the driver arity are read from it).
    targets : torch.Tensor
        The pattern to generate, shape ``(batch, T, feedback_dim)``.  Its
        feature dimension **must** equal the model's feedback input dimension
        (and the readout's ``out_features``) — the autoregression contract,
        since the readout output is fed back as the next feedback.  ``T`` is
        the number of closed-loop FORCE steps.
    forgetting : float, optional
        Overrides the readout's forgetting factor :math:`\lambda \in (0, 1]`
        for this pass.  ``None`` (default) keeps the readout's configured
        value.  FORCE is classically run with ``1.0`` (growing window).
    delta : float, optional
        Overrides the readout's :math:`P_0 = I / \delta` prior for this pass.
        ``None`` (default) keeps the readout's configured value.  A moderate
        ``delta`` (around ``1.0``) is the standard FORCE choice; a very small
        ``delta`` (huge initial ``P``) makes the early closed-loop updates
        explode.
    initial_feedback : torch.Tensor, optional
        Feedback of shape ``(batch, 1, feedback_dim)`` seeding the first
        closed-loop step.  If ``None`` (default) the last warmup target slice
        is used (or a zero seed when the warmup is empty).
    reset : bool, default=True
        If ``True`` (default), reset the reservoir states before the
        teacher-forced warmup.  Set ``False`` to continue from the current
        (already warmed) reservoir state.

    Returns
    -------
    torch.Tensor
        The closed-loop training outputs of shape ``(batch, T, feedback_dim)``
        — the network's own generated signal at each FORCE step, *after* that
        step's RLS update.  By the end of training these track ``targets``
        closely; free-run generation with :meth:`ESNModel.forecast` continues
        the pattern from the final reservoir state.

    Raises
    ------
    TypeError
        If the model's readout is not an :class:`RLSReadout` (FORCE needs the
        online RLS ``partial_fit`` machinery).
    ValueError
        If the model has more than one readout (FORCE trains a single
        generator readout), if ``warmup_inputs`` is empty, if ``targets`` is
        not 3-D, if the readout's ``out_features`` (and hence the fed-back
        output dimension) does not match ``targets``'s feature dimension (the
        autoregression contract), if the driver arity of ``warmup_inputs``
        does not match the model, or if ``initial_feedback`` has the wrong
        shape.

    Notes
    -----
    - **Single-readout, self-generating only.**  FORCE here trains exactly one
      readout whose output is fed straight back as the model's feedback, so
      the readout must be an ``RLSReadout`` and its ``out_features`` must equal
      the feedback dimension.  Multi-readout or driven-generation setups are
      out of scope and rejected with a clear error.
    - **Driven models.**  If the model has driver inputs, pass the driver
      warmup slices in ``warmup_inputs``; during the closed-loop phase the
      drivers are held at their last warmup value (FORCE classically targets
      autonomous generation, not driven mapping).  A feedback-only model is
      the common case.
    - **Reservoir regime.**  FORCE works best with a *chaotic* reservoir
      (spectral radius ``> 1``, e.g. ``1.5``); the online readout tames the
      chaos into the target pattern.  See the module example.

    Examples
    --------
    >>> from resdag.models import classic_esn
    >>> from resdag import RLSReadout
    >>> from resdag.training import ESNTrainer
    >>> model = classic_esn(  # doctest: +SKIP
    ...     600, feedback_size=1, output_size=1, spectral_radius=1.5,
    ...     readout=RLSReadout, readout_name="output", readout_bias=False,
    ... )
    >>> trainer = ESNTrainer(model)  # doctest: +SKIP
    >>> trainer.fit_force(  # doctest: +SKIP
    ...     warmup_inputs=(target[:, :200, :],),
    ...     targets=target[:, 200:, :],
    ...     delta=1.0,
    ... )
    >>> generated = model.forecast(target, horizon=1000)  # doctest: +SKIP

    See Also
    --------
    fit : Open-loop teacher-forced algebraic readout fitting.
    resdag.layers.RLSReadout : The online RLS readout FORCE drives.
    resdag.core.ESNModel.forecast : Free-run generation after FORCE training.
    """
    if len(warmup_inputs) == 0:
        raise ValueError(
            "fit_force() requires at least one warmup input (the feedback tensor), "
            "even if it is length-0 in time; the model's dtype/device/batch and "
            "driver arity are read from it."
        )
    if targets.dim() != 3:
        raise ValueError(
            f"fit_force() targets must be 3-D (batch, T, feedback_dim); "
            f"got shape {tuple(targets.shape)}."
        )

    # FORCE trains exactly one online RLS readout whose output is the
    # generated signal fed straight back as feedback.
    readouts = self._discover_readouts()
    if len(readouts) != 1:
        names = [name for name, _ in readouts]
        raise ValueError(
            f"fit_force() trains a single self-generating readout, but the model has "
            f"{len(readouts)} readout(s): {names}. FORCE feeds one readout's output "
            f"straight back as the model feedback; multi-readout generation is out of "
            f"scope."
        )
    readout_name, readout = readouts[0]
    if not isinstance(readout, RLSReadout):
        raise TypeError(
            f"fit_force() requires the readout ('{readout_name}') to be an RLSReadout "
            f"(the online recursive-least-squares readout whose partial_fit drives the "
            f"FORCE update), but it is a {type(readout).__name__}. Rebuild the model "
            f"with readout=RLSReadout."
        )

    feedback_dim = targets.shape[-1]
    if readout.out_features != feedback_dim:
        raise ValueError(
            f"fit_force(): the readout ('{readout_name}') has out_features="
            f"{readout.out_features}, but targets have feedback_dim={feedback_dim}. "
            f"FORCE feeds the readout output back as the next feedback, so the readout "
            f"output dimension must equal the feedback (target) dimension."
        )

    num_drivers = len(warmup_inputs) - 1
    expected_drivers = len(self.model.inputs) - 1
    if num_drivers != expected_drivers:
        raise ValueError(
            f"fit_force(): the model has {expected_drivers} driver input(s) but "
            f"warmup_inputs carries {num_drivers}. Pass (feedback, driver1, ...) matching "
            f"the model's inputs."
        )

    # Driven models hold each driver at its last warmup value during the
    # closed-loop phase, so an empty warmup leaves no value to hold.
    if num_drivers and warmup_inputs[0].shape[1] == 0:
        raise ValueError(
            "fit_force(): a driven model needs a non-empty warmup so each driver has a "
            "last value to hold during the autonomous closed-loop phase; the warmup "
            "feedback has 0 timesteps. Provide at least one warmup step."
        )

    batch_size = warmup_inputs[0].shape[0]
    device = warmup_inputs[0].device
    dtype = warmup_inputs[0].dtype

    if initial_feedback is not None:
        expected = (batch_size, 1, feedback_dim)
        if initial_feedback.dim() != 3 or tuple(initial_feedback.shape) != expected:
            raise ValueError(
                f"fit_force(): initial_feedback must have shape {expected} "
                f"(batch, 1, feedback_dim); got {tuple(initial_feedback.shape)}."
            )

    # Apply per-pass overrides, then reset the RLS recursion (P = I/delta,
    # zero weights) so a re-run trains from scratch.
    if forgetting is not None:
        if not 0.0 < forgetting <= 1.0:
            raise ValueError(
                f"fit_force(): forgetting (lambda) must be in (0, 1], got {forgetting}."
            )
        readout.forgetting = forgetting
    if delta is not None:
        if delta <= 0.0:
            raise ValueError(f"fit_force(): delta must be positive, got {delta}.")
        readout.delta = delta
    readout.reset_state()

    # Capture the readout's input feature vector each step via a pre-hook.
    # The normal graph forward (unlike the flat forecast engine) fires hooks,
    # so this sees exactly what the readout consumes — reservoir state,
    # concatenated input, etc. — regardless of the architecture upstream.
    captured: dict[str, torch.Tensor] = {}

    def _capture_hook(_module: torch.nn.Module, args: tuple[torch.Tensor, ...]) -> None:
        captured["features"] = args[0]

    handle = readout.register_forward_pre_hook(_capture_hook)

    total_steps = targets.shape[1]
    outputs = torch.empty(batch_size, total_steps, feedback_dim, dtype=dtype, device=device)

    # Drivers are held at their last warmup value during the closed-loop
    # phase (FORCE generates autonomously). Feedback-only models skip this.
    held_drivers: tuple[torch.Tensor, ...] = ()
    try:
        with torch.no_grad():
            # Phase 1: teacher-forced warmup, driving the reservoir open-loop
            # with the true signal. ``RLSReadout.forward`` is plain linear
            # (no fitted-guard, weights zero after reset_state), so it flows a
            # harmless zero downstream while the reservoir synchronises — the
            # warmup only needs the reservoir state, not the readout output.
            if warmup_inputs[0].shape[1] > 0:
                self.model.warmup(*warmup_inputs, reset=reset)
                seed = warmup_inputs[0][:, -1:, :]
            elif reset:
                self.model.reset_reservoirs()
                seed = torch.zeros(batch_size, 1, feedback_dim, dtype=dtype, device=device)
            else:
                seed = torch.zeros(batch_size, 1, feedback_dim, dtype=dtype, device=device)

            if num_drivers:
                held_drivers = tuple(d[:, -1:, :] for d in warmup_inputs[1:])

            if initial_feedback is not None:
                z = initial_feedback
            else:
                z = seed

            # Phase 2: closed-loop FORCE. Each step drives the model with the
            # network's own feedback, reads the readout feature vector, applies
            # one RLS update toward the target, and feeds the *updated* output
            # back into the next step.
            for t in range(total_steps):
                d = targets[:, t : t + 1, :]
                self.model(z, *held_drivers)
                features = captured["features"]
                readout.partial_fit(features, d)
                z = readout(features)
                outputs[:, t, :] = z.squeeze(1)
    finally:
        handle.remove()

    return outputs