code: def create_embed_rl_dataset(
      env: EnvWrapper,
      memory_network: MemoryNetwork,
      num_state_embed_transitions: int,
      batch_size: int,
      seq_len: int,
      hidden_dim: int,
      use_gpu: bool,
  ):
      assert isinstance(env.action_space, gym.spaces.Discrete)
      assert isinstance(env.observation_space, gym.spaces.Box)
      assert len(env.observation_space.shape) == 1
      logger.info("Starting to create embedded RL Dataset!")
  
      # seqlen+1 because MDNRNN embeds the first seq_len steps and then
      # the embedded state will be concatenated with the last step
      # Ie.. (o1,o2,...,on) -> RNN -> h1,h2,...,hn
      # and we set s_{n+1} = [o_{n+1}, h_n]
      # pyre-fixme[45]: Cannot instantiate abstract class `StateEmbedEnvironment`.
      embed_env = StateEmbedEnvironment(
          gym_env=env,
          mdnrnn=memory_network,
          max_embed_seq_len=seq_len + 1,
      )
      # now create a filled replay buffer of embeddings
      # new obs shape dim = state_dim + hidden_dim
      embed_rb = ReplayBuffer(
          replay_capacity=num_state_embed_transitions,
          batch_size=batch_size,
          stack_size=1,
      )
      random_policy = make_random_policy_for_env(env)
      agent = Agent.create_for_env(env, policy=random_policy)
      fill_replay_buffer(
          env=embed_env,
          replay_buffer=embed_rb,
          desired_size=num_state_embed_transitions,
          agent=agent,
      )
      batch = embed_rb.sample_transition_batch(
          batch_size=num_state_embed_transitions
      )
      state_min = min(batch.state.min(), batch.next_state.min()).item()
      state_max = max(batch.state.max(), batch.next_state.max()).item()
      logger.info(
          f"Finished making embed dataset with size {embed_rb.size}, "
          f"min {state_min}, max {state_max}"
      )
      return embed_rb, state_min, state_max
  
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 5 for calling a global function with example usage: create_embed_rl_dataset(env, memory_network, num_state_embed_transitions, batch_size, seq_len, hidden_dim, use_gpu) and returns: embed_rb,  state_min,  state_max