[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?