Learning rate scheduler confusion

[opcode81]

The way in which learning rate schedulers are parametrised has led to confusion.
Particularly torch's LambdaLR is used throughout Tianshou. While we use it correctly (as far as I can tell), the terminology that is used can confuse users (and has confused a user in #1157). The reason for the confusion is that the learning rate scheduler has a different notion of "epoch" than Tianshou itself does:

For the learning rate scheduler, if its step method is called without the explicit epoch argument (which we do), the "epoch" is simply a counter of the number of times the step function has been called in the past. Since we call the step function once in every training step, the maximum "epoch" for the scheduler is the total number of training steps (which can be computed as ceil(step_per_epoch / step_per_collect) * num_epochs). We thus correctly parametrise the scheduler as follows (e.g. in atari_ppo.py):

tianshou/examples/atari/atari_ppo.py

Lines 128 to 132 in bd74273

	if args.lr_decay:
	# decay learning rate to 0 linearly
	max_update_num = np.ceil(args.step_per_epoch / args.step_per_collect) * args.epoch
	lr_scheduler = LambdaLR(optim, lr_lambda=lambda epoch: 1 - epoch / max_update_num)

A user thinking it is the Tianshou "epoch" semantics (like the reporter of #1157), will be tempted to do this instead:

lr_scheduler = LambdaLR(optim, lr_lambda=lambda epoch: 1 - epoch / args.epoch)

The fact that we use the term epoch in our code for parametrising schedulers is the main issue here.

Note that if a user defines the scheduler as above, because the number of training steps is much larger than the number of epochs, the lambda factor becomes negative. This is not caught in LambdaLR and effectively causes the gradient to be reversed, resulting in fatal learning failure.

I would suggest to:

• Replace the term epoch with the term update_num or training_step (also in the high-level API)
• Try to catch errors where the learning rate becomes negative.
• Improve documentation (where appropriate) in order to avoid this kind of confusion

[MischaPanch]

Loosely related: #1198