chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,2 @@
|
||||
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import copy
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from utils.logger import logger
|
||||
|
||||
from agents.diffusion import Diffusion
|
||||
from agents.model import MLP
|
||||
|
||||
|
||||
class Diffusion_BC(object):
|
||||
def __init__(self,
|
||||
state_dim,
|
||||
action_dim,
|
||||
max_action,
|
||||
device,
|
||||
discount,
|
||||
tau,
|
||||
beta_schedule='linear',
|
||||
n_timesteps=100,
|
||||
lr=2e-4,
|
||||
):
|
||||
|
||||
self.model = MLP(state_dim=state_dim, action_dim=action_dim, device=device)
|
||||
self.actor = Diffusion(state_dim=state_dim, action_dim=action_dim, model=self.model, max_action=max_action,
|
||||
beta_schedule=beta_schedule, n_timesteps=n_timesteps,
|
||||
).to(device)
|
||||
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)
|
||||
|
||||
self.max_action = max_action
|
||||
self.action_dim = action_dim
|
||||
self.discount = discount
|
||||
self.tau = tau
|
||||
self.device = device
|
||||
|
||||
def train(self, replay_buffer, iterations, batch_size=100, log_writer=None):
|
||||
|
||||
metric = {'bc_loss': [], 'ql_loss': [], 'actor_loss': [], 'critic_loss': []}
|
||||
for _ in range(iterations):
|
||||
# Sample replay buffer / batch
|
||||
state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)
|
||||
|
||||
loss = self.actor.loss(action, state)
|
||||
|
||||
self.actor_optimizer.zero_grad()
|
||||
loss.backward()
|
||||
self.actor_optimizer.step()
|
||||
|
||||
metric['actor_loss'].append(0.)
|
||||
metric['bc_loss'].append(loss.item())
|
||||
metric['ql_loss'].append(0.)
|
||||
metric['critic_loss'].append(0.)
|
||||
|
||||
return metric
|
||||
|
||||
def sample_action(self, state):
|
||||
state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
|
||||
with torch.no_grad():
|
||||
action = self.actor.sample(state)
|
||||
return action.cpu().data.numpy().flatten()
|
||||
|
||||
def save_model(self, dir, id=None):
|
||||
if id is not None:
|
||||
torch.save(self.actor.state_dict(), f'{dir}/actor_{id}.pth')
|
||||
else:
|
||||
torch.save(self.actor.state_dict(), f'{dir}/actor.pth')
|
||||
|
||||
def load_model(self, dir, id=None):
|
||||
if id is not None:
|
||||
self.actor.load_state_dict(torch.load(f'{dir}/actor_{id}.pth'))
|
||||
else:
|
||||
self.actor.load_state_dict(torch.load(f'{dir}/actor.pth'))
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import copy
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
from agents.helpers import (cosine_beta_schedule,
|
||||
linear_beta_schedule,
|
||||
vp_beta_schedule,
|
||||
extract,
|
||||
Losses)
|
||||
from utils.utils import Progress, Silent
|
||||
|
||||
|
||||
class Diffusion(nn.Module):
|
||||
def __init__(self, state_dim, action_dim, model, max_action,
|
||||
beta_schedule='linear', n_timesteps=100,
|
||||
loss_type='l2', clip_denoised=True, predict_epsilon=True):
|
||||
super(Diffusion, self).__init__()
|
||||
|
||||
self.state_dim = state_dim
|
||||
self.action_dim = action_dim
|
||||
self.max_action = max_action
|
||||
self.model = model
|
||||
|
||||
if beta_schedule == 'linear':
|
||||
betas = linear_beta_schedule(n_timesteps)
|
||||
elif beta_schedule == 'cosine':
|
||||
betas = cosine_beta_schedule(n_timesteps)
|
||||
elif beta_schedule == 'vp':
|
||||
betas = vp_beta_schedule(n_timesteps)
|
||||
|
||||
alphas = 1. - betas
|
||||
alphas_cumprod = torch.cumprod(alphas, axis=0)
|
||||
alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])
|
||||
|
||||
self.n_timesteps = int(n_timesteps)
|
||||
self.clip_denoised = clip_denoised
|
||||
self.predict_epsilon = predict_epsilon
|
||||
|
||||
self.register_buffer('betas', betas)
|
||||
self.register_buffer('alphas_cumprod', alphas_cumprod)
|
||||
self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
|
||||
|
||||
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||
self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
|
||||
self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
|
||||
self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
|
||||
self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
|
||||
self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1))
|
||||
|
||||
# calculations for posterior q(x_{t-1} | x_t, x_0)
|
||||
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
|
||||
self.register_buffer('posterior_variance', posterior_variance)
|
||||
|
||||
## log calculation clipped because the posterior variance
|
||||
## is 0 at the beginning of the diffusion chain
|
||||
self.register_buffer('posterior_log_variance_clipped',
|
||||
torch.log(torch.clamp(posterior_variance, min=1e-20)))
|
||||
self.register_buffer('posterior_mean_coef1',
|
||||
betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
|
||||
self.register_buffer('posterior_mean_coef2',
|
||||
(1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))
|
||||
|
||||
self.loss_fn = Losses[loss_type]()
|
||||
|
||||
# ------------------------------------------ sampling ------------------------------------------#
|
||||
|
||||
def predict_start_from_noise(self, x_t, t, noise):
|
||||
'''
|
||||
if self.predict_epsilon, model output is (scaled) noise;
|
||||
otherwise, model predicts x0 directly
|
||||
'''
|
||||
if self.predict_epsilon:
|
||||
return (
|
||||
extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
|
||||
extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
|
||||
)
|
||||
else:
|
||||
return noise
|
||||
|
||||
def q_posterior(self, x_start, x_t, t):
|
||||
posterior_mean = (
|
||||
extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
|
||||
extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
|
||||
)
|
||||
posterior_variance = extract(self.posterior_variance, t, x_t.shape)
|
||||
posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
|
||||
return posterior_mean, posterior_variance, posterior_log_variance_clipped
|
||||
|
||||
def p_mean_variance(self, x, t, s):
|
||||
x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, t, s))
|
||||
|
||||
if self.clip_denoised:
|
||||
x_recon.clamp_(-self.max_action, self.max_action)
|
||||
else:
|
||||
assert RuntimeError()
|
||||
|
||||
model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
|
||||
return model_mean, posterior_variance, posterior_log_variance
|
||||
|
||||
# @torch.no_grad()
|
||||
def p_sample(self, x, t, s):
|
||||
b, *_, device = *x.shape, x.device
|
||||
model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, s=s)
|
||||
noise = torch.randn_like(x)
|
||||
# no noise when t == 0
|
||||
nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
|
||||
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
|
||||
|
||||
# @torch.no_grad()
|
||||
def p_sample_loop(self, state, shape, verbose=False, return_diffusion=False):
|
||||
device = self.betas.device
|
||||
|
||||
batch_size = shape[0]
|
||||
x = torch.randn(shape, device=device)
|
||||
|
||||
if return_diffusion: diffusion = [x]
|
||||
|
||||
progress = Progress(self.n_timesteps) if verbose else Silent()
|
||||
for i in reversed(range(0, self.n_timesteps)):
|
||||
timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
|
||||
x = self.p_sample(x, timesteps, state)
|
||||
|
||||
progress.update({'t': i})
|
||||
|
||||
if return_diffusion: diffusion.append(x)
|
||||
|
||||
progress.close()
|
||||
|
||||
if return_diffusion:
|
||||
return x, torch.stack(diffusion, dim=1)
|
||||
else:
|
||||
return x
|
||||
|
||||
# @torch.no_grad()
|
||||
def sample(self, state, *args, **kwargs):
|
||||
batch_size = state.shape[0]
|
||||
shape = (batch_size, self.action_dim)
|
||||
action = self.p_sample_loop(state, shape, *args, **kwargs)
|
||||
return action.clamp_(-self.max_action, self.max_action)
|
||||
|
||||
# ------------------------------------------ training ------------------------------------------#
|
||||
|
||||
def q_sample(self, x_start, t, noise=None):
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x_start)
|
||||
|
||||
sample = (
|
||||
extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
|
||||
extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
|
||||
)
|
||||
|
||||
return sample
|
||||
|
||||
def p_losses(self, x_start, state, t, weights=1.0):
|
||||
noise = torch.randn_like(x_start)
|
||||
|
||||
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
|
||||
|
||||
x_recon = self.model(x_noisy, t, state)
|
||||
|
||||
assert noise.shape == x_recon.shape
|
||||
|
||||
if self.predict_epsilon:
|
||||
loss = self.loss_fn(x_recon, noise, weights)
|
||||
else:
|
||||
loss = self.loss_fn(x_recon, x_start, weights)
|
||||
|
||||
return loss
|
||||
|
||||
def loss(self, x, state, weights=1.0):
|
||||
batch_size = len(x)
|
||||
t = torch.randint(0, self.n_timesteps, (batch_size,), device=x.device).long()
|
||||
return self.p_losses(x, state, t, weights)
|
||||
|
||||
def forward(self, state, *args, **kwargs):
|
||||
return self.sample(state, *args, **kwargs)
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import math
|
||||
import time
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class SinusoidalPosEmb(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
device = x.device
|
||||
half_dim = self.dim // 2
|
||||
emb = math.log(10000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
|
||||
emb = x[:, None] * emb[None, :]
|
||||
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
|
||||
return emb
|
||||
|
||||
#-----------------------------------------------------------------------------#
|
||||
#---------------------------------- sampling ---------------------------------#
|
||||
#-----------------------------------------------------------------------------#
|
||||
|
||||
|
||||
def extract(a, t, x_shape):
|
||||
b, *_ = t.shape
|
||||
out = a.gather(-1, t)
|
||||
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
|
||||
|
||||
|
||||
def cosine_beta_schedule(timesteps, s=0.008, dtype=torch.float32):
|
||||
"""
|
||||
cosine schedule
|
||||
as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
|
||||
"""
|
||||
steps = timesteps + 1
|
||||
x = np.linspace(0, steps, steps)
|
||||
alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
|
||||
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
||||
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
|
||||
betas_clipped = np.clip(betas, a_min=0, a_max=0.999)
|
||||
return torch.tensor(betas_clipped, dtype=dtype)
|
||||
|
||||
|
||||
def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=2e-2, dtype=torch.float32):
|
||||
betas = np.linspace(
|
||||
beta_start, beta_end, timesteps
|
||||
)
|
||||
return torch.tensor(betas, dtype=dtype)
|
||||
|
||||
|
||||
def vp_beta_schedule(timesteps, dtype=torch.float32):
|
||||
t = np.arange(1, timesteps + 1)
|
||||
T = timesteps
|
||||
b_max = 10.
|
||||
b_min = 0.1
|
||||
alpha = np.exp(-b_min / T - 0.5 * (b_max - b_min) * (2 * t - 1) / T ** 2)
|
||||
betas = 1 - alpha
|
||||
return torch.tensor(betas, dtype=dtype)
|
||||
|
||||
#-----------------------------------------------------------------------------#
|
||||
#---------------------------------- losses -----------------------------------#
|
||||
#-----------------------------------------------------------------------------#
|
||||
|
||||
class WeightedLoss(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, pred, targ, weights=1.0):
|
||||
'''
|
||||
pred, targ : tensor [ batch_size x action_dim ]
|
||||
'''
|
||||
loss = self._loss(pred, targ)
|
||||
weighted_loss = (loss * weights).mean()
|
||||
return weighted_loss
|
||||
|
||||
class WeightedL1(WeightedLoss):
|
||||
|
||||
def _loss(self, pred, targ):
|
||||
return torch.abs(pred - targ)
|
||||
|
||||
class WeightedL2(WeightedLoss):
|
||||
|
||||
def _loss(self, pred, targ):
|
||||
return F.mse_loss(pred, targ, reduction='none')
|
||||
|
||||
|
||||
Losses = {
|
||||
'l1': WeightedL1,
|
||||
'l2': WeightedL2,
|
||||
}
|
||||
|
||||
|
||||
class EMA():
|
||||
'''
|
||||
empirical moving average
|
||||
'''
|
||||
def __init__(self, beta):
|
||||
super().__init__()
|
||||
self.beta = beta
|
||||
|
||||
def update_model_average(self, ma_model, current_model):
|
||||
for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
|
||||
old_weight, up_weight = ma_params.data, current_params.data
|
||||
ma_params.data = self.update_average(old_weight, up_weight)
|
||||
|
||||
def update_average(self, old, new):
|
||||
if old is None:
|
||||
return new
|
||||
return old * self.beta + (1 - self.beta) * new
|
||||
@@ -0,0 +1,50 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from agents.helpers import SinusoidalPosEmb
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""
|
||||
MLP Model
|
||||
"""
|
||||
def __init__(self,
|
||||
state_dim,
|
||||
action_dim,
|
||||
device,
|
||||
t_dim=16):
|
||||
|
||||
super(MLP, self).__init__()
|
||||
self.device = device
|
||||
|
||||
self.time_mlp = nn.Sequential(
|
||||
SinusoidalPosEmb(t_dim),
|
||||
nn.Linear(t_dim, t_dim * 2),
|
||||
nn.Mish(),
|
||||
nn.Linear(t_dim * 2, t_dim),
|
||||
)
|
||||
|
||||
input_dim = state_dim + action_dim + t_dim
|
||||
self.mid_layer = nn.Sequential(nn.Linear(input_dim, 256),
|
||||
nn.Mish(),
|
||||
nn.Linear(256, 256),
|
||||
nn.Mish(),
|
||||
nn.Linear(256, 256),
|
||||
nn.Mish())
|
||||
|
||||
self.final_layer = nn.Linear(256, action_dim)
|
||||
|
||||
def forward(self, x, time, state):
|
||||
|
||||
t = self.time_mlp(time)
|
||||
x = torch.cat([x, t, state], dim=1)
|
||||
x = self.mid_layer(x)
|
||||
|
||||
return self.final_layer(x)
|
||||
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import copy
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.optim.lr_scheduler import CosineAnnealingLR
|
||||
from utils.logger import logger
|
||||
|
||||
from agents.diffusion import Diffusion
|
||||
from agents.model import MLP
|
||||
from agents.helpers import EMA
|
||||
|
||||
|
||||
class Critic(nn.Module):
|
||||
def __init__(self, state_dim, action_dim, hidden_dim=256):
|
||||
super(Critic, self).__init__()
|
||||
self.q1_model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
|
||||
nn.Mish(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.Mish(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.Mish(),
|
||||
nn.Linear(hidden_dim, 1))
|
||||
|
||||
self.q2_model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
|
||||
nn.Mish(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.Mish(),
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
nn.Mish(),
|
||||
nn.Linear(hidden_dim, 1))
|
||||
|
||||
def forward(self, state, action):
|
||||
x = torch.cat([state, action], dim=-1)
|
||||
return self.q1_model(x), self.q2_model(x)
|
||||
|
||||
def q1(self, state, action):
|
||||
x = torch.cat([state, action], dim=-1)
|
||||
return self.q1_model(x)
|
||||
|
||||
def q_min(self, state, action):
|
||||
q1, q2 = self.forward(state, action)
|
||||
return torch.min(q1, q2)
|
||||
|
||||
|
||||
class Diffusion_QL(object):
|
||||
def __init__(self,
|
||||
state_dim,
|
||||
action_dim,
|
||||
max_action,
|
||||
device,
|
||||
discount,
|
||||
tau,
|
||||
max_q_backup=False,
|
||||
eta=1.0,
|
||||
beta_schedule='linear',
|
||||
n_timesteps=100,
|
||||
ema_decay=0.995,
|
||||
step_start_ema=1000,
|
||||
update_ema_every=5,
|
||||
lr=3e-4,
|
||||
lr_decay=False,
|
||||
lr_maxt=1000,
|
||||
grad_norm=1.0,
|
||||
):
|
||||
|
||||
self.model = MLP(state_dim=state_dim, action_dim=action_dim, device=device)
|
||||
|
||||
self.actor = Diffusion(state_dim=state_dim, action_dim=action_dim, model=self.model, max_action=max_action,
|
||||
beta_schedule=beta_schedule, n_timesteps=n_timesteps,).to(device)
|
||||
self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)
|
||||
|
||||
self.lr_decay = lr_decay
|
||||
self.grad_norm = grad_norm
|
||||
|
||||
self.step = 0
|
||||
self.step_start_ema = step_start_ema
|
||||
self.ema = EMA(ema_decay)
|
||||
self.ema_model = copy.deepcopy(self.actor)
|
||||
self.update_ema_every = update_ema_every
|
||||
|
||||
self.critic = Critic(state_dim, action_dim).to(device)
|
||||
self.critic_target = copy.deepcopy(self.critic)
|
||||
self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=3e-4)
|
||||
|
||||
if lr_decay:
|
||||
self.actor_lr_scheduler = CosineAnnealingLR(self.actor_optimizer, T_max=lr_maxt, eta_min=0.)
|
||||
self.critic_lr_scheduler = CosineAnnealingLR(self.critic_optimizer, T_max=lr_maxt, eta_min=0.)
|
||||
|
||||
self.state_dim = state_dim
|
||||
self.max_action = max_action
|
||||
self.action_dim = action_dim
|
||||
self.discount = discount
|
||||
self.tau = tau
|
||||
self.eta = eta # q_learning weight
|
||||
self.device = device
|
||||
self.max_q_backup = max_q_backup
|
||||
|
||||
def step_ema(self):
|
||||
if self.step < self.step_start_ema:
|
||||
return
|
||||
self.ema.update_model_average(self.ema_model, self.actor)
|
||||
|
||||
def train(self, replay_buffer, iterations, batch_size=100, log_writer=None):
|
||||
|
||||
metric = {'bc_loss': [], 'ql_loss': [], 'actor_loss': [], 'critic_loss': []}
|
||||
for _ in range(iterations):
|
||||
# Sample replay buffer / batch
|
||||
state, action, next_state, reward, not_done = replay_buffer.sample(batch_size)
|
||||
|
||||
""" Q Training """
|
||||
current_q1, current_q2 = self.critic(state, action)
|
||||
|
||||
if self.max_q_backup:
|
||||
next_state_rpt = torch.repeat_interleave(next_state, repeats=10, dim=0)
|
||||
next_action_rpt = self.ema_model(next_state_rpt)
|
||||
target_q1, target_q2 = self.critic_target(next_state_rpt, next_action_rpt)
|
||||
target_q1 = target_q1.view(batch_size, 10).max(dim=1, keepdim=True)[0]
|
||||
target_q2 = target_q2.view(batch_size, 10).max(dim=1, keepdim=True)[0]
|
||||
target_q = torch.min(target_q1, target_q2)
|
||||
else:
|
||||
next_action = self.ema_model(next_state)
|
||||
target_q1, target_q2 = self.critic_target(next_state, next_action)
|
||||
target_q = torch.min(target_q1, target_q2)
|
||||
|
||||
target_q = (reward + not_done * self.discount * target_q).detach()
|
||||
|
||||
critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q)
|
||||
|
||||
self.critic_optimizer.zero_grad()
|
||||
critic_loss.backward()
|
||||
if self.grad_norm > 0:
|
||||
critic_grad_norms = nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=self.grad_norm, norm_type=2)
|
||||
self.critic_optimizer.step()
|
||||
|
||||
""" Policy Training """
|
||||
bc_loss = self.actor.loss(action, state)
|
||||
new_action = self.actor(state)
|
||||
|
||||
q1_new_action, q2_new_action = self.critic(state, new_action)
|
||||
if np.random.uniform() > 0.5:
|
||||
q_loss = - q1_new_action.mean() / q2_new_action.abs().mean().detach()
|
||||
else:
|
||||
q_loss = - q2_new_action.mean() / q1_new_action.abs().mean().detach()
|
||||
actor_loss = bc_loss + self.eta * q_loss
|
||||
|
||||
self.actor_optimizer.zero_grad()
|
||||
actor_loss.backward()
|
||||
if self.grad_norm > 0:
|
||||
actor_grad_norms = nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=self.grad_norm, norm_type=2)
|
||||
self.actor_optimizer.step()
|
||||
|
||||
|
||||
""" Step Target network """
|
||||
if self.step % self.update_ema_every == 0:
|
||||
self.step_ema()
|
||||
|
||||
for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
|
||||
target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
|
||||
|
||||
self.step += 1
|
||||
|
||||
""" Log """
|
||||
if log_writer is not None:
|
||||
if self.grad_norm > 0:
|
||||
log_writer.add_scalar('Actor Grad Norm', actor_grad_norms.max().item(), self.step)
|
||||
log_writer.add_scalar('Critic Grad Norm', critic_grad_norms.max().item(), self.step)
|
||||
log_writer.add_scalar('BC Loss', bc_loss.item(), self.step)
|
||||
log_writer.add_scalar('QL Loss', q_loss.item(), self.step)
|
||||
log_writer.add_scalar('Critic Loss', critic_loss.item(), self.step)
|
||||
log_writer.add_scalar('Target_Q Mean', target_q.mean().item(), self.step)
|
||||
|
||||
metric['actor_loss'].append(actor_loss.item())
|
||||
metric['bc_loss'].append(bc_loss.item())
|
||||
metric['ql_loss'].append(q_loss.item())
|
||||
metric['critic_loss'].append(critic_loss.item())
|
||||
|
||||
if self.lr_decay:
|
||||
self.actor_lr_scheduler.step()
|
||||
self.critic_lr_scheduler.step()
|
||||
|
||||
return metric
|
||||
|
||||
def sample_action(self, state):
|
||||
state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
|
||||
state_rpt = torch.repeat_interleave(state, repeats=50, dim=0)
|
||||
with torch.no_grad():
|
||||
action = self.actor.sample(state_rpt)
|
||||
q_value = self.critic_target.q_min(state_rpt, action).flatten()
|
||||
idx = torch.multinomial(F.softmax(q_value, dim=0), 1)
|
||||
return action[idx].cpu().data.numpy().flatten()
|
||||
|
||||
def save_model(self, dir, id=None):
|
||||
if id is not None:
|
||||
torch.save(self.actor.state_dict(), f'{dir}/actor_{id}.pth')
|
||||
torch.save(self.critic.state_dict(), f'{dir}/critic_{id}.pth')
|
||||
else:
|
||||
torch.save(self.actor.state_dict(), f'{dir}/actor.pth')
|
||||
torch.save(self.critic.state_dict(), f'{dir}/critic.pth')
|
||||
|
||||
def load_model(self, dir, id=None):
|
||||
if id is not None:
|
||||
self.actor.load_state_dict(torch.load(f'{dir}/actor_{id}.pth'))
|
||||
self.critic.load_state_dict(torch.load(f'{dir}/critic_{id}.pth'))
|
||||
else:
|
||||
self.actor.load_state_dict(torch.load(f'{dir}/actor.pth'))
|
||||
self.critic.load_state_dict(torch.load(f'{dir}/critic.pth'))
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,98 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Author: SilverWings
|
||||
GitHub: https://github.com/silverwingsbot
|
||||
|
||||
Simple example: Run a trained Diffusion_QL model in easycarla
|
||||
"""
|
||||
|
||||
import gym
|
||||
import easycarla
|
||||
import numpy as np
|
||||
import torch
|
||||
import os
|
||||
from agents.ql_diffusion import Diffusion_QL
|
||||
|
||||
# ===================== Helper Functions =====================
|
||||
def convert_obs_dict_to_vector(obs_dict):
|
||||
"""Convert observation dictionary to a flattened state vector."""
|
||||
return np.concatenate([
|
||||
obs_dict['ego_state'], # 9 dimensions
|
||||
obs_dict['lane_info'], # 2 dimensions
|
||||
obs_dict['lidar'], # 240 dimensions
|
||||
obs_dict['nearby_vehicles'], # 20 dimensions
|
||||
obs_dict['waypoints'] # 36 dimensions
|
||||
]).astype(np.float32)
|
||||
|
||||
# ===================== Environment Configuration =====================
|
||||
carla_params = {
|
||||
'number_of_vehicles': 100,
|
||||
'number_of_walkers': 0,
|
||||
'dt': 0.1, # time interval between two frames
|
||||
'ego_vehicle_filter': 'vehicle.tesla.model3', # filter for defining ego vehicle
|
||||
'surrounding_vehicle_spawned_randomly': True, # Whether surrounding vehicles are spawned randomly (True) or set manually (False)
|
||||
'port': 2000, # connection port
|
||||
'town': 'Town03', # which town to simulate
|
||||
'max_time_episode': 1000, # maximum timesteps per episode
|
||||
'max_waypoints': 12, # maximum number of waypoints
|
||||
'visualize_waypoints': True, # Whether to visualize waypoints (default: True)
|
||||
'desired_speed': 8, # desired speed (m/s)
|
||||
'max_ego_spawn_times': 200, # maximum times to spawn ego vehicle
|
||||
'view_mode' : 'top', # 'top' for bird's-eye view, 'follow' for third-person view
|
||||
'traffic': 'off', # 'on' for normal traffic lights, 'off' for always green and frozen
|
||||
'lidar_max_range': 50.0, # Maximum LIDAR perception range (meters)
|
||||
'max_nearby_vehicles': 5, # Maximum number of nearby vehicles to observe
|
||||
}
|
||||
|
||||
# ===================== Initialize Environment =====================
|
||||
env = gym.make('carla-v0', params=carla_params)
|
||||
|
||||
# ===================== Initialize Model =====================
|
||||
state_dim = 307
|
||||
action_dim = 3
|
||||
max_action = 1.0
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
model = Diffusion_QL(
|
||||
state_dim=state_dim,
|
||||
action_dim=action_dim,
|
||||
max_action=max_action,
|
||||
device=device,
|
||||
discount=0.99,
|
||||
tau=0.005,
|
||||
eta=0.01,
|
||||
beta_schedule='vp',
|
||||
n_timesteps=5
|
||||
)
|
||||
|
||||
# ===================== Load Pretrained Model =====================
|
||||
model_id = 200 # Model checkpoint ID to load
|
||||
save_path = './params_dql' # Model checkpoint directory
|
||||
model.load_model(save_path, id=model_id)
|
||||
print(f"Successfully loaded model ID {model_id}")
|
||||
|
||||
# ===================== Run One Episode =====================
|
||||
obs = env.reset()
|
||||
done = False
|
||||
step = 0
|
||||
episode_reward = 0.0
|
||||
|
||||
while not done:
|
||||
obs_vec = convert_obs_dict_to_vector(obs)
|
||||
action = model.sample_action(obs_vec)
|
||||
|
||||
try:
|
||||
next_obs, reward, cost, done, info = env.step(action)
|
||||
except Exception as e:
|
||||
print(f"[Error] Carla step failed: {e}")
|
||||
obs = env.reset()
|
||||
continue
|
||||
|
||||
obs = next_obs
|
||||
episode_reward += reward
|
||||
step += 1
|
||||
|
||||
# Optional: add a delay for better visualization
|
||||
# time.sleep(0.05)
|
||||
|
||||
print(f"Episode finished. Total reward: {episode_reward:.2f}, Total steps: {step}")
|
||||
@@ -0,0 +1,55 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class Data_Sampler(object):
|
||||
def __init__(self, data, device, reward_tune='no'):
|
||||
self.state = torch.from_numpy(data['observations']).float()
|
||||
self.action = torch.from_numpy(data['actions']).float()
|
||||
self.next_state = torch.from_numpy(data['next_observations']).float()
|
||||
reward = torch.from_numpy(data['rewards']).reshape(-1, 1).float()
|
||||
self.not_done = 1. - torch.from_numpy(data['dones']).reshape(-1, 1).float()
|
||||
|
||||
self.size = self.state.shape[0]
|
||||
self.state_dim = self.state.shape[1]
|
||||
self.action_dim = self.action.shape[1]
|
||||
self.device = device
|
||||
|
||||
# 保留原 reward 调整逻辑
|
||||
if reward_tune == 'normalize':
|
||||
reward = (reward - reward.mean()) / reward.std()
|
||||
elif reward_tune == 'iql_antmaze':
|
||||
reward = reward - 1.0
|
||||
elif reward_tune == 'iql_locomotion':
|
||||
reward = iql_normalize(reward, self.not_done)
|
||||
elif reward_tune == 'cql_antmaze':
|
||||
reward = (reward - 0.5) * 4.0
|
||||
elif reward_tune == 'antmaze':
|
||||
reward = (reward - 0.25) * 2.0
|
||||
|
||||
self.reward = reward
|
||||
|
||||
def sample(self, batch_size):
|
||||
ind = torch.randint(0, self.size, size=(batch_size,))
|
||||
return (
|
||||
self.state[ind].to(self.device),
|
||||
self.action[ind].to(self.device),
|
||||
self.next_state[ind].to(self.device),
|
||||
self.reward[ind].to(self.device),
|
||||
self.not_done[ind].to(self.device)
|
||||
)
|
||||
|
||||
def iql_normalize(reward, not_done):
|
||||
trajs_rt = []
|
||||
episode_return = 0.0
|
||||
for i in range(len(reward)):
|
||||
episode_return += reward[i]
|
||||
if not not_done[i]:
|
||||
trajs_rt.append(episode_return)
|
||||
episode_return = 0.0
|
||||
rt_max, rt_min = torch.max(torch.tensor(trajs_rt)), torch.min(torch.tensor(trajs_rt))
|
||||
reward /= (rt_max - rt_min)
|
||||
reward *= 1000.
|
||||
return reward
|
||||
@@ -0,0 +1,493 @@
|
||||
"""
|
||||
Based on rllab's logger.
|
||||
|
||||
https://github.com/rll/rllab
|
||||
"""
|
||||
from enum import Enum
|
||||
from contextlib import contextmanager
|
||||
import numpy as np
|
||||
import os
|
||||
import os.path as osp
|
||||
import sys
|
||||
import datetime
|
||||
import dateutil.tz
|
||||
import csv
|
||||
import json
|
||||
import pickle
|
||||
import errno
|
||||
from collections import OrderedDict
|
||||
from numbers import Number
|
||||
import os
|
||||
|
||||
from tabulate import tabulate
|
||||
import dateutil.tz
|
||||
import os.path as osp
|
||||
|
||||
def dict_to_safe_json(d):
|
||||
"""
|
||||
Convert each value in the dictionary into a JSON'able primitive.
|
||||
:param d:
|
||||
:return:
|
||||
"""
|
||||
new_d = {}
|
||||
for key, item in d.items():
|
||||
if safe_json(item):
|
||||
new_d[key] = item
|
||||
else:
|
||||
if isinstance(item, dict):
|
||||
new_d[key] = dict_to_safe_json(item)
|
||||
else:
|
||||
new_d[key] = str(item)
|
||||
return new_d
|
||||
|
||||
|
||||
def safe_json(data):
|
||||
if data is None:
|
||||
return True
|
||||
elif isinstance(data, (bool, int, float)):
|
||||
return True
|
||||
elif isinstance(data, (tuple, list)):
|
||||
return all(safe_json(x) for x in data)
|
||||
elif isinstance(data, dict):
|
||||
return all(isinstance(k, str) and safe_json(v) for k, v in data.items())
|
||||
return False
|
||||
|
||||
def create_exp_name(exp_prefix, exp_id=0, seed=0):
|
||||
"""
|
||||
Create a semi-unique experiment name that has a timestamp
|
||||
:param exp_prefix:
|
||||
:param exp_id:
|
||||
:return:
|
||||
"""
|
||||
now = datetime.datetime.now(dateutil.tz.tzlocal())
|
||||
timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')
|
||||
return "%s_%s_%04d--s-%d" % (exp_prefix, timestamp, exp_id, seed)
|
||||
|
||||
def create_log_dir(
|
||||
exp_prefix,
|
||||
exp_id=0,
|
||||
seed=0,
|
||||
base_log_dir=None,
|
||||
include_exp_prefix_sub_dir=True,
|
||||
):
|
||||
"""
|
||||
Creates and returns a unique log directory.
|
||||
:param exp_prefix: All experiments with this prefix will have log
|
||||
directories be under this directory.
|
||||
:param exp_id: The number of the specific experiment run within this
|
||||
experiment.
|
||||
:param base_log_dir: The directory where all log should be saved.
|
||||
:return:
|
||||
"""
|
||||
exp_name = create_exp_name(exp_prefix, exp_id=exp_id,
|
||||
seed=seed)
|
||||
if base_log_dir is None:
|
||||
base_log_dir = './data'
|
||||
if include_exp_prefix_sub_dir:
|
||||
log_dir = osp.join(base_log_dir, exp_prefix.replace("_", "-"), exp_name)
|
||||
else:
|
||||
log_dir = osp.join(base_log_dir, exp_name)
|
||||
if osp.exists(log_dir):
|
||||
print("WARNING: Log directory already exists {}".format(log_dir), flush=True)
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
return log_dir
|
||||
|
||||
|
||||
def setup_logger(
|
||||
exp_prefix="default",
|
||||
variant=None,
|
||||
text_log_file="debug.log",
|
||||
variant_log_file="variant.json",
|
||||
tabular_log_file="progress.csv",
|
||||
snapshot_mode="last",
|
||||
snapshot_gap=1,
|
||||
log_tabular_only=False,
|
||||
log_dir=None,
|
||||
git_infos=None,
|
||||
script_name=None,
|
||||
**create_log_dir_kwargs
|
||||
):
|
||||
"""
|
||||
Set up logger to have some reasonable default settings.
|
||||
Will save log output to
|
||||
based_log_dir/exp_prefix/exp_name.
|
||||
exp_name will be auto-generated to be unique.
|
||||
If log_dir is specified, then that directory is used as the output dir.
|
||||
:param exp_prefix: The sub-directory for this specific experiment.
|
||||
:param variant:
|
||||
:param text_log_file:
|
||||
:param variant_log_file:
|
||||
:param tabular_log_file:
|
||||
:param snapshot_mode:
|
||||
:param log_tabular_only:
|
||||
:param snapshot_gap:
|
||||
:param log_dir:
|
||||
:param git_infos:
|
||||
:param script_name: If set, save the script name to this.
|
||||
:return:
|
||||
"""
|
||||
first_time = log_dir is None
|
||||
if first_time:
|
||||
log_dir = create_log_dir(exp_prefix, **create_log_dir_kwargs)
|
||||
|
||||
if variant is not None:
|
||||
logger.log("Variant:")
|
||||
logger.log(json.dumps(dict_to_safe_json(variant), indent=2))
|
||||
variant_log_path = osp.join(log_dir, variant_log_file)
|
||||
logger.log_variant(variant_log_path, variant)
|
||||
|
||||
tabular_log_path = osp.join(log_dir, tabular_log_file)
|
||||
text_log_path = osp.join(log_dir, text_log_file)
|
||||
|
||||
logger.add_text_output(text_log_path)
|
||||
if first_time:
|
||||
logger.add_tabular_output(tabular_log_path)
|
||||
else:
|
||||
logger._add_output(tabular_log_path, logger._tabular_outputs,
|
||||
logger._tabular_fds, mode='a')
|
||||
for tabular_fd in logger._tabular_fds:
|
||||
logger._tabular_header_written.add(tabular_fd)
|
||||
logger.set_snapshot_dir(log_dir)
|
||||
logger.set_snapshot_mode(snapshot_mode)
|
||||
logger.set_snapshot_gap(snapshot_gap)
|
||||
logger.set_log_tabular_only(log_tabular_only)
|
||||
exp_name = log_dir.split("/")[-1]
|
||||
logger.push_prefix("[%s] " % exp_name)
|
||||
|
||||
if script_name is not None:
|
||||
with open(osp.join(log_dir, "script_name.txt"), "w") as f:
|
||||
f.write(script_name)
|
||||
return log_dir
|
||||
|
||||
|
||||
def create_stats_ordered_dict(
|
||||
name,
|
||||
data,
|
||||
stat_prefix=None,
|
||||
always_show_all_stats=True,
|
||||
exclude_max_min=False,
|
||||
):
|
||||
if stat_prefix is not None:
|
||||
name = "{}{}".format(stat_prefix, name)
|
||||
if isinstance(data, Number):
|
||||
return OrderedDict({name: data})
|
||||
|
||||
if len(data) == 0:
|
||||
return OrderedDict()
|
||||
|
||||
if isinstance(data, tuple):
|
||||
ordered_dict = OrderedDict()
|
||||
for number, d in enumerate(data):
|
||||
sub_dict = create_stats_ordered_dict(
|
||||
"{0}_{1}".format(name, number),
|
||||
d,
|
||||
)
|
||||
ordered_dict.update(sub_dict)
|
||||
return ordered_dict
|
||||
|
||||
if isinstance(data, list):
|
||||
try:
|
||||
iter(data[0])
|
||||
except TypeError:
|
||||
pass
|
||||
else:
|
||||
data = np.concatenate(data)
|
||||
|
||||
if (isinstance(data, np.ndarray) and data.size == 1
|
||||
and not always_show_all_stats):
|
||||
return OrderedDict({name: float(data)})
|
||||
|
||||
stats = OrderedDict([
|
||||
(name + ' Mean', np.mean(data)),
|
||||
(name + ' Std', np.std(data)),
|
||||
])
|
||||
if not exclude_max_min:
|
||||
stats[name + ' Max'] = np.max(data)
|
||||
stats[name + ' Min'] = np.min(data)
|
||||
return stats
|
||||
|
||||
|
||||
class TerminalTablePrinter(object):
|
||||
def __init__(self):
|
||||
self.headers = None
|
||||
self.tabulars = []
|
||||
|
||||
def print_tabular(self, new_tabular):
|
||||
if self.headers is None:
|
||||
self.headers = [x[0] for x in new_tabular]
|
||||
else:
|
||||
assert len(self.headers) == len(new_tabular)
|
||||
self.tabulars.append([x[1] for x in new_tabular])
|
||||
self.refresh()
|
||||
|
||||
def refresh(self):
|
||||
import os
|
||||
rows, columns = os.popen('stty size', 'r').read().split()
|
||||
tabulars = self.tabulars[-(int(rows) - 3):]
|
||||
sys.stdout.write("\x1b[2J\x1b[H")
|
||||
sys.stdout.write(tabulate(tabulars, self.headers))
|
||||
sys.stdout.write("\n")
|
||||
|
||||
|
||||
class MyEncoder(json.JSONEncoder):
|
||||
def default(self, o):
|
||||
if isinstance(o, type):
|
||||
return {'$class': o.__module__ + "." + o.__name__}
|
||||
elif isinstance(o, Enum):
|
||||
return {
|
||||
'$enum': o.__module__ + "." + o.__class__.__name__ + '.' + o.name
|
||||
}
|
||||
elif callable(o):
|
||||
return {
|
||||
'$function': o.__module__ + "." + o.__name__
|
||||
}
|
||||
return json.JSONEncoder.default(self, o)
|
||||
|
||||
|
||||
def mkdir_p(path):
|
||||
try:
|
||||
os.makedirs(path)
|
||||
except OSError as exc: # Python >2.5
|
||||
if exc.errno == errno.EEXIST and os.path.isdir(path):
|
||||
pass
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
class Logger(object):
|
||||
def __init__(self):
|
||||
self._prefixes = []
|
||||
self._prefix_str = ''
|
||||
|
||||
self._tabular_prefixes = []
|
||||
self._tabular_prefix_str = ''
|
||||
|
||||
self._tabular = []
|
||||
|
||||
self._text_outputs = []
|
||||
self._tabular_outputs = []
|
||||
|
||||
self._text_fds = {}
|
||||
self._tabular_fds = {}
|
||||
self._tabular_header_written = set()
|
||||
|
||||
self._snapshot_dir = None
|
||||
self._snapshot_mode = 'all'
|
||||
self._snapshot_gap = 1
|
||||
|
||||
self._log_tabular_only = False
|
||||
self._header_printed = False
|
||||
self.table_printer = TerminalTablePrinter()
|
||||
|
||||
def reset(self):
|
||||
self.__init__()
|
||||
|
||||
def _add_output(self, file_name, arr, fds, mode='a'):
|
||||
if file_name not in arr:
|
||||
mkdir_p(os.path.dirname(file_name))
|
||||
arr.append(file_name)
|
||||
fds[file_name] = open(file_name, mode)
|
||||
|
||||
def _remove_output(self, file_name, arr, fds):
|
||||
if file_name in arr:
|
||||
fds[file_name].close()
|
||||
del fds[file_name]
|
||||
arr.remove(file_name)
|
||||
|
||||
def push_prefix(self, prefix):
|
||||
self._prefixes.append(prefix)
|
||||
self._prefix_str = ''.join(self._prefixes)
|
||||
|
||||
def add_text_output(self, file_name):
|
||||
self._add_output(file_name, self._text_outputs, self._text_fds,
|
||||
mode='a')
|
||||
|
||||
def remove_text_output(self, file_name):
|
||||
self._remove_output(file_name, self._text_outputs, self._text_fds)
|
||||
|
||||
def add_tabular_output(self, file_name, relative_to_snapshot_dir=False):
|
||||
if relative_to_snapshot_dir:
|
||||
file_name = osp.join(self._snapshot_dir, file_name)
|
||||
self._add_output(file_name, self._tabular_outputs, self._tabular_fds,
|
||||
mode='w')
|
||||
|
||||
def remove_tabular_output(self, file_name, relative_to_snapshot_dir=False):
|
||||
if relative_to_snapshot_dir:
|
||||
file_name = osp.join(self._snapshot_dir, file_name)
|
||||
if self._tabular_fds[file_name] in self._tabular_header_written:
|
||||
self._tabular_header_written.remove(self._tabular_fds[file_name])
|
||||
self._remove_output(file_name, self._tabular_outputs, self._tabular_fds)
|
||||
|
||||
def set_snapshot_dir(self, dir_name):
|
||||
self._snapshot_dir = dir_name
|
||||
|
||||
def get_snapshot_dir(self, ):
|
||||
return self._snapshot_dir
|
||||
|
||||
def get_snapshot_mode(self, ):
|
||||
return self._snapshot_mode
|
||||
|
||||
def set_snapshot_mode(self, mode):
|
||||
self._snapshot_mode = mode
|
||||
|
||||
def get_snapshot_gap(self, ):
|
||||
return self._snapshot_gap
|
||||
|
||||
def set_snapshot_gap(self, gap):
|
||||
self._snapshot_gap = gap
|
||||
|
||||
def set_log_tabular_only(self, log_tabular_only):
|
||||
self._log_tabular_only = log_tabular_only
|
||||
|
||||
def get_log_tabular_only(self, ):
|
||||
return self._log_tabular_only
|
||||
|
||||
def log(self, s, with_prefix=True, with_timestamp=True):
|
||||
out = s
|
||||
if with_prefix:
|
||||
out = self._prefix_str + out
|
||||
if with_timestamp:
|
||||
now = datetime.datetime.now(dateutil.tz.tzlocal())
|
||||
timestamp = now.strftime('%y-%m-%d.%H:%M') # :%S
|
||||
out = "%s|%s" % (timestamp, out)
|
||||
if not self._log_tabular_only:
|
||||
# Also log to stdout
|
||||
print(out, flush=True)
|
||||
for fd in list(self._text_fds.values()):
|
||||
fd.write(out + '\n')
|
||||
fd.flush()
|
||||
sys.stdout.flush()
|
||||
|
||||
def record_tabular(self, key, val):
|
||||
self._tabular.append((self._tabular_prefix_str + str(key), str(val)))
|
||||
|
||||
def record_dict(self, d, prefix=None):
|
||||
if prefix is not None:
|
||||
self.push_tabular_prefix(prefix)
|
||||
for k, v in d.items():
|
||||
self.record_tabular(k, v)
|
||||
if prefix is not None:
|
||||
self.pop_tabular_prefix()
|
||||
|
||||
def push_tabular_prefix(self, key):
|
||||
self._tabular_prefixes.append(key)
|
||||
self._tabular_prefix_str = ''.join(self._tabular_prefixes)
|
||||
|
||||
def pop_tabular_prefix(self, ):
|
||||
del self._tabular_prefixes[-1]
|
||||
self._tabular_prefix_str = ''.join(self._tabular_prefixes)
|
||||
|
||||
def save_extra_data(self, data, file_name='extra_data.pkl', mode='joblib'):
|
||||
"""
|
||||
Data saved here will always override the last entry
|
||||
|
||||
:param data: Something pickle'able.
|
||||
"""
|
||||
file_name = osp.join(self._snapshot_dir, file_name)
|
||||
if mode == 'joblib':
|
||||
import joblib
|
||||
joblib.dump(data, file_name, compress=3)
|
||||
elif mode == 'pickle':
|
||||
pickle.dump(data, open(file_name, "wb"))
|
||||
else:
|
||||
raise ValueError("Invalid mode: {}".format(mode))
|
||||
return file_name
|
||||
|
||||
def get_table_dict(self, ):
|
||||
return dict(self._tabular)
|
||||
|
||||
def get_table_key_set(self, ):
|
||||
return set(key for key, value in self._tabular)
|
||||
|
||||
@contextmanager
|
||||
def prefix(self, key):
|
||||
self.push_prefix(key)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.pop_prefix()
|
||||
|
||||
@contextmanager
|
||||
def tabular_prefix(self, key):
|
||||
self.push_tabular_prefix(key)
|
||||
yield
|
||||
self.pop_tabular_prefix()
|
||||
|
||||
def log_variant(self, log_file, variant_data):
|
||||
mkdir_p(os.path.dirname(log_file))
|
||||
with open(log_file, "w") as f:
|
||||
json.dump(variant_data, f, indent=2, sort_keys=True, cls=MyEncoder)
|
||||
|
||||
def record_tabular_misc_stat(self, key, values, placement='back'):
|
||||
if placement == 'front':
|
||||
prefix = ""
|
||||
suffix = key
|
||||
else:
|
||||
prefix = key
|
||||
suffix = ""
|
||||
if len(values) > 0:
|
||||
self.record_tabular(prefix + "Average" + suffix, np.average(values))
|
||||
self.record_tabular(prefix + "Std" + suffix, np.std(values))
|
||||
self.record_tabular(prefix + "Median" + suffix, np.median(values))
|
||||
self.record_tabular(prefix + "Min" + suffix, np.min(values))
|
||||
self.record_tabular(prefix + "Max" + suffix, np.max(values))
|
||||
else:
|
||||
self.record_tabular(prefix + "Average" + suffix, np.nan)
|
||||
self.record_tabular(prefix + "Std" + suffix, np.nan)
|
||||
self.record_tabular(prefix + "Median" + suffix, np.nan)
|
||||
self.record_tabular(prefix + "Min" + suffix, np.nan)
|
||||
self.record_tabular(prefix + "Max" + suffix, np.nan)
|
||||
|
||||
def dump_tabular(self, *args, **kwargs):
|
||||
wh = kwargs.pop("write_header", None)
|
||||
if len(self._tabular) > 0:
|
||||
if self._log_tabular_only:
|
||||
self.table_printer.print_tabular(self._tabular)
|
||||
else:
|
||||
for line in tabulate(self._tabular).split('\n'):
|
||||
self.log(line, *args, **kwargs)
|
||||
tabular_dict = dict(self._tabular)
|
||||
# Also write to the csv files
|
||||
# This assumes that the keys in each iteration won't change!
|
||||
for tabular_fd in list(self._tabular_fds.values()):
|
||||
writer = csv.DictWriter(tabular_fd,
|
||||
fieldnames=list(tabular_dict.keys()))
|
||||
if wh or (
|
||||
wh is None and tabular_fd not in self._tabular_header_written):
|
||||
writer.writeheader()
|
||||
self._tabular_header_written.add(tabular_fd)
|
||||
writer.writerow(tabular_dict)
|
||||
tabular_fd.flush()
|
||||
del self._tabular[:]
|
||||
|
||||
def pop_prefix(self, ):
|
||||
del self._prefixes[-1]
|
||||
self._prefix_str = ''.join(self._prefixes)
|
||||
|
||||
def save_itr_params(self, itr, params):
|
||||
if self._snapshot_dir:
|
||||
if self._snapshot_mode == 'all':
|
||||
file_name = osp.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
|
||||
pickle.dump(params, open(file_name, "wb"))
|
||||
elif self._snapshot_mode == 'last':
|
||||
# override previous params
|
||||
file_name = osp.join(self._snapshot_dir, 'params.pkl')
|
||||
pickle.dump(params, open(file_name, "wb"))
|
||||
elif self._snapshot_mode == "gap":
|
||||
if itr % self._snapshot_gap == 0:
|
||||
file_name = osp.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
|
||||
pickle.dump(params, open(file_name, "wb"))
|
||||
elif self._snapshot_mode == "gap_and_last":
|
||||
if itr % self._snapshot_gap == 0:
|
||||
file_name = osp.join(self._snapshot_dir, 'itr_%d.pkl' % itr)
|
||||
pickle.dump(params, open(file_name, "wb"))
|
||||
file_name = osp.join(self._snapshot_dir, 'params.pkl')
|
||||
pickle.dump(params, open(file_name, "wb"))
|
||||
elif self._snapshot_mode == 'none':
|
||||
pass
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
logger = Logger()
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def soft_update_from_to(source, target, tau):
|
||||
for target_param, param in zip(target.parameters(), source.parameters()):
|
||||
target_param.data.copy_(
|
||||
target_param.data * (1.0 - tau) + param.data * tau
|
||||
)
|
||||
|
||||
|
||||
def copy_model_params_from_to(source, target):
|
||||
for target_param, param in zip(target.parameters(), source.parameters()):
|
||||
target_param.data.copy_(param.data)
|
||||
|
||||
|
||||
def fanin_init(tensor, scale=1):
|
||||
size = tensor.size()
|
||||
if len(size) == 2:
|
||||
fan_in = size[0]
|
||||
elif len(size) > 2:
|
||||
fan_in = np.prod(size[1:])
|
||||
else:
|
||||
raise Exception("Shape must be have dimension at least 2.")
|
||||
bound = scale / np.sqrt(fan_in)
|
||||
return tensor.data.uniform_(-bound, bound)
|
||||
|
||||
|
||||
def orthogonal_init(tensor, gain=0.01):
|
||||
torch.nn.init.orthogonal_(tensor, gain=gain)
|
||||
|
||||
|
||||
def fanin_init_weights_like(tensor):
|
||||
size = tensor.size()
|
||||
if len(size) == 2:
|
||||
fan_in = size[0]
|
||||
elif len(size) > 2:
|
||||
fan_in = np.prod(size[1:])
|
||||
else:
|
||||
raise Exception("Shape must be have dimension at least 2.")
|
||||
bound = 1. / np.sqrt(fan_in)
|
||||
new_tensor = torch.FloatTensor(tensor.size())
|
||||
new_tensor.uniform_(-bound, bound)
|
||||
return new_tensor
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
# Copyright 2022 Twitter, Inc and Zhendong Wang.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import time
|
||||
import math
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
def print_banner(s, separator="-", num_star=60):
|
||||
print(separator * num_star, flush=True)
|
||||
print(s, flush=True)
|
||||
print(separator * num_star, flush=True)
|
||||
|
||||
|
||||
class Progress:
|
||||
|
||||
def __init__(self, total, name='Progress', ncol=3, max_length=20, indent=0, line_width=100, speed_update_freq=100):
|
||||
self.total = total
|
||||
self.name = name
|
||||
self.ncol = ncol
|
||||
self.max_length = max_length
|
||||
self.indent = indent
|
||||
self.line_width = line_width
|
||||
self._speed_update_freq = speed_update_freq
|
||||
|
||||
self._step = 0
|
||||
self._prev_line = '\033[F'
|
||||
self._clear_line = ' ' * self.line_width
|
||||
|
||||
self._pbar_size = self.ncol * self.max_length
|
||||
self._complete_pbar = '#' * self._pbar_size
|
||||
self._incomplete_pbar = ' ' * self._pbar_size
|
||||
|
||||
self.lines = ['']
|
||||
self.fraction = '{} / {}'.format(0, self.total)
|
||||
|
||||
self.resume()
|
||||
|
||||
def update(self, description, n=1):
|
||||
self._step += n
|
||||
if self._step % self._speed_update_freq == 0:
|
||||
self._time0 = time.time()
|
||||
self._step0 = self._step
|
||||
self.set_description(description)
|
||||
|
||||
def resume(self):
|
||||
self._skip_lines = 1
|
||||
print('\n', end='')
|
||||
self._time0 = time.time()
|
||||
self._step0 = self._step
|
||||
|
||||
def pause(self):
|
||||
self._clear()
|
||||
self._skip_lines = 1
|
||||
|
||||
def set_description(self, params=[]):
|
||||
|
||||
if type(params) == dict:
|
||||
params = sorted([
|
||||
(key, val)
|
||||
for key, val in params.items()
|
||||
])
|
||||
|
||||
############
|
||||
# Position #
|
||||
############
|
||||
self._clear()
|
||||
|
||||
###########
|
||||
# Percent #
|
||||
###########
|
||||
percent, fraction = self._format_percent(self._step, self.total)
|
||||
self.fraction = fraction
|
||||
|
||||
#########
|
||||
# Speed #
|
||||
#########
|
||||
speed = self._format_speed(self._step)
|
||||
|
||||
##########
|
||||
# Params #
|
||||
##########
|
||||
num_params = len(params)
|
||||
nrow = math.ceil(num_params / self.ncol)
|
||||
params_split = self._chunk(params, self.ncol)
|
||||
params_string, lines = self._format(params_split)
|
||||
self.lines = lines
|
||||
|
||||
description = '{} | {}{}'.format(percent, speed, params_string)
|
||||
print(description)
|
||||
self._skip_lines = nrow + 1
|
||||
|
||||
def append_description(self, descr):
|
||||
self.lines.append(descr)
|
||||
|
||||
def _clear(self):
|
||||
position = self._prev_line * self._skip_lines
|
||||
empty = '\n'.join([self._clear_line for _ in range(self._skip_lines)])
|
||||
print(position, end='')
|
||||
print(empty)
|
||||
print(position, end='')
|
||||
|
||||
def _format_percent(self, n, total):
|
||||
if total:
|
||||
percent = n / float(total)
|
||||
|
||||
complete_entries = int(percent * self._pbar_size)
|
||||
incomplete_entries = self._pbar_size - complete_entries
|
||||
|
||||
pbar = self._complete_pbar[:complete_entries] + self._incomplete_pbar[:incomplete_entries]
|
||||
fraction = '{} / {}'.format(n, total)
|
||||
string = '{} [{}] {:3d}%'.format(fraction, pbar, int(percent * 100))
|
||||
else:
|
||||
fraction = '{}'.format(n)
|
||||
string = '{} iterations'.format(n)
|
||||
return string, fraction
|
||||
|
||||
def _format_speed(self, n):
|
||||
num_steps = n - self._step0
|
||||
t = time.time() - self._time0
|
||||
speed = num_steps / t
|
||||
string = '{:.1f} Hz'.format(speed)
|
||||
if num_steps > 0:
|
||||
self._speed = string
|
||||
return string
|
||||
|
||||
def _chunk(self, l, n):
|
||||
return [l[i:i + n] for i in range(0, len(l), n)]
|
||||
|
||||
def _format(self, chunks):
|
||||
lines = [self._format_chunk(chunk) for chunk in chunks]
|
||||
lines.insert(0, '')
|
||||
padding = '\n' + ' ' * self.indent
|
||||
string = padding.join(lines)
|
||||
return string, lines
|
||||
|
||||
def _format_chunk(self, chunk):
|
||||
line = ' | '.join([self._format_param(param) for param in chunk])
|
||||
return line
|
||||
|
||||
def _format_param(self, param):
|
||||
k, v = param
|
||||
return '{} : {}'.format(k, v)[:self.max_length]
|
||||
|
||||
def stamp(self):
|
||||
if self.lines != ['']:
|
||||
params = ' | '.join(self.lines)
|
||||
string = '[ {} ] {}{} | {}'.format(self.name, self.fraction, params, self._speed)
|
||||
self._clear()
|
||||
print(string, end='\n')
|
||||
self._skip_lines = 1
|
||||
else:
|
||||
self._clear()
|
||||
self._skip_lines = 0
|
||||
|
||||
def close(self):
|
||||
self.pause()
|
||||
|
||||
|
||||
class Silent:
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def __getattr__(self, attr):
|
||||
return lambda *args: None
|
||||
|
||||
|
||||
class EarlyStopping(object):
|
||||
def __init__(self, tolerance=5, min_delta=0):
|
||||
self.tolerance = tolerance
|
||||
self.min_delta = min_delta
|
||||
self.counter = 0
|
||||
self.early_stop = False
|
||||
|
||||
def __call__(self, train_loss, validation_loss):
|
||||
if (validation_loss - train_loss) > self.min_delta:
|
||||
self.counter += 1
|
||||
if self.counter >= self.tolerance:
|
||||
return True
|
||||
else:
|
||||
self.counter = 0
|
||||
return False
|
||||
Reference in New Issue
Block a user