code: def test_split_preserves_number_of_rollouts(self):
      batch_size = 2
      env, _, _, num_dones = self.init_batch_and_play(
          TEST_ENV_NAME,
          steps_per_epoch=20,
          generate_data=True,
          batch_size=batch_size,
      )
  
      num_rollouts_after_split = sum(
          len(env.current_epoch_rollouts(split))
          for split in self.splits
      )
      # After the end of epoch all environments are reset, which increases number
      # of rollouts by batch size. Number of rollouts could be increased by one
      # in case a rollout is broken on a boundary between the dataset splits.
      self.assertGreaterEqual(
          num_rollouts_after_split, num_dones + batch_size
      )
      self.assertLessEqual(
          num_rollouts_after_split, num_dones + batch_size + 1
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import shutil
  import gym
  from gym.spaces import Box
  from gym.spaces import Discrete
  import numpy
  from tensor2tensor.data_generators import gym_env
  from tensor2tensor.data_generators import problem
  from tensor2tensor.rl.gym_utils import make_gym_env
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling the GymEnvTest obj's underlying member method code with example usage: obj.test_split_preserves_number_of_rollouts() without return types