Trainer

The trainer is the orchestration layer. Three entry points, ordered by abstraction:

1. train_from_config(config) — high-level. Pass a populated TrainConfig, get back (state, history). Honors checkpointing, logging, early stopping, optimizer switching, warm start, replay buffer. This is what the CLI calls under the hood. Use this from agent code.
2. train(model_name, episodes, ...) — backward-compat wrapper that builds a TrainConfig from positional args and delegates.
3. create_train_state(...) + make_train_step(...) — low-level. Use when you need to drive the training loop yourself (custom outer loop, distributed setup, hand-coded LR schedule, …). The single @jax.jit boundary is around train_step.

The "rollout + minibatch sweep" cycle is shared across all five optimizer families (STANDARD, PCGRAD, MAO, LBFGS, GN); only the per-batch grad step differs (dispatched at construction time, before JIT).

For the full surface and example call patterns, see Training entry points in REFERENCE.md.

deqn_jax.training.trainer

Main training loop for DEQN-JAX.

Key design: single JIT boundary around entire train_step for maximum performance. Three step variants dispatched at construction time (before JIT):

• STANDARD: normal jax.grad + opt.update(grads, state, params)
• MAO: jax.jacrev(per_eq_loss_vector) -> per-equation Jacobian -> mao.update(eq_jac, state, params)
• LBFGS: optax.lbfgs (GradientTransformationExtraArgs) -- needs value + value_fn for line search

train_from_config

train_from_config(config) -> Tuple[Any, Dict[str, list]]

Train from a TrainConfig object.

This is the primary entry point for config-driven training. Supports checkpoint resume, mid-training optimizer switching, and grouped TensorBoard logging.

Parameters:

Name	Type	Description	Default
config		TrainConfig instance	required

Returns:

Type	Description
Tuple[Any, Dict[str, list]]	Tuple of (trained_params, history_dict)

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

