def METHOD_NAME(self):
    env = MyEnv(batch_size=3)
    alg1 = create_algorithm(env)
    iter_num = 50
    for _ in range(iter_num):
        alg1.train_iter()
    time_step = common.get_initial_time_step(env)
    state = alg1.get_initial_predict_state(env.batch_size)
    policy_step = alg1.rollout_step(time_step, state)
    action_dist = alf.nest.find_field(policy_step.info,
                                      'action_distribution')[0]
    logits = action_dist.log_prob(torch.arange(3).reshape(3, 1))
    print("logits: ", logits)
    self.assertTrue(torch.all(logits[1, :] > logits[0, :]))
    self.assertTrue(torch.all(logits[1, :] > logits[2, :]))