def METHOD_NAME(self, reward_dim):
    env_class = PolicyUnittestEnv
    learning_rate = 1e-1
    iterations = 20
    batch_size = 100
    steps_per_episode = 13
    env = env_class(batch_size, steps_per_episode, reward_dim=reward_dim)
    eval_env = env_class(
        batch_size, steps_per_episode, reward_dim=reward_dim)
    algorithm = create_algorithm(env, learning_rate=learning_rate)
    env.reset()
    eval_env.reset()
    for i in range(iterations):
        algorithm.train_iter()
        eval_env.reset()
        eval_time_step = unroll(eval_env, algorithm, steps_per_episode - 1)
        logging.info("%d reward=%f", i,
                     float(eval_time_step.reward.mean()))
    self.assertAlmostEqual(
        1.0, float(eval_time_step.reward.mean()), delta=1e-1)