def play_env(self,
               env=None,
               nsteps=100,
               base_env_name=None,
               batch_size=5,
               reward_range=None):
    """Creates `GymEnvProblem` with the given arguments and plays it randomly.

    Args:
      env: optional env.
      nsteps: plays the env randomly for nsteps.
      base_env_name: passed to GymEnvProblem's init.
      batch_size: passed to GymEnvProblem's init.
      reward_range: passed to GymEnvProblem's init.

    Returns:
      tuple of gym_env_problem, number of trajectories done,
      number of trajectories done in the last step.
    """

    if env is None:
      env = gym_env_problem.GymEnvProblem(
          base_env_name=base_env_name,
          batch_size=batch_size,
          reward_range=reward_range)
      # Usually done by a registered subclass, we do this manually in the test.
      env.name = base_env_name

    # Reset all environments.
    env.reset()

    # Play for some steps to generate data.
    num_dones = 0
    num_dones_in_last_step = 0
    for _ in range(nsteps):
      # Sample actions.
      actions = np.stack([env.action_space.sample() for _ in range(batch_size)])
      # Step through it.
      _, _, dones, _ = env.step(actions)
      # Get the indices where we are done ...
      done_indices = env_problem_utils.done_indices(dones)
      # ... and reset those.
      env.reset(indices=done_indices)
      # count the number of dones we got, in this step and overall.
      num_dones_in_last_step = sum(dones)
      num_dones += num_dones_in_last_step

    return env, num_dones, num_dones_in_last_step
Creates `GymEnvProblem` with the given arguments and plays it randomly.

    Args:
      env: optional env.
      nsteps: plays the env randomly for nsteps.
      base_env_name: passed to GymEnvProblem's init.
      batch_size: passed to GymEnvProblem's init.
      reward_range: passed to GymEnvProblem's init.

    Returns:
      tuple of gym_env_problem, number of trajectories done,
      number of trajectories done in the last step.
    
tensor2tensor/envs/gym_env_problem_test.py