code: def calculate_feature_sensitivity(
      env: EnvWrapper,
      trainer: MDNRNNTrainer,
      use_gpu: bool,
      test_batch: rlt.MemoryNetworkInput,
  ):
      assert isinstance(env.action_space, gym.spaces.Discrete)
      assert isinstance(env.observation_space, gym.spaces.Box)
      assert len(env.observation_space.shape) == 1
      state_dim = env.observation_space.shape[0]
      feature_sensitivity_evaluator = FeatureSensitivityEvaluator(
          trainer,
          state_feature_num=state_dim,
          sorted_state_feature_start_indices=list(range(state_dim)),
      )
      feature_sensitivity_vector = (
          feature_sensitivity_evaluator.evaluate(test_batch)[
              "feature_sensitivity"
          ]
      )
      feature_sensitivity_map = {}
      for i in range(state_dim):
          feature_sensitivity_map[
              "state" + str(i)
          ] = feature_sensitivity_vector[i].item()
          print(
              "state {}, feature sensitivity: {}".format(
                  i, feature_sensitivity_vector[i].item()
              )
          )
      return feature_sensitivity_map
  
imports: import logging
  import os
  import unittest
  from typing import Dict, List, Optional
  import gym
  import numpy
  import reagent.core.types
  import torch
  from reagent.evaluation.world_model_evaluator import FeatureImportanceEvaluator, FeatureSensitivityEvaluator
  from reagent.gym.agents.agent import Agent
  from reagent.gym.envs import EnvWrapper, Gym
  from reagent.gym.envs.pomdp.state_embed_env import StateEmbedEnvironment
  from reagent.gym.policies.random_policies import make_random_policy_for_env
  from reagent.gym.preprocessors import make_replay_buffer_trainer_preprocessor
  from reagent.gym.runners.gymrunner import evaluate_for_n_episodes
  from reagent.gym.utils import build_normalizer, fill_replay_buffer
  from reagent.model_managers.union import ModelManager__Union
  from reagent.models.world_model import MemoryNetwork
  from reagent.replay_memory.circular_replay_buffer import ReplayBuffer
  from reagent.test.base.horizon_test_base import HorizonTestBase
  from reagent.training.world_model.mdnrnn_trainer import MDNRNNTrainer
  from reagent.workflow.types import RewardOptions
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: calculate_feature_sensitivity(env, trainer, use_gpu, test_batch) and returns: feature_sensitivity_map