def train_from_config(config) -> Tuple[Any, Dict[str, list]]:
    """Train from a TrainConfig object.

    This is the primary entry point for config-driven training.
    Supports checkpoint resume, mid-training optimizer switching,
    and grouped TensorBoard logging.

    Args:
        config: TrainConfig instance

    Returns:
        Tuple of (trained_params, history_dict)
    """
    _validate_train_config(config)
    model, n_equations = _resolve_model_for_training(config)

    fp64 = jnp.zeros(1).dtype == jnp.float64
    hidden_sizes = config.network.hidden_sizes
    key = jax.random.PRNGKey(config.seed)

    # ---- Build LR schedule helper for logging ----
    from deqn_jax.optimizers.registry import _build_lr_schedule

    # When a schedule is active, the optimizer is created with lr=1.0.
    # The actual LR is passed as a dynamic scalar to train_step each episode.
    has_schedule = config.optimizer.lr_schedule != "constant"
    if has_schedule:
        effective_opt_cfg = config.optimizer.model_copy(
            update={"learning_rate": 1.0, "lr_schedule": "constant"}
        )
    else:
        effective_opt_cfg = config.optimizer

    state, opt, kind, start_episode, total_for_schedule = _build_initial_state(
        config,
        model,
        key,
        n_equations,
        effective_opt_cfg,
    )

    # ---- Metric logger ----
    wandb_config = config.to_dict() if config.wandb_project else None
    logger = create_logger(
        tensorboard_dir=config.tensorboard_dir,
        wandb_project=config.wandb_project,
        wandb_config=wandb_config,
    )

    # ---- Print header ----
    if config.verbose:
        _print_header(
            model_spec=model,
            optimizer=config.optimizer.name,
            learning_rate=config.optimizer.learning_rate,
            hidden_sizes=hidden_sizes,
            n_params=_count_params(state.params),
            batch_size=config.batch_size,
            mc_samples=config.mc_samples,
            warm_start=config.warm_start,
            grad_clip=config.optimizer.grad_clip,
            loss_reweight=config.loss_reweight,
            fp64=fp64,
            lr_schedule=config.optimizer.lr_schedule,
            lr_warmup=config.optimizer.lr_warmup,
            lr_min_factor=config.optimizer.lr_min_factor,
            net_type=getattr(config.network, "type", "mlp")
            if config.network
            else "mlp",
            history_len=get_history_len(state.params),
        )

    # Build LR schedule function for computing per-episode LR (None if constant)
    lr_schedule_fn = None
    if has_schedule:
        lr_schedule_fn = _build_lr_schedule(config.optimizer, total_for_schedule)

    # ---- Pre-compute quadrature nodes (if using Gauss-Hermite) ----
    quad_nodes_jax = None
    quad_weights_jax = None
    exp_type = config.expectation_type
    if exp_type in ("quadrature", "gh", "gauss_hermite"):
        n_qp = config.n_quadrature_points
        quad = gauss_hermite_nd(n_qp, model.n_shocks)
        if quad is not None:
            quad_nodes_jax = jnp.array(quad[0])
            quad_weights_jax = jnp.array(quad[1])
            if config.verbose:
                print(
                    f"  Quadrature: {n_qp}^{model.n_shocks} = {quad[0].shape[0]} nodes (Gauss-Hermite)"
                )
        else:
            n_total = n_qp**model.n_shocks
            if config.verbose:
                print(
                    f"  Quadrature: {n_total} nodes exceeds limit, falling back to MC"
                )

    # ---- Determine history length from network (Python-level, before JIT) ----
    history_len = get_history_len(state.params)

    # ---- Shock mask ----
    if config.shock_mask is not None and config.verbose:
        shock_names = (
            model.shock_names
            if model.shock_names
            else tuple(f"shock_{i}" for i in range(model.n_shocks))
        )
        active = [n for n, m in zip(shock_names, config.shock_mask) if m > 0]
        zeroed = [n for n, m in zip(shock_names, config.shock_mask) if m == 0]
        print(f"  Shock mask: active={active}, zeroed={zeroed}")

    custom_loss_fn = _build_custom_loss_fn(config, model, history_len)

    # ---- Target network setup ----
    use_target = config.target_update_every > 0
    if use_target:
        state = state._replace(target_params=state.params)
        if config.verbose:
            print(
                f"  Target network: update every {config.target_update_every} episodes"
                f" (tau={config.target_tau})"
            )

    # ---- Create JIT-compiled train step ----
    gradient_surgery = config.gradient_surgery
    train_step = make_train_step(
        model,
        opt,
        config.episode_length,
        config.mc_samples,
        config.batch_size,
        loss_reweight=config.loss_reweight,
        reweight_alpha=config.reweight_alpha,
        kind=kind,
        gradient_surgery=gradient_surgery,
        grad_clip=config.optimizer.grad_clip,
        quad_nodes=quad_nodes_jax,
        quad_weights=quad_weights_jax,
        history_len=history_len,
        compute_loss_fn=custom_loss_fn,
        ss_reset_frac=config.ss_reset_frac,
        use_target_network=use_target,
        n_epochs_per_rollout=config.n_epochs_per_rollout,
        n_minibatches_per_epoch=config.n_minibatches_per_epoch,
        initialize_each_episode=config.initialize_each_episode,
        sorted_within_batch=config.sorted_within_batch,
        replay_cfg=config.replay_buffer,
    )

    if (
        config.verbose
        and kind == OptimizerKind.STANDARD
        and gradient_surgery != "pcgrad"
    ):
        # Compute and report the effective schedule so users can see what
        # the trainer is actually doing per outer iteration.
        ep_samples = config.episode_length * config.batch_size
        mbs_avail = max(1, ep_samples // config.batch_size)
        mbs_this_epoch = (
            min(config.n_minibatches_per_epoch, mbs_avail)
            if config.n_minibatches_per_epoch is not None
            else mbs_avail
        )
        updates_per_cycle = config.n_epochs_per_rollout * mbs_this_epoch
        print(
            f"  Schedule: 1 rollout ({config.episode_length}×{config.batch_size}="
            f"{ep_samples} states) → {config.n_epochs_per_rollout} epoch(s) × "
            f"{mbs_this_epoch} minibatch(es) of {config.batch_size} "
            f"= {updates_per_cycle} grad updates/cycle "
            f"({updates_per_cycle * config.episodes} total over "
            f"{config.episodes} cycles)"
        )

    return _run_training_loop(
        config=config,
        model=model,
        state=state,
        opt=opt,
        kind=kind,
        gradient_surgery=gradient_surgery,
        train_step=train_step,
        lr_schedule_fn=lr_schedule_fn,
        quad_nodes_jax=quad_nodes_jax,
        quad_weights_jax=quad_weights_jax,
        history_len=history_len,
        custom_loss_fn=custom_loss_fn,
        use_target=use_target,
        n_equations=n_equations,
        start_episode=start_episode,
        logger=logger,
    )

train

train(
    model_name: str,
    episodes: int = 1000,
    hidden_sizes: Tuple[int, ...] = (64, 64),
    learning_rate: float = 0.001,
    batch_size: int = 64,
    episode_length: int = 100,
    mc_samples: int = 5,
    optimizer: str = "adam",
    warm_start: bool = False,
    seed: int = 42,
    log_every: int = 100,
    verbose: bool = True,
    grad_clip: Optional[float] = None,
    loss_weights: Optional[List[float]] = None,
    loss_reweight: str = "none",
    reweight_alpha: float = 0.9,
    tensorboard_dir: Optional[str] = None,
    wandb_project: Optional[str] = None,
    checkpoint_dir: Optional[str] = None,
    checkpoint_every: Optional[int] = None,
) -> Tuple[Any, Dict[str, list]]

Train DEQN model (backward-compatible wrapper).

Builds a TrainConfig and delegates to train_from_config().

Preserves the pre-sweep per-cycle training budget (one gradient step per cycle) to avoid surprising legacy callers. New code should prefer train_from_config(TrainConfig(...)) which defaults to the full rollout+minibatch-sweep schedule matching reference DEQN.

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

def train(
    model_name: str,
    episodes: int = 1000,
    hidden_sizes: Tuple[int, ...] = (64, 64),
    learning_rate: float = 1e-3,
    batch_size: int = 64,
    episode_length: int = 100,
    mc_samples: int = 5,
    optimizer: str = "adam",
    warm_start: bool = False,
    seed: int = 42,
    log_every: int = 100,
    verbose: bool = True,
    grad_clip: Optional[float] = None,
    loss_weights: Optional[List[float]] = None,
    loss_reweight: str = "none",
    reweight_alpha: float = 0.9,
    tensorboard_dir: Optional[str] = None,
    wandb_project: Optional[str] = None,
    checkpoint_dir: Optional[str] = None,
    checkpoint_every: Optional[int] = None,
) -> Tuple[Any, Dict[str, list]]:
    """Train DEQN model (backward-compatible wrapper).

    Builds a TrainConfig and delegates to train_from_config().

    Preserves the pre-sweep per-cycle training budget (one gradient step
    per cycle) to avoid surprising legacy callers. New code should prefer
    ``train_from_config(TrainConfig(...))`` which defaults to the full
    rollout+minibatch-sweep schedule matching reference DEQN.
    """
    from deqn_jax.config import NetworkConfig, OptimizerConfig, TrainConfig

    config = TrainConfig(
        model=model_name,
        episodes=episodes,
        batch_size=batch_size,
        episode_length=episode_length,
        mc_samples=mc_samples,
        seed=seed,
        network=NetworkConfig(hidden_sizes=hidden_sizes),
        optimizer=OptimizerConfig(
            name=optimizer,
            learning_rate=learning_rate,
            grad_clip=grad_clip,
        ),
        warm_start=warm_start,
        loss_weights=list(loss_weights) if loss_weights is not None else None,
        loss_reweight=loss_reweight,
        reweight_alpha=reweight_alpha,
        log_every=log_every,
        verbose=verbose,
        tensorboard_dir=tensorboard_dir,
        wandb_project=wandb_project,
        checkpoint_dir=checkpoint_dir,
        checkpoint_every=checkpoint_every,
        n_minibatches_per_epoch=1,
    )

    return train_from_config(config)