def evaluate(self, tdp: TrainingBatch): """ Calculate state feature sensitivity due to actions: randomly permutating actions and see how much the prediction of next state feature deviates. """ self.trainer.mdnrnn.mdnrnn.eval() batch_size, seq_len, state_dim = tdp.training_input.next_state.size() state_feature_num = self.feature_extractor.state_feature_num feature_sensitivity = torch.zeros(state_feature_num) mdnrnn_input = tdp.training_input mdnrnn_output = self.trainer.mdnrnn(mdnrnn_input) predicted_next_state_means = mdnrnn_output.mus shuffled_mdnrnn_input = MemoryNetworkInput( state=tdp.training_input.state, # shuffle the actions action=FeatureVector(float_features=tdp.training_input.action. float_features[torch.randperm(batch_size)]), next_state=tdp.training_input.next_state, reward=tdp.training_input.reward, not_terminal=tdp.training_input.not_terminal, ) shuffled_mdnrnn_output = self.trainer.mdnrnn(shuffled_mdnrnn_input) shuffled_predicted_next_state_means = shuffled_mdnrnn_output.mus assert (predicted_next_state_means.size() == shuffled_predicted_next_state_means.size() == (batch_size, seq_len, self.trainer.params.num_gaussians, state_dim)) state_feature_boundaries = ( self.feature_extractor.sorted_state_feature_start_indices + [state_dim]) for i in range(state_feature_num): boundary_start, boundary_end = ( state_feature_boundaries[i], state_feature_boundaries[i + 1], ) abs_diff = torch.mean( torch.sum( torch.abs( shuffled_predicted_next_state_means[:, :, :, boundary_start: boundary_end] - predicted_next_state_means[:, :, :, boundary_start:boundary_end] ), dim=3, )) feature_sensitivity[i] = abs_diff.cpu().detach().item() self.trainer.mdnrnn.mdnrnn.train() logger.info("**** Debug tool feature sensitivity ****: {}".format( feature_sensitivity)) return {"feature_sensitivity": feature_sensitivity.numpy()}
def handle(self, tdp: TrainingBatch) -> None: batch_size, _, _ = tdp.training_input.next_state.size() tdp = TrainingBatch( training_input=MemoryNetworkInput( state=tdp.training_input.state, action=tdp.training_input.action, # shuffle the data next_state=tdp.training_input.next_state[torch.randperm( batch_size)], reward=tdp.training_input.reward[torch.randperm(batch_size)], not_terminal=tdp.training_input.not_terminal[torch.randperm( batch_size)], ), extras=ExtraData(), ) losses = self.trainer_or_evaluator.train(tdp, batch_first=True) self.results.append(losses)
def evaluate(self, tdp: TrainingBatch): """ Calculate feature importance: setting each state/action feature to the mean value and observe loss increase. """ self.trainer.mdnrnn.mdnrnn.eval() state_features = tdp.training_input.state.float_features action_features = tdp.training_input.action.float_features # type: ignore batch_size, seq_len, state_dim = state_features.size() # type: ignore action_dim = action_features.size()[2] # type: ignore action_feature_num = self.action_feature_num state_feature_num = self.state_feature_num feature_importance = torch.zeros(action_feature_num + state_feature_num) orig_losses = self.trainer.get_loss(tdp, state_dim=state_dim, batch_first=True) orig_loss = orig_losses["loss"].cpu().detach().item() del orig_losses action_feature_boundaries = self.sorted_action_feature_start_indices + [ action_dim ] state_feature_boundaries = self.sorted_state_feature_start_indices + [state_dim] for i in range(action_feature_num): action_features = tdp.training_input.action.float_features.reshape( # type: ignore (batch_size * seq_len, action_dim) ).data.clone() # if actions are discrete, an action's feature importance is the loss # increase due to setting all actions to this action if self.discrete_action: assert action_dim == action_feature_num action_vec = torch.zeros(action_dim) action_vec[i] = 1 action_features[:] = action_vec # type: ignore # if actions are continuous, an action's feature importance is the loss # increase due to masking this action feature to its mean value else: boundary_start, boundary_end = ( action_feature_boundaries[i], action_feature_boundaries[i + 1], ) action_features[ # type: ignore :, boundary_start:boundary_end ] = self.compute_median_feature_value( # type: ignore action_features[:, boundary_start:boundary_end] # type: ignore ) action_features = action_features.reshape( # type: ignore (batch_size, seq_len, action_dim) ) # type: ignore new_tdp = TrainingBatch( training_input=MemoryNetworkInput( # type: ignore state=tdp.training_input.state, action=FeatureVector( # type: ignore float_features=action_features ), # type: ignore next_state=tdp.training_input.next_state, reward=tdp.training_input.reward, not_terminal=tdp.training_input.not_terminal, ), extras=ExtraData(), ) losses = self.trainer.get_loss( new_tdp, state_dim=state_dim, batch_first=True ) feature_importance[i] = losses["loss"].cpu().detach().item() - orig_loss del losses for i in range(state_feature_num): state_features = tdp.training_input.state.float_features.reshape( # type: ignore (batch_size * seq_len, state_dim) ).data.clone() boundary_start, boundary_end = ( state_feature_boundaries[i], state_feature_boundaries[i + 1], ) state_features[ # type: ignore :, boundary_start:boundary_end ] = self.compute_median_feature_value( state_features[:, boundary_start:boundary_end] # type: ignore ) state_features = state_features.reshape( # type: ignore (batch_size, seq_len, state_dim) ) # type: ignore new_tdp = TrainingBatch( training_input=MemoryNetworkInput( # type: ignore state=FeatureVector(float_features=state_features), # type: ignore action=tdp.training_input.action, next_state=tdp.training_input.next_state, reward=tdp.training_input.reward, not_terminal=tdp.training_input.not_terminal, ), extras=ExtraData(), ) losses = self.trainer.get_loss( new_tdp, state_dim=state_dim, batch_first=True ) feature_importance[i + action_feature_num] = ( losses["loss"].cpu().detach().item() - orig_loss ) del losses self.trainer.mdnrnn.mdnrnn.train() logger.info( "**** Debug tool feature importance ****: {}".format(feature_importance) ) return {"feature_loss_increase": feature_importance.numpy()}
def evaluate(self, tdp: TrainingBatch): """ Calculate feature importance: setting each state/action feature to the mean value and observe loss increase. """ self.trainer.mdnrnn.mdnrnn.eval() state_features = tdp.training_input.state.float_features action_features = tdp.training_input.action.float_features batch_size, seq_len, state_dim = state_features.size() action_dim = action_features.size()[2] action_feature_num = self.feature_extractor.action_feature_num state_feature_num = self.feature_extractor.state_feature_num feature_importance = torch.zeros(action_feature_num + state_feature_num) orig_losses = self.trainer.get_loss(tdp, state_dim=state_dim, batch_first=True) orig_loss = orig_losses["loss"].cpu().detach().item() del orig_losses action_feature_boundaries = ( self.feature_extractor.sorted_action_feature_start_indices + [action_dim]) state_feature_boundaries = ( self.feature_extractor.sorted_state_feature_start_indices + [state_dim]) for i in range(action_feature_num): action_features = tdp.training_input.action.float_features.reshape( (batch_size * seq_len, action_dim)).data.clone() boundary_start, boundary_end = ( action_feature_boundaries[i], action_feature_boundaries[i + 1], ) action_features[:, boundary_start: boundary_end] = action_features[:, boundary_start: boundary_end].mean( dim=0) action_features = action_features.reshape( (batch_size, seq_len, action_dim)) new_tdp = TrainingBatch( training_input=MemoryNetworkInput( state=tdp.training_input.state, action=FeatureVector(float_features=action_features), next_state=tdp.training_input.next_state, reward=tdp.training_input.reward, not_terminal=tdp.training_input.not_terminal, ), extras=ExtraData(), ) losses = self.trainer.get_loss(new_tdp, state_dim=state_dim, batch_first=True) feature_importance[i] = losses["loss"].cpu().detach().item( ) - orig_loss del losses for i in range(state_feature_num): state_features = tdp.training_input.state.float_features.reshape( (batch_size * seq_len, state_dim)).data.clone() boundary_start, boundary_end = ( state_feature_boundaries[i], state_feature_boundaries[i + 1], ) state_features[:, boundary_start: boundary_end] = state_features[:, boundary_start: boundary_end].mean( dim=0) state_features = state_features.reshape( (batch_size, seq_len, state_dim)) new_tdp = TrainingBatch( training_input=MemoryNetworkInput( state=FeatureVector(float_features=state_features), action=tdp.training_input.action, next_state=tdp.training_input.next_state, reward=tdp.training_input.reward, not_terminal=tdp.training_input.not_terminal, ), extras=ExtraData(), ) losses = self.trainer.get_loss(new_tdp, state_dim=state_dim, batch_first=True) feature_importance[i + action_feature_num] = ( losses["loss"].cpu().detach().item() - orig_loss) del losses self.trainer.mdnrnn.mdnrnn.train() logger.info("**** Debug tool feature importance ****: {}".format( feature_importance)) return {"feature_loss_increase": feature_importance.numpy()}