[myscience]

A DCL estimator saved with the sklearn backend restores its encoder on load but leaves its dynamics model freshly initialized. Solver checkpoints omitted the dynamics parameters, so embeddings could match after a round trip while predicted latent transitions changed.

This patch saves and restores the dynamics state for both minibatch and full-batch DCL solvers through a shared mixin. It adds round-trip and resumed-training regression tests for linear and orthogonal-linear dynamics, plus documentation of historical checkpoint behavior.

Reproducer

import numpy as np
import torch
from cebra.integrations.sklearn.cebra import CEBRA, DCL

torch.manual_seed(4)
X = np.random.default_rng(4).normal(size=(100, 5)).astype("float32")
model = DCL(
    device="cpu", output_dimension=3, max_iterations=2,
    batch_size=8, batch_size_negatives=11,
).fit(X)
z = torch.randn(7, 3)
before = model.dynamics_model_(z).detach().clone()
model.save("dcl.pt", backend="sklearn")
restored = CEBRA.load("dcl.pt", backend="sklearn")

# The encoder round trip already works.
np.testing.assert_allclose(model.transform(X), restored.transform(X))
# Fails before this patch; passes with it.
torch.testing.assert_close(restored.dynamics_model_(z), before)

The failure was independently reproduced on upstream commit b7e4b9e.

Checkpoint compatibility

Historical DCL checkpoints lack the saved dynamics weights. Strict loading raises an explanatory error for those checkpoints. Explicit low-level solver.load_state_dict(state, strict=False) loads the available components, warns, and retains the initialized dynamics. The change is confined to the DCL solvers.

Validation

• Full required CPU suite: 3,187 passed, 102 skipped, zero failures; 6,965 dataset-dependent cases deselected.
• Regression coverage includes both DCL batch modes, linear and orthogonal-linear dynamics, encoder/dynamics round trips, resumed training, and strict/non-strict handling of incomplete checkpoints.
• Formatting/import ordering, spelling, docstring coverage, citation validation, binary checks, package build, and an installed-wheel checkpoint smoke test passed.
• Local environment: Python 3.10.12, PyTorch 2.5.1, scikit-learn 1.7.2, NumPy 2.2.6, CPU.

I have signed the Contributor License Agreement.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing. Review the Copilot code review billing settings.

[stes]

This feels like a patch to an upstream bug; the dynamics model should be part of the state dict (?). Did we forget an ``nn.Module'' somewhere?

[myscience]

Yeah, I think that's a fair point. From what I understand the issue is that Solver has it's own state_dict method that saves self.model, while the DCL puts the dynamics module in self.reference_model = Sequential(model, dynamics), so the issue is that the custom saving never saves reference_model (DCL still needs model as a separate attribute to compute the embeddings and the pos/neg pairs etc...). I guess one clean option would be to add a custom state_dict and load_state_dict methods to AuxiliaryVariableSolver (which DCL inherits) that save and load reference_model. Otherwise one modifies the current method of the base Solver to handle the optional reference_model if present... what do you think @stes ?