release
This commit is contained in:
allenzren
@@ -0,0 +1,87 @@
|
||||
"""
|
||||
To balance actor and critic losses, the rewards are divided through by the standard deviation of a rolling discounted sum of the rewards (without subtracting and re-adding the mean).
|
||||
|
||||
Code is based on: https://github.com/openai/phasic-policy-gradient/blob/master/phasic_policy_gradient/reward_normalizer.py
|
||||
|
||||
Reference: https://arxiv.org/pdf/2005.12729.pdf
|
||||
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class RunningMeanStd:
|
||||
def __init__(
|
||||
self,
|
||||
epsilon=1e-4, # initial count (with mean=0 ,var=1)
|
||||
shape=(), # unbatched shape of data, shape[0] is the batch size
|
||||
):
|
||||
super().__init__()
|
||||
self.mean = np.zeros(shape)
|
||||
self.var = np.ones(shape)
|
||||
self.count = epsilon
|
||||
|
||||
def update(self, x):
|
||||
batch_mean = np.mean(x, axis=0)
|
||||
batch_var = np.var(x, axis=0)
|
||||
batch_count = x.shape[0]
|
||||
self.update_from_moments(batch_mean, batch_var, batch_count)
|
||||
|
||||
def update_from_moments(self, batch_mean, batch_var, batch_count):
|
||||
delta = batch_mean - self.mean
|
||||
tot_count = self.count + batch_count
|
||||
|
||||
self.mean = self.mean + delta * batch_count / tot_count
|
||||
m_a = self.var * self.count
|
||||
m_b = batch_var * batch_count
|
||||
M2 = m_a + m_b + delta**2 * self.count * batch_count / tot_count
|
||||
self.var = M2 / (tot_count - 1)
|
||||
self.count = tot_count
|
||||
|
||||
|
||||
class RunningRewardScaler:
|
||||
"""
|
||||
Pseudocode can be found in https://arxiv.org/pdf/1811.02553.pdf
|
||||
section 9.3 (which is based on our Baselines code, haha)
|
||||
Motivation is that we'd rather normalize the returns = sum of future rewards,
|
||||
but we haven't seen the future yet. So we assume that the time-reversed rewards
|
||||
have similar statistics to the rewards, and normalize the time-reversed rewards.
|
||||
"""
|
||||
|
||||
def __init__(self, num_envs, cliprew=10.0, gamma=0.99, epsilon=1e-8, per_env=False):
|
||||
ret_rms_shape = (num_envs,) if per_env else ()
|
||||
self.ret_rms = RunningMeanStd(shape=ret_rms_shape)
|
||||
self.cliprew = cliprew
|
||||
self.ret = np.zeros(num_envs)
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
self.per_env = per_env
|
||||
|
||||
def __call__(self, reward, first):
|
||||
rets = backward_discounted_sum(
|
||||
prevret=self.ret, reward=reward, first=first, gamma=self.gamma
|
||||
)
|
||||
self.ret = rets[:, -1]
|
||||
self.ret_rms.update(rets if self.per_env else rets.reshape(-1))
|
||||
return self.transform(reward)
|
||||
|
||||
def transform(self, reward):
|
||||
return np.clip(
|
||||
reward / np.sqrt(self.ret_rms.var + self.epsilon),
|
||||
-self.cliprew,
|
||||
self.cliprew,
|
||||
)
|
||||
|
||||
|
||||
def backward_discounted_sum(
|
||||
prevret, # value predictions
|
||||
reward, # reward
|
||||
first, # mark beginning of episodes"
|
||||
gamma, # discount
|
||||
):
|
||||
assert first.ndim == 2
|
||||
_, nstep = reward.shape
|
||||
ret = np.zeros_like(reward)
|
||||
for t in range(nstep):
|
||||
prevret = ret[:, t] = reward[:, t] + (1 - first[:, t]) * gamma * prevret
|
||||
return ret
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2022 Naoki Katsura
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
"""
|
||||
|
||||
# From https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup
|
||||
|
||||
import math
|
||||
import torch
|
||||
from torch.optim.lr_scheduler import _LRScheduler
|
||||
|
||||
|
||||
class CosineAnnealingWarmupRestarts(_LRScheduler):
|
||||
"""
|
||||
optimizer (Optimizer): Wrapped optimizer.
|
||||
first_cycle_steps (int): First cycle step size.
|
||||
cycle_mult(float): Cycle steps magnification. Default: -1.
|
||||
max_lr(float): First cycle's max learning rate. Default: 0.1.
|
||||
min_lr(float): Min learning rate. Default: 0.001.
|
||||
warmup_steps(int): Linear warmup step size. Default: 0.
|
||||
gamma(float): Decrease rate of max learning rate by cycle. Default: 1.
|
||||
last_epoch (int): The index of last epoch. Default: -1.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
first_cycle_steps: int,
|
||||
cycle_mult: float = 1.0,
|
||||
max_lr: float = 0.1,
|
||||
min_lr: float = 0.001,
|
||||
warmup_steps: int = 0,
|
||||
gamma: float = 1.0,
|
||||
last_epoch: int = -1,
|
||||
):
|
||||
assert warmup_steps < first_cycle_steps
|
||||
|
||||
self.first_cycle_steps = first_cycle_steps # first cycle step size
|
||||
self.cycle_mult = cycle_mult # cycle steps magnification
|
||||
self.base_max_lr = max_lr # first max learning rate
|
||||
self.max_lr = max_lr # max learning rate in the current cycle
|
||||
self.min_lr = min_lr # min learning rate
|
||||
self.warmup_steps = warmup_steps # warmup step size
|
||||
self.gamma = gamma # decrease rate of max learning rate by cycle
|
||||
|
||||
self.cur_cycle_steps = first_cycle_steps # first cycle step size
|
||||
self.cycle = 0 # cycle count
|
||||
self.step_in_cycle = last_epoch # step size of the current cycle
|
||||
|
||||
super(CosineAnnealingWarmupRestarts, self).__init__(optimizer, last_epoch)
|
||||
|
||||
# set learning rate min_lr
|
||||
self.init_lr()
|
||||
|
||||
def init_lr(self):
|
||||
self.base_lrs = []
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group["lr"] = self.min_lr
|
||||
self.base_lrs.append(self.min_lr)
|
||||
|
||||
def get_lr(self):
|
||||
if self.step_in_cycle == -1:
|
||||
return self.base_lrs
|
||||
elif self.step_in_cycle < self.warmup_steps:
|
||||
return [
|
||||
(self.max_lr - base_lr) * self.step_in_cycle / self.warmup_steps
|
||||
+ base_lr
|
||||
for base_lr in self.base_lrs
|
||||
]
|
||||
else:
|
||||
return [
|
||||
base_lr
|
||||
+ (self.max_lr - base_lr)
|
||||
* (
|
||||
1
|
||||
+ math.cos(
|
||||
math.pi
|
||||
* (self.step_in_cycle - self.warmup_steps)
|
||||
/ (self.cur_cycle_steps - self.warmup_steps)
|
||||
)
|
||||
)
|
||||
/ 2
|
||||
for base_lr in self.base_lrs
|
||||
]
|
||||
|
||||
def step(self, epoch=None):
|
||||
if epoch is None:
|
||||
epoch = self.last_epoch + 1
|
||||
self.step_in_cycle = self.step_in_cycle + 1
|
||||
if self.step_in_cycle >= self.cur_cycle_steps:
|
||||
self.cycle += 1
|
||||
self.step_in_cycle = self.step_in_cycle - self.cur_cycle_steps
|
||||
self.cur_cycle_steps = (
|
||||
int((self.cur_cycle_steps - self.warmup_steps) * self.cycle_mult)
|
||||
+ self.warmup_steps
|
||||
)
|
||||
else:
|
||||
if epoch >= self.first_cycle_steps:
|
||||
if self.cycle_mult == 1.0:
|
||||
self.step_in_cycle = epoch % self.first_cycle_steps
|
||||
self.cycle = epoch // self.first_cycle_steps
|
||||
else:
|
||||
n = int(
|
||||
math.log(
|
||||
(
|
||||
epoch / self.first_cycle_steps * (self.cycle_mult - 1)
|
||||
+ 1
|
||||
),
|
||||
self.cycle_mult,
|
||||
)
|
||||
)
|
||||
self.cycle = n
|
||||
self.step_in_cycle = epoch - int(
|
||||
self.first_cycle_steps
|
||||
* (self.cycle_mult**n - 1)
|
||||
/ (self.cycle_mult - 1)
|
||||
)
|
||||
self.cur_cycle_steps = self.first_cycle_steps * self.cycle_mult ** (
|
||||
n
|
||||
)
|
||||
else:
|
||||
self.cur_cycle_steps = self.first_cycle_steps
|
||||
self.step_in_cycle = epoch
|
||||
|
||||
self.max_lr = self.base_max_lr * (self.gamma**self.cycle)
|
||||
self.last_epoch = math.floor(epoch)
|
||||
for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
|
||||
param_group["lr"] = lr
|
||||
@@ -0,0 +1,19 @@
|
||||
"""
|
||||
Simple timer from https://github.com/jannerm/diffuser/blob/main/diffuser/utils/timer.py
|
||||
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
|
||||
class Timer:
|
||||
|
||||
def __init__(self):
|
||||
self._start = time.time()
|
||||
|
||||
def __call__(self, reset=True):
|
||||
now = time.time()
|
||||
diff = now - self._start
|
||||
if reset:
|
||||
self._start = now
|
||||
return diff
|
||||
Reference in New Issue
Block a user