Training: Optax & Evosax Loops

HybridModel.fit and HybridModel.fit_ensemble are the entry points. Objective says what is minimised; OptaxTrainingConfig or EvosaxTrainingConfig says how to search. train_with_optax and train_with_evosax are the functions fit dispatches to: they take the model first, then fit's arguments, and return a FitResult.

Quick links

• Objective
• OptaxTrainingConfig
• TrainingPhase
• train_with_optax
• EvosaxTrainingConfig
• train_with_evosax
• register_algorithm

---

Objective

from jaxhybridmodels.objective import Objective  ·  also re-exported as jaxhybridmodels.Objective

Objective(
    loss: 'Callable[..., Array] | str' = 'mse',
    channel_weights: 'Mapping[str, float] | Sequence[float] | None' = None,
    penalty_weight: 'float' = 0.0,
    penalty_points: 'Mapping[str, ArrayLike] | None' = None,
    penalty_fn: 'Callable[..., Array] | None' = None,
    trajectory_penalty_fn: 'Callable[..., Array] | None' = None,
    trajectory_penalty_weight: 'float' = 0.0,
) -> None

The training objective: data loss plus optional penalties.

Keyword-only. Objective() is plain MSE over every channel with no penalty.

Attributes

Field	Type	Description
loss	`Callable	str`
channel_weights	`Mapping[str, float]	Sequence[float]
penalty_weight	float	Weight on the bound-saturation penalty, relative to the per-bucket-averaged data term. 0.0 disables it. TrainingPhase.penalty_weight overrides it for one Optax phase.
penalty_points	`Mapping[str, ArrayLike]	None`
penalty_fn	`Callable	None`
trajectory_penalty_fn	`Callable	None`
trajectory_penalty_weight	float	Weight on trajectory_penalty_fn. 0.0 disables it even when a function is set; a non-zero weight without a function raises.

Source

Objective.extra_penalty_points()

extra_penalty_points(self, predictors: 'Any') -> 'tuple[Array, ...]'

penalty_points as one array per BoundedPredictor, in traversal order.

Predictors with no entry get an empty [0, n_inputs] array. An unknown path raises, naming the valid ones.

Source

Objective.loss_fn()

loss_fn(
    self,
    output_channel_names: 'Sequence[str]',
) -> Callable[[Array, BucketPayload], Array]

The resolved (pred_obs, bp) -> scalar for a dataset with these channels.

Channel names in channel_weights are matched here, so an unknown name raises before anything compiles.

Source

---

OptaxTrainingConfig

from jaxhybridmodels.training import OptaxTrainingConfig  ·  also re-exported as jaxhybridmodels.OptaxTrainingConfig

OptaxTrainingConfig(
    steps: 'InitVar[int | None]' = None,
    lr: 'InitVar[float | optax.Schedule | None]' = None,
    optimizer: 'InitVar[OptimizerSpec | None]' = None,
    phases: 'tuple[TrainingPhase, ...]' = (),
    tournament_attempts: 'int' = 1,
    tournament_steps: 'int' = 0,
    tournament_lr: 'float' = 0.0001,
    patience: 'int' = 0,
    restore_best: 'bool' = True,
    verbose: 'bool' = True,
) -> None

Search settings for train_with_optax and Optax HybridModel.fit.

What is minimised (loss, channel weights, penalties) lives in Objective, passed separately as objective=. This config says how to minimise it.

A single phase is spelled with steps= (plus optional lr= and optimizer=). A schedule is spelled with phases=, a tuple of TrainingPhase. The two forms do not mix; either way the config stores phases::

OptaxTrainingConfig(steps=500, lr=1e-3)
OptaxTrainingConfig(steps=500, lr=optax.cosine_decay_schedule(1e-2, 500))
OptaxTrainingConfig(
    phases=(
        TrainingPhase(steps=200, lr=3e-2, horizon=0.5),
        TrainingPhase(steps=200, lr=5e-3, reset_optimizer=True),
    )
)

Attributes

Field	Type	Description
phases	tuple[TrainingPhase, ...]	The schedule, run in order. See TrainingPhase for the per-phase settings.
tournament_attempts, tournament_steps	int	The tournament runs only when tournament_steps > 0 and tournament_attempts > 1. Each attempt re-initialises the predictors, trains for tournament_steps steps, and is scored on the data term plus any configured trajectory penalty by a forward-only pass; the bound penalty remains excluded. The lowest score wins. An attempt that raises a diffrax error or a non-finite loss is dropped and the next key tried. If all fail, the original predictors are used and a RuntimeWarning is raised.
tournament_lr	float	Learning rate for the tournament's short bursts, independent of the phase learning rates.
patience	int	Consecutive steps without a new best tracked loss before the current phase stops early. The tracked loss is the data term plus any trajectory penalty, excluding the fixed-point bound penalty. Counted within a phase and reset at every phase boundary, so a plateau at the end of one phase cannot kill the next before its new learning rate acts. 0 disables early stopping.
restore_best	bool	Return the predictors from the lowest tracked-loss step instead of the last one. The running minimum resets whenever the horizon changes, since tracked losses over different horizons are not comparable and the shortest horizon would otherwise always own the minimum.
verbose	bool	Selects RichTrainingUI over SilentUI when ui=None. An explicit ui=... argument always wins.

Source

---

TrainingPhase

from jaxhybridmodels.training import TrainingPhase  ·  also re-exported as jaxhybridmodels.TrainingPhase

TrainingPhase(
    steps: 'int',
    lr: 'float | optax.Schedule' = 0.001,
    optimizer: 'OptimizerSpec' = 'adamw',
    horizon: 'float' = 1.0,
    reset_optimizer: 'bool' = False,
    penalty_weight: 'float | None' = None,
) -> None

One contiguous block of Optax steps with its own schedule settings.

Attributes

Field	Type	Description
steps	int	Step budget for the phase. At least 1.
lr	`float	optax.Schedule`
optimizer	OptimizerSpec	A registered name ("adamw", "adabelief"), a factory taking learning_rate and returning an optax.GradientTransformation (e.g. optax.adamw, or lambda learning_rate: optax.chain( optax.clip_by_global_norm(1.0), optax.adamw(learning_rate))), or a ready-made transformation instance. Names and factories are wrapped in optax.inject_hyperparams, so a phase boundary or a schedule can move the learning rate without a rebuild. A raw instance cannot be re-hyperparametrised, so a phase that changes lr with a raw instance must set reset_optimizer. A phase that changes the optimiser must set it too, because optimiser state belongs to the optimiser that built it.
horizon	float	Fraction of each experiment's timeline the loss looks at, in (0, 1]. It masks the loss, never the integration: the solver still runs the full trajectory, and only the first horizon of the observation times is scored, which stops a long-horizon divergence from drowning the gradient. Being a runtime mask rather than a shape change, a phase boundary costs no recompile.
reset_optimizer	bool	Rebuild the optimiser and discard its state at the start of this phase. Set it when switching optimiser, and when a horizon change has made the accumulated momentum wrong.
penalty_weight	`float	None`

Source

---

train_with_optax()

from jaxhybridmodels.training import train_with_optax

train_with_optax(
    model: 'HybridModel',
    data: 'Data',
    config: 'OptaxTrainingConfig',
    key: 'Array',
    objective: 'Objective | None' = None,
    freeze: 'Sequence[FreezeSpec] | FreezeSpec | None' = None,
    trainable: 'Any' = None,
    ui: 'TrainingUI | None' = None,
    validation: 'Data | None' = None,
) -> FitResult

Train model's predictors with Optax; what model.fit runs for an Optax config.

Runs the phases described by config, optionally preceded by a tournament. See the module docstring for what a step, a phase, and the tournament are, and OptaxTrainingConfig for the fields. Every keyword means what it does for HybridModel.fit>, and data and validation may be a Dataset or a list of experiments.

history.loss is the raw per-step loss, measured before that step's update, concatenated across phases. It can go up. It is the data term plus any trajectory penalty; the bound penalty is in history.penalty, so a ramping bound weight does not move the loss series. Evosax records best-so-far instead, so the two series mean different things.

The returned model comes from the lowest-loss step when config.restore_best=True, else the final step. That minimum resets whenever the phase horizon changes, so the returned model always comes from the last horizon trained on.

Source

---

EvosaxTrainingConfig

from jaxhybridmodels.training import EvosaxTrainingConfig  ·  also re-exported as jaxhybridmodels.EvosaxTrainingConfig

EvosaxTrainingConfig(
    algorithm: 'str' = 'CMA_ES',
    population_size: 'int' = 64,
    steps: 'int' = 100,
    init: "Literal['warm', 'uniform_box', 'lhs_box']" = 'warm',
    init_box_extent: 'float' = 2.0,
    sigma_init: 'float' = 0.1,
    log_every: 'int' = 1,
    verbose: 'bool' = True,
) -> None

Search settings for train_with_evosax and Evosax HybridModel.fit.

What is minimised lives in Objective, passed separately as objective=. Unlike OptaxTrainingConfig there are no phases: the run is a flat loop, so every field is a scalar.

Attributes

Field	Type	Description
algorithm		Evosax strategy name, one of the keys of ALGORITHM_REGISTRY: "CMA_ES" (default), "Sep_CMA_ES", "SimpleES". Add another strategy with register_algorithm.
population_size		Individuals per generation, evaluated in parallel through vmap.
steps		Number of generations, run in sequence. Named like the Optax config's steps: one generation is one Evosax step.
init		Initial-population scheme, one of "warm", "uniform_box", "lhs_box". The module docstring explains the difference.
init_box_extent		Half-width of the box for "uniform_box" and "lhs_box". Ignored by "warm".
sigma_init		Initial CMA-ES step size. Used directly by "warm" and as the prior step size for the box-init modes. CMA-ES adapts it after the first tell.
log_every		UI heartbeat cadence. Honoured only by the Rich UIs; the silent and recording UIs see every generation.
verbose		Selects RichEvosaxUI over SilentUI when ui=None. An explicit ui=... argument always wins.

Source

---

train_with_evosax()

from jaxhybridmodels.training import train_with_evosax

train_with_evosax(
    model: 'HybridModel',
    data: 'Data',
    config: 'EvosaxTrainingConfig',
    key: 'Array',
    objective: 'Objective | None' = None,
    freeze: 'Sequence[FreezeSpec] | FreezeSpec | None' = None,
    trainable: 'Any' = None,
    ui: 'EvosaxUI | None' = None,
    validation: 'Data | None' = None,
) -> FitResult

Train model's predictors with an evolutionary strategy; fit for an Evosax config.

Runs config.steps generations over a population of config.population_size candidates, taking no gradient. See the module docstring for the JIT boundary and the initial-population modes. Every keyword means what it does for HybridModel.fit>, and data and validation may be a Dataset or a list of experiments.

Each individual's fitness is the bucket-averaged data loss plus any trajectory penalty plus objective.penalty_weight times the bound penalty, since evosax ranks by one scalar.

history.loss is the best fitness so far after each generation, so it is monotone non-increasing; Optax records the raw per-step loss instead, so the two series mean different things. Without validation the returned model is rebuilt from the lowest-loss candidate seen in any generation, the warm-up evaluation of the input included. With it, each generation's best is scored on the validation set and the returned model is the one that scored lowest there.

Source

---

register_algorithm()

from jaxhybridmodels.training.evosax import register_algorithm

register_algorithm(name: 'str', cls: 'type') -> 'None'

Register an evosax strategy class under name for algorithm=.

After registration, EvosaxTrainingConfig(algorithm=name) builds that strategy. Re-registering an existing name overwrites without warning. Mirrors register_solver.

Source