Runs the training loop for PPO, with fixed policy and value nets.

def training_loop(env=None,
                  env_name="CartPole-v0",
                  epochs=EPOCHS,
                  policy_net_fun=None,
                  value_net_fun=None,
                  policy_and_value_net_fun=None,
                  policy_optimizer_fun=None,
                  value_optimizer_fun=None,
                  policy_and_value_optimizer_fun=None,
                  batch_size=BATCH_TRAJECTORIES,
                  num_optimizer_steps=NUM_OPTIMIZER_STEPS,
                  print_every_optimizer_steps=PRINT_EVERY_OPTIMIZER_STEP,
                  boundary=20,
                  max_timestep=None,
                  random_seed=None,
                  gamma=GAMMA,
                  lambda_=LAMBDA,
                  epsilon=EPSILON,
                  c1=1.0,
                  c2=0.01):
  """Runs the training loop for PPO, with fixed policy and value nets."""
  jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)

  value_losses = []
  ppo_objective = []
  combined_losses = []
  average_rewards = []

  env = env if env is not None else gym.make(env_name)

  # Batch Observations Shape = [-1, -1] + OBS, because we will eventually call
  # policy and value networks on shape [B, T] +_OBS
  batch_observations_shape = (-1, -1) + env.observation_space.shape

  assert isinstance(env.action_space, gym.spaces.Discrete)
  num_actions = env.action_space.n

  policy_and_value_net_params, policy_and_value_net_apply = None, None
  policy_and_value_opt_state, policy_and_value_opt_update = None, None
  policy_net_params, policy_net_apply = None, None
  value_net_params, value_net_apply = None, None
  if policy_and_value_net_fun is not None:
    jax_rng_key, subkey = jax_random.split(jax_rng_key)

    # Initialize the policy and value network.
    policy_and_value_net_params, policy_and_value_net_apply = (
        policy_and_value_net_fun(subkey, batch_observations_shape, num_actions))

    # Initialize the optimizers.
    policy_and_value_opt_state, policy_and_value_opt_update = (
        policy_and_value_optimizer_fun(policy_and_value_net_params))
  else:
    # Initialize the policy and value functions.
    assert policy_net_fun and value_net_fun
    jax_rng_key, key1, key2 = jax_random.split(jax_rng_key, num=3)

    policy_net_params, policy_net_apply = policy_net_fun(
        key1, batch_observations_shape, num_actions)
    value_net_params, value_net_apply = value_net_fun(key2,
                                                      batch_observations_shape,
                                                      num_actions)

    # Initialize the optimizers.
    ppo_opt_state, ppo_opt_update = policy_optimizer_fun(policy_net_params)
    value_opt_state, value_opt_update = value_optimizer_fun(value_net_params)

  # A function that will call the appropriate policy function with parameters.
  def get_policy_output(observations):
    if policy_net_apply is not None:
      assert policy_net_params
      return policy_net_apply(observations, policy_net_params)

    assert policy_and_value_net_apply and policy_and_value_net_params
    policy_predictions, unused_value_predictions = policy_and_value_net_apply(
        observations, policy_and_value_net_params)
    return policy_predictions

  for i in range(epochs):
    t = time.time()
    t0 = t
    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
    trajs = collect_trajectories(
        env,
        policy_fun=get_policy_output,
        num_trajectories=batch_size,
        policy=POLICY,
        max_timestep=max_timestep,
        epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.

    avg_reward = float(sum(np.sum(traj[2]) for traj in trajs)) / len(trajs)
    max_reward = max(np.sum(traj[2]) for traj in trajs)
    min_reward = min(np.sum(traj[2]) for traj in trajs)
    average_rewards.append(avg_reward)

    logging.vlog(1, "Rewards average=[%0.2f], max=[%0.2f], min=[%0.2f]",
                 avg_reward, max_reward, min_reward)
    logging.vlog(1, "Collecting trajectories took %0.2f msec.", get_time(t))
    logging.vlog(1,
                 "Trajectory Length average=[%0.2f], max=[%0.2f], min=[%0.2f]",
                 float(sum(len(traj[0]) for traj in trajs)) / len(trajs),
                 max(len(traj[0]) for traj in trajs),
                 min(len(traj[0]) for traj in trajs))

    t = time.time()
    (_, reward_mask, padded_observations, padded_actions,
     padded_rewards) = pad_trajectories(trajs, boundary=boundary)

    logging.vlog(1, "Padding trajectories took %0.2f msec.", get_time(t))
    logging.vlog(1, "Padded Observations' shape [%s]",
                 str(padded_observations.shape))
    logging.vlog(1, "Padded Actions' shape [%s]", str(padded_actions.shape))
    logging.vlog(1, "Padded Rewards' shape [%s]", str(padded_rewards.shape))

    # Some assertions.
    B, T = padded_actions.shape  # pylint: disable=invalid-name
    assert (B, T) == padded_rewards.shape
    assert (B, T) == reward_mask.shape
    assert (B, T + 1) == padded_observations.shape[:2]
    assert (B, T + 1) + env.observation_space.shape == padded_observations.shape

    # Linear annealing from 0.1 to 0.0
    epsilon_schedule = epsilon if epochs == 1 else epsilon * (1.0 -
                                                              (i /
                                                               (epochs - 1)))

    # Compute value and ppo losses.
    cur_value_loss, cur_ppo_loss, cur_combined_loss = None, None, None
    if policy_and_value_net_apply is not None:
      t = time.time()
      cur_combined_loss, cur_ppo_loss, cur_value_loss, _ = (
          combined_loss(
              policy_and_value_net_params,
              policy_and_value_net_params,
              policy_and_value_net_apply,
              padded_observations,
              padded_actions,
              padded_rewards,
              reward_mask,
              gamma=gamma,
              lambda_=lambda_,
              epsilon=epsilon_schedule,
              c1=c1,
              c2=c2))
      logging.vlog(
          1, "Calculating P&V loss [%10.2f(%10.2f, %10.2f)] took %0.2f msec.",
          cur_combined_loss, cur_value_loss, cur_ppo_loss, get_time(t))
    else:
      t = time.time()
      cur_value_loss = value_loss(
          value_net_apply,
          value_net_params,
          padded_observations,
          padded_rewards,
          reward_mask,
          gamma=gamma)

      logging.vlog(1, "Calculating value loss took %0.2f msec.", get_time(t))

      t = time.time()
      cur_ppo_loss = ppo_loss(
          policy_net_apply,
          policy_net_params,
          policy_net_params,
          value_net_apply,
          value_net_params,
          padded_observations,
          padded_actions,
          padded_rewards,
          reward_mask,
          gamma=gamma,
          lambda_=lambda_,
          epsilon=epsilon_schedule)
      logging.vlog(1, "Calculating PPO loss took %0.2f msec.", get_time(t))

    value_losses.append(cur_value_loss)
    ppo_objective.append(-1.0 * cur_ppo_loss)
    combined_losses.append(cur_combined_loss)

    if policy_and_value_net_apply:
      logging.vlog(1, "Policy and Value Optimization")
      t1 = time.time()
      for j in range(num_optimizer_steps):
        t = time.time()
        # Update the optimizer state.
        policy_and_value_opt_state = policy_and_value_opt_step(
            j,
            policy_and_value_opt_state,
            policy_and_value_opt_update,
            policy_and_value_net_apply,
            policy_and_value_net_params,
            padded_observations,
            padded_actions,
            padded_rewards,
            reward_mask,
            c1=c1,
            c2=c2,
            gamma=gamma,
            lambda_=lambda_,
            epsilon=epsilon_schedule)
        t2 = time.time()
        # Get the new params.
        new_policy_and_value_net_params = trax_opt.get_params(
            policy_and_value_opt_state)
        if ((j + 1) %
            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
          # Compute and log the loss.
          (loss_combined, loss_ppo, loss_value, unused_entropy_bonus) = (
              combined_loss(
                  new_policy_and_value_net_params,
                  policy_and_value_net_params,  # old params
                  policy_and_value_net_apply,
                  padded_observations,
                  padded_actions,
                  padded_rewards,
                  reward_mask,
                  gamma=gamma,
                  lambda_=lambda_,
                  epsilon=epsilon_schedule,
                  c1=c1,
                  c2=c2))
          logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                       get_time(t, t2))
          logging.vlog(
              1,
              "Combined Loss(value, ppo) [%10.2f] -> [%10.2f(%10.2f,%10.2f)]",
              cur_combined_loss, loss_combined, loss_value, loss_ppo)
        # Update the params.
        policy_and_value_net_params = new_policy_and_value_net_params

      logging.vlog(
          1, "Total PPO loss reduction [%0.2f]%%",
          (100 *
           (cur_combined_loss - loss_combined) / np.abs(cur_combined_loss)))

      logging.info(
          "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], Combined"
          " Loss(value, ppo) [%10.2f(%10.2f,%10.2f)], took [%10.2f msec]",
          i, min_reward, max_reward, avg_reward, loss_combined, loss_value,
          loss_ppo, get_time(t1))
    else:
      # Run optimizers.
      logging.vlog(1, "PPO Optimization")
      t1 = time.time()

      for j in range(num_optimizer_steps):
        t = time.time()
        # Update the optimizer state.
        ppo_opt_state = ppo_opt_step(
            j,
            ppo_opt_state,
            ppo_opt_update,
            policy_net_apply,
            policy_net_params,
            value_net_apply,
            value_net_params,
            padded_observations,
            padded_actions,
            padded_rewards,
            reward_mask,
            gamma=gamma,
            lambda_=lambda_,
            epsilon=epsilon_schedule,
        )
        t2 = time.time()
        # Get the new params.
        new_policy_net_params = trax_opt.get_params(ppo_opt_state)
        if ((j + 1) %
            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
          new_ppo_loss = ppo_loss(
              policy_net_apply,
              new_policy_net_params,
              policy_net_params,
              value_net_apply,
              value_net_params,
              padded_observations,
              padded_actions,
              padded_rewards,
              reward_mask,
              gamma=gamma,
              lambda_=lambda_,
              epsilon=epsilon_schedule,
          )
          logging.vlog(1, "One PPO grad desc took: %0.2f msec", get_time(t, t2))
          logging.vlog(1, "PPO loss [%10.2f] -> [%10.2f]", cur_ppo_loss,
                       new_ppo_loss)
        # Update the params.
        policy_net_params = new_policy_net_params

      logging.vlog(1, "Total PPO loss reduction [%0.2f]%%",
                   (100 * (cur_ppo_loss - new_ppo_loss) / np.abs(cur_ppo_loss)))

      logging.vlog(1, "Value Optimization")

      for j in range(num_optimizer_steps):
        t = time.time()
        value_opt_state = value_opt_step(
            j,
            value_opt_state,
            value_opt_update,
            value_net_apply,
            padded_observations,
            padded_rewards,
            reward_mask,
            gamma=gamma)
        t2 = time.time()
        value_net_params = trax_opt.get_params(value_opt_state)
        if ((j + 1) %
            print_every_optimizer_steps == 0) or (j == num_optimizer_steps - 1):
          new_value_loss = value_loss(
              value_net_apply,
              value_net_params,
              padded_observations,
              padded_rewards,
              reward_mask,
              gamma=gamma)
          logging.vlog(1, "One value grad desc took: %0.2f msec",
                       get_time(t, t2))
          logging.vlog(1, "Value loss [%10.2f] -> [%10.2f]", cur_value_loss,
                       new_value_loss)
      logging.vlog(1, "Total value loss reduction [%0.2f]%%",
                   (100 *
                    (cur_value_loss - new_value_loss) / np.abs(cur_value_loss)))

      logging.vlog(1, "Grad desc took %0.2f msec", get_time(t1))

      # Set the optimized params to new params.
      policy_net_params = trax_opt.get_params(ppo_opt_state)
      value_net_params = trax_opt.get_params(value_opt_state)

      logging.info(
          "Epoch [% 6d], Reward[min, max, avg] [%10.2f,%10.2f,%10.2f], "
          "ppo loss [%10.2f], value loss [%10.2f], took [%10.2f msec]", i,
          min_reward, max_reward, avg_reward, new_ppo_loss, new_value_loss,
          get_time(t0))

  # Log the parameters, just for the sake of it.
  if policy_net_params:
    log_params(policy_net_params, "policy_net_params")
  if value_net_params:
    log_params(value_net_params, "value_net_params")
  if policy_and_value_net_params:
    log_params(policy_and_value_net_params, "policy_and_value_net_params")

  if value_losses:
    logging.vlog(1, "value_losses: %s", np.stack(value_losses))
  if ppo_objective:
    logging.vlog(1, "ppo_objective: %s", np.stack(ppo_objective))
  if average_rewards:
    logging.vlog(1, "average_rewards: %s", average_rewards)

  return ((policy_net_params, value_net_params), average_rewards,
          np.stack(value_losses), np.stack(ppo_objective))