class DQN(Policy):
def __init__(self, is_train=False, dataset='Multiwoz'):
with open(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'config.json'), 'r') as f:
cfg = json.load(f)
self.save_dir = os.path.join(
os.path.dirname(os.path.abspath(__file__)), cfg['save_dir'])
self.save_per_epoch = cfg['save_per_epoch']
self.training_iter = cfg['training_iter']
self.training_batch_iter = cfg['training_batch_iter']
self.batch_size = cfg['batch_size']
self.gamma = cfg['gamma']
self.is_train = is_train
if is_train:
init_logging_handler(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
cfg['log_dir']))
# construct multiwoz vector
if dataset == 'Multiwoz':
voc_file = os.path.join(root_dir, 'data/multiwoz/sys_da_voc.txt')
voc_opp_file = os.path.join(root_dir,
'data/multiwoz/usr_da_voc.txt')
self.vector = MultiWozVector(voc_file,
voc_opp_file,
composite_actions=True,
vocab_size=cfg['vocab_size'])
#replay memory
self.memory = MemoryReplay(cfg['memory_size'])
self.net = EpsilonGreedyPolicy(self.vector.state_dim, cfg['hv_dim'],
self.vector.da_dim,
cfg['epsilon_spec']).to(device=DEVICE)
self.target_net = copy.deepcopy(self.net)
self.online_net = self.target_net
self.eval_net = self.target_net
if is_train:
self.net_optim = optim.Adam(self.net.parameters(), lr=cfg['lr'])
self.loss_fn = nn.MSELoss()
def update_memory(self, sample):
self.memory.append(sample)
def predict(self, state):
"""
Predict an system action given state.
Args:
state (dict): Dialog state. Please refer to util/state.py
Returns:
action : System act, with the form of (act_type, {slot_name_1: value_1, slot_name_2, value_2, ...})
"""
s_vec = torch.Tensor(self.vector.state_vectorize(state))
a = self.net.select_action(s_vec.to(device=DEVICE))
action = self.vector.action_devectorize(a.numpy())
state['system_action'] = action
return action
def init_session(self):
"""
Restore after one session
"""
self.memory.reset()
def calc_q_loss(self, batch):
'''Compute the Q value loss using predicted and target Q values from the appropriate networks'''
s = torch.from_numpy(np.stack(batch.state)).to(device=DEVICE)
a = torch.from_numpy(np.stack(batch.action)).to(device=DEVICE)
r = torch.from_numpy(np.stack(batch.reward)).to(device=DEVICE)
next_s = torch.from_numpy(np.stack(batch.next_state)).to(device=DEVICE)
mask = torch.Tensor(np.stack(batch.mask)).to(device=DEVICE)
q_preds = self.net(s)
with torch.no_grad():
# Use online_net to select actions in next state
online_next_q_preds = self.online_net(next_s)
# Use eval_net to calculate next_q_preds for actions chosen by online_net
next_q_preds = self.eval_net(next_s)
act_q_preds = q_preds.gather(
-1,
a.argmax(-1).long().unsqueeze(-1)).squeeze(-1)
online_actions = online_next_q_preds.argmax(dim=-1, keepdim=True)
max_next_q_preds = next_q_preds.gather(-1, online_actions).squeeze(-1)
max_q_targets = r + self.gamma * mask * max_next_q_preds
q_loss = self.loss_fn(act_q_preds, max_q_targets)
return q_loss
def update(self, epoch):
total_loss = 0.
for i in range(self.training_iter):
round_loss = 0.
# 1. batch a sample from memory
batch = self.memory.get_batch(batch_size=self.batch_size)
for _ in range(self.training_batch_iter):
# 2. calculate the Q loss
loss = self.calc_q_loss(batch)
# 3. make a optimization step
self.net_optim.zero_grad()
loss.backward()
self.net_optim.step()
round_loss += loss.item()
logging.debug(
'<<dialog policy dqn>> epoch {}, iteration {}, loss {}'.format(
epoch, i, round_loss / self.training_batch_iter))
total_loss += round_loss
total_loss /= (self.training_batch_iter * self.training_iter)
logging.debug('<<dialog policy dqn>> epoch {}, total_loss {}'.format(
epoch, total_loss))
# update the epsilon value
self.net.update_epsilon(epoch)
# update the target network
self.target_net.load_state_dict(self.net.state_dict())
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch)
def save(self, directory, epoch):
if not os.path.exists(directory):
os.makedirs(directory)
torch.save(self.net.state_dict(),
directory + '/' + str(epoch) + '_dqn.pol.mdl')
logging.info(
'<<dialog policy>> epoch {}: saved network to mdl'.format(epoch))
def load(self, filename):
dqn_mdl_candidates = [
filename + '.dqn.mdl',
os.path.join(os.path.dirname(os.path.abspath(__file__)),
filename + '.dqn.mdl'),
]
for dqn_mdl in dqn_mdl_candidates:
if os.path.exists(dqn_mdl):
self.net.load_state_dict(
torch.load(dqn_mdl, map_location=DEVICE))
self.target_net.load_state_dict(
torch.load(dqn_mdl, map_location=DEVICE))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
dqn_mdl))
break
class PG_generator(Policy):
def __init__(self, is_train=False, dataset='Multiwoz'):
with open("/home/raliegh/图片/ConvLab-2/convlab2/policy/pg/config.json",
'r') as f:
cfg = json.load(f)
self.save_dir = os.path.join(
os.path.dirname(os.path.abspath(__file__)), cfg['save_dir'])
self.save_per_epoch = cfg['save_per_epoch']
self.update_round = cfg['update_round']
self.optim_batchsz = cfg['batchsz']
self.gamma = cfg['gamma']
self.is_train = is_train
if is_train:
init_logging_handler(cfg['log_dir'])
# load vocabulary
if dataset == 'Multiwoz':
voc_file = os.path.join(root_dir, 'data/multiwoz/sys_da_voc.txt')
voc_opp_file = os.path.join(root_dir,
'data/multiwoz/usr_da_voc.txt')
self.vector = MultiWozVector(voc_file, voc_opp_file)
self.policy = MultiDiscretePolicy(
self.vector.state_dim, cfg['h_dim'],
self.vector.da_dim).to(device=DEVICE)
# self.policy = MultiDiscretePolicy(self.vector.state_dim, cfg['h_dim'], self.vector.da_dim).to(device=DEVICE)
if is_train:
self.policy_optim = optim.RMSprop(self.policy.parameters(),
lr=cfg['lr'])
# load_best model from the web.
self.load(
"/home/raliegh/图片/ConvLab-2/convlab2/policy/pg/save/best/best_pg_from_web.pol.mdl"
)
def predict(self, state):
"""
Predict an system action given state.
Args:
state (tensor): Dialog state. Please refer to util/state.py
Returns:
action : System act, with the form of (act_type, {slot_name_1: value_1, slot_name_2, value_2, ...})
"""
with torch.no_grad():
s_vec = torch.Tensor(self.vector.state_vectorize(state))
a = self.policy.select_action(s_vec.to(device=DEVICE),
self.is_train).cpu()
# print(a)
# action = self.vector.action_devectorize(a.detach().numpy())
# state['system_action'] = action
return a
def load(self, filename):
policy_mdl_candidates = [
filename, filename + '.pol.mdl', filename + '_pg.pol.mdl',
os.path.join(os.path.dirname(os.path.abspath(__file__)), filename),
os.path.join(os.path.dirname(os.path.abspath(__file__)),
filename + '.pol.mdl'),
os.path.join(os.path.dirname(os.path.abspath(__file__)),
filename + '_pg.pol.mdl')
]
for policy_mdl in policy_mdl_candidates:
if os.path.exists(policy_mdl):
self.policy.load_state_dict(
torch.load(policy_mdl, map_location=DEVICE))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
policy_mdl))
break
class DQfD(Policy):
def __init__(self, train=True):
# load configuration file
with open(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'config.json'), 'r') as f:
cfg = json.load(f)
self.gamma = cfg['gamma']
self.epsilon_init = cfg['epsilon_init']
self.epsilon_final = cfg['epsilon_final']
self.istrain = train
if self.istrain:
self.epsilon = self.epsilon_init
else:
self.epsilon = self.epsilon_final
self.epsilon_degrade_period = cfg['epsilon_degrade_period']
self.tau = cfg['tau']
self.action_number = cfg[
'action_number'] # total number of actions considered
init_logging_handler(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
cfg['log_dir']))
# load action mapping file
action_map_file = os.path.join(root_dir,
'convlab2/policy/act_500_list.txt')
_, self.ind2act_dict = read_action_map(action_map_file)
# load vector for MultiWoz 2.1
voc_file = os.path.join(root_dir, 'data/multiwoz/sys_da_voc.txt')
voc_opp_file = os.path.join(root_dir, 'data/multiwoz/usr_da_voc.txt')
self.vector = MultiWozVector(voc_file, voc_opp_file)
# build Q network
# current Q network to be trained
self.Q = DuelDQN(self.vector.state_dim, cfg['h_dim'],
self.action_number).to(device=DEVICE)
# target Q network
self.target_Q = DuelDQN(self.vector.state_dim, cfg['h_dim'],
self.action_number).to(device=DEVICE)
self.target_Q.load_state_dict(self.Q.state_dict())
# define optimizer
# self.optimizer = RAdam(self.Q.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
self.optimizer = optim.Adam(self.Q.parameters(),
lr=cfg['lr'],
weight_decay=cfg['weight_decay'])
self.scheduler = StepLR(self.optimizer,
step_size=cfg['lr_decay_step'],
gamma=cfg['lr_decay'])
self.min_lr = cfg['min_lr']
# loss function
self.criterion = torch.nn.MSELoss()
def predict(self, state):
"""Predict an system action and its index given state."""
s_vec = torch.Tensor(self.vector.state_vectorize(state))
if state['user_action'] == [['bye', 'general', 'none', 'none']]:
action = [['bye', 'general', 'none', 'none']]
else:
a, a_ind = self.Q.select_action(s_vec.to(device=DEVICE),
self.epsilon, self.ind2act_dict,
self.istrain)
action = self.vector.action_devectorize(a)
state['system_action'] = action
return action
def predict_ind(self, state):
"""Predict an system action and its index given state."""
s_vec = torch.Tensor(self.vector.state_vectorize(state))
if state['user_action'] == [['bye', 'general', 'none', 'none']]:
action = [['bye', 'general', 'none', 'none']]
a_ind = 489
else:
a, a_ind = self.Q.select_action(s_vec.to(device=DEVICE),
self.epsilon, self.ind2act_dict,
self.istrain)
action = self.vector.action_devectorize(a)
state['system_action'] = action
return action, a_ind
def init_session(self):
"""Restore after one session"""
pass
def aux_loss(self, s, a, candidate_a_ind, expert_label):
"""compute auxiliary loss given batch of states, actions and expert labels"""
# only keep those expert demonstrations by setting expert label to 1
s_exp = s[np.where(expert_label == 1)[0]]
a_exp = a[np.where(expert_label == 1)[0]]
candidate_a_ind_exp = candidate_a_ind[np.where(expert_label == 1)[0]]
# if there exist expert demonstration in current batch
if s_exp.size(0) > 0:
# compute q value predictions for states for each action
q_all = self.Q(s_exp)
# only when agent take the same action as the expert does, the act_diff term(i.e. l(a_e,a)) is 0
act_diff = q_all.new_full(q_all.size(), self.tau)
for row, exp_ind in enumerate(candidate_a_ind_exp):
act_diff[row, exp_ind] = 0
# compute aux_loss = max(Q(s, a) + l(a_e, a)) - Q(s, a_e)
q_max_act = (q_all + act_diff).max(dim=1)[0]
q_exp_act = q_all.gather(-1, a_exp.unsqueeze(1)).squeeze(-1)
aux_loss = (q_max_act - q_exp_act).sum() / s_exp.size(0)
else:
aux_loss = 0
return aux_loss
def update_net(self):
"""update target network by copying parameters from online network"""
self.target_Q.load_state_dict(self.Q.state_dict())
def compute_loss(self, s, a, r, s_next, mask, expert_label,
candidate_a_ind):
"""compute loss for batch"""
# q value predictions for current state for each action
q_preds = self.Q(s)
with torch.no_grad():
# online net for action selection in next state
online_next_q_preds = self.Q(s_next)
# target net for q value predicting in the next state
next_q_preds = self.target_Q(s_next)
# select q value predictions for corresponding actions
act_q_preds = q_preds.gather(-1, a.unsqueeze(1)).squeeze(-1)
# use online net to choose action for the next state
online_actions = online_next_q_preds.argmax(dim=-1, keepdim=True)
# use target net to predict the corresponding value
max_next_q_preds = next_q_preds.gather(-1, online_actions).squeeze(-1)
# compute target q values
max_q_targets = r + self.gamma * max_next_q_preds * mask
# q loss
q_loss = self.criterion(act_q_preds, max_q_targets)
# auxiliary loss
aux_loss_term = self.aux_loss(s, a, candidate_a_ind, expert_label)
# total loss
loss = q_loss + aux_loss_term
return loss
def update(self, loss):
"""update online network"""
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def save(self, directory, epoch):
"""save model to directory"""
if not os.path.exists(directory):
os.makedirs(directory)
torch.save(self.Q.state_dict(),
directory + '/' + str(epoch) + '_dqn.mdl')
logging.info(
'<<dialog policy>> epoch {}: saved network to mdl'.format(epoch))
def load(self, filename):
"""load model"""
dqn_mdl_candidates = [
filename + '_dqn.mdl',
os.path.join(os.path.dirname(os.path.abspath(__file__)),
filename + '_dqn.mdl')
]
for dqn_mdl in dqn_mdl_candidates:
if os.path.exists(dqn_mdl):
self.Q.load_state_dict(torch.load(dqn_mdl,
map_location=DEVICE))
self.target_Q.load_state_dict(
torch.load(dqn_mdl, map_location=DEVICE))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
dqn_mdl))
break
class GDPL(Policy):
def __init__(self, is_train=False, dataset='Multiwoz'):
with open(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'config.json'), 'r') as f:
cfg = json.load(f)
self.save_dir = os.path.join(
os.path.dirname(os.path.abspath(__file__)), cfg['save_dir'])
self.save_per_epoch = cfg['save_per_epoch']
self.update_round = cfg['update_round']
self.optim_batchsz = cfg['batchsz']
self.gamma = cfg['gamma']
self.epsilon = cfg['epsilon']
self.tau = cfg['tau']
self.is_train = is_train
if is_train:
init_logging_handler(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
cfg['log_dir']))
# construct policy and value network
if dataset == 'Multiwoz':
voc_file = os.path.join(root_dir, 'data/multiwoz/sys_da_voc.txt')
voc_opp_file = os.path.join(root_dir,
'data/multiwoz/usr_da_voc.txt')
self.vector = MultiWozVector(voc_file, voc_opp_file)
self.policy = MultiDiscretePolicy(
self.vector.state_dim, cfg['h_dim'],
self.vector.da_dim).to(device=DEVICE)
self.value = Value(self.vector.state_dim,
cfg['hv_dim']).to(device=DEVICE)
if is_train:
self.policy_optim = optim.RMSprop(self.policy.parameters(),
lr=cfg['lr'])
self.value_optim = optim.Adam(self.value.parameters(),
lr=cfg['lr'])
def predict(self, state):
"""
Predict an system action given state.
Args:
state (dict): Dialog state. Please refer to util/state.py
Returns:
action : System act, with the form of (act_type, {slot_name_1: value_1, slot_name_2, value_2, ...})
"""
s_vec = torch.Tensor(self.vector.state_vectorize(state))
a = self.policy.select_action(s_vec.to(device=DEVICE),
self.is_train).cpu()
return self.vector.action_devectorize(a.numpy())
def init_session(self):
"""
Restore after one session
"""
pass
def est_adv(self, r, v, mask):
"""
we save a trajectory in continuous space and it reaches the ending of current trajectory when mask=0.
:param r: reward, Tensor, [b]
:param v: estimated value, Tensor, [b]
:param mask: indicates ending for 0 otherwise 1, Tensor, [b]
:return: A(s, a), V-target(s), both Tensor
"""
batchsz = v.size(0)
# v_target is worked out by Bellman equation.
v_target = torch.Tensor(batchsz).to(device=DEVICE)
delta = torch.Tensor(batchsz).to(device=DEVICE)
A_sa = torch.Tensor(batchsz).to(device=DEVICE)
prev_v_target = 0
prev_v = 0
prev_A_sa = 0
for t in reversed(range(batchsz)):
# mask here indicates a end of trajectory
# this value will be treated as the target value of value network.
# mask = 0 means the immediate reward is the real V(s) since it's end of trajectory.
# formula: V(s_t) = r_t + gamma * V(s_t+1)
v_target[t] = r[t] + self.gamma * prev_v_target * mask[t]
# please refer to : https://arxiv.org/abs/1506.02438
# for generalized adavantage estimation
# formula: delta(s_t) = r_t + gamma * V(s_t+1) - V(s_t)
delta[t] = r[t] + self.gamma * prev_v * mask[t] - v[t]
# formula: A(s, a) = delta(s_t) + gamma * lamda * A(s_t+1, a_t+1)
# here use symbol tau as lambda, but original paper uses symbol lambda.
A_sa[t] = delta[t] + self.gamma * self.tau * prev_A_sa * mask[t]
# update previous
prev_v_target = v_target[t]
prev_v = v[t]
prev_A_sa = A_sa[t]
# normalize A_sa
A_sa = (A_sa - A_sa.mean()) / A_sa.std()
return A_sa, v_target
def update(self, epoch, batchsz, s, a, next_s, mask, rewarder):
# update reward estimator
rewarder.update_irl((s, a, next_s), batchsz, epoch)
# get estimated V(s) and PI_old(s, a)
# actually, PI_old(s, a) can be saved when interacting with env, so as to save the time of one forward elapsed
# v: [b, 1] => [b]
v = self.value(s).squeeze(-1).detach()
log_pi_old_sa = self.policy.get_log_prob(s, a).detach()
# estimate advantage and v_target according to GAE and Bellman Equation
r = rewarder.estimate(s, a, next_s, log_pi_old_sa).detach()
A_sa, v_target = self.est_adv(r, v, mask)
for i in range(self.update_round):
# 1. shuffle current batch
perm = torch.randperm(batchsz)
# shuffle the variable for mutliple optimize
v_target_shuf, A_sa_shuf, s_shuf, a_shuf, log_pi_old_sa_shuf = v_target[perm], A_sa[perm], s[perm], a[perm], \
log_pi_old_sa[perm]
# 2. get mini-batch for optimizing
optim_chunk_num = int(np.ceil(batchsz / self.optim_batchsz))
# chunk the optim_batch for total batch
v_target_shuf, A_sa_shuf, s_shuf, a_shuf, log_pi_old_sa_shuf = torch.chunk(v_target_shuf, optim_chunk_num), \
torch.chunk(A_sa_shuf, optim_chunk_num), \
torch.chunk(s_shuf, optim_chunk_num), \
torch.chunk(a_shuf, optim_chunk_num), \
torch.chunk(log_pi_old_sa_shuf,
optim_chunk_num)
# 3. iterate all mini-batch to optimize
policy_loss, value_loss = 0., 0.
for v_target_b, A_sa_b, s_b, a_b, log_pi_old_sa_b in zip(
v_target_shuf, A_sa_shuf, s_shuf, a_shuf,
log_pi_old_sa_shuf):
# print('optim:', batchsz, v_target_b.size(), A_sa_b.size(), s_b.size(), a_b.size(), log_pi_old_sa_b.size())
# 1. update value network
self.value_optim.zero_grad()
v_b = self.value(s_b).squeeze(-1)
loss = (v_b - v_target_b).pow(2).mean()
value_loss += loss.item()
# backprop
loss.backward()
# nn.utils.clip_grad_norm(self.value.parameters(), 4)
self.value_optim.step()
# 2. update policy network by clipping
self.policy_optim.zero_grad()
# [b, 1]
log_pi_sa = self.policy.get_log_prob(s_b, a_b)
# ratio = exp(log_Pi(a|s) - log_Pi_old(a|s)) = Pi(a|s) / Pi_old(a|s)
# we use log_pi for stability of numerical operation
# [b, 1] => [b]
ratio = (log_pi_sa - log_pi_old_sa_b).exp().squeeze(-1)
surrogate1 = ratio * A_sa_b
surrogate2 = torch.clamp(ratio, 1 - self.epsilon,
1 + self.epsilon) * A_sa_b
# this is element-wise comparing.
# we add negative symbol to convert gradient ascent to gradient descent
surrogate = -torch.min(surrogate1, surrogate2).mean()
policy_loss += surrogate.item()
# backprop
surrogate.backward()
# gradient clipping, for stability
torch.nn.utils.clip_grad_norm(self.policy.parameters(), 10)
# self.lock.acquire() # retain lock to update weights
self.policy_optim.step()
# self.lock.release() # release lock
value_loss /= optim_chunk_num
policy_loss /= optim_chunk_num
logging.debug(
'<<dialog policy ppo>> epoch {}, iteration {}, value, loss {}'.
format(epoch, i, value_loss))
logging.debug(
'<<dialog policy ppo>> epoch {}, iteration {}, policy, loss {}'
.format(epoch, i, policy_loss))
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch)
def save(self, directory, epoch):
if not os.path.exists(directory):
os.makedirs(directory)
torch.save(self.value.state_dict(),
directory + '/' + str(epoch) + '_ppo.val.mdl')
torch.save(self.policy.state_dict(),
directory + '/' + str(epoch) + '_ppo.pol.mdl')
logging.info(
'<<dialog policy>> epoch {}: saved network to mdl'.format(epoch))
def load(self, filename):
value_mdl = os.path.join(os.path.dirname(os.path.abspath(__file__)),
filename + '_ppo.val.mdl')
if os.path.exists(value_mdl):
self.value.load_state_dict(torch.load(value_mdl))
print('<<dialog policy>> loaded checkpoint from file: {}'.format(
value_mdl))
policy_mdl = os.path.join(os.path.dirname(os.path.abspath(__file__)),
filename + '_ppo.pol.mdl')
if os.path.exists(policy_mdl):
self.policy.load_state_dict(torch.load(policy_mdl))
print('<<dialog policy>> loaded checkpoint from file: {}'.format(
policy_mdl))
class PG(Policy):
def __init__(self, is_train=False, dataset='Multiwoz'):
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config.json'), 'r') as f:
cfg = json.load(f)
self.save_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), cfg['save_dir'])
self.save_per_epoch = cfg['save_per_epoch']
self.update_round = cfg['update_round']
self.optim_batchsz = cfg['batchsz']
self.gamma = cfg['gamma']
self.is_train = is_train
if is_train:
init_logging_handler(cfg['log_dir'])
if dataset == 'Multiwoz':
voc_file = os.path.join(root_dir, 'data/multiwoz/sys_da_voc.txt')
voc_opp_file = os.path.join(root_dir, 'data/multiwoz/usr_da_voc.txt')
self.vector = MultiWozVector(voc_file, voc_opp_file)
self.policy = MultiDiscretePolicy(self.vector.state_dim, cfg['h_dim'], self.vector.da_dim).to(device=DEVICE)
# self.policy = MultiDiscretePolicy(self.vector.state_dim, cfg['h_dim'], self.vector.da_dim).to(device=DEVICE)
if is_train:
self.policy_optim = optim.RMSprop(self.policy.parameters(), lr=cfg['lr'])
def predict(self, state):
"""
Predict an system action given state.
Args:
state (dict): Dialog state. Please refer to util/state.py
Returns:
action : System act, with the form of (act_type, {slot_name_1: value_1, slot_name_2, value_2, ...})
"""
s_vec = torch.Tensor(self.vector.state_vectorize(state))
a = self.policy.select_action(s_vec.to(device=DEVICE), self.is_train).cpu()
action = self.vector.action_devectorize(a.numpy())
state['system_action'] = action
return action
def init_session(self):
"""
Restore after one session
"""
pass
def est_return(self, r, mask):
"""
we save a trajectory in continuous space and it reaches the ending of current trajectory when mask=0.
:param r: reward, Tensor, [b]
:param mask: indicates ending for 0 otherwise 1, Tensor, [b]
:return: V-target(s), Tensor
"""
batchsz = r.size(0)
# v_target is worked out by Bellman equation.
v_target = torch.Tensor(batchsz).to(device=DEVICE)
prev_v_target = 0
for t in reversed(range(batchsz)):
# mask here indicates a end of trajectory
# this value will be treated as the target value of value network.
# mask = 0 means the immediate reward is the real V(s) since it's end of trajectory.
# formula: V(s_t) = r_t + gamma * V(s_t+1)
v_target[t] = r[t] + self.gamma * prev_v_target * mask[t]
# update previous
prev_v_target = v_target[t]
return v_target
def update(self, epoch, batchsz, s, a, r, mask):
v_target = self.est_return(r, mask)
for i in range(self.update_round):
# 1. shuffle current batch
perm = torch.randperm(batchsz)
# shuffle the variable for mutliple optimize
v_target_shuf, s_shuf, a_shuf = v_target[perm], s[perm], a[perm]
# 2. get mini-batch for optimizing
optim_chunk_num = int(np.ceil(batchsz / self.optim_batchsz))
# chunk the optim_batch for total batch
v_target_shuf, s_shuf, a_shuf = torch.chunk(v_target_shuf, optim_chunk_num), \
torch.chunk(s_shuf, optim_chunk_num), \
torch.chunk(a_shuf, optim_chunk_num)
# 3. iterate all mini-batch to optimize
policy_loss = 0.
for v_target_b, s_b, a_b in zip(v_target_shuf, s_shuf, a_shuf):
# print('optim:', batchsz, v_target_b.size(), A_sa_b.size(), s_b.size(), a_b.size(), log_pi_old_sa_b.size())
# update policy network by clipping
self.policy_optim.zero_grad()
# [b, 1]
log_pi_sa = self.policy.get_log_prob(s_b, a_b)
# ratio = exp(log_Pi(a|s) - log_Pi_old(a|s)) = Pi(a|s) / Pi_old(a|s)
# we use log_pi for stability of numerical operation
# [b, 1] => [b]
# this is element-wise comparing.
# we add negative symbol to convert gradient ascent to gradient descent
surrogate = - (log_pi_sa * v_target_b).mean()
policy_loss += surrogate.item()
# backprop
surrogate.backward()
for p in self.policy.parameters():
p.grad[p.grad != p.grad] = 0.0
# gradient clipping, for stability
torch.nn.utils.clip_grad_norm(self.policy.parameters(), 10)
# self.lock.acquire() # retain lock to update weights
self.policy_optim.step()
# self.lock.release() # release lock
policy_loss /= optim_chunk_num
logging.debug('<<dialog policy pg>> epoch {}, iteration {}, policy, loss {}'.format(epoch, i, policy_loss))
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch)
def save(self, directory, epoch):
if not os.path.exists(directory):
os.makedirs(directory)
torch.save(self.policy.state_dict(), directory + '/' + str(epoch) + '_pg.pol.mdl')
logging.info('<<dialog policy>> epoch {}: saved network to mdl'.format(epoch))
def load(self, filename):
policy_mdl_candidates = [
filename,
filename + '.pol.mdl',
filename + '_pg.pol.mdl',
os.path.join(os.path.dirname(os.path.abspath(__file__)), filename),
os.path.join(os.path.dirname(os.path.abspath(__file__)), filename + '.pol.mdl'),
os.path.join(os.path.dirname(os.path.abspath(__file__)), filename + '_pg.pol.mdl')
]
for policy_mdl in policy_mdl_candidates:
if os.path.exists(policy_mdl):
self.policy.load_state_dict(torch.load(policy_mdl, map_location=DEVICE))
logging.info('<<dialog policy>> loaded checkpoint from file: {}'.format(policy_mdl))
break
def load_from_pretrained(self, archive_file, model_file, filename):
if not os.path.isfile(archive_file):
if not model_file:
raise Exception("No model for PG Policy is specified!")
archive_file = cached_path(model_file)
model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'save')
if not os.path.exists(model_dir):
os.mkdir(model_dir)
if not os.path.exists(os.path.join(model_dir, 'best_pg.pol.mdl')):
archive = zipfile.ZipFile(archive_file, 'r')
archive.extractall(model_dir)
policy_mdl = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename + '_pg.pol.mdl')
if os.path.exists(policy_mdl):
self.policy.load_state_dict(torch.load(policy_mdl, map_location=DEVICE))
logging.info('<<dialog policy>> loaded checkpoint from file: {}'.format(policy_mdl))
@classmethod
def from_pretrained(cls,
archive_file="",
model_file="https://convlab.blob.core.windows.net/convlab-2/pg_policy_multiwoz.zip"):
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config.json'), 'r') as f:
cfg = json.load(f)
model = cls()
model.load_from_pretrained(archive_file, model_file, cfg['load'])
return model