code: def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
      """Load T2TGymEnv with data from one epoch.
  
      Args:
        hparams: hparams.
        data_dir: data directory.
        which_epoch_data: data from which epoch to load.
  
      Returns:
        env.
      """
      t2t_env = rl_utils.setup_env(
          hparams,
          batch_size=hparams.real_batch_size,
          max_num_noops=hparams.max_num_noops,
      )
      # Load data.
      if which_epoch_data is not None:
          if which_epoch_data == "last":
              which_epoch_data = infer_last_epoch_num(data_dir)
          assert isinstance(which_epoch_data, int), "{}".format(
              type(which_epoch_data)
          )
          t2t_env.start_new_epoch(which_epoch_data, data_dir)
      else:
          t2t_env.start_new_epoch(-999)
      return t2t_env
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import os
  import re
  import gym
  import numpy
  import six
  from tensor2tensor.models.research.rl import get_policy
  from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
  from tensor2tensor.rl import rl_utils
  from tensor2tensor.rl.envs.simulated_batch_gym_env import FlatBatchEnv
  from tensor2tensor.utils import hparam
  from tensor2tensor.utils import trainer_lib
  from tensor2tensor.utils.misc_utils import camelcase_to_snakecase
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: setup_and_load_epoch(hparams, data_dir, which_epoch_data) and returns: t2t_env