Hi! A small fix to the scheduler params.

Summary

The LinearWarmupCosineAnnealingLR scheduler config in train.py only specifies {"type": "LinearWarmupCosineAnnealingLR"} without the required warmup_steps and max_steps arguments.

This causes training to crash immediately for all tasks (not just PushT) with:

TypeError: LinearWarmupCosineAnnealingLR.__init__() missing 2 required positional arguments: 'warmup_steps' and 'max_steps'

Root Cause

stable_pretraining's create_scheduler has a fallback that attempts to auto-infer these parameters from module.trainer.estimated_stepping_batches.

However, at the time configure_optimizers is called by Lightning, the trainer has not yet estimated stepping batches (the dataloader iterator hasn't been created), so estimated_stepping_batches returns None.

The auto-inference code filters out None values, which means max_steps is silently dropped, and the scheduler constructor fails with the missing argument error.

Full traceback:

File ".../stable_pretraining/module.py", line 634, in configure_optimizers
    scheduler = create_scheduler(opt, sched_config, module=self)
File ".../stable_pretraining/optim/lr_scheduler.py", line 193, in create_scheduler
    return fn(optimizer, **params)
TypeError: LinearWarmupCosineAnnealingLR.__init__() missing 2 required positional arguments: 'warmup_steps' and 'max_steps'

Changes

In train.py, compute max_steps and warmup_steps explicitly from the dataloader length and max_epochs, then pass them directly to the scheduler config:

• max_steps = len(train_dataloader) * max_epochs (total training steps)
• warmup_steps = 1% of max_steps (linear warmup phase)
• Also passes warmup_start_lr=0.0 and eta_min=0.0 for completeness
• Changes scheduler interval from "epoch" to "step" to match the step-level scheduling

This removes the dependency on the auto-inference path entirely, making it robust across stable_pretraining versions.

Testing

• Verified that python train.py data=pusht launches training successfully and progresses through epochs (tested on macOS with MPS backend, stable_pretraining==0.1.6, stable_worldmodel==0.0.6)
• The scheduler initializes without error and learning rate warmup + cosine annealing behaves as expected