code: def test_split_preserves_number_of_frames(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_frames = sum(
          len(rollout)
          for split in self.splits
          for rollout in env.current_epoch_rollouts(split)
      )
      # There are 3 frames in every rollout: the initial one and two returned by
      # step(). Additionally there are batch_size observations coming from final
      # reset at the end of epoch.
      self.assertEqual(num_frames, 3 * num_dones + batch_size)
  
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 6 for calling the GymEnvTest obj's underlying member method code with example usage: obj.test_split_preserves_number_of_frames() without return types