|
|
from abc import ABC, abstractmethod |
|
|
|
|
|
import numpy as np |
|
|
import torch |
|
|
from torch.optim.lr_scheduler import LambdaLR, LRScheduler |
|
|
|
|
|
from tianshou.highlevel.config import SamplingConfig |
|
|
from tianshou.utils.string import ToStringMixin |
|
|
|
|
|
|
|
|
class LRSchedulerFactory(ToStringMixin, ABC): |
|
|
"""Factory for the creation of a learning rate scheduler.""" |
|
|
|
|
|
@abstractmethod |
|
|
def create_scheduler(self, optim: torch.optim.Optimizer) -> LRScheduler: |
|
|
pass |
|
|
|
|
|
|
|
|
class LRSchedulerFactoryLinear(LRSchedulerFactory): |
|
|
def __init__(self, sampling_config: SamplingConfig): |
|
|
self.sampling_config = sampling_config |
|
|
|
|
|
def create_scheduler(self, optim: torch.optim.Optimizer) -> LRScheduler: |
|
|
return LambdaLR(optim, lr_lambda=self._LRLambda(self.sampling_config).compute) |
|
|
|
|
|
class _LRLambda: |
|
|
def __init__(self, sampling_config: SamplingConfig): |
|
|
self.max_update_num = ( |
|
|
np.ceil(sampling_config.step_per_epoch / sampling_config.step_per_collect) |
|
|
* sampling_config.num_epochs |
|
|
) |
|
|
|
|
|
def compute(self, epoch: int) -> float: |
|
|
return 1.0 - epoch / self.max_update_num |
|
|
|
