# coding=utf-8
# Copyright 2019 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""PPO in JAX.

Notation:

B, scalar  - batch size
T, scalar  - number of time-steps in a trajectory, or the value of the padded
             time-step dimension.
OBS, tuple - shape of a singular observation from the environment.
             Ex: For CartPole-v0 this is (4,) and Pong-v0 it's (210, 160, 3)
A, scalar  - Number of actions, assuming a discrete space.

Policy and Value function signatures:

Policy            Function :: [B, T] + OBS ->  [B, T, A]
Value             Function :: [B, T] + OBS ->  [B, T, 1]
Policy and Value  Function :: [B, T] + OBS -> ([B, T, A], [B, T, 1])

i.e. the policy net should take a batch of *trajectories* and at each time-step
in each batch deliver a probability distribution over actions.

NOTE: It doesn't return logits, rather the expectation is that it returns
log-probabilities instead.

NOTE: The policy and value functions need to take care to not take into account
future time-steps while deciding the actions (or value) for the current
time-step.

Policy and Value Function produces a tuple of the expected output of a policy
function and a value function.

"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import os
import pickle
import time

from absl import logging
import gym
from jax import grad
from jax import jit
from jax import lax
from jax import numpy as np
from jax import random as jax_random
import numpy as onp
from tensor2tensor.envs import env_problem
from tensor2tensor.envs import env_problem_utils
from tensor2tensor.trax import jaxboard
from tensor2tensor.trax import layers
from tensor2tensor.trax import optimizers as trax_opt
from tensor2tensor.trax import trax
from tensorflow.io import gfile

DEBUG_LOGGING = False
GAMMA = 0.99
LAMBDA = 0.95
EPSILON = 0.1
EPOCHS = 50  # 100
NUM_OPTIMIZER_STEPS = 100
PRINT_EVERY_OPTIMIZER_STEP = 20
BATCH_TRAJECTORIES = 32


def policy_and_value_net(rng_key,
                         batch_observations_shape,
                         num_actions,
                         bottom_layers_fn=None,
                         two_towers=True):
  """A policy and value net function."""

  # Layers.

  # Now, with the current logits, one head computes action probabilities and the
  # other computes the value function.
  # NOTE: The LogSoftmax instead of the Softmax because of numerical stability.

  net = None
  if not two_towers:
    tower = [] if bottom_layers_fn is None else bottom_layers_fn()
    tower.extend([
        layers.Branch(
            layers.Serial(layers.Dense(num_actions), layers.LogSoftmax()),
            layers.Dense(1))
    ])
    net = layers.Serial(*tower)
  else:
    tower1 = [] if bottom_layers_fn is None else bottom_layers_fn()
    tower2 = [] if bottom_layers_fn is None else bottom_layers_fn()

    tower1.extend([layers.Dense(num_actions), layers.LogSoftmax()])
    tower2.extend([layers.Dense(1)])

    net = layers.Branch(
        layers.Serial(*tower1),
        layers.Serial(*tower2),
    )
  assert net
  return net.initialize(batch_observations_shape, rng_key), net


def optimizer_fun(net_params, step_size=1e-3):
  opt = trax_opt.Adam(step_size=step_size, b1=0.9, b2=0.999, eps=1e-08)
  opt_init = lambda x: (x, opt.tree_init(x))
  opt_update = lambda i, g, s: opt.tree_update(i, g, s[0], s[1])
  get_params = lambda x: x[0]
  opt_state = opt_init(net_params)
  return opt_state, opt_update, get_params


# Should this be collect 'n' trajectories, or
# Run the env for 'n' steps and take completed trajectories, or
# Any other option?
# TODO(afrozm): Replace this with EnvProblem?
def collect_trajectories(env,
                         policy_fun,
                         num_trajectories=1,
                         policy=env_problem_utils.CATEGORICAL_SAMPLING,
                         max_timestep=None,
                         boundary=20,
                         epsilon=0.1,
                         reset=True,
                         rng=None):
  """Collect trajectories with the given policy net and behaviour.

  Args:
    env: A gym env interface, for now this is not-batched.
    policy_fun: observations(B,T+1) -> log-probabs(B,T+1, A) callable.
    num_trajectories: int, number of trajectories.
    policy: string, "greedy", "epsilon-greedy", or "categorical-sampling" i.e.
      how to use the policy_fun to return an action.
    max_timestep: int or None, the index of the maximum time-step at which we
      return the trajectory, None for ending a trajectory only when env returns
      done.
    boundary: int, boundary for padding, used in EnvProblem envs.
    epsilon: float, the epsilon for `epsilon-greedy` policy.
    reset: bool, true if we want to reset the envs. The envs are also reset if
      max_max_timestep is None or < 0
    rng: jax rng, splittable.

  Returns:
    A tuple (trajectory, number of trajectories that are done)
    trajectory: list of (observation, action, reward) tuples, where each element
    `i` is a tuple of numpy arrays with shapes as follows:
    observation[i] = (B, T_i + 1)
    action[i] = (B, T_i)
    reward[i] = (B, T_i)
  """

  assert isinstance(env, env_problem.EnvProblem)
  # This is an env_problem, run its collect function.
  return env_problem_utils.play_env_problem_with_policy(
      env,
      policy_fun,
      num_trajectories=num_trajectories,
      max_timestep=max_timestep,
      boundary=boundary,
      policy_sampling=policy,
      eps=epsilon,
      reset=reset,
      rng=rng)


# This function can probably be simplified, ask how?
# Can we do something much simpler than lax.pad, maybe np.pad?
# Others?


def get_padding_value(dtype):
  """Returns the padding value given a dtype."""
  padding_value = None
  if dtype == np.uint8:
    padding_value = np.uint8(0)
  elif dtype == np.uint16:
    padding_value = np.uint16(0)
  elif dtype == np.float32 or dtype == np.float64:
    padding_value = 0.0
  else:
    padding_value = 0
  assert padding_value is not None
  return padding_value


