Learning rate does not decay in examples/vizdoom_ppo.py when using ICM wrapper

In order to test if the learning rate decay was actually working, I added a train_fn to the example which tracks the ppo policy learning rate. The git diff is here (can be applied with git apply): tianshou_vizdoom_ppo_no_extend.txt

I ran (call this run_grey) the code with the following command (only to make the training resource / time consumption smaller and ensure ICM gets used): python vizdoom_ppo.py --buffer_size 10000 --epoch 5 --icm_lr_scale 10 --task "D2_navigation"

The primary change (adding the training_fn to track the learning rates in the nested ppo policy and the top level icm policy) is copied below:

    def train_fn(curr_epoch, env_step):
        if ppo.lr_schedulers:
            lr = ppo.lr_schedulers[0].get_last_lr()[0]
            logger.writer.add_scalar("training/lr", lr, global_step=env_step)
        else:
            logger.writer.add_scalar("training/no_ppo_lr", 0, global_step=env_step)
        if algorithm.lr_schedulers:
            for idx, scheduler in enumerate(algorithm.lr_schedulers):
                lr = scheduler.get_last_lr()[0]
                logger.writer.add_scalar(f"training/algorithm_lr_{idx}", lr, global_step=env_step)
        else:
            logger.writer.add_scalar("training/no_algorithm_lr", 0, global_step=env_step)

My hypothesis is that the scheduler attached to the ppo policy never gets stepped because:

1. the top level ICM policy itself does not have any lr_schedulers - it does not check for lr_schedulers in wrapped algorithms, and it's update method calls update_with_batch for the wrapped algorithm (in OnPolicyWrapperAlgorithm in algorithm_base.py) rather than update
2. the lr_schedulers are only stepped in the _update method (in Algorithm class in algorithm_base.py)

And based on my hypothesis I added the wrapped_algorithm's scheduler to the top level (ICM) policy:

algorithm.lr_schedulers.extend(algorithm.wrapped_algorithm.lr_schedulers)

The git patch with this change is here: tianshou_vizdoom_ppo_with_extend.txt And ran the code again with the same command (call this run_pink).

And now in tensorboard I can see the learning rate decreasing.

Note: In the attached images grey colour is for the run_grey run, and the pink colour is for the run_pink run.

According to the train_fn code:

1. if the ppo algorithm itself does not have a learning rate scheduler attached, a training/no_ppo_lr graph should be plotted - no such graph is plotted in either run - meaning in both runs there is a scheduler attached to the ppo algorithm - and correspondingly the training/lr plot shows up in both runs: As we can see, in pink, the learning rate decays, correctly from 2e-5 to 0 over the course of training, while in grey the learning rate stays at the initial value of 2e-5.
2. If the top level algorithm (ICM wrapper) does not have any learning rate schedulers attached then a plot of training/no_algorithm_lr with value of 0 should show up. Which it does only for grey:
3. If the top level algorithm (ICM wrapper) does have a learning rate scheduler attached, a graph of training/algorithm_lr_{index} should show up where index is the index of the attached scheduler. In run_pink since I have explicitly added the scheduler of the ppo algorithm to the top level ICM algorithm, it does show up, and has the same values as the pink plot in training/lr - meaning it decays from 2e-5 to 0 - since the same optimizer is being tracked here:

Also, the issue only happens when ICM is enabled (icm_lr_scale > 0). When ICM is not enabled, the learning rate decays without issue.

Perhaps the fix could be to check if ICM has a wrapped algorithm and if that wrapped algorithm has a learning rate scheduler attached, and add the wrapped algorithm's learning rate scheduler to the ICM on policy wrapper's list of learning rate schedulers - like the .extend I did above, but I am not certain this is the idiomatic fix.