in Experiments/PolicyManagers.py [0:0]
def evaluate_loglikelihoods(self, sample_traj, sample_action_seq, concatenated_traj, latent_z_indices, latent_b):
# Initialize both loglikelihoods to 0.
subpolicy_loglikelihood = 0.
latent_loglikelihood = 0.
# Need to assemble inputs first - returns a Torch CUDA Tensor.
# This doesn't need to take in actions, because we can evaluate for all actions then select.
assembled_inputs, subpolicy_inputs, padded_action_seq = self.assemble_inputs(concatenated_traj, latent_z_indices, latent_b, sample_action_seq, self.conditional_information)
###########################
# Compute learnt subpolicy loglikelihood.
###########################
learnt_subpolicy_loglikelihoods, entropy = self.policy_network.forward(subpolicy_inputs, padded_action_seq)
# Clip values. # Comment this out to remove clipping.
learnt_subpolicy_loglikelihoods = torch.clamp(learnt_subpolicy_loglikelihoods,min=self.args.subpolicy_clamp_value)
# Multiplying the likelihoods with the subpolicy ratio before summing.
learnt_subpolicy_loglikelihoods = self.args.subpolicy_ratio*learnt_subpolicy_loglikelihoods
# Summing until penultimate timestep.
# learnt_subpolicy_loglikelihood = learnt_subpolicy_loglikelihoods[:-1].sum()
# TAKING AVERAGE HERE AS WELL.
learnt_subpolicy_loglikelihood = learnt_subpolicy_loglikelihoods[:-1].mean()
###########################
# Compute Latent policy loglikelihood values.
###########################
# Whether to clone assembled_inputs based on the phase of training.
# In phase one it doesn't matter if we use the clone or not, because we never use latent policy loss.
# So just clone anyway.
# For now, ignore phase 3. This prevents gradients from going into the variational policy from the latent policy.
assembled_inputs_copy = assembled_inputs.clone().detach()
latent_z_copy = latent_z_indices.clone().detach()
# Consideration for later:
# if self.training_phase==3:
# Don't clone.
if self.args.discrete_z:
# Return discrete probabilities from latent policy network.
latent_z_logprobabilities, latent_b_logprobabilities, latent_b_probabilities, latent_z_probabilities = self.latent_policy.forward(assembled_inputs_copy)
# # Selects first option for variable = 1, second option for variable = 0.
# Use this to check if latent_z elements are equal:
diff_val = (1-(latent_z_indices==latent_z_indices.roll(1,0))[1:]).to(device).float()
# We rolled latent_z, we didn't roll diff. This works because latent_b is always guaranteed to be 1 in the first timestep, so it doesn't matter what's in diff_val[0].
diff_val = diff_val.roll(1,0)
# Selects first option for variable = 1, second option for variable = 0.
latent_z_temporal_logprobabilities = torch.where(latent_b[:-1].byte(), latent_z_logprobabilities[range(len(sample_traj)-1),latent_z_indices[:-1].long()], -self.lambda_likelihood_penalty*diff_val)
latent_z_logprobability = latent_z_temporal_logprobabilities.mean()
else:
# If not, we need to evaluate the latent probabilties of latent_z_indices under latent_policy.
latent_b_logprobabilities, latent_b_probabilities, latent_distributions = self.latent_policy.forward(assembled_inputs_copy, self.epsilon)
# Evalute loglikelihood of latent z vectors under the latent policy's distributions.
latent_z_logprobabilities = latent_distributions.log_prob(latent_z_copy.unsqueeze(1))
# Multiply logprobabilities by the latent policy ratio.
latent_z_temporal_logprobabilities = latent_z_logprobabilities[:-1]*self.args.latentpolicy_ratio
latent_z_logprobability = latent_z_temporal_logprobabilities.mean()
latent_z_probabilities = None
# LATENT LOGLIKELIHOOD is defined as:
# = \sum_{t=1}^T \log p(\zeta_t | \tau_{1:t}, \zeta_{1:t-1})
# = \sum_{t=1}^T \log { \phi_t(b_t)} + \log { 1[b_t==1] \eta_t(h_t|s_{1:t}) + 1[b_t==0] 1[z_t==z_{t-1}] }
# Adding log probabilities of termination (of whether it terminated or not), till penultimate step.
latent_b_temporal_logprobabilities = latent_b_logprobabilities[range(len(sample_traj)-1),latent_b[:-1].long()]
latent_b_logprobability = latent_b_temporal_logprobabilities.mean()
latent_loglikelihood += latent_b_logprobability
latent_loglikelihood += latent_z_logprobability
# DON'T CLAMP, JUST MULTIPLY BY SUITABLE RATIO! Probably use the same lat_z_wt and lat_b_wt ratios from the losses.
latent_temporal_loglikelihoods = self.args.lat_b_wt*latent_b_temporal_logprobabilities + self.args.lat_z_wt*latent_z_temporal_logprobabilities.squeeze(1)
##################################################
#### Manage merging likelihoods for REINFORCE ####
##################################################
if self.training_phase==1:
temporal_loglikelihoods = learnt_subpolicy_loglikelihoods[:-1].squeeze(1)
elif self.training_phase==2 or self.training_phase==3:
# temporal_loglikelihoods = learnt_subpolicy_loglikelihoods[:-1].squeeze(1) + self.args.temporal_latentpolicy_ratio*latent_temporal_loglikelihoods
temporal_loglikelihoods = learnt_subpolicy_loglikelihoods[:-1].squeeze(1)
if self.args.debug:
if self.iter%self.args.debug==0:
print("Embedding in the Evaluate Likelihoods Function.")
embed()
return None, None, None, latent_loglikelihood, \
latent_b_logprobabilities, latent_z_logprobabilities, latent_b_probabilities, latent_z_probabilities, \
latent_z_logprobability, latent_b_logprobability, learnt_subpolicy_loglikelihood, learnt_subpolicy_loglikelihoods, temporal_loglikelihoods