# TODO(afrozm): Use np.pad instead and make jittable?
def pad_trajectories(trajectories, boundary=20):
  """Pad trajectories to a bucket length that is a multiple of boundary.

  Args:
    trajectories: list[(observation, actions, rewards)], where each observation
      is shaped (t+1,) + OBS and actions & rewards are shaped (t,), with the
      length of the list being B (batch size).
    boundary: int, bucket length, the actions and rewards are padded to integer
      multiples of boundary.

  Returns:
    tuple: (padding lengths, reward_mask, padded_observations, padded_actions,
        padded_rewards) where padded_observations is shaped (B, T+1) + OBS and
        padded_actions, padded_rewards & reward_mask are shaped (B, T).
        Where T is max(t) rounded up to an integer multiple of boundary.
        padded_length is how much padding we've added and
        reward_mask is 1s for actual rewards and 0s for the padding.
  """

  # Let's compute max(t) over all trajectories.
  t_max = max(r.shape[0] for (_, _, r) in trajectories)

  # t_max is rounded to the next multiple of `boundary`
  boundary = int(boundary)
  bucket_length = boundary * int(np.ceil(float(t_max) / boundary))

  # So all obs will be padded to t_max + 1 and actions and rewards to t_max.
  padded_observations = []
  padded_actions = []
  padded_rewards = []
  padded_lengths = []
  reward_masks = []
  for (o, a, r) in trajectories:
    # Determine the amount to pad, this holds true for obs, actions and rewards.
    num_to_pad = bucket_length + 1 - o.shape[0]
    padded_lengths.append(num_to_pad)
    if num_to_pad == 0:
      padded_observations.append(o)
      padded_actions.append(a)
      padded_rewards.append(r)
      reward_masks.append(onp.ones_like(r, dtype=np.int32))
      continue

    # First pad observations.
    padding_config = [(0, num_to_pad, 0)]
    for _ in range(o.ndim - 1):
      padding_config.append((0, 0, 0))
    padding_config = tuple(padding_config)

    padding_value = get_padding_value(o.dtype)
    action_padding_value = get_padding_value(a.dtype)
    reward_padding_value = get_padding_value(r.dtype)

    padded_obs = lax.pad(o, padding_value, padding_config)
    padded_observations.append(padded_obs)

    # Now pad actions and rewards.
    assert a.ndim == 1 and r.ndim == 1
    padding_config = ((0, num_to_pad, 0),)

    padded_action = lax.pad(a, action_padding_value, padding_config)
    padded_actions.append(padded_action)
    padded_reward = lax.pad(r, reward_padding_value, padding_config)
    padded_rewards.append(padded_reward)

    # Also create the mask to use later.
    reward_mask = onp.ones_like(r, dtype=np.int32)
    reward_masks.append(lax.pad(reward_mask, 0, padding_config))

  return padded_lengths, np.stack(reward_masks), np.stack(
      padded_observations), np.stack(padded_actions), np.stack(padded_rewards)


# TODO(afrozm): JAX-ify this, this is too slow for pong.
def rewards_to_go(rewards, mask, gamma=0.99):
  r"""Computes rewards to go.

  Reward to go is defined as follows, the discounted reward that we have to
  yet collect, going forward from this point, i.e.:

  r2g_t = \sum_{l=0}^{\infty} (\gamma^{l} * reward_{t+l})

  Args:
    rewards: np.ndarray of shape (B, T) of rewards.
    mask: np.ndarray of shape (B, T) of mask for the rewards.
    gamma: float, discount factor.

  Returns:
    rewards to go, np.ndarray of shape (B, T).
  """
  B, T = rewards.shape  # pylint: disable=invalid-name,unused-variable

  masked_rewards = rewards * mask  # (B, T)

  # We use the following recurrence relation, derived from the equation above:
  #
  # r2g[t+1] = (r2g[t] - r[t]) / gamma
  #
  # This means we'll need to calculate r2g[0] first and then r2g[1] and so on ..
  #
  # **However** this leads to overflows for long sequences: r2g[t] - r[t] > 0
  # and gamma < 1.0, so the division keeps increasing.
  #
  # So we just run the recurrence in reverse, i.e.
  #
  # r2g[t] = r[t] + (gamma*r2g[t+1])
  #
  # This is much better, but might have lost updates since the (small) rewards
  # at earlier time-steps may get added to a (very?) large sum.

  # Compute r2g_{T-1} at the start and then compute backwards in time.
  r2gs = [masked_rewards[:, -1]]

  # Go from T-2 down to 0.
  for t in reversed(range(T - 1)):
    r2gs.append(masked_rewards[:, t] + (gamma * r2gs[-1]))

  # The list should have length T.
  assert T == len(r2gs)

  # First we stack them in the correct way to make it (B, T), but these are
  # still from newest (T-1) to oldest (0), so then we flip it on time axis.
  return np.flip(np.stack(r2gs, axis=1), axis=1)


@jit
def value_loss_given_predictions(value_prediction,
                                 rewards,
                                 reward_mask,
                                 gamma=0.99,
                                 epsilon=0.2,
                                 value_prediction_old=None):
  """Computes the value loss given the prediction of the value function.

  Args:
    value_prediction: np.ndarray of shape (B, T+1, 1)
    rewards: np.ndarray of shape (B, T) of rewards.
    reward_mask: np.ndarray of shape (B, T), the mask over rewards.
    gamma: float, discount factor.
    epsilon: float, clip-fraction, used if value_value_prediction_old isn't None
    value_prediction_old: np.ndarray of shape (B, T+1, 1) of value predictions
      using the old parameters. If provided, we incorporate this in the loss as
      well. This is from the OpenAI baselines implementation.

  Returns:
    The average L2 value loss, averaged over instances where reward_mask is 1.
  """

  B, T = rewards.shape  # pylint: disable=invalid-name
  assert (B, T) == reward_mask.shape
  assert (B, T + 1, 1) == value_prediction.shape

  value_prediction = np.squeeze(value_prediction, axis=2)  # (B, T+1)
  value_prediction = value_prediction[:, :-1] * reward_mask  # (B, T)
  r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
  loss = (value_prediction - r2g)**2

  # From the baselines implementation.
  if value_prediction_old is not None:
    value_prediction_old = np.squeeze(value_prediction_old, axis=2)  # (B, T+1)
    value_prediction_old = value_prediction_old[:, :-1] * reward_mask  # (B, T)

    v_clipped = value_prediction_old + np.clip(
        value_prediction - value_prediction_old, -epsilon, epsilon)
    v_clipped_loss = (v_clipped - r2g)**2
    loss = np.maximum(v_clipped_loss, loss)

  # Take an average on only the points where mask != 0.
  return np.sum(loss) / np.sum(reward_mask)


# TODO(afrozm): JAX-ify this, this is too slow for pong.
def deltas(predicted_values, rewards, mask, gamma=0.99):
  r"""Computes TD-residuals from V(s) and rewards.

  Where a `delta`, i.e. a td-residual is defined as:

  delta_{b,t} = r_{b,t} + \gamma * v_{b,t+1} - v_{b,t}.

  Args:
    predicted_values: ndarray of shape (B, T+1). NOTE: Expects axis 2 was
      squeezed. These represent V(s_bt) for b < B and t < T+1
    rewards: ndarray of shape (B, T) of rewards.
    mask: ndarray of shape (B, T) of mask for rewards.
    gamma: float, discount factor.

  Returns:
    ndarray of shape (B, T) of one-step TD-residuals.
  """

  # `d`s are basically one-step TD residuals.
  d = []
  _, T = rewards.shape  # pylint: disable=invalid-name
  for t in range(T):
    d.append(rewards[:, t] + (gamma * predicted_values[:, t + 1]) -
             predicted_values[:, t])

  return np.array(d).T * mask


def gae_advantages(td_deltas, mask, lambda_=0.95, gamma=0.99):
  r"""Computes the GAE advantages given the one step TD-residuals.

  The formula for a GAE advantage estimator is as follows:

  A_{bt} = \sum_{l=0}^{\infty}(\gamma * \lambda)^{l}(\delta_{b,t+l}).

  Internally we just call rewards_to_go, since it is the same computation.

  Args:
    td_deltas: np.ndarray of shape (B, T) of one step TD-residuals.
    mask: np.ndarray of shape (B, T) of mask for the residuals. It maybe the
      case that the `td_deltas` are already masked correctly since they are
      produced by `deltas(...)`
    lambda_: float, lambda parameter for GAE estimators.
    gamma: float, lambda parameter for GAE estimators.

  Returns:
    GAE advantage estimates.
  """

  return rewards_to_go(td_deltas, mask, lambda_ * gamma)


