code: def train_mdnrnn_and_compute_feature_stats(
      env_name: str,
      model: ModelManager__Union,
      num_train_transitions: int,
      num_test_transitions: int,
      seq_len: int,
      batch_size: int,
      num_train_epochs: int,
      use_gpu: bool,
      saved_mdnrnn_path: Optional[str] = None,
  ):
      """Train MDNRNN Memory Network and compute feature importance/sensitivity."""
      env: gym.Env = Gym(env_name=env_name)
      env.seed(SEED)
  
      manager = model.value
      trainer = manager.build_trainer(
          use_gpu=use_gpu,
          reward_options=RewardOptions(),
          normalization_data_map=build_normalizer(env),
      )
  
      device = "cuda" if use_gpu else "cpu"
      # pyre-fixme[6]: Expected `device` for 2nd param but got `str`.
      trainer_preprocessor = make_replay_buffer_trainer_preprocessor(
          trainer, device, env
      )
      test_replay_buffer = ReplayBuffer(
          replay_capacity=num_test_transitions,
          batch_size=batch_size,
          stack_size=seq_len,
          return_everything_as_stack=True,
      )
      random_policy = make_random_policy_for_env(env)
      agent = Agent.create_for_env(env, policy=random_policy)
      fill_replay_buffer(
          env, test_replay_buffer, num_test_transitions, agent
      )
  
      if saved_mdnrnn_path is None:
          # train from scratch
          trainer = train_mdnrnn(
              env=env,
              trainer=trainer,
              trainer_preprocessor=trainer_preprocessor,
              num_train_transitions=num_train_transitions,
              seq_len=seq_len,
              batch_size=batch_size,
              num_train_epochs=num_train_epochs,
              test_replay_buffer=test_replay_buffer,
          )
      else:
          # load a pretrained model, and just evaluate it
          trainer.memory_network.mdnrnn.load_state_dict(
              torch.load(saved_mdnrnn_path)
          )
  
      with torch.no_grad():
          trainer.memory_network.mdnrnn.eval()
          test_batch = test_replay_buffer.sample_transition_batch(
              batch_size=test_replay_buffer.size
          )
          preprocessed_test_batch = trainer_preprocessor(test_batch)
          feature_importance = calculate_feature_importance(
              env=env,
              trainer=trainer,
              use_gpu=use_gpu,
              test_batch=preprocessed_test_batch,
          )
  
          feature_sensitivity = calculate_feature_sensitivity(
              env=env,
              trainer=trainer,
              use_gpu=use_gpu,
              test_batch=preprocessed_test_batch,
          )
  
          trainer.memory_network.mdnrnn.train()
      return feature_importance, feature_sensitivity
  
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 4 for calling a global function with example usage: train_mdnrnn_and_compute_feature_stats(env_name, model, num_train_transitions, num_test_transitions, seq_len, batch_size, num_train_epochs, use_gpu, saved_mdnrnn_path) and returns: feature_importance,  feature_sensitivity