class_docstr: Evaluate state feature sensitivity caused by varying actions
class_docstr_tok: ["Evaluate","state","feature","sensitivity","caused","by","varying","actions"]
docstr: Calculate state feature sensitivity due to actions:
  randomly permutating actions and see how much the prediction of next
  state feature deviates.
docstr_tok: ["Calculate","state","feature","sensitivity","due","to","actions",":","randomly","permutating","actions","and","see","how","much","the","prediction","of","next","state","feature","deviates","."]
code: def evaluate(self, batch: MemoryNetworkInput):
      """Calculate state feature sensitivity due to actions:
      randomly permutating actions and see how much the prediction of next
      state feature deviates."""
      assert isinstance(batch, MemoryNetworkInput)
  
      self.trainer.memory_network.mdnrnn.eval()
  
      (
          seq_len,
          batch_size,
          state_dim,
      ) = batch.next_state.float_features.size()
      state_feature_num = self.state_feature_num
      feature_sensitivity = torch.zeros(state_feature_num)
  
      # the input of world_model has seq-len as the first dimension
      mdnrnn_output = self.trainer.memory_network(
          batch.state, batch.action
      )
      predicted_next_state_means = mdnrnn_output.mus
  
      shuffled_mdnrnn_output = self.trainer.memory_network(
          batch.state,
          # shuffle the actions
          FeatureData(
              batch.action.float_features[
                  :, torch.randperm(batch_size), :
              ]
          ),
      )
      shuffled_predicted_next_state_means = shuffled_mdnrnn_output.mus
  
      assert (
          predicted_next_state_means.size()
          == shuffled_predicted_next_state_means.size()
          == (
              seq_len,
              batch_size,
              self.trainer.params.num_gaussians,
              state_dim,
          )
      )
  
      state_feature_boundaries = (
          self.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.memory_network.mdnrnn.train()
      logger.info(
          "**** Debug tool feature sensitivity ****: {}".format(
              feature_sensitivity
          )
      )
      return {"feature_sensitivity": feature_sensitivity.numpy()}
  
file_path: ReAgent/reagent/evaluation/world_model_evaluator.py
imports: import logging
  from typing import Dict, List
  import torch
  from reagent.core.types import FeatureData, MemoryNetworkInput
  from reagent.training.world_model.mdnrnn_trainer import MDNRNNTrainer
label_desc: Use image node_id 2 for calling the FeatureSensitivityEvaluator obj's underlying member method code with example usage: obj.evaluate(batch) and returns: dict