DQN playing BreakoutNoFrameskip-v4 from https://github.com/sgoodfriend/rl-algo-impls/tree/0511de345b17175b7cf1ea706c3e05981f11761c
d30b0db
from torch.optim import Optimizer | |
from typing import Callable | |
Schedule = Callable[[float], float] | |
def linear_schedule( | |
start_val: float, end_val: float, end_fraction: float = 1.0 | |
) -> Schedule: | |
def func(progress_fraction: float) -> float: | |
if progress_fraction >= end_fraction: | |
return end_val | |
else: | |
return start_val + (end_val - start_val) * progress_fraction / end_fraction | |
return func | |
def constant_schedule(val: float) -> Schedule: | |
return lambda f: val | |
def spike_schedule( | |
max_value: float, | |
start_fraction: float = 1e-2, | |
end_fraction: float = 1e-4, | |
peak_progress: float = 0.1, | |
) -> Schedule: | |
assert 0 < peak_progress < 1 | |
def func(progress_fraction: float) -> float: | |
if progress_fraction < peak_progress: | |
fraction = ( | |
start_fraction | |
+ (1 - start_fraction) * progress_fraction / peak_progress | |
) | |
else: | |
fraction = 1 + (end_fraction - 1) * (progress_fraction - peak_progress) / ( | |
1 - peak_progress | |
) | |
return max_value * fraction | |
return func | |
def schedule(name: str, start_val: float) -> Schedule: | |
if name == "linear": | |
return linear_schedule(start_val, 0) | |
elif name == "none": | |
return constant_schedule(start_val) | |
elif name == "spike": | |
return spike_schedule(start_val) | |
else: | |
raise ValueError(f"Schedule {name} not supported") | |
def update_learning_rate(optimizer: Optimizer, learning_rate: float) -> None: | |
for param_group in optimizer.param_groups: | |
param_group["lr"] = learning_rate | |