def getImage(srl_model, state, device):
"""
Gets an image by using the decoder of a SRL model
(when available)
:param srl_model: (Pytorch model)
:param state: ([float]) the state vector from latent space
:param device: (pytorch device)
:return: ([float])
"""
with th.no_grad():
state = th.from_numpy(np.array(state).reshape(1, -1)).float()
state = state.to(device)
net_out = srl_model.decode(state)
img = detachToNumpy(net_out)[0].T
img = deNormalize(img, mode="tf")
return img[:, :, ::-1]
def learn(self, images_path, actions, rewards, episode_starts):
"""
Learn a state representation
:param images_path: (numpy 1D array)
:param actions: (np.ndarray)
:param rewards: (numpy 1D array)
:param episode_starts: (numpy 1D array) boolean array
the ith index is True if one episode starts at this frame
:return: (np.ndarray) the learned states for the given observations
"""
print("\nYour are using the following weights for the losses:")
pprint(self.losses_weights_dict)
# PREPARE DATA -------------------------------------------------------------------------------------------------
# here, we organize the data into minibatches
# and find pairs for the respective loss terms (for robotics priors only)
num_samples = images_path.shape[0] - 1 # number of samples
# indices for all time steps where the episode continues
indices = np.array([i for i in range(num_samples) if not episode_starts[i + 1]], dtype='int64')
np.random.shuffle(indices)
# split indices into minibatches. minibatchlist is a list of lists; each
# list is the id of the observation preserved through the training
minibatchlist = [np.array(sorted(indices[start_idx:start_idx + self.batch_size]))
for start_idx in range(0, len(indices) - self.batch_size + 1, self.batch_size)]
test_minibatchlist = DataLoader.createTestMinibatchList(len(images_path), MAX_BATCH_SIZE_GPU)
# Number of minibatches used for validation:
n_val_batches = np.round(VALIDATION_SIZE * len(minibatchlist)).astype(np.int64)
val_indices = np.random.permutation(len(minibatchlist))[:n_val_batches]
# Print some info
print("{} minibatches for training, {} samples".format(len(minibatchlist) - n_val_batches,
(len(minibatchlist) - n_val_batches) * BATCH_SIZE))
print("{} minibatches for validation, {} samples".format(n_val_batches, n_val_batches * BATCH_SIZE))
assert n_val_batches > 0, "Not enough sample to create a validation set"
# Stats about actions
if not self.continuous_action:
print('Discrete action space:')
action_set = set(actions)
n_actions = int(np.max(actions) + 1)
print("{} unique actions / {} actions".format(len(action_set), n_actions))
n_pairs_per_action = np.zeros(n_actions, dtype=np.int64)
n_obs_per_action = np.zeros(n_actions, dtype=np.int64)
for i in range(n_actions):
n_obs_per_action[i] = np.sum(actions == i)
print("Number of observations per action")
print(n_obs_per_action)
else:
print('Continuous action space:')
print('Action dimension: {}'.format(self.dim_action))
dissimilar_pairs, same_actions_pairs = None, None
if not self.no_priors:
if self.continuous_action:
print('This option (priors) doesnt support continuous action space for now !')
dissimilar_pairs, same_actions_pairs = findPriorsPairs(self.batch_size, minibatchlist, actions, rewards,
n_actions, n_pairs_per_action)
if self.use_vae and self.perceptual_similarity_loss and self.path_to_dae is not None:
self.denoiser = SRLModules(state_dim=self.state_dim_dae, action_dim=self.dim_action,
model_type="custom_cnn",
cuda=self.cuda, losses=["dae"])
self.denoiser.load_state_dict(th.load(self.path_to_dae))
self.denoiser.eval()
self.denoiser = self.denoiser.to(self.device)
for param in self.denoiser.parameters():
param.requires_grad = False
if self.episode_prior:
idx_to_episode = {idx: episode_idx for idx, episode_idx in enumerate(np.cumsum(episode_starts))}
minibatch_episodes = [[idx_to_episode[i] for i in minibatch] for minibatch in minibatchlist]
data_loader = DataLoader(minibatchlist, images_path, n_workers=N_WORKERS, multi_view=self.multi_view,
use_triplets=self.use_triplets, is_training=True, apply_occlusion=self.use_dae,
occlusion_percentage=self.occlusion_percentage)
test_data_loader = DataLoader(test_minibatchlist, images_path, n_workers=N_WORKERS, multi_view=self.multi_view,
use_triplets=self.use_triplets, max_queue_len=1, is_training=False,
apply_occlusion=self.use_dae, occlusion_percentage=self.occlusion_percentage)
# TRAINING -----------------------------------------------------------------------------------------------------
loss_history = defaultdict(list)
loss_manager = LossManager(self.model, loss_history)
best_error = np.inf
best_model_path = "{}/srl_model.pth".format(self.log_folder)
start_time = time.time()
# Random features, we don't need to train a model
if len(self.losses) == 1 and self.losses[0] == 'random':
global N_EPOCHS
N_EPOCHS = 0
printYellow("Skipping training because using random features")
th.save(self.model.state_dict(), best_model_path)
for epoch in range(N_EPOCHS):
# In each epoch, we do a full pass over the training data:
epoch_loss, epoch_batches = 0, 0
val_loss = 0
pbar = tqdm(total=len(minibatchlist))
for minibatch_num, (minibatch_idx, obs, next_obs, noisy_obs, next_noisy_obs) in enumerate(data_loader):
validation_mode = minibatch_idx in val_indices
if validation_mode:
self.model.eval()
else:
self.model.train()
if self.use_dae:
noisy_obs = noisy_obs.to(self.device)
next_noisy_obs = next_noisy_obs.to(self.device)
obs, next_obs = obs.to(self.device), next_obs.to(self.device)
self.optimizer.zero_grad()
loss_manager.resetLosses()
decoded_obs, decoded_next_obs = None, None
states_denoiser = None
states_denoiser_predicted = None
next_states_denoiser = None
next_states_denoiser_predicted = None
# Predict states given observations as in Time Contrastive Network (Triplet Loss) [Sermanet et al.]
if self.use_triplets:
states, positive_states, negative_states = self.model.forwardTriplets(obs[:, :3:, :, :],
obs[:, 3:6, :, :],
obs[:, 6:, :, :])
next_states, next_positive_states, next_negative_states = self.model.forwardTriplets(
next_obs[:, :3:, :, :],
next_obs[:, 3:6, :, :],
next_obs[:, 6:, :, :])
elif self.use_autoencoder:
(states, decoded_obs), (next_states, decoded_next_obs) = self.model(obs), self.model(next_obs)
elif self.use_dae:
(states, decoded_obs), (next_states, decoded_next_obs) = \
self.model(noisy_obs), self.model(next_noisy_obs)
elif self.use_vae:
(decoded_obs, mu, logvar), (next_decoded_obs, next_mu, next_logvar) = self.model(obs), \
self.model(next_obs)
states, next_states = self.model.getStates(obs), self.model.getStates(next_obs)
if self.perceptual_similarity_loss:
# Predictions for the perceptual similarity loss as in DARLA
# https://arxiv.org/pdf/1707.08475.pdf
(states_denoiser, decoded_obs_denoiser), (next_states_denoiser, decoded_next_obs_denoiser) = \
self.denoiser(obs), self.denoiser(next_obs)
(states_denoiser_predicted, decoded_obs_denoiser_predicted) = self.denoiser(decoded_obs)
(next_states_denoiser_predicted,
decoded_next_obs_denoiser_predicted) = self.denoiser(next_decoded_obs)
else:
states, next_states = self.model(obs), self.model(next_obs)
# Actions associated to the observations of the current minibatch
actions_st = actions[minibatchlist[minibatch_idx]]
if not self.continuous_action:
# Discrete actions, rearrange action to have n_minibatch ligns and one column, containing the int action
actions_st = th.from_numpy(actions_st).view(-1, 1).requires_grad_(False).to(self.device)
else:
# Continuous actions, rearrange action to have n_minibatch ligns and dim_action columns
actions_st = th.from_numpy(actions_st).view(-1, self.dim_action).requires_grad_(False).to(self.device)
# L1 regularization
if self.losses_weights_dict['l1_reg'] > 0:
l1Loss(loss_manager.reg_params, self.losses_weights_dict['l1_reg'], loss_manager)
if self.losses_weights_dict['l2_reg'] > 0:
l2Loss(loss_manager.reg_params, self.losses_weights_dict['l2_reg'], loss_manager)
if not self.no_priors:
if self.n_actions == np.inf:
print('This option (priors) doesnt support continuous action space for now !')
roboticPriorsLoss(states, next_states, minibatch_idx=minibatch_idx,
dissimilar_pairs=dissimilar_pairs, same_actions_pairs=same_actions_pairs,
weight=self.losses_weights_dict['priors'], loss_manager=loss_manager)
# TODO change here to classic call (forward and backward)
if self.use_forward_loss:
next_states_pred = self.model.forwardModel(states, actions_st)
forwardModelLoss(next_states_pred, next_states,
weight=self.losses_weights_dict['forward'],
loss_manager=loss_manager)
if self.use_inverse_loss:
actions_pred = self.model.inverseModel(states, next_states)
inverseModelLoss(actions_pred, actions_st, weight=self.losses_weights_dict['inverse'],
loss_manager=loss_manager, continuous_action=self.continuous_action)
if self.use_reward_loss:
rewards_st = rewards[minibatchlist[minibatch_idx]].copy()
# Removing negative reward
rewards_st[rewards_st == -1] = 0
rewards_st = th.from_numpy(rewards_st).to(self.device)
rewards_pred = self.model.rewardModel(states, next_states)
rewardModelLoss(rewards_pred, rewards_st.long(), weight=self.losses_weights_dict['reward'],
loss_manager=loss_manager)
if self.use_autoencoder or self.use_dae:
loss_type = "dae" if self.use_dae else "autoencoder"
autoEncoderLoss(obs, decoded_obs, next_obs, decoded_next_obs,
weight=self.losses_weights_dict[loss_type], loss_manager=loss_manager)
if self.use_vae:
kullbackLeiblerLoss(mu, next_mu, logvar, next_logvar, loss_manager=loss_manager, beta=self.beta)
if self.perceptual_similarity_loss:
perceptualSimilarityLoss(states_denoiser, states_denoiser_predicted, next_states_denoiser,
next_states_denoiser_predicted,
weight=self.losses_weights_dict['perceptual'],
loss_manager=loss_manager)
else:
generationLoss(decoded_obs, next_decoded_obs, obs, next_obs,
weight=self.losses_weights_dict['vae'], loss_manager=loss_manager)
if self.reward_prior:
rewards_st = rewards[minibatchlist[minibatch_idx]]
rewards_st = th.from_numpy(rewards_st).float().view(-1, 1).to(self.device)
rewardPriorLoss(states, rewards_st, weight=self.losses_weights_dict['reward-prior'],
loss_manager=loss_manager)
if self.episode_prior:
episodePriorLoss(minibatch_idx, minibatch_episodes, states, self.discriminator,
BALANCED_SAMPLING, weight=self.losses_weights_dict['episode-prior'],
loss_manager=loss_manager)
if self.use_triplets:
tripletLoss(states, positive_states, negative_states, weight=self.losses_weights_dict['triplet'],
loss_manager=loss_manager, alpha=0.2)
# Compute weighted average of losses
loss_manager.updateLossHistory()
loss = loss_manager.computeTotalLoss()
# We have to call backward in both train/val
# to avoid memory error
loss.backward()
if validation_mode:
val_loss += loss.item()
# We do not optimize on validation data
# so optimizer.step() is not called
else:
self.optimizer.step()
epoch_loss += loss.item()
epoch_batches += 1
pbar.update(1)
pbar.close()
train_loss = epoch_loss / float(epoch_batches)
val_loss /= float(n_val_batches)
# Even if loss_history is modified by LossManager
# we make it explicit
loss_history = loss_manager.loss_history
loss_history['train_loss'].append(train_loss)
loss_history['val_loss'].append(val_loss)
for key in loss_history.keys():
if key in ['train_loss', 'val_loss']:
continue
loss_history[key][-1] /= epoch_batches
if epoch + 1 < N_EPOCHS:
loss_history[key].append(0)
# Save best model
if val_loss < best_error:
best_error = val_loss
th.save(self.model.state_dict(), best_model_path)
if np.isnan(train_loss):
printRed("NaN Loss, consider increasing NOISE_STD in the gaussian noise layer")
sys.exit(NAN_ERROR)
# Then we print the results for this epoch:
if (epoch + 1) % EPOCH_FLAG == 0:
print("Epoch {:3}/{}, train_loss:{:.4f} val_loss:{:.4f}".format(epoch + 1, N_EPOCHS, train_loss,
val_loss))
print("{:.2f}s/epoch".format((time.time() - start_time) / (epoch + 1)))
if DISPLAY_PLOTS:
with th.no_grad():
self.model.eval()
# Optionally plot the current state space
plotRepresentation(self.predStatesWithDataLoader(test_data_loader), rewards,
add_colorbar=epoch == 0,
name="Learned State Representation (Training Data)")
if self.use_autoencoder or self.use_vae or self.use_dae:
# Plot Reconstructed Image
if obs[0].shape[0] == 3: # RGB
plotImage(deNormalize(detachToNumpy(obs[0])), "Input Image (Train)")
if self.use_dae:
plotImage(deNormalize(detachToNumpy(noisy_obs[0])), "Noisy Input Image (Train)")
if self.perceptual_similarity_loss:
plotImage(deNormalize(detachToNumpy(decoded_obs_denoiser[0])),
"Reconstructed Image DAE")
plotImage(deNormalize(detachToNumpy(decoded_obs_denoiser_predicted[0])),
"Reconstructed Image predicted DAE")
plotImage(deNormalize(detachToNumpy(decoded_obs[0])), "Reconstructed Image")
elif obs[0].shape[0] % 3 == 0: # Multi-RGB
for k in range(obs[0].shape[0] // 3):
plotImage(deNormalize(detachToNumpy(obs[0][k * 3:(k + 1) * 3, :, :]), "image_net"),
"Input Image {} (Train)".format(k + 1))
if self.use_dae:
plotImage(deNormalize(detachToNumpy(noisy_obs[0][k * 3:(k + 1) * 3, :, :])),
"Noisy Input Image (Train)".format(k + 1))
if self.perceptual_similarity_loss:
plotImage(deNormalize(
detachToNumpy(decoded_obs_denoiser[0][k * 3:(k + 1) * 3, :, :])),
"Reconstructed Image DAE")
plotImage(deNormalize(
detachToNumpy(decoded_obs_denoiser_predicted[0][k * 3:(k + 1) * 3, :, :])),
"Reconstructed Image predicted DAE")
plotImage(deNormalize(detachToNumpy(decoded_obs[0][k * 3:(k + 1) * 3, :, :])),
"Reconstructed Image {}".format(k + 1))
if DISPLAY_PLOTS:
plt.close("Learned State Representation (Training Data)")
# Load best model before predicting states
self.model.load_state_dict(th.load(best_model_path))
print("Predicting states for all the observations...")
# return predicted states for training observations
self.model.eval()
with th.no_grad():
pred_states = self.predStatesWithDataLoader(test_data_loader)
pairs_loss_weight = [k for k in zip(loss_manager.names, loss_manager.weights)]
return loss_history, pred_states, pairs_loss_weight