def gpcf(self, z_goal, h):
#Compares possible next state as a result of each action to goal.
#Chooses the action which will reduce the loss the most
output = []
for action in range(len(h)):
if self.hiddengoals:
if self.Forward_model == 'D':
output.append(self.LSEloss(h[action], z_goal))
else:
#compares the goal with the next state using the NLL loss of the latent representation of the next hidden state
#h[action][4] = next_hidden_mu, h[action][5] = next_hidden_sigma, logpi = -1
output.append(
gmm_loss(
z_goal, h[action][4], h[action][5],
torch.tensor([-1.0], dtype=torch.float32).to(
self.device)) / 33)
else:
if self.Forward_model == 'D':
output.append(criterion(h[action], z_goal).item())
else:
#compares the goal with the next state using the NLL loss of the latent representation of the next hidden state
#h[action][0] = mu, h[action][1] = sigma, logpi = logpi
output.append(
gmm_loss(z_goal, h[action][0], h[action][1],
h[action][2]) / 33)
return output.index(min(output))
def get_loss(self, latent_obs, action, reward, terminal, latent_next_obs,
include_reward: bool):
""" Compute losses.
The loss that is computed is:
(GMMLoss(latent_next_obs, GMMPredicted) + MSE(reward, predicted_reward) +
BCE(terminal, logit_terminal)) / (LSIZE + 2)
The LSIZE + 2 factor is here to counteract the fact that the GMMLoss scales
approximately linearily with LSIZE. All losses are averaged both on the
batch and the sequence dimensions (the two first dimensions).
:args latent_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:args action: (BSIZE, SEQ_LEN, ASIZE) torch tensor
:args reward: (BSIZE, SEQ_LEN) torch tensor
:args latent_next_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:returns: dictionary of losses, containing the gmm, the mse, the bce and
the averaged loss.
"""
mus, sigmas, logpi, rs, ds = self.mdrnnBIG(action, latent_obs)
gmm = gmm_loss(latent_next_obs, mus, sigmas, logpi)
bce = F.binary_cross_entropy_with_logits(ds, terminal)
if include_reward:
mse = F.mse_loss(rs, reward)
scale = LSIZE + 2
else:
mse = 0
scale = LSIZE + 1
loss = (gmm + bce + mse) / scale
return dict(gmm=gmm, bce=bce, mse=mse, loss=loss)
def get_loss(latent_obs, action, reward, terminal, latent_next_obs):
""" Compute losses.
The loss that is computed is:
(GMMLoss(latent_next_obs, GMMPredicted) + MSE(reward, predicted_reward) +
BCE(terminal, logit_terminal)) / (LSIZE + 2)
The LSIZE + 2 factor is here to counteract the fact that the GMMLoss scales
approximately linearily with LSIZE. All losses are averaged both on the
batch and the sequence dimensions (the two first dimensions).
:args latent_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:args action: (BSIZE, SEQ_LEN, ASIZE) torch tensor
:args reward: (BSIZE, SEQ_LEN) torch tensor
:args latent_next_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:returns: dictionary of losses, containing the gmm, the mse, the bce and
the averaged loss.
"""
latent_obs, action,\
reward, terminal,\
latent_next_obs = [arr.transpose(1, 0)
for arr in [latent_obs, action,
reward, terminal,
latent_next_obs]]
mus, sigmas, logpi, rs, ds = mdrnn(action, latent_obs)
gmm = gmm_loss(latent_next_obs, mus, sigmas, logpi)
bce = f.binary_cross_entropy_with_logits(ds, terminal)
mse = f.mse_loss(rs, reward)
loss = (gmm + bce + mse) / (LSIZE + 2)
return dict(gmm=gmm, bce=bce, mse=mse, loss=loss)
def get_loss(latent_obs, action, reward, terminal, latent_next_obs):
""" Compute losses.
The loss that is computed is:
(GMMLoss(latent_next_obs, GMMPredicted) + MSE(reward, predicted_reward) +
BCE(terminal, logit_terminal)) / (LSIZE + 2)
The LSIZE + 2 factor is here to counteract the fact that the GMMLoss scales
approximately linearily with LSIZE. All losses are averaged both on the
batch and the sequence dimensions (the two first dimensions).
:args latent_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:args action: (BSIZE, SEQ_LEN, ASIZE) torch tensor
:args reward: (BSIZE, SEQ_LEN) torch tensor
:args latent_next_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:returns: dictionary of losses, containing the gmm, the mse, the bce and
the averaged loss.
"""
# transpose such that seq_len is the first dimension
latent_obs, action,\
reward, terminal,\
latent_next_obs = [arr.transpose(1, 0)
for arr in [latent_obs, action,
reward, terminal,
latent_next_obs]]
mus, sigmas, logpi, rs, ds = mdrnn(action, latent_obs)
gmm = gmm_loss(latent_next_obs, mus, sigmas, logpi)
bce = f.binary_cross_entropy_with_logits(ds, terminal)
mse = f.mse_loss(rs, reward)
loss = (gmm + bce + mse) / (LSIZE + 2)
return dict(gmm=gmm, bce=bce, mse=mse, loss=loss)
def mdrnn_exp_reward(self, latent_obs, action, reward, latent_next_obs,
hidden):
""" # REMOVE TERMINAL
Compute losses.
The loss that is computed is:
(GMMLoss(latent_next_obs, GMMPredicted) + MSE(reward, predicted_reward) +
BCE(terminal, logit_terminal)) / (LSIZE + 2)
The LSIZE + 2 factor is here to counteract the fact that the GMMLoss scales
approximately linearily with LSIZE. All losses are averaged both on the
batch and the sequence dimensions (the two first dimensions).
:args latent_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:args action: (BSIZE, SEQ_LEN, ASIZE) torch tensor
:args reward: (BSIZE, SEQ_LEN) torch tensor
:args latent_next_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:returns: dictionary of losses, containing the gmm, the mse, the bce and
the averaged loss.
"""
mus, sigmas, logpi, rs, ds, next_hidden = self.mdrnn(
action, latent_obs, hidden)
gmm = gmm_loss(latent_next_obs, mus, sigmas, logpi)
# bce = f.binary_cross_entropy_with_logits(ds, terminal)
mse = f.mse_loss(rs, reward)
loss = (gmm + mse) / (LSIZE + 2)
return loss.squeeze().cpu().numpy()
def test_gmm_loss_my(self):
# seq_len x batch_size x gaussian_size x feature_size
# 1 x 1 x 2 x 2
mus = torch.Tensor([[
[[0., 0.], [6., 6.]],
]])
sigmas = torch.Tensor([[
[[2., 2.], [2., 2.]],
]])
# seq_len x batch_size x gaussian_size
pi = torch.Tensor([[[.5, .5]]])
logpi = torch.log(pi)
# seq_len x batch_size x feature_size
batch = torch.Tensor([[[3., 3.]]])
gl = gmm_loss(batch, mus, sigmas, logpi)
# first component, first dimension
n11 = Normal(mus[0, 0, 0, 0], sigmas[0, 0, 0, 0])
# first component, second dimension
n12 = Normal(mus[0, 0, 0, 1], sigmas[0, 0, 0, 1])
p1 = pi[0, 0, 0] * torch.exp(n11.log_prob(batch[0, 0, 0])) * torch.exp(
n12.log_prob(batch[0, 0, 1]))
# second component, first dimension
n21 = Normal(mus[0, 0, 1, 0], sigmas[0, 0, 1, 0])
# second component, second dimension
n22 = Normal(mus[0, 0, 1, 1], sigmas[0, 0, 1, 1])
p2 = pi[0, 0, 1] * torch.exp(n21.log_prob(batch[0, 0, 0])) * torch.exp(
n22.log_prob(batch[0, 0, 1]))
print("gmm loss={}, p1={}, p2={}, p1+p2={}, -log(p1+p2)={}".format(
gl, p1, p2, p1 + p2, -torch.log(p1 + p2)))
assert -torch.log(p1 + p2) == gl
print()
def test_gmm_loss(self):
""" Test case 1 """
n_samples = 10000
means = torch.Tensor([[0., 0.],
[1., 1.],
[-1., 1.]])
stds = torch.Tensor([[.03, .05],
[.02, .1],
[.1, .03]])
pi = torch.Tensor([.2, .3, .5])
cat_dist = Categorical(pi)
indices = cat_dist.sample((n_samples,)).long()
rands = torch.randn(n_samples, 2)
samples = means[indices] + rands * stds[indices]
class _model(nn.Module):
def __init__(self, gaussians):
super().__init__()
self.means = nn.Parameter(torch.Tensor(1, gaussians, 2).normal_())
self.pre_stds = nn.Parameter(torch.Tensor(1, gaussians, 2).normal_())
self.pi = nn.Parameter(torch.Tensor(1, gaussians).normal_())
def forward(self, *inputs):
return self.means, torch.exp(self.pre_stds), f.softmax(self.pi, dim=1)
model = _model(3)
optimizer = torch.optim.Adam(model.parameters())
iterations = 100000
log_step = iterations // 10
pbar = tqdm(total=iterations)
cum_loss = 0
for i in range(iterations):
batch = samples[torch.LongTensor(128).random_(0, n_samples)]
m, s, p = model.forward()
loss = gmm_loss(batch, m, s, p)
cum_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.set_postfix_str("avg_loss={:10.6f}".format(
cum_loss / (i + 1)))
pbar.update(1)
if i % log_step == log_step - 1:
print(m)
print(s)
print(p)
def test_gmm_loss(self):
""" Test case 1 """
n_samples = 10000
means = torch.Tensor([[0., 0.], [1., 1.], [-1., 1.]])
stds = torch.Tensor([[.03, .05], [.02, .1], [.1, .03]])
pi = torch.Tensor([.2, .3, .5])
cat_dist = Categorical(pi)
indices = cat_dist.sample((n_samples, )).long()
rands = torch.randn(n_samples, 2)
samples = means[indices] + rands * stds[indices]
class _model(nn.Module):
def __init__(self, gaussians):
super().__init__()
self.means = nn.Parameter(
torch.Tensor(1, gaussians, 2).normal_())
self.pre_stds = nn.Parameter(
torch.Tensor(1, gaussians, 2).normal_())
self.pi = nn.Parameter(torch.Tensor(1, gaussians).normal_())
def forward(self, *inputs):
return self.means, torch.exp(self.pre_stds), f.softmax(self.pi,
dim=1)
model = _model(3)
optimizer = torch.optim.Adam(model.parameters())
iterations = 100000
log_step = iterations // 10
pbar = tqdm(total=iterations)
cum_loss = 0
for i in range(iterations):
batch = samples[torch.LongTensor(128).random_(0, n_samples)]
m, s, p = model.forward()
loss = gmm_loss(batch, m, s, p)
cum_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
pbar.set_postfix_str("avg_loss={:10.6f}".format(cum_loss /
(i + 1)))
pbar.update(1)
if i % log_step == log_step - 1:
print(m)
print(s)
print(p)
def mdn_rnn_loss_function(latent_obs, action, reward, terminal,
latent_next_obs, mdn_rnn_prediction):
latent_obs, action,\
reward, terminal,\
latent_next_obs = [arr.transpose(1, 0)
for arr in [latent_obs, action,
reward, terminal,
latent_next_obs]]
mus, sigmas, logpi, rs, ds = mdn_rnn_prediction
gmm = gmm_loss(latent_next_obs, mus, sigmas, logpi)
bce = F.binary_cross_entropy_with_logits(ds, terminal)
mse = F.binary_cross_entropy_with_logits(rs, reward)
scale = LSIZE + 2
loss = (gmm + bce + mse) / scale
return dict(gmm=gmm, bce=bce, mse=mse, loss=loss)
def mdn_rnn_loss_function(latent_obs, action, reward, terminal,
latent_next_obs, mdn_rnn_prediction):
""" Compute losses.
The loss that is computed is:
(GMMLoss(latent_next_obs, GMMPredicted) + MSE(reward, predicted_reward) +
BCE(terminal, logit_terminal)) / (LSIZE + 2)
The LSIZE + 2 factor is here to counteract the fact that the GMMLoss scales
approximately linearily with LSIZE. All losses are averaged both on the
batch and the sequence dimensions (the two first dimensions).
:args latent_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:args action: (BSIZE, SEQ_LEN, ASIZE) torch tensor
:args reward: (BSIZE, SEQ_LEN) torch tensor
:args latent_next_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:returns: dictionary of losses, containing the gmm, the mse, the bce and
the averaged loss.
"""
(latent_obs, lengths), (action, _),\
(reward, _), (terminal, _),\
(latent_next_obs, _) = [
pad_packed_sequence(a)
for a in [latent_obs, action,
reward, terminal,
latent_next_obs]]
(mus, _), (sigmas, _), (logpi, _), (rs, _), (ds, _) = [
pad_packed_sequence(o) for o in mdn_rnn_prediction
]
gmm_losses = []
bce_losses = []
for b, length in enumerate(lengths):
gmm_losses.append(
gmm_loss(latent_next_obs[:length, b], mus[:length, b],
sigmas[:length, b], logpi[:length, b]))
bce_losses.append(
F.binary_cross_entropy_with_logits(ds[:length, b],
terminal[:length, b]))
gmm = torch.mean(torch.stack(gmm_losses))
bce = torch.mean(torch.stack(bce_losses))
loss = (gmm + bce) / (LSIZE + 1)
return dict(gmm=gmm, bce=bce, loss=loss)
def get_loss(input, output, train):
""" Compute losses.
The loss that is computed is:
(GMMLoss(latent_next_obs, GMMPredicted) + MSE(reward, predicted_reward) +
BCE(terminal, logit_terminal)) / (LSIZE + 2)
The LSIZE + 2 factor is here to counteract the fact that the GMMLoss scales
approximately linearily with LSIZE. All losses are averaged both on the
batch and the sequence dimensions (the two first dimensions).
:args latent_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:args action: (BSIZE, SEQ_LEN, ASIZE) torch tensor
:args reward: (BSIZE, SEQ_LEN) torch tensor
:args latent_next_obs: (BSIZE, SEQ_LEN, LSIZE) torch tensor
:returns: dictionary of losses, containing the gmm, the mse, the bce and
the averaged loss.
"""
mu, sigma, pi = mdrnn(input, train)
gmm = gmm_loss(mu, sigma, pi, output)
return gmm
def rollout(self, params, render=False):
""" Executes rollouts for number of goals
"""
# copy params into the controller
if params is not None:
load_parameters(params, self.fmodel)
optimizer = optim.Adam(params=self.fmodel.parameters(), lr=0.0001)
MDRNNoptimizer = torch.optim.RMSprop(self.mdrnnBIG.parameters(),
lr=1e-3,
alpha=.9)
#MDRNNoptimizer.load_state_dict(self.rnn_state["optimizer"])
VAEOptimizer = optim.Adam(self.vae.parameters())
VAEOptimizer.load_state_dict(self.vae_state["optimizer"])
HiddenVAEOptimizer = optim.Adam(self.HiddenVAE.parameters())
HiddenVAEOptimizer.load_state_dict(self.hiddenvae_state["optimizer"])
zstate_list = []
self.env.seed(1337)
#self.env.reset()
obs = self.env.reset()
obs = obs['image']
expl_rate = 0.4
hidden = [torch.zeros(1, RSIZE).to(self.device) for _ in range(2)]
_, latent_mu, logsigma, z = self.tolatent(obs)
i = 0
#Bootstrapping, collect 100 initial states for goal space
while True:
action = random.randrange(6)
_, hidden, z, zh = self.transform_obs_hidden(obs, hidden, action)
_, _, hidden_latent = self.tohiddenlatent(hidden)
obs, exreward, done, _ = self.env.step(action)
obs = obs['image']
if self.hiddengoals:
zstate_list.append(
np.array(hidden_latent.cpu().detach().numpy())
) #if we use pure hidden
else:
zstate_list.append(np.array(
z.cpu().detach().numpy())) #if we use latent_space
i += 1
if render:
self.env.render()
if i > self.time_limit:
break
s = obs
loss_list = []
WM_loss = []
VAE_loss_per_rollout = []
hiddenvae_loss_per_rollout = []
rollout_reward = []
visitationcount = []
exreward_per_rollout = []
visitationarray = np.zeros((25, 25))
final_loss = []
#Goal Exploration
for c in range(self.number_goals):
#reset env, uncomment below if necessary to reset agent in enviroinment after every episode
'''
self.env.seed(1337)
self.env.reset()
s = self.env.reset()
#reset obs and hidden state
s = s['image']
_,_,_,z = self.tolatent(s)
hidden = [
torch.zeros(1, RSIZE).to(self.device) for _ in range(2)]
'''
print('Goal Number', c)
zstate_list = np.array(zstate_list)
zstate_list = zstate_list.squeeze(1)
kde = scipy.stats.gaussian_kde(zstate_list.T)
z_goal = sampling_method(
kde) #sample goal from goal space using KDE
z_goal = torch.tensor([z_goal], dtype=torch.float32).to(
self.device) #controller requires both as tensors
if not self.hiddengoals:
z_goal_obs = self.vae.decoder(z_goal)
z_goal_obs = z_goal_obs.reshape(7, 7, 3)
z_goal_obs = np.array(z_goal_obs.cpu().detach())
plt9 = plt.figure('Zgoal')
plt.cla()
sn.heatmap(z_goal_obs[:, :, 0],
cmap='Reds',
annot=True,
cbar=False).invert_yaxis()
total_hiddenvae_loss = 0
total_vae_loss = 0
total_reward = 0
total_exreward = 0
total_loss = 0
goal_loss = []
scur_rollout = []
snext_rollout = []
r_rollout = []
d_rollout = []
act_rollout = []
zstate_list = zstate_list[:, np.newaxis, :]
zstate_list = zstate_list.tolist()
for goalattempts in range(100):
if visitationarray[self.env.agent_pos[0],
self.env.agent_pos[1]] == 0:
visitationarray[self.env.agent_pos[0],
self.env.agent_pos[1]] += 1
h = []
for a in range(6):
if self.Forward_model == 'D':
h.append(
self.fmodel(
z.detach(), hidden[0].detach(),
torch.tensor([[a]], dtype=torch.float32).to(
self.device)))
else:
#Perform a prediction of next state for every action. Add to a list spo comparison with goal can occur
z, hmus, hsigmas, hlogpi, zh, next_hidden, next_hidden_latent, next_hidden_mu, next_hidden_sigma = self.predict_next(
s, hidden, a)
h.append([
hmus, hsigmas, hlogpi, next_hidden_latent,
next_hidden_mu, next_hidden_sigma
])
if expl_rate > random.random():
m = random.randrange(6)
else:
#choose action which will bring us closer to goal
m = self.gpcf(z_goal, h)
z, hmus, hsigmas, hlogpi, zh, hidden, hidden_latent, hidden_mu, hidden_sigma = self.predict_next(
s, hidden, m
) #gets mean, standard deviation and pi, next latent of prediction of next latent obs
if not self.hiddengoals:
if self.Forward_model == 'D':
predicted_next_obs = self.vae.decoder(h[m])
predicted_next_obs = predicted_next_obs.reshape(
7, 7, 3)
p = np.array(predicted_next_obs.cpu().detach())
else:
predicted_next_obs = self.vae.decoder(zh)
predicted_next_obs = predicted_next_obs.reshape(
7, 7, 3)
p = np.array(predicted_next_obs.cpu().detach())
else:
predicted_next_obs = self.vae.decoder(zh)
predicted_next_obs = predicted_next_obs.reshape(7, 7, 3)
p = np.array(predicted_next_obs.cpu().detach())
#Show predicted next observation
if render:
plt5 = plt.figure('Predicted obs')
plt.cla()
sn.heatmap(p[:, :, 0], cmap='Reds', annot=True,
cbar=False).invert_yaxis()
s, exreward, _, _ = self.env.step(
m
) #perform action , get next observation and external reward if any
total_exreward += exreward
s = s['image']
recons, next_mu, next_logsigma, next_z = self.tolatent(
s) #transform observation to latent representation
if self.hiddengoals:
reconhidden, hiddenmu, hiddenlogsigma = self.HiddenVAE(
hidden[0].detach()
) #transoform hidden state into latent representation if using goals in world model
#Show actual observation
if render:
plt6 = plt.figure('Actual obs')
plt.cla()
sn.heatmap(s[:, :, 0], cmap='Reds', annot=True,
cbar=False).invert_yaxis()
#Collect information for training World Model
scur_rollout.append(np.array(z.cpu().detach()))
snext_rollout.append(np.array(next_z.cpu().detach()))
r_rollout.append([0.0])
act_rollout.append([[np.float(m)]])
d_rollout.append([0.0])
if render:
self.env.render()
if self.hiddengoals:
hiddenvae_loss = self.VAEloss(reconhidden,
hidden[0].detach(), hiddenmu,
hiddenlogsigma)
total_hiddenvae_loss += hiddenvae_loss
VAE_loss = self.VAEloss(
recons,
torch.tensor(s.flatten(),
dtype=torch.float32).unsqueeze(0).to(
self.device), next_mu, next_logsigma)
total_vae_loss += VAE_loss
#Curiosity reward is how far the next state was from the prediction
Curiosityreward = gmm_loss(next_z.detach(), hmus, hsigmas,
hlogpi) / 33
#Uncomment below if requiring to add only completely new hidden states to goal space
'''
if Curiosityreward > 1.29: #only add this to the goal space if it was new: this promotes sampling goals which we are unsure about
if self.hiddengoals:
zstate_list.append(np.array(hidden_latent.cpu().detach().numpy()))#if we use pure hidden
else:
zstate_list.append(np.array(z.cpu().detach().numpy()))#if we use latent_space
'''
#add all states to goal space
if self.hiddengoals:
zstate_list.append(
np.array(hidden_latent.cpu().detach().numpy())
) #if we use pure hidden
else:
zstate_list.append(np.array(
z.cpu().detach().numpy())) #if we use latent_space
#if forward model is a linear layer then there are vastly different loss functions. This performs badly so is not recommended to use
if self.Forward_model == 'D':
if self.hiddengoals:
goal_loss.append(
self.LSEloss(hidden_latent, z_goal)
) #how far away the achieved step is from the goal
floss = self.LSEloss(
h[m], hidden_latent.detach()
) #difference between forward model prediction and next hidden
else:
goal_loss.append(
criterion(next_z.detach(), z_goal).item()
) #how far away the achieved step is from the goal
floss = criterion(
h[m], next_z.detach()
) #difference between forward model prediction and next latent
else:
if self.hiddengoals:
goal_loss.append(
gmm_loss(
z_goal, hidden_mu, hidden_sigma,
torch.tensor([-1.0], dtype=torch.float32).to(
self.device)) / 33
) #how far away the achieved step is from the goal
floss = Curiosityreward #difference between forward model prediction and next hidden
else:
goal_loss.append(
gmm_loss(
z_goal, next_mu, next_logsigma.exp(),
torch.tensor([-1.0], dtype=torch.float32).to(
self.device)) / 33)
floss = Curiosityreward
total_loss += floss
#train forward model D if necessary
if self.Forward_model == 'D':
optimizer.zero_grad()
floss.backward()
optimizer.step()
#To see what goals look like at lowest observed distance throughout testing
if goal_loss[-1] < 1.5:
'''
plt84 = plt.figure('Actual obs')
plt.cla()
sn.heatmap(s[:,:,0],cmap = 'Reds', annot=True,cbar = False).invert_yaxis()
plt85 = plt.figure('Zgoal')
plt.cla()
sn.heatmap(z_goal_obs[:,:,0],cmap = 'Reds', annot=True,cbar = False).invert_yaxis()
plt.show()
'''
reward = 4.0 #this reward is more of a place holder
else:
reward = 0.0
if self.curiosityreward:
reward = reward + Curiosityreward
total_reward += reward
final_loss.append(goal_loss[-1])
#Using every single observation, action, next observation,terminality condition and reward seen in the episode, get the loss of the world model
mdrnnlosses = self.get_loss(
torch.tensor(scur_rollout).to(self.device),
torch.tensor(act_rollout).to(self.device),
torch.tensor(r_rollout).to(self.device),
torch.tensor(d_rollout).to(self.device),
torch.tensor(snext_rollout).to(self.device),
include_reward=False)
#train world model
MDRNNoptimizer.zero_grad()
mdrnnlosses['loss'].backward()
MDRNNoptimizer.step()
WM_loss.append(
mdrnnlosses['loss']) #append to world model loss graph
#train VAE and HiddenVAE if representation learning is not static
if not self.static:
VAE_loss_per_rollout.append(
total_vae_loss / (goalattempts + 1)
) #average VAE loss metric when non static representations are being used
VAEOptimizer.zero_grad()
VAE_loss_per_rollout[-1].backward()
VAEOptimizer.step()
if self.hiddengoals:
hiddenvae_loss_per_rollout.append(
total_hiddenvae_loss / (goalattempts + 1)
) #average HiddenVAE loss metric when non static representations of hiddens states are being used
HiddenVAEOptimizer.zero_grad()
hiddenvae_loss_per_rollout[-1].backward()
HiddenVAEOptimizer.step()
if goalattempts % 10 == 0: #every 10 goals update the MDRNN cell for use in predicting the next state
self.mdrnn.load_state_dict(self.mdrnnBIG.state_dict())
loss_list.append(total_loss / (goalattempts + 1))
rollout_reward.append(total_reward)
visitationcount.append(np.sum(visitationarray))
exreward_per_rollout.append(total_exreward)
plot1 = plt.figure('Average Forward model loss')
plt.plot(loss_list)
plt7 = plt.figure('WM_loss')
plt.plot(WM_loss)
plt4 = plt.figure('Distance to goal per step')
plt.cla()
plt.plot(goal_loss)
rolloutrewardplot = plt.figure('Reward per rollout')
plt.plot(rollout_reward)
if not self.static:
vaerolloutplot = plt.figure('VAE loss per rollout')
plt.plot(VAE_loss_per_rollout)
if self.hiddengoals:
hiddenvaerolloutplot = plt.figure('HiddenVAE loss per rollout')
plt.plot(hiddenvae_loss_per_rollout)
plt8 = plt.figure('Visitation')
plt.plot(visitationcount)
pltexreward = plt.figure('Extrinsic Reward per rollout')
plt.plot(exreward_per_rollout)
pltgoalloss = plt.figure('Final Goal Loss per Episode')
plt.plot(final_loss)
plt.show()
input('stop')
def test_mdrnn_learning(self):
num_epochs = 300
num_episodes = 400
batch_size = 200
action_dim = 2
seq_len = 5
state_dim = 2
simulated_num_gaussian = 2
mdrnn_num_gaussian = 2
simulated_hidden_size = 3
mdrnn_hidden_size = 10
mdrnn_hidden_layer = 1
adam_lr = 0.01
cur_state_mem = numpy.zeros((num_episodes, seq_len, state_dim))
next_state_mem = numpy.zeros((num_episodes, seq_len, state_dim))
action_mem = numpy.zeros((num_episodes, seq_len, action_dim))
reward_mem = numpy.zeros((num_episodes, seq_len))
terminal_mem = numpy.zeros((num_episodes, seq_len))
next_mus_mem = numpy.zeros(
(num_episodes, seq_len, simulated_num_gaussian, state_dim))
swm = SimulatedWorldModel(
action_dim=action_dim,
state_dim=state_dim,
num_gaussian=simulated_num_gaussian,
lstm_num_layer=1,
lstm_hidden_dim=simulated_hidden_size,
)
actions = torch.eye(action_dim)
for e in range(num_episodes):
swm.init_hidden(batch_size=1)
next_state = torch.randn((1, 1, state_dim))
for s in range(seq_len):
cur_state = next_state
action = torch.tensor(
actions[numpy.random.randint(action_dim)]).view(
1, 1, action_dim)
next_mus, reward = swm(action, cur_state)
terminal = 0
if s == seq_len - 1:
terminal = 1
next_pi = torch.ones(
simulated_num_gaussian) / simulated_num_gaussian
index = Categorical(next_pi).sample((1, )).long().item()
next_state = next_mus[0, 0, index].view(1, 1, state_dim)
print(
"{} cur_state: {}, action: {}, next_state: {}, reward: {}, terminal: {}"
.format(e, cur_state, action, next_state, reward,
terminal))
print("next_pi: {}, sampled index: {}".format(next_pi, index))
print("next_mus:", next_mus, "\n")
cur_state_mem[e, s, :] = cur_state.detach().numpy()
action_mem[e, s, :] = action.numpy()
reward_mem[e, s] = reward.detach().numpy()
terminal_mem[e, s] = terminal
next_state_mem[e, s, :] = next_state.detach().numpy()
next_mus_mem[e, s, :, :] = next_mus.detach().numpy()
mdrnn = MDRNN(
latents=state_dim,
actions=action_dim,
gaussians=mdrnn_num_gaussian,
hiddens=mdrnn_hidden_size,
layers=mdrnn_hidden_layer,
)
mdrnn.train()
optimizer = torch.optim.Adam(mdrnn.parameters(), lr=adam_lr)
num_batch = num_episodes // batch_size
earlystopping = EarlyStopping('min', patience=30)
cum_loss = []
cum_gmm = []
cum_bce = []
cum_mse = []
for e in range(num_epochs):
for i in range(0, num_batch):
mdrnn.init_hidden(batch_size=batch_size)
optimizer.zero_grad()
sample_indices = numpy.random.randint(num_episodes,
size=batch_size)
obs, action, reward, terminal, next_obs = \
cur_state_mem[sample_indices], \
action_mem[sample_indices], \
reward_mem[sample_indices], \
terminal_mem[sample_indices], \
next_state_mem[sample_indices]
obs, action, reward, terminal, next_obs = \
torch.tensor(obs, dtype=torch.float), \
torch.tensor(action, dtype=torch.float), \
torch.tensor(reward, dtype=torch.float), \
torch.tensor(terminal, dtype=torch.float), \
torch.tensor(next_obs, dtype=torch.float)
print("learning at epoch {} step {} best score {} counter {}".
format(e, i, earlystopping.best,
earlystopping.num_bad_epochs))
losses = self.get_loss(obs, action, reward, terminal, next_obs,
state_dim, mdrnn)
losses['loss'].backward()
optimizer.step()
cum_loss += [losses['loss'].item()]
cum_gmm += [losses['gmm'].item()]
cum_bce += [losses['bce'].item()]
cum_mse += [losses['mse'].item()]
print(
"loss={loss:10.6f} bce={bce:10.6f} gmm={gmm:10.6f} mse={mse:10.6f}"
.format(
loss=losses['loss'],
bce=losses['bce'],
gmm=losses['gmm'],
mse=losses['mse'],
))
print(
"cum loss={loss:10.6f} cum bce={bce:10.6f} cum gmm={gmm:10.6f} cum mse={mse:10.6f}"
.format(
loss=numpy.mean(cum_loss),
bce=numpy.mean(cum_bce),
gmm=numpy.mean(cum_gmm),
mse=numpy.mean(cum_mse),
))
print()
earlystopping.step(numpy.mean(cum_loss[-num_batch:]))
if numpy.mean(cum_loss[-num_batch:]) < -3. and earlystopping.stop:
break
assert numpy.mean(cum_loss[-num_batch:]) < -3.
sample_indices = [0]
mdrnn.init_hidden(batch_size=len(sample_indices))
mdrnn.eval()
obs, action, reward, terminal, next_obs = \
cur_state_mem[sample_indices], \
action_mem[sample_indices], \
reward_mem[sample_indices], \
terminal_mem[sample_indices], \
next_state_mem[sample_indices]
obs, action, reward, terminal, next_obs = \
torch.tensor(obs, dtype=torch.float), \
torch.tensor(action, dtype=torch.float), \
torch.tensor(reward, dtype=torch.float), \
torch.tensor(terminal, dtype=torch.float), \
torch.tensor(next_obs, dtype=torch.float)
transpose_obs, transpose_action, transpose_reward, transpose_terminal, transpose_next_obs = \
self.transpose(obs, action, reward, terminal, next_obs)
mus, sigmas, logpi, rs, ds = mdrnn(transpose_action, transpose_obs)
pi = torch.exp(logpi)
gl = gmm_loss(transpose_next_obs, mus, sigmas, logpi)
print(gl)
print()