def chosen_probabs(probab_observations, actions):
  """Picks out the probabilities of the actions along batch and time-steps.

  Args:
    probab_observations: ndarray of shape `[B, T+1, A]`, where
      probab_observations[b, t, i] contains the log-probability of action = i at
      the t^th time-step in the b^th trajectory.
    actions: ndarray of shape `[B, T]`, with each entry in [0, A) denoting which
      action was chosen in the b^th trajectory's t^th time-step.

  Returns:
    `[B, T]` ndarray with the log-probabilities of the chosen actions.
  """
  B, T = actions.shape  # pylint: disable=invalid-name
  assert (B, T + 1) == probab_observations.shape[:2]
  return probab_observations[np.arange(B)[:, None], np.arange(T), actions]


def compute_probab_ratios(p_new, p_old, actions, reward_mask):
  """Computes the probability ratios for each time-step in a trajectory.

  Args:
    p_new: ndarray of shape [B, T+1, A] of the log-probabilities that the policy
      network assigns to all the actions at each time-step in each batch using
      the old parameters.
    p_old: ndarray of shape [B, T+1, A], same as above, but using old policy
      network parameters.
    actions: ndarray of shape [B, T] where each element is from [0, A).
    reward_mask: ndarray of shape [B, T] masking over probabilities.

  Returns:
    probab_ratios: ndarray of shape [B, T], where
    probab_ratios_{b,t} = p_new_{b,t,action_{b,t}} / p_old_{b,t,action_{b,t}}
  """

  B, T = actions.shape  # pylint: disable=invalid-name
  assert (B, T + 1) == p_old.shape[:2]
  assert (B, T + 1) == p_new.shape[:2]

  logp_old = chosen_probabs(p_old, actions)
  logp_new = chosen_probabs(p_new, actions)

  assert (B, T) == logp_old.shape
  assert (B, T) == logp_new.shape

  # Since these are log-probabilities, we just subtract them.
  probab_ratios = np.exp(logp_new - logp_old) * reward_mask
  assert (B, T) == probab_ratios.shape
  return probab_ratios


def clipped_probab_ratios(probab_ratios, epsilon=0.2):
  return np.clip(probab_ratios, 1 - epsilon, 1 + epsilon)


def clipped_objective(probab_ratios, advantages, reward_mask, epsilon=0.2):
  return np.minimum(
      probab_ratios * advantages,
      clipped_probab_ratios(probab_ratios, epsilon=epsilon) *
      advantages) * reward_mask


@jit
def ppo_loss_given_predictions(log_probab_actions_new,
                               log_probab_actions_old,
                               value_predictions_old,
                               padded_actions,
                               padded_rewards,
                               reward_mask,
                               gamma=0.99,
                               lambda_=0.95,
                               epsilon=0.2):
  """PPO objective, with an eventual minus sign, given predictions."""
  B, T = padded_rewards.shape  # pylint: disable=invalid-name
  assert (B, T) == padded_actions.shape
  assert (B, T) == reward_mask.shape

  _, _, A = log_probab_actions_old.shape  # pylint: disable=invalid-name
  assert (B, T + 1, 1) == value_predictions_old.shape
  assert (B, T + 1, A) == log_probab_actions_old.shape
  assert (B, T + 1, A) == log_probab_actions_new.shape

  # (B, T)
  td_deltas = deltas(
      np.squeeze(value_predictions_old, axis=2),  # (B, T+1)
      padded_rewards,
      reward_mask,
      gamma=gamma)

  # (B, T)
  advantages = gae_advantages(
      td_deltas, reward_mask, lambda_=lambda_, gamma=gamma)

  # Normalize the advantages.
  advantages = (advantages - np.mean(advantages)) / np.std(advantages)

  # (B, T)
  ratios = compute_probab_ratios(log_probab_actions_new, log_probab_actions_old,
                                 padded_actions, reward_mask)
  assert (B, T) == ratios.shape

  # (B, T)
  objective = clipped_objective(
      ratios, advantages, reward_mask, epsilon=epsilon)
  assert (B, T) == objective.shape

  # ()
  average_objective = np.sum(objective) / np.sum(reward_mask)

  # Loss is negative objective.
  return -average_objective


@jit
def combined_loss_given_predictions(log_probab_actions_new,
                                    log_probab_actions_old,
                                    value_prediction_new,
                                    value_prediction_old,
                                    padded_actions,
                                    padded_rewards,
                                    reward_mask,
                                    gamma=0.99,
                                    lambda_=0.95,
                                    epsilon=0.2,
                                    c1=1.0,
                                    c2=0.01):
  """Computes the combined (clipped loss + value loss) given predictions."""
  loss_value = value_loss_given_predictions(
      value_prediction_new,
      padded_rewards,
      reward_mask,
      gamma=gamma,
      value_prediction_old=value_prediction_old,
      epsilon=epsilon)
  loss_ppo = ppo_loss_given_predictions(
      log_probab_actions_new,
      log_probab_actions_old,
      value_prediction_old,
      padded_actions,
      padded_rewards,
      reward_mask,
      gamma=gamma,
      lambda_=lambda_,
      epsilon=epsilon)
  entropy_bonus = masked_entropy(log_probab_actions_new, reward_mask)
  return (loss_ppo + (c1 * loss_value) - (c2 * entropy_bonus), loss_ppo,
          loss_value, entropy_bonus)


@functools.partial(jit, static_argnums=(3,))
def combined_loss(new_params,
                  log_probab_actions_old,
                  value_predictions_old,
                  policy_and_value_net_apply,
                  padded_observations,
                  padded_actions,
                  padded_rewards,
                  reward_mask,
                  gamma=0.99,
                  lambda_=0.95,
                  epsilon=0.2,
                  c1=1.0,
                  c2=0.01,
                  rng=None):
  """Computes the combined (clipped loss + value loss) given observations."""
  log_probab_actions_new, value_predictions_new = policy_and_value_net_apply(
      padded_observations, new_params, rng=rng)

  # (combined_loss, ppo_loss, value_loss, entropy_bonus)
  return combined_loss_given_predictions(
      log_probab_actions_new,
      log_probab_actions_old,
      value_predictions_new,
      value_predictions_old,
      padded_actions,
      padded_rewards,
      reward_mask,
      gamma=gamma,
      lambda_=lambda_,
      epsilon=epsilon,
      c1=c1,
      c2=c2)


@functools.partial(jit, static_argnums=(2, 3, 4))
def policy_and_value_opt_step(i,
                              opt_state,
                              opt_update,
                              get_params,
                              policy_and_value_net_apply,
                              log_probab_actions_old,
                              value_predictions_old,
                              padded_observations,
                              padded_actions,
                              padded_rewards,
                              reward_mask,
                              c1=1.0,
                              c2=0.01,
                              gamma=0.99,
                              lambda_=0.95,
                              epsilon=0.1,
                              rng=None):
  """Policy and Value optimizer step."""

  # Combined loss function given the new params.
  def policy_and_value_loss(params):
    """Returns the combined loss given just parameters."""
    (loss, _, _, _) = combined_loss(
        params,
        log_probab_actions_old,
        value_predictions_old,
        policy_and_value_net_apply,
        padded_observations,
        padded_actions,
        padded_rewards,
        reward_mask,
        c1=c1,
        c2=c2,
        gamma=gamma,
        lambda_=lambda_,
        epsilon=epsilon,
        rng=rng)
    return loss

  new_params = get_params(opt_state)
  g = grad(policy_and_value_loss)(new_params)
  # TODO(afrozm): Maybe clip gradients?
  return opt_update(i, g, opt_state)


