chc.train

Stability tier: stable.

System identification: fit the learned residual to trajectory data (learn what physics misses).

Trains only the residual parameters (the known mechanism is frozen) to minimise one-step or multi-step (rollout) prediction error via autodiff through the RK4 step + Optax. This is the "hybrid learns the unknown part" step; the known dynamics are never re-learned.

one_step_mse(model, xs, us, x_next, dt)

Mean squared one-step prediction error of model over a batch of transitions.

fit_residual(model, data, dt, steps=2000, lr=0.01)

Fit the residual of model to one-step transitions data = {x, u, x_next}.

Returns the trained model (same frozen known, updated residual) and the loss history.

rollout_mse(model, x0s, us, xs, dt)

Mean squared multi-step rollout error over a batch of trajectories.

x0s (N, n), us (N, H, m), xs (N, H+1, n): realised states incl. the initial one.

fit_residual_multistep(model, data, dt, steps=1500, lr=0.01)

Fit the residual to trajectories by minimising multi-step rollout error.

data = {x0, us, xs} with shapes (N, n), (N, H, m), (N, H+1, n). Unlike one-step fitting this directly penalises rollout drift. Returns the trained model and loss history.