def METHOD_NAME(self, action_type):
    batch_size = 100
    steps_per_episode = 13
    env = PolicyUnittestEnv(
        batch_size, steps_per_episode, action_type=action_type)
    time_step = env.reset()
    for _ in range(10):
        for s in range(steps_per_episode):
            if s == 0:
                step_type = StepType.FIRST
                discount = 1.0
            elif s == steps_per_episode - 1:
                step_type = StepType.LAST
                discount = 0.0
            else:
                step_type = StepType.MID
                discount = 1.0
            if s == 0:
                reward = torch.zeros(batch_size)
            else:
                reward = (action == prev_observation.to(torch.int64)).to(
                    torch.float32)
                reward = reward.reshape(batch_size)
            self.assertEqual(time_step.step_type,
                             torch.full([batch_size], step_type))
            self.assertEqual(time_step.reward, reward)
            self.assertEqual(time_step.discount,
                             torch.full([batch_size], discount))
            action = torch.randint(0, 2, (batch_size, 1))
            prev_observation = time_step.observation
            time_step = env.step(action)