code: def test_play_env_problem_with_policy(self):
      env = gym_env_problem.GymEnvProblem(
          base_env_name="CartPole-v0",
          batch_size=2,
          reward_range=(-1, 1),
      )
  
      # Let's make sure that at-most 4 observations come to the policy function.
      len_history_for_policy = 4
  
      def policy_fun(observations, lengths, state=None, rng=None):
          del lengths
          b = observations.shape[0]
          # Assert that observations from time-step len_history_for_policy onwards
          # are zeros.
          self.assertTrue(
              np.all(observations[:, len_history_for_policy:, ...] == 0)
          )
          self.assertFalse(
              np.all(observations[:, :len_history_for_policy, ...] == 0)
          )
          a = env.action_space.n
          p = np.random.uniform(size=(b, 1, a))
          p = np.exp(p)
          p = p / np.sum(p, axis=-1, keepdims=True)
          return np.log(p), np.mean(p, axis=-1), state, rng
  
      max_timestep = 15
      num_trajectories = 2
      (
          trajectories,
          _,
          _,
          _,
      ) = env_problem_utils.play_env_problem_with_policy(
          env,
          policy_fun,
          num_trajectories=num_trajectories,
          max_timestep=max_timestep,
          len_history_for_policy=len_history_for_policy,
      )
  
      self.assertEqual(num_trajectories, len(trajectories))
  
      # Check shapes within trajectories.
      traj = trajectories[0]
      T = traj[1].shape[0]  # pylint: disable=invalid-name
      self.assertEqual((T + 1, 4), traj[0].shape)  # (4,) is OBS
      self.assertEqual((T,), traj[2].shape)
      self.assertEqual(T, len(traj[4]["log_prob_actions"]))
      self.assertEqual(T, len(traj[4]["value_predictions"]))
      self.assertLessEqual(T, max_timestep)
  
      traj = trajectories[1]
      T = traj[1].shape[0]  # pylint: disable=invalid-name
      self.assertEqual((T + 1, 4), traj[0].shape)
      self.assertEqual((T,), traj[2].shape)
      self.assertEqual(T, len(traj[4]["log_prob_actions"]))
      self.assertEqual(T, len(traj[4]["value_predictions"]))
      self.assertLessEqual(T, max_timestep)
  
file_path: tensor2tensor/tensor2tensor/envs/env_problem_utils_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.envs import env_problem_utils
  from tensor2tensor.envs import gym_env_problem
  from tensor2tensor.envs import tic_tac_toe_env
  from tensor2tensor.envs import tic_tac_toe_env_problem
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the EnvProblemUtilsTest obj's underlying member method code with example usage: obj.test_play_env_problem_with_policy() and returns: np,  np,  state,  rng