def get_time(t1, t2=None):
  if t2 is None:
    t2 = time.time()
  return round((t2 - t1) * 1000, 2)


def approximate_kl(log_prob_new, log_prob_old, mask):
  """Computes the approximate KL divergence between the old and new log-probs.

  Args:
    log_prob_new: (B, T+1, A) log probs new
    log_prob_old: (B, T+1, A) log probs old
    mask: (B, T)

  Returns:
    Approximate KL.
  """
  diff = log_prob_old - log_prob_new
  # Cut the last time-step out.
  diff = diff[:, :-1]
  # Mask out the irrelevant part.
  diff *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
  # Average on non-masked part.
  return np.sum(diff) / np.sum(mask)


def masked_entropy(log_probs, mask):
  """Computes the entropy for the given log-probs.

  Args:
    log_probs: (B, T+1, A) log probs
    mask: (B, T) mask.

  Returns:
    Entropy.
  """
  # Cut the last time-step out.
  lp = log_probs[:, :-1]
  # Mask out the irrelevant part.
  lp *= mask[:, :, np.newaxis]  # make mask (B, T, 1)
  p = np.exp(lp) * mask[:, :, np.newaxis]  # (B, T, 1)
  # Average on non-masked part and take negative.
  return -(np.sum(lp * p) / np.sum(mask))


def evaluate_policy(eval_env,
                    get_predictions,
                    boundary,
                    max_timestep=20000,
                    rng=None):
  """Evaluate the policy."""

  avg_rewards = {}
  for policy in [
      env_problem_utils.CATEGORICAL_SAMPLING, env_problem_utils.GUMBEL_SAMPLING,
      env_problem_utils.EPSILON_GREEDY
  ]:
    trajs, _ = env_problem_utils.play_env_problem_with_policy(
        eval_env,
        get_predictions,
        boundary=boundary,
        max_timestep=max_timestep,
        reset=True,
        policy_sampling=policy,
        rng=rng)
    avg_rewards[policy] = float(sum(
        np.sum(traj[2]) for traj in trajs)) / len(trajs)
  return avg_rewards


def maybe_restore_params(output_dir, policy_and_value_net_params):
  """Maybe restore the params from the checkpoint dir.

  Args:
    output_dir: Directory where saved model checkpoints are stored.
    policy_and_value_net_params: Default params, returned if model is'nt found.

  Returns:
    triple (restore (bool), params, iter(int)) where iter is the epoch from
    which we restored the params, 0 is restore = False.
  """
  model_files = gfile.glob(os.path.join(output_dir, "model-??????.pkl"))
  if not model_files:
    return False, policy_and_value_net_params, 0

  model_file = sorted(model_files)[-1]
  model_file_basename = os.path.basename(model_file)  # model-??????.pkl
  i = int(filter(str.isdigit, model_file_basename))
  with gfile.GFile(model_file, "rb") as f:
    policy_and_value_net_params = pickle.load(f)
  return True, policy_and_value_net_params, i


def training_loop(
    env=None,
    epochs=EPOCHS,
    policy_and_value_net_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,
    target_kl=0.01,
    boundary=20,
    max_timestep=None,
    max_timestep_eval=20000,
    random_seed=None,
    gamma=GAMMA,
    lambda_=LAMBDA,
    epsilon=EPSILON,
    c1=1.0,
    c2=0.01,
    output_dir=None,
    eval_every_n=1000,
    eval_env=None,
    done_frac_for_policy_save=0.5,
    enable_early_stopping=True,
    env_name=None,
):
  """Runs the training loop for PPO, with fixed policy and value nets."""
  assert env
  assert output_dir
  assert env_name

  gfile.makedirs(output_dir)

  # Create summary writers and history.
  train_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "train"))
  timing_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "timing"))
  eval_sw = jaxboard.SummaryWriter(os.path.join(output_dir, "eval"))

  train_sw.text("env_name", env_name)
  timing_sw.text("env_name", env_name)
  eval_sw.text("env_name", env_name)

  jax_rng_key = trax.get_random_number_generator_and_set_seed(random_seed)

  # 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

  jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)

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

  # Maybe restore the policy params. If there is nothing to restore, then
  # iteration = 0 and policy_and_value_net_params are returned as is.
  restore, policy_and_value_net_params, iteration = (
      maybe_restore_params(output_dir, policy_and_value_net_params))

  if restore:
    logging.info("Restored parameters from iteration [%d]", iteration)
    # We should start from the next iteration.
    iteration += 1

  policy_and_value_net_apply = jit(policy_and_value_net_apply)

  # Initialize the optimizers.
  policy_and_value_optimizer = (
      policy_and_value_optimizer_fun(policy_and_value_net_params))
  (policy_and_value_opt_state, policy_and_value_opt_update,
   policy_and_value_get_params) = policy_and_value_optimizer

  num_trajectories_done = 0
  last_saved_at = 0

  logging.info("Starting the PPO training loop.")
  for i in range(iteration, epochs):
    epoch_start_time = time.time()

    # Params we'll use to collect the trajectories.
    policy_and_value_net_params = policy_and_value_get_params(
        policy_and_value_opt_state)

    # A function to get the policy and value predictions.
    def get_predictions(observations, rng=None):
      """Returns log-probs, value predictions and key back."""
      key, key1 = jax_random.split(rng, num=2)

      log_probs, value_preds = policy_and_value_net_apply(
          observations, policy_and_value_net_params, rng=key1)

      return log_probs, value_preds, key

    # Evaluate the policy.
    policy_eval_start_time = time.time()
    if ((i + 1) % eval_every_n == 0) or (i == epochs - 1):
      jax_rng_key, key = jax_random.split(jax_rng_key, num=2)

      logging.vlog(1, "Epoch [% 6d] evaluating policy.", i)

      avg_reward = evaluate_policy(
          eval_env,
          get_predictions,
          boundary,
          max_timestep=max_timestep_eval,
          rng=key)
      for k, v in avg_reward.items():
        eval_sw.scalar("eval/mean_reward/%s" % k, v, step=i)
        logging.info("Epoch [% 6d] Policy Evaluation [%s] = %10.2f", i, k, v)
    policy_eval_time = get_time(policy_eval_start_time)

    trajectory_collection_start_time = time.time()
    logging.vlog(1, "Epoch [% 6d] collecting trajectories.", i)
    jax_rng_key, key = jax_random.split(jax_rng_key)
    trajs, num_done = collect_trajectories(
        env,
        policy_fun=get_predictions,
        num_trajectories=batch_size,
        max_timestep=max_timestep,
        boundary=boundary,
        rng=key,
        reset=(i == 0) or restore,
        epsilon=(10.0 / (i + 10.0)))  # this is a different epsilon.
    trajectory_collection_time = get_time(trajectory_collection_start_time)

    logging.vlog(1, "Collecting trajectories took %0.2f msec.",
                 trajectory_collection_time)

    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)

    train_sw.scalar("train/mean_reward", avg_reward, step=i)

    logging.vlog(1, "Rewards avg=[%0.2f], max=[%0.2f], min=[%0.2f], all=%s",
                 avg_reward, max_reward, min_reward,
                 [float(np.sum(traj[2])) for traj in trajs])

    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))
    logging.vlog(2, "Trajectory Lengths: %s", [len(traj[0]) for traj in trajs])

    padding_start_time = time.time()
    (_, reward_mask, padded_observations, padded_actions,
     padded_rewards) = pad_trajectories(
         trajs, boundary=boundary)
    padding_time = get_time(padding_start_time)

    logging.vlog(1, "Padding trajectories took %0.2f msec.",
                 get_time(padding_start_time))
    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))

    # Calculate log-probabilities and value predictions of the trajectories.
    # We'll pass these to the loss functions so as to not get recomputed.

    # NOTE:
    # There is a slight problem here, if the policy network contains
    # stochasticity in the log-probabilities (ex: dropout), then calculating
    # these again here is not going to be correct and should be done in the
    # collect function.

    log_prob_recompute_start_time = time.time()
    jax_rng_key, key = jax_random.split(jax_rng_key)
    log_probabs_traj, value_predictions_traj, _ = get_predictions(
        padded_observations, rng=key)
    log_prob_recompute_time = get_time(log_prob_recompute_start_time)

    # 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)))

    # Constant epsilon.
    epsilon_schedule = epsilon

    # Compute value and ppo losses.
    jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
    logging.vlog(2, "Starting to compute P&V loss.")
    loss_compute_start_time = time.time()
    cur_combined_loss, cur_ppo_loss, cur_value_loss, entropy_bonus = (
        combined_loss(
            policy_and_value_net_params,
            log_probabs_traj,
            value_predictions_traj,
            policy_and_value_net_apply,
            padded_observations,
            padded_actions,
            padded_rewards,
            reward_mask,
            gamma=gamma,
            lambda_=lambda_,
            epsilon=epsilon_schedule,
            c1=c1,
            c2=c2,
            rng=key1))
    loss_compute_time = get_time(loss_compute_start_time)
    logging.vlog(
        1,
        "Calculating P&V loss [%10.2f(%10.2f, %10.2f, %10.2f)] took %0.2f msec.",
        cur_combined_loss, cur_value_loss, cur_ppo_loss, entropy_bonus,
        get_time(loss_compute_start_time))

    jax_rng_key, key1 = jax_random.split(jax_rng_key, num=2)
    logging.vlog(1, "Policy and Value Optimization")
    optimization_start_time = time.time()
    keys = jax_random.split(key1, num=num_optimizer_steps)
    for j in range(num_optimizer_steps):
      k1, k2, k3 = jax_random.split(keys[j], num=3)
      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_get_params,
          policy_and_value_net_apply,
          log_probabs_traj,
          value_predictions_traj,
          padded_observations,
          padded_actions,
          padded_rewards,
          reward_mask,
          c1=c1,
          c2=c2,
          gamma=gamma,
          lambda_=lambda_,
          epsilon=epsilon_schedule,
          rng=k1)

      # Compute the approx KL for early stopping.
      new_policy_and_value_net_params = policy_and_value_get_params(
          policy_and_value_opt_state)

      log_probab_actions_new, _ = policy_and_value_net_apply(
          padded_observations, new_policy_and_value_net_params, rng=k2)

      approx_kl = approximate_kl(log_probab_actions_new, log_probabs_traj,
                                 reward_mask)

      early_stopping = enable_early_stopping and approx_kl > 1.5 * target_kl
      if early_stopping:
        logging.vlog(
            1, "Early stopping policy and value optimization at iter: %d, "
            "with approx_kl: %0.2f", j, approx_kl)
        # We don't return right-away, we want the below to execute on the last
        # iteration.

      t2 = time.time()
      if (((j + 1) % print_every_optimizer_steps == 0) or
          (j == num_optimizer_steps - 1) or early_stopping):
        # Compute and log the loss.
        (loss_combined, loss_ppo, loss_value, entropy_bonus) = (
            combined_loss(
                new_policy_and_value_net_params,
                log_probabs_traj,
                value_predictions_traj,
                policy_and_value_net_apply,
                padded_observations,
                padded_actions,
                padded_rewards,
                reward_mask,
                gamma=gamma,
                lambda_=lambda_,
                epsilon=epsilon_schedule,
                c1=c1,
                c2=c2,
                rng=k3))
        logging.vlog(1, "One Policy and Value grad desc took: %0.2f msec",
                     get_time(t, t2))
        logging.vlog(
            1, "Combined Loss(value, ppo, entropy_bonus) [%10.2f] ->"
            " [%10.2f(%10.2f,%10.2f,%10.2f)]", cur_combined_loss, loss_combined,
            loss_value, loss_ppo, entropy_bonus)

      if early_stopping:
        break

    optimization_time = get_time(optimization_start_time)

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

    # Save parameters every time we see the end of at least a fraction of batch
    # number of trajectories that are done (not completed -- completed includes
    # truncated and done).
    # Also don't save too frequently, enforce a minimum gap.
    # Or if this is the last iteration.
    policy_save_start_time = time.time()
    num_trajectories_done += num_done
    if (((num_trajectories_done >= done_frac_for_policy_save * batch_size)
         and (i - last_saved_at > eval_every_n)) or (i == epochs - 1)):
      logging.vlog(1, "Epoch [% 6d] saving model.", i)
      params_file = os.path.join(output_dir, "model-%06d.pkl" % i)
      with gfile.GFile(params_file, "wb") as f:
        pickle.dump(policy_and_value_net_params, f)
      # Reset this number.
      num_trajectories_done = 0
      last_saved_at = i
    policy_save_time = get_time(policy_save_start_time)

    epoch_time = get_time(epoch_start_time)

    logging.info(
        "Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined"
        " Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]", i, min_reward,
        max_reward, avg_reward, loss_combined, loss_value, loss_ppo,
        entropy_bonus)

    timing_dict = {
        "epoch": epoch_time,
        "policy_eval": policy_eval_time,
        "trajectory_collection": trajectory_collection_time,
        "padding": padding_time,
        "log_prob_recompute": log_prob_recompute_time,
        "loss_compute": loss_compute_time,
        "optimization": optimization_time,
        "policy_save": policy_save_time,
    }

    for k, v in timing_dict.items():
      timing_sw.scalar("timing/%s" % k, v, step=i)

    max_key_len = max(len(k) for k in timing_dict)
    timing_info_list = [
        "%s : % 10.2f" % (k.rjust(max_key_len + 1), v)
        for k, v in sorted(timing_dict.items())
    ]
    logging.info("Epoch [% 6d], Timings: \n%s", i, "\n".join(timing_info_list))

    # Reset restore.
    restore = False

    # Flush summary writers once in a while.
    if (i+1) % 1000 == 0 or i == epochs - 1:
      train_sw.flush()
      timing_sw.flush()
      eval_sw.flush()

[((576, 1, 576, 44), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((611, 1, 611, 49), 'functools.partial', 'functools.partial', (), '', False, 'import functools\n'), ((121, 8, 121, 71), 'tensor2tensor.trax.optimizers.Adam', 'trax_opt.Adam', (), '', True, 'from tensor2tensor.trax import optimizers as trax_opt\n'), ((170, 9, 179, 14), 'tensor2tensor.envs.env_problem_utils.play_env_problem_with_policy', 'env_problem_utils.play_env_problem_with_policy', (), '', False, 'from tensor2tensor.envs import env_problem_utils\n'), ((355, 21, 355, 57), 'jax.numpy.squeeze', 'np.squeeze', (), '', True, 'from jax import numpy as np\n'), ((479, 9, 479, 57), 'jax.numpy.clip', 'np.clip', ({(479, 17, 479, 30): 'probab_ratios', (479, 32, 479, 43): '(1 - epsilon)', (479, 45, 479, 56): '(1 + epsilon)'}, {}), '(probab_ratios, 1 - epsilon, 1 + epsilon)', True, 'from jax import numpy as np\n'), ((743, 24, 743, 52), 'os.path.basename', 'os.path.basename', ({(743, 41, 743, 51): 'model_file'}, {}), '(model_file)', False, 'import os\n'), ((780, 2, 780, 28), 'tensorflow.io.gfile.makedirs', 'gfile.makedirs', ({(780, 17, 780, 27): 'output_dir'}, {}), '(output_dir)', False, 'from tensorflow.io import gfile\n'), ((791, 16, 791, 74), 'tensor2tensor.trax.trax.get_random_number_generator_and_set_seed', 'trax.get_random_number_generator_and_set_seed', ({(791, 62, 791, 73): 'random_seed'}, {}), '(random_seed)', False, 'from tensor2tensor.trax import trax\n'), ((800, 22, 800, 58), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((816, 31, 816, 62), 'jax.jit', 'jit', ({(816, 35, 816, 61): 'policy_and_value_net_apply'}, {}), '(policy_and_value_net_apply)', False, 'from jax import jit\n'), ((827, 2, 827, 49), 'absl.logging.info', 'logging.info', ({(827, 15, 827, 48): '"""Starting the PPO training loop."""'}, {}), "('Starting the PPO training loop.')", False, 'from absl import logging\n'), ((104, 10, 104, 31), 'tensor2tensor.trax.layers.Serial', 'layers.Serial', ({(104, 24, 104, 30): '*tower'}, {}), '(*tower)', False, 'from tensor2tensor.trax import layers\n'), ((191, 20, 191, 31), 'jax.numpy.uint8', 'np.uint8', ({(191, 29, 191, 30): '0'}, {}), '(0)', True, 'from jax import numpy as np\n'), ((256, 17, 256, 58), 'jax.lax.pad', 'lax.pad', ({(256, 25, 256, 26): 'o', (256, 28, 256, 41): 'padding_value', (256, 43, 256, 57): 'padding_config'}, {}), '(o, padding_value, padding_config)', False, 'from jax import lax\n'), ((263, 20, 263, 68), 'jax.lax.pad', 'lax.pad', ({(263, 28, 263, 29): 'a', (263, 31, 263, 51): 'action_padding_value', (263, 53, 263, 67): 'padding_config'}, {}), '(a, action_padding_value, padding_config)', False, 'from jax import lax\n'), ((265, 20, 265, 68), 'jax.lax.pad', 'lax.pad', ({(265, 28, 265, 29): 'r', (265, 31, 265, 51): 'reward_padding_value', (265, 53, 265, 67): 'padding_config'}, {}), '(r, reward_padding_value, padding_config)', False, 'from jax import lax\n'), ((269, 18, 269, 50), 'numpy.ones_like', 'onp.ones_like', (), '', True, 'import numpy as onp\n'), ((272, 25, 272, 47), 'jax.numpy.stack', 'np.stack', ({(272, 34, 272, 46): 'reward_masks'}, {}), '(reward_masks)', True, 'from jax import numpy as np\n'), ((272, 49, 273, 26), 'jax.numpy.stack', 'np.stack', ({(273, 6, 273, 25): 'padded_observations'}, {}), '(padded_observations)', True, 'from jax import numpy as np\n'), ((273, 28, 273, 52), 'jax.numpy.stack', 'np.stack', ({(273, 37, 273, 51): 'padded_actions'}, {}), '(padded_actions)', True, 'from jax import numpy as np\n'), ((273, 54, 273, 78), 'jax.numpy.stack', 'np.stack', ({(273, 63, 273, 77): 'padded_rewards'}, {}), '(padded_rewards)', True, 'from jax import numpy as np\n'), ((325, 17, 325, 39), 'jax.numpy.stack', 'np.stack', (), '', True, 'from jax import numpy as np\n'), ((362, 27, 362, 67), 'jax.numpy.squeeze', 'np.squeeze', (), '', True, 'from jax import numpy as np\n'), ((368, 11, 368, 43), 'jax.numpy.maximum', 'np.maximum', ({(368, 22, 368, 36): 'v_clipped_loss', (368, 38, 368, 42): 'loss'}, {}), '(v_clipped_loss, loss)', True, 'from jax import numpy as np\n'), ((371, 9, 371, 21), 'jax.numpy.sum', 'np.sum', ({(371, 16, 371, 20): 'loss'}, {}), '(loss)', True, 'from jax import numpy as np\n'), ((371, 24, 371, 43), 'jax.numpy.sum', 'np.sum', ({(371, 31, 371, 42): 'reward_mask'}, {}), '(reward_mask)', True, 'from jax import numpy as np\n'), ((473, 18, 473, 45), 'jax.numpy.exp', 'np.exp', ({(473, 25, 473, 44): '(logp_new - logp_old)'}, {}), '(logp_new - logp_old)', True, 'from jax import numpy as np\n'), ((511, 6, 511, 47), 'jax.numpy.squeeze', 'np.squeeze', (), '', True, 'from jax import numpy as np\n'), ((521, 52, 521, 70), 'jax.numpy.std', 'np.std', ({(521, 59, 521, 69): 'advantages'}, {}), '(advantages)', True, 'from jax import numpy as np\n'), ((534, 22, 534, 39), 'jax.numpy.sum', 'np.sum', ({(534, 29, 534, 38): 'objective'}, {}), '(objective)', True, 'from jax import numpy as np\n'), ((534, 42, 534, 61), 'jax.numpy.sum', 'np.sum', ({(534, 49, 534, 60): 'reward_mask'}, {}), '(reward_mask)', True, 'from jax import numpy as np\n'), ((652, 6, 652, 33), 'jax.grad', 'grad', ({(652, 11, 652, 32): 'policy_and_value_loss'}, {}), '(policy_and_value_loss)', False, 'from jax import grad\n'), ((659, 9, 659, 20), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((680, 9, 680, 21), 'jax.numpy.sum', 'np.sum', ({(680, 16, 680, 20): 'diff'}, {}), '(diff)', True, 'from jax import numpy as np\n'), ((680, 24, 680, 36), 'jax.numpy.sum', 'np.sum', ({(680, 31, 680, 35): 'mask'}, {}), '(mask)', True, 'from jax import numpy as np\n'), ((697, 6, 697, 16), 'jax.numpy.exp', 'np.exp', ({(697, 13, 697, 15): 'lp'}, {}), '(lp)', True, 'from jax import numpy as np\n'), ((714, 15, 721, 16), 'tensor2tensor.envs.env_problem_utils.play_env_problem_with_policy', 'env_problem_utils.play_env_problem_with_policy', (), '', False, 'from tensor2tensor.envs import env_problem_utils\n'), ((738, 27, 738, 71), 'os.path.join', 'os.path.join', ({(738, 40, 738, 50): 'output_dir', (738, 52, 738, 70): '"""model-??????.pkl"""'}, {}), "(output_dir, 'model-??????.pkl')", False, 'import os\n'), ((745, 7, 745, 36), 'tensorflow.io.gfile.GFile', 'gfile.GFile', ({(745, 19, 745, 29): 'model_file', (745, 31, 745, 35): '"""rb"""'}, {}), "(model_file, 'rb')", False, 'from tensorflow.io import gfile\n'), ((746, 34, 746, 48), 'pickle.load', 'pickle.load', ({(746, 46, 746, 47): 'f'}, {}), '(f)', False, 'import pickle\n'), ((783, 36, 783, 69), 'os.path.join', 'os.path.join', ({(783, 49, 783, 59): 'output_dir', (783, 61, 783, 68): '"""train"""'}, {}), "(output_dir, 'train')", False, 'import os\n'), ((784, 37, 784, 71), 'os.path.join', 'os.path.join', ({(784, 50, 784, 60): 'output_dir', (784, 62, 784, 70): '"""timing"""'}, {}), "(output_dir, 'timing')", False, 'import os\n'), ((785, 35, 785, 67), 'os.path.join', 'os.path.join', ({(785, 48, 785, 58): 'output_dir', (785, 60, 785, 66): '"""eval"""'}, {}), "(output_dir, 'eval')", False, 'import os\n'), ((812, 4, 812, 70), 'absl.logging.info', 'logging.info', ({(812, 17, 812, 58): '"""Restored parameters from iteration [%d]"""', (812, 60, 812, 69): 'iteration'}, {}), "('Restored parameters from iteration [%d]', iteration)", False, 'from absl import logging\n'), ((829, 23, 829, 34), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((846, 29, 846, 40), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((863, 39, 863, 50), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((864, 4, 864, 63), 'absl.logging.vlog', 'logging.vlog', ({(864, 17, 864, 18): '(1)', (864, 20, 864, 59): '"""Epoch [% 6d] collecting trajectories."""', (864, 61, 864, 62): 'i'}, {}), "(1, 'Epoch [% 6d] collecting trajectories.', i)", False, 'from absl import logging\n'), ((865, 23, 865, 52), 'jax.random.split', 'jax_random.split', ({(865, 40, 865, 51): 'jax_rng_key'}, {}), '(jax_rng_key)', True, 'from jax import random as jax_random\n'), ((877, 4, 878, 44), 'absl.logging.vlog', 'logging.vlog', ({(877, 17, 877, 18): '(1)', (877, 20, 877, 62): '"""Collecting trajectories took %0.2f msec."""', (878, 17, 878, 43): 'trajectory_collection_time'}, {}), "(1, 'Collecting trajectories took %0.2f msec.',\n    trajectory_collection_time)", False, 'from absl import logging\n'), ((897, 25, 897, 36), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((919, 36, 919, 47), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((920, 23, 920, 52), 'jax.random.split', 'jax_random.split', ({(920, 40, 920, 51): 'jax_rng_key'}, {}), '(jax_rng_key)', True, 'from jax import random as jax_random\n'), ((941, 24, 941, 60), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((942, 4, 942, 52), 'absl.logging.vlog', 'logging.vlog', ({(942, 17, 942, 18): '(2)', (942, 20, 942, 51): '"""Starting to compute P&V loss."""'}, {}), "(2, 'Starting to compute P&V loss.')", False, 'from absl import logging\n'), ((943, 30, 943, 41), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((967, 24, 967, 60), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((968, 4, 968, 52), 'absl.logging.vlog', 'logging.vlog', ({(968, 17, 968, 18): '(1)', (968, 20, 968, 51): '"""Policy and Value Optimization"""'}, {}), "(1, 'Policy and Value Optimization')", False, 'from absl import logging\n'), ((969, 30, 969, 41), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((970, 11, 970, 58), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((1053, 29, 1053, 40), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((1068, 4, 1072, 22), 'absl.logging.info', 'logging.info', ({(1069, 8, 1070, 63): '"""Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]"""', (1070, 65, 1070, 66): 'i', (1070, 68, 1070, 78): 'min_reward', (1071, 8, 1071, 18): 'max_reward', (1071, 20, 1071, 30): 'avg_reward', (1071, 32, 1071, 45): 'loss_combined', (1071, 47, 1071, 57): 'loss_value', (1071, 59, 1071, 67): 'loss_ppo', (1072, 8, 1072, 21): 'entropy_bonus'}, {}), "(\n    'Epoch [% 6d], Reward[min, max, avg] [%5.2f,%5.2f,%5.2f], Combined Loss(value, ppo, entropy) [%2.5f(%2.5f,%2.5f,%2.5f)]'\n    , i, min_reward, max_reward, avg_reward, loss_combined, loss_value,\n    loss_ppo, entropy_bonus)", False, 'from absl import logging\n'), ((113, 8, 113, 30), 'tensor2tensor.trax.layers.Serial', 'layers.Serial', ({(113, 22, 113, 29): '*tower1'}, {}), '(*tower1)', False, 'from tensor2tensor.trax import layers\n'), ((114, 8, 114, 30), 'tensor2tensor.trax.layers.Serial', 'layers.Serial', ({(114, 22, 114, 29): '*tower2'}, {}), '(*tower2)', False, 'from tensor2tensor.trax import layers\n'), ((193, 20, 193, 32), 'jax.numpy.uint16', 'np.uint16', ({(193, 30, 193, 31): '0'}, {}), '(0)', True, 'from jax import numpy as np\n'), ((270, 24, 270, 63), 'jax.lax.pad', 'lax.pad', ({(270, 32, 270, 43): 'reward_mask', (270, 45, 270, 46): '(0)', (270, 48, 270, 62): 'padding_config'}, {}), '(reward_mask, 0, padding_config)', False, 'from jax import lax\n'), ((365, 39, 366, 67), 'jax.numpy.clip', 'np.clip', ({(366, 8, 366, 47): '(value_prediction - value_prediction_old)', (366, 49, 366, 57): '(-epsilon)', (366, 59, 366, 66): 'epsilon'}, {}), '(value_prediction - value_prediction_old, -epsilon, epsilon)', True, 'from jax import numpy as np\n'), ((400, 9, 400, 20), 'jax.numpy.array', 'np.array', ({(400, 18, 400, 19): 'd'}, {}), '(d)', True, 'from jax import numpy as np\n'), ((521, 29, 521, 48), 'jax.numpy.mean', 'np.mean', ({(521, 37, 521, 47): 'advantages'}, {}), '(advantages)', True, 'from jax import numpy as np\n'), ((699, 11, 699, 25), 'jax.numpy.sum', 'np.sum', ({(699, 18, 699, 24): '(lp * p)'}, {}), '(lp * p)', True, 'from jax import numpy as np\n'), ((699, 28, 699, 40), 'jax.numpy.sum', 'np.sum', ({(699, 35, 699, 39): 'mask'}, {}), '(mask)', True, 'from jax import numpy as np\n'), ((838, 18, 838, 46), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((848, 25, 848, 61), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((850, 6, 850, 59), 'absl.logging.vlog', 'logging.vlog', ({(850, 19, 850, 20): '(1)', (850, 22, 850, 55): '"""Epoch [% 6d] evaluating policy."""', (850, 57, 850, 58): 'i'}, {}), "(1, 'Epoch [% 6d] evaluating policy.', i)", False, 'from absl import logging\n'), ((972, 19, 972, 51), 'jax.random.split', 'jax_random.split', (), '', True, 'from jax import random as jax_random\n'), ((973, 10, 973, 21), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((1012, 11, 1012, 22), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((1057, 6, 1057, 54), 'absl.logging.vlog', 'logging.vlog', ({(1057, 19, 1057, 20): '(1)', (1057, 22, 1057, 50): '"""Epoch [% 6d] saving model."""', (1057, 52, 1057, 53): 'i'}, {}), "(1, 'Epoch [% 6d] saving model.', i)", False, 'from absl import logging\n'), ((1058, 20, 1058, 66), 'os.path.join', 'os.path.join', ({(1058, 33, 1058, 43): 'output_dir', (1058, 45, 1058, 65): "'model-%06d.pkl' % i"}, {}), "(output_dir, 'model-%06d.pkl' % i)", False, 'import os\n'), ((109, 19, 109, 44), 'tensor2tensor.trax.layers.Dense', 'layers.Dense', ({(109, 32, 109, 43): 'num_actions'}, {}), '(num_actions)', False, 'from tensor2tensor.trax import layers\n'), ((109, 46, 109, 65), 'tensor2tensor.trax.layers.LogSoftmax', 'layers.LogSoftmax', ({}, {}), '()', False, 'from tensor2tensor.trax import layers\n'), ((110, 19, 110, 34), 'tensor2tensor.trax.layers.Dense', 'layers.Dense', ({(110, 32, 110, 33): '(1)'}, {}), '(1)', False, 'from tensor2tensor.trax import layers\n'), ((243, 26, 243, 58), 'numpy.ones_like', 'onp.ones_like', (), '', True, 'import numpy as onp\n'), ((442, 52, 442, 64), 'jax.numpy.arange', 'np.arange', ({(442, 62, 442, 63): 'T'}, {}), '(T)', True, 'from jax import numpy as np\n'), ((860, 8, 860, 77), 'absl.logging.info', 'logging.info', ({(860, 21, 860, 67): '"""Epoch [% 6d] Policy Evaluation [%s] = %10.2f"""', (860, 69, 860, 70): 'i', (860, 72, 860, 73): 'k', (860, 75, 860, 76): 'v'}, {}), "('Epoch [% 6d] Policy Evaluation [%s] = %10.2f', i, k, v)", False, 'from absl import logging\n'), ((881, 21, 881, 36), 'jax.numpy.sum', 'np.sum', ({(881, 28, 881, 35): 'traj[2]'}, {}), '(traj[2])', True, 'from jax import numpy as np\n'), ((882, 21, 882, 36), 'jax.numpy.sum', 'np.sum', ({(882, 28, 882, 35): 'traj[2]'}, {}), '(traj[2])', True, 'from jax import numpy as np\n'), ((1006, 8, 1008, 50), 'absl.logging.vlog', 'logging.vlog', ({(1007, 12, 1007, 13): '(1)', (1007, 15, 1008, 35): '"""Early stopping policy and value optimization at iter: %d, with approx_kl: %0.2f"""', (1008, 37, 1008, 38): 'j', (1008, 40, 1008, 49): 'approx_kl'}, {}), "(1,\n    'Early stopping policy and value optimization at iter: %d, with approx_kl: %0.2f'\n    , j, approx_kl)", False, 'from absl import logging\n'), ((1034, 8, 1037, 48), 'absl.logging.vlog', 'logging.vlog', ({(1035, 12, 1035, 13): '(1)', (1035, 15, 1036, 45): '"""Combined Loss(value, ppo, entropy_bonus) [%10.2f] -> [%10.2f(%10.2f,%10.2f,%10.2f)]"""', (1036, 47, 1036, 64): 'cur_combined_loss', (1036, 66, 1036, 79): 'loss_combined', (1037, 12, 1037, 22): 'loss_value', (1037, 24, 1037, 32): 'loss_ppo', (1037, 34, 1037, 47): 'entropy_bonus'}, {}), "(1,\n    'Combined Loss(value, ppo, entropy_bonus) [%10.2f] -> [%10.2f(%10.2f,%10.2f,%10.2f)]'\n    , cur_combined_loss, loss_combined, loss_value, loss_ppo, entropy_bonus)", False, 'from absl import logging\n'), ((1046, 53, 1046, 78), 'jax.numpy.abs', 'np.abs', ({(1046, 60, 1046, 77): 'cur_combined_loss'}, {}), '(cur_combined_loss)', True, 'from jax import numpy as np\n'), ((1059, 11, 1059, 41), 'tensorflow.io.gfile.GFile', 'gfile.GFile', ({(1059, 23, 1059, 34): 'params_file', (1059, 36, 1059, 40): '"""wb"""'}, {}), "(params_file, 'wb')", False, 'from tensorflow.io import gfile\n'), ((1060, 8, 1060, 51), 'pickle.dump', 'pickle.dump', ({(1060, 20, 1060, 47): 'policy_and_value_net_params', (1060, 49, 1060, 50): 'f'}, {}), '(policy_and_value_net_params, f)', False, 'import pickle\n'), ((102, 12, 102, 27), 'tensor2tensor.trax.layers.Dense', 'layers.Dense', ({(102, 25, 102, 26): '(1)'}, {}), '(1)', False, 'from tensor2tensor.trax import layers\n'), ((442, 29, 442, 41), 'jax.numpy.arange', 'np.arange', ({(442, 39, 442, 40): 'B'}, {}), '(B)', True, 'from jax import numpy as np\n'), ((888, 24, 888, 39), 'jax.numpy.sum', 'np.sum', ({(888, 31, 888, 38): 'traj[2]'}, {}), '(traj[2])', True, 'from jax import numpy as np\n'), ((101, 26, 101, 51), 'tensor2tensor.trax.layers.Dense', 'layers.Dense', ({(101, 39, 101, 50): 'num_actions'}, {}), '(num_actions)', False, 'from tensor2tensor.trax import layers\n'), ((101, 53, 101, 72), 'tensor2tensor.trax.layers.LogSoftmax', 'layers.LogSoftmax', ({}, {}), '()', False, 'from tensor2tensor.trax import layers\n'), ((723, 8, 723, 23), 'jax.numpy.sum', 'np.sum', ({(723, 15, 723, 22): 'traj[2]'}, {}), '(traj[2])', True, 'from jax import numpy as np\n'), ((880, 27, 880, 42), 'jax.numpy.sum', 'np.sum', ({(880, 34, 880, 41): 'traj[2]'}, {}), '(traj[2])', True, 'from jax import numpy as np\n')]