Bug: `init_for_epoch(None)` silently skips schedule initialization for constant schedules

[jpdunc23]

Note that this isn't triggered in our usual Trainer.train_one_epoch pipeline, but could be encountered during direct use of TrainStepper.train_on_batch.

When a TimeLengthSchedule has no milestones (e.g., created via from_constant(n)), calling init_for_epoch(None) fails to initialize the sampler because self._epoch starts as None and the cache check self._epoch == epoch short-circuits. sample() then returns n_data_steps instead of the configured constant, silently ignoring the user's n_forward_steps setting.

A secondary consequence: after a successful init_for_epoch(k) with a valid integer epoch, a subsequent init_for_epoch(None) crashes with AssertionError because the cache check no longer short-circuits but the code unconditionally asserts epoch is not None.

[mcgibbon]

Good catch, we should add a test and fix this case. Probably either by making sure the init code gets force-called (add a force flag?) during init itself, or adding a lazy flag to know if we've initialized at all and force-init if not.