code: def create_string_game_data(
      dataset_size=10000,
      training_data_ratio=0.9,
      filter_short_sequence=False,
  ):
      SEQ_LEN = 6
      NUM_ACTION = 2
      NUM_MDP_PER_BATCH = 5
  
      env = Gym(env_name="StringGame-v0", set_max_steps=SEQ_LEN)
      df = create_df_from_replay_buffer(
          env=env,
          problem_domain=ProblemDomain.DISCRETE_ACTION,
          desired_size=dataset_size,
          multi_steps=None,
          ds="2020-10-10",
      )
  
      if filter_short_sequence:
          batch_size = NUM_MDP_PER_BATCH
          time_diff = torch.ones(SEQ_LEN, batch_size)
          valid_step = (
              SEQ_LEN
              * torch.ones(batch_size, dtype=torch.int64)[:, None]
          )
          not_terminal = torch.Tensor(
              [0 if i == SEQ_LEN - 1 else 1 for i in range(SEQ_LEN)]
          )
          not_terminal = torch.transpose(
              not_terminal.tile(NUM_MDP_PER_BATCH, 1), 0, 1
          )
      else:
          batch_size = NUM_MDP_PER_BATCH * SEQ_LEN
          time_diff = torch.ones(SEQ_LEN, batch_size)
          valid_step = torch.arange(SEQ_LEN, 0, -1).tile(
              NUM_MDP_PER_BATCH
          )[:, None]
          not_terminal = torch.transpose(
              torch.tril(
                  torch.ones(SEQ_LEN, SEQ_LEN), diagonal=-1
              ).tile(NUM_MDP_PER_BATCH, 1),
              0,
              1,
          )
  
      num_batches = int(dataset_size / SEQ_LEN / NUM_MDP_PER_BATCH)
      batches = [None for _ in range(num_batches)]
      batch_count, batch_seq_count = 0, 0
      batch_reward = torch.zeros(SEQ_LEN, batch_size)
      batch_action = torch.zeros(SEQ_LEN, batch_size, NUM_ACTION)
      batch_state = torch.zeros(SEQ_LEN, batch_size, NUM_ACTION)
      for mdp_id in sorted(set(df.mdp_id)):
          mdp = df[df["mdp_id"] == mdp_id].sort_values(
              "sequence_number", ascending=True
          )
          if len(mdp) != SEQ_LEN:
              continue
  
          all_step_reward = torch.Tensor(list(mdp["reward"]))
          all_step_state = torch.Tensor(
              [list(s.values()) for s in mdp["state_features"]]
          )
          all_step_action = torch.zeros_like(all_step_state)
          all_step_action[
              torch.arange(SEQ_LEN), [int(a) for a in mdp["action"]]
          ] = 1.0
  
          for j in range(SEQ_LEN):
              if filter_short_sequence and j > 0:
                  break
  
              reward = torch.zeros_like(all_step_reward)
              reward[: SEQ_LEN - j] = all_step_reward[-(SEQ_LEN - j) :]
              batch_reward[:, batch_seq_count] = reward
  
              state = torch.zeros_like(all_step_state)
              state[: SEQ_LEN - j] = all_step_state[-(SEQ_LEN - j) :]
              batch_state[:, batch_seq_count] = state
  
              action = torch.zeros_like(all_step_action)
              action[: SEQ_LEN - j] = all_step_action[-(SEQ_LEN - j) :]
              batch_action[:, batch_seq_count] = action
  
              batch_seq_count += 1
  
          if batch_seq_count == batch_size:
              batches[batch_count] = rlt.MemoryNetworkInput(
                  reward=batch_reward,
                  action=rlt.FeatureData(float_features=batch_action),
                  state=rlt.FeatureData(float_features=batch_state),
                  next_state=rlt.FeatureData(
                      float_features=torch.zeros_like(batch_state)
                  ),  # fake, not used anyway
                  not_terminal=not_terminal,
                  time_diff=time_diff,
                  valid_step=valid_step,
                  step=None,
              )
              batch_count += 1
              batch_seq_count = 0
              batch_reward = torch.zeros_like(batch_reward)
              batch_action = torch.zeros_like(batch_action)
              batch_state = torch.zeros_like(batch_state)
      assert batch_count == num_batches
  
      num_training_batches = int(training_data_ratio * num_batches)
      training_data = DataLoader(
          batches[:num_training_batches], collate_fn=lambda x: x[0]
      )
      eval_data = DataLoader(
          batches[num_training_batches:], collate_fn=lambda x: x[0]
      )
      return training_data, eval_data
  
imports: import logging
  import os
  import random
  import unittest
  from typing import Optional
  import numpy
  import pytorch_lightning
  import torch
  import torch.nn
  from parameterized import parameterized
  from reagent.core import types
  from reagent.core.parameters import NormalizationData, NormalizationParameters, ProblemDomain, Seq2RewardTrainerParameters
  from reagent.gym.envs import Gym
  from reagent.gym.utils import create_df_from_replay_buffer
  from reagent.models.seq2reward_model import Seq2RewardNetwork
  from reagent.net_builder.value.fully_connected import FullyConnected
  from reagent.prediction.predictor_wrapper import FAKE_STATE_ID_LIST_FEATURES, FAKE_STATE_ID_SCORE_LIST_FEATURES, Seq2RewardPlanShortSeqWithPreprocessor, Seq2RewardWithPreprocessor
  from reagent.preprocessing.identify_types import DO_NOT_PREPROCESS
  from reagent.preprocessing.preprocessor import Preprocessor
  from reagent.training.utils import gen_permutations
  from reagent.training.world_model.compress_model_trainer import CompressModelTrainer
  from reagent.training.world_model.seq2reward_trainer import get_Q, Seq2RewardTrainer
  from torch.utils.data import DataLoader
label_desc: Use image node_id 1 for calling a global function with example usage: create_string_game_data(dataset_size, training_data_ratio, filter_short_sequence) and returns: training_data,  eval_data