code: def train_mdnrnn_and_train_on_embedded_env(
      env_name: str,
      embedding_model: ModelManager__Union,
      num_embedding_train_transitions: int,
      seq_len: int,
      batch_size: int,
      num_embedding_train_epochs: int,
      train_model: ModelManager__Union,
      num_state_embed_transitions: int,
      num_agent_train_epochs: int,
      num_agent_eval_epochs: int,
      use_gpu: bool,
      passing_score_bar: float,
      # pyre-fixme[9]: saved_mdnrnn_path has type `str`; used as `None`.
      saved_mdnrnn_path: str = None,
  ):
      """Train an agent on embedded states by the MDNRNN."""
      env = Gym(env_name=env_name)
      env.seed(SEED)
  
      embedding_manager = embedding_model.value
      embedding_trainer = embedding_manager.build_trainer(
          use_gpu=use_gpu,
          reward_options=RewardOptions(),
          normalization_data_map=build_normalizer(env),
      )
  
      device = "cuda" if use_gpu else "cpu"
      embedding_trainer_preprocessor = make_replay_buffer_trainer_preprocessor(
          embedding_trainer,
          # pyre-fixme[6]: Expected `device` for 2nd param but got `str`.
          device,
          env,
      )
      if saved_mdnrnn_path is None:
          # train from scratch
          embedding_trainer = train_mdnrnn(
              env=env,
              trainer=embedding_trainer,
              trainer_preprocessor=embedding_trainer_preprocessor,
              num_train_transitions=num_embedding_train_transitions,
              seq_len=seq_len,
              batch_size=batch_size,
              num_train_epochs=num_embedding_train_epochs,
          )
      else:
          # load a pretrained model, and just evaluate it
          embedding_trainer.memory_network.mdnrnn.load_state_dict(
              torch.load(saved_mdnrnn_path)
          )
  
      # create embedding dataset
      embed_rb, state_min, state_max = create_embed_rl_dataset(
          env=env,
          memory_network=embedding_trainer.memory_network,
          num_state_embed_transitions=num_state_embed_transitions,
          batch_size=batch_size,
          seq_len=seq_len,
          hidden_dim=embedding_trainer.params.hidden_size,
          use_gpu=use_gpu,
      )
      embed_env = StateEmbedEnvironment(
          gym_env=env,
          mdnrnn=embedding_trainer.memory_network,
          max_embed_seq_len=seq_len,
          state_min_value=state_min,
          state_max_value=state_max,
      )
      agent_manager = train_model.value
      agent_trainer = agent_manager.build_trainer(
          use_gpu=use_gpu,
          reward_options=RewardOptions(),
          # pyre-fixme[6]: Expected `EnvWrapper` for 1st param but got
          #  `StateEmbedEnvironment`.
          normalization_data_map=build_normalizer(embed_env),
      )
      device = "cuda" if use_gpu else "cpu"
      agent_trainer_preprocessor = make_replay_buffer_trainer_preprocessor(
          agent_trainer,
          # pyre-fixme[6]: Expected `device` for 2nd param but got `str`.
          device,
          env,
      )
      num_batch_per_epoch = embed_rb.size // batch_size
      # FIXME: This has to be wrapped in dataloader
      for epoch in range(num_agent_train_epochs):
          for _ in tqdm(
              range(num_batch_per_epoch), desc=f"epoch {epoch}"
          ):
              batch = embed_rb.sample_transition_batch(
                  batch_size=batch_size
              )
              preprocessed_batch = agent_trainer_preprocessor(batch)
              # FIXME: This should be fitted with Lightning's trainer
              agent_trainer.train(preprocessed_batch)
  
      # evaluate model
      rewards = []
      policy = agent_manager.create_policy(agent_trainer, serving=False)
      # pyre-fixme[6]: Expected `EnvWrapper` for 1st param but got
      #  `StateEmbedEnvironment`.
      agent = Agent.create_for_env(
          embed_env, policy=policy, device=device
      )
      # num_processes=1 needed to avoid workers from dying on CircleCI tests
      rewards = evaluate_for_n_episodes(
          n=num_agent_eval_epochs,
          # pyre-fixme[6]: Expected `EnvWrapper` for 2nd param but got
          #  `StateEmbedEnvironment`.
          env=embed_env,
          agent=agent,
          num_processes=1,
      )
      assert (
          np.mean(rewards) >= passing_score_bar
      ), f"average reward doesn't pass our bar {passing_score_bar}"
      return rewards
  
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 6 for calling a global function with example usage: train_mdnrnn_and_train_on_embedded_env(env_name, embedding_model, num_embedding_train_transitions, seq_len, batch_size, num_embedding_train_epochs, train_model, num_state_embed_transitions, num_agent_train_epochs, num_agent_eval_epochs, use_gpu, passing_score_bar, saved_mdnrnn_path) and returns: rewards