text: # Copyright 2021 The Flax 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.
  
  """Library file which executes the PPO training."""
  
  import functools
  from typing import Tuple, List
  import jax
  import jax.random
  import jax.numpy as jnp
  import numpy as np
  import flax
  
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  import ml_collections
  
  import agent
  import models
  import test_episodes
  
  @jax.jit
  @functools.partial(jax.vmap, in_axes=(1, 1, 1, None, None), out_axes=1)
  def gae_advantages(
      rewards: np.ndarray,
      terminal_masks: np.ndarray,
      values: np.ndarray,
      discount: float,
      gae_param: float):
    """Use Generalized Advantage Estimation (GAE) to compute advantages.
  
    As defined by eqs. (11-12) in PPO paper arXiv: 1707.06347. Implementation uses
    key observation that A_{t} = delta_t + gamma*lambda*A_{t+1}.
  
    Args:
      rewards: array shaped (actor_steps, num_agents), rewards from the game
      terminal_masks: array shaped (actor_steps, num_agents), zeros for terminal
                      and ones for non-terminal states
      values: array shaped (actor_steps, num_agents), values estimated by critic
      discount: RL discount usually denoted with gamma
      gae_param: GAE parameter usually denoted with lambda
  
    Returns:
      advantages: calculated advantages shaped (actor_steps, num_agents)
    """
    assert rewards.shape[0] + 1 == values.shape[0], ('One more value needed; Eq. '
                                                     '(12) in PPO paper requires '
                                                     'V(s_{t+1}) for delta_t')
    advantages = []
    gae = 0.
    for t in reversed(range(len(rewards))):
      # Masks used to set next state value to 0 for terminal states.
      value_diff = discount * values[t + 1] * terminal_masks[t] - values[t]
      delta = rewards[t] + value_diff
      # Masks[t] used to ensure that values before and after a terminal state
      # are independent of each other.
      gae = delta + discount * gae_param * terminal_masks[t] * gae
      advantages.append(gae)
    advantages = advantages[::-1]
    return jnp.array(advantages)
  
  @functools.partial(jax.jit, static_argnums=1)
  def loss_fn(
      params: flax.core.frozen_dict.FrozenDict,
      module: models.ActorCritic,
      minibatch: Tuple,
      clip_param: float,
      vf_coeff: float,
      entropy_coeff: float):
    """Evaluate the loss function.
  
    Compute loss as a sum of three components: the negative of the PPO clipped
    surrogate objective, the value function loss and the negative of the entropy
    bonus.
  
    Args:
      params: the parameters of the actor-critic model
      module: the actor-critic model
      minibatch: Tuple of five elements forming one experience batch:
                 states: shape (batch_size, 84, 84, 4)
                 actions: shape (batch_size, 84, 84, 4)
                 old_log_probs: shape (batch_size,)
                 returns: shape (batch_size,)
                 advantages: shape (batch_size,)
      clip_param: the PPO clipping parameter used to clamp ratios in loss function
      vf_coeff: weighs value function loss in total loss
      entropy_coeff: weighs entropy bonus in the total loss
  
    Returns:
      loss: the PPO loss, scalar quantity
    """
    states, actions, old_log_probs, returns, advantages = minibatch
    log_probs, values = agent.policy_action(params, module, states)
    values = values[:, 0]  # Convert shapes: (batch, 1) to (batch, ).
    probs = jnp.exp(log_probs)
  
    value_loss = jnp.mean(jnp.square(returns - values), axis=0)
  
    entropy = jnp.sum(-probs*log_probs, axis=1).mean()
  
    log_probs_act_taken = jax.vmap(lambda lp, a: lp[a])(log_probs, actions)
    ratios = jnp.exp(log_probs_act_taken - old_log_probs)
    # Advantage normalization (following the OpenAI baselines).
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
    PG_loss = ratios * advantages
    clipped_loss = advantages * jax.lax.clamp(1. - clip_param, ratios,
                                              1. + clip_param)
    PPO_loss = -jnp.mean(jnp.minimum(PG_loss, clipped_loss), axis=0)
  
    return PPO_loss + vf_coeff*value_loss - entropy_coeff*entropy
  
  @functools.partial(jax.jit, static_argnums=(0,7))
  def train_step(
      module: models.ActorCritic,
      optimizer: flax.optim.base.Optimizer,
      trajectories: Tuple,
      clip_param: float,
      vf_coeff: float,
      entropy_coeff: float,
      lr: float,
      batch_size: int):
    """Compilable train step.
  
    Runs an entire epoch of training (i.e. the loop over minibatches within
    an epoch is included here for performance reasons).
  
    Args:
      module: the actor-critic model
      optimizer: optimizer for the actor-critic model
      trajectories: Tuple of the following five elements forming the experience:
                    states: shape (steps_per_agent*num_agents, 84, 84, 4)
                    actions: shape (steps_per_agent*num_agents, 84, 84, 4)
                    old_log_probs: shape (steps_per_agent*num_agents, )
                    returns: shape (steps_per_agent*num_agents, )
                    advantages: (steps_per_agent*num_agents, )
      clip_param: the PPO clipping parameter used to clamp ratios in loss function
      vf_coeff: weighs value function loss in total loss
      entropy_coeff: weighs entropy bonus in the total loss
      lr: learning rate, varies between optimization steps
          if decaying_lr_and_clip_param is set to true
      batch_size: the minibatch size, static argument
  
    Returns:
      optimizer: new optimizer after the parameters update
      loss: loss summed over training steps
    """
    iterations = trajectories[0].shape[0] // batch_size
    trajectories = jax.tree_map(
        lambda x: x.reshape((iterations, batch_size) + x.shape[1:]), trajectories)
    loss = 0.
    for batch in zip(*trajectories):
      grad_fn = jax.value_and_grad(loss_fn)
      l, grad = grad_fn(optimizer.target, module, batch, clip_param, vf_coeff,
                        entropy_coeff)
      loss += l
      optimizer = optimizer.apply_gradient(grad, learning_rate=lr)
    return optimizer, loss
  
  def get_experience(
      params: flax.core.frozen_dict.FrozenDict,
      module: models.ActorCritic,
      simulators: List[agent.RemoteSimulator],
      steps_per_actor: int):
    """Collect experience from agents.
  
    Runs `steps_per_actor` time steps of the game for each of the `simulators`.
    """
    all_experience = []
    # Range up to steps_per_actor + 1 to get one more value needed for GAE.
    for _ in range(steps_per_actor + 1):
      states = []
      for sim in simulators:
        state = sim.conn.recv()
        states.append(state)
      states = np.concatenate(states, axis=0)
      log_probs, values = agent.policy_action(params, module, states)
      log_probs, values = jax.device_get((log_probs, values))
      probs = np.exp(np.array(log_probs))
      for i, sim in enumerate(simulators):
        probabilities = probs[i]
        action = np.random.choice(probs.shape[1], p=probabilities)
        sim.conn.send(action)
      experiences = []
      for i, sim in enumerate(simulators):
        state, action, reward, done = sim.conn.recv()
        value = values[i, 0]
        log_prob = log_probs[i][action]
        sample = agent.ExpTuple(state, action, reward, value, log_prob, done)
        experiences.append(sample)
      all_experience.append(experiences)
    return all_experience
  
  def process_experience(
      experience: List[List[agent.ExpTuple]],
      actor_steps: int,
      num_agents: int,
      gamma: float,
      lambda_: float):
    """Process experience for training, including advantage estimation.
  
    Args:
      experience: collected from agents in the form of nested lists/namedtuple
      actor_steps: number of steps each agent has completed
      num_agents: number of agents that collected experience
      gamma: dicount parameter
      lambda_: GAE parameter
  
    Returns:
      trajectories: trajectories readily accessible for `train_step()` function
    """
    obs_shape = (84, 84, 4)
    exp_dims = (actor_steps, num_agents)
    values_dims = (actor_steps + 1, num_agents)
    states = np.zeros(exp_dims + obs_shape, dtype=np.float32)
    actions = np.zeros(exp_dims, dtype=np.int32)
    rewards = np.zeros(exp_dims, dtype=np.float32)
    values = np.zeros(values_dims, dtype=np.float32)
    log_probs = np.zeros(exp_dims, dtype=np.float32)
    dones = np.zeros(exp_dims, dtype=np.float32)
  
    for t in range(len(experience) - 1):  # experience[-1] only for next_values
      for agent_id, exp_agent in enumerate(experience[t]):
        states[t, agent_id, ...] = exp_agent.state
        actions[t, agent_id] = exp_agent.action
        rewards[t, agent_id] = exp_agent.reward
        values[t, agent_id] = exp_agent.value
        log_probs[t, agent_id] = exp_agent.log_prob
        # Dones need to be 0 for terminal states.
        dones[t, agent_id] = float(not exp_agent.done)
    for a in range(num_agents):
      values[-1, a] = experience[-1][a].value
    advantages = gae_advantages(rewards, dones, values, gamma, lambda_)
    returns = advantages + values[:-1, :]
    # After preprocessing, concatenate data from all agents.
    trajectories = (states, actions, log_probs, returns, advantages)
    trajectory_len = num_agents * actor_steps
    trajectories = tuple(map(
        lambda x: np.reshape(x, (trajectory_len,) + x.shape[2:]), trajectories))
    return trajectories
  
  def train(
      module: models.ActorCritic,
      optimizer: flax.optim.base.Optimizer,
      config: ml_collections.ConfigDict,
      model_dir: str):
    """Main training loop.
  
    Args:
      module: the actor-critic model
      optimizer: optimizer for the actor-critic model
      config: object holding hyperparameters and the training information
      model_dir: path to dictionary where checkpoints and logging info are stored
  
    Returns:
      optimizer: the trained optimizer
    """
    game = config.game + 'NoFrameskip-v4'
    simulators = [agent.RemoteSimulator(game)
                  for _ in range(config.num_agents)]
    summary_writer = tensorboard.SummaryWriter(model_dir)
    summary_writer.hparams(dict(config))
    loop_steps = config.total_frames // (config.num_agents * config.actor_steps)
    log_frequency = 40
    checkpoint_frequency = 500
  
  
    for s in range(loop_steps):
      # Bookkeeping and testing.
      if s % log_frequency == 0:
        score = test_episodes.policy_test(1, module, optimizer.target, game)
        frames = s * config.num_agents * config.actor_steps
        summary_writer.scalar('game_score', score, frames)
        print(f'Step {s}:\nframes seen {frames}\nscore {score}\n\n')
      if s % checkpoint_frequency == 0:
        checkpoints.save_checkpoint(model_dir, optimizer, s)
  
      # Core training code.
      alpha = 1. - s/loop_steps if config.decaying_lr_and_clip_param else 1.
      all_experiences = get_experience(
          optimizer.target, module, simulators, config.actor_steps)
      trajectories = process_experience(
          all_experiences, config.actor_steps, config.num_agents, config.gamma,
          config.lambda_)
      lr = config.learning_rate * alpha
      clip_param = config.clip_param * alpha
      for e in range(config.num_epochs):
        permutation = np.random.permutation(
            config.num_agents * config.actor_steps)
        trajectories = tuple(map(lambda x: x[permutation], trajectories))
        optimizer, loss = train_step(
            module, optimizer, trajectories, clip_param, config.vf_coeff,
            config.entropy_coeff, lr, config.batch_size)
    return optimizer
  