text: # Copyright 2021 The Brax 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.
  
  """Proximal policy optimization training.
  
  See: https://arxiv.org/pdf/1707.06347.pdf
  """
  
  import functools
  import time
  from typing import Any, Callable, Dict, Optional, List
  
  from absl import logging
  import flax
  import jax
  import jax.numpy as jnp
  from brax import envs
  from brax.training import distribution
  from brax.training import env
  from brax.training import networks
  from brax.training import normalization
  from brax.physics.base import take
  
  
  def compute_gae(truncation: jnp.ndarray,
                  rewards: jnp.ndarray,
                  values: jnp.ndarray,
                  bootstrap_value: jnp.ndarray,
                  lambda_: float = 1.0,
                  discount: float = 0.99):
    r"""Calculates the Generalized Advantage Estimation (GAE).
  
    Args:
      truncation: A float32 tensor of shape [T, B] with truncation signal.
      rewards: A float32 tensor of shape [T, B] containing rewards generated by
        following the behaviour policy.
      values: A float32 tensor of shape [T, B] with the value function estimates
        wrt. the target policy.
      bootstrap_value: A float32 of shape [B] with the value function estimate at
        time T.
      lambda_: Mix between 1-step (lambda_=0) and n-step (lambda_=1). Defaults to
        lambda_=1.
      discount: TD discount.
  
    Returns:
      A float32 tensor of shape [T, B]. Can be used as target to
        train a baseline (V(x_t) - vs_t)^2.
      A float32 tensor of shape [T, B] of advantages.
    """
    truncation_mask = 1 - truncation
    # Append bootstrapped value to get [v1, ..., v_t+1]
    values_t_plus_1 = jnp.concatenate(
        [values[1:], jnp.expand_dims(bootstrap_value, 0)], axis=0)
    deltas = rewards + discount * values_t_plus_1 - values
    deltas *= truncation_mask
  
    acc = jnp.zeros_like(bootstrap_value)
    vs_minus_v_xs = []
  
    def compute_vs_minus_v_xs(carry, target_t):
      lambda_, acc = carry
      truncation_mask, delta = target_t
      acc = delta + discount * truncation_mask * lambda_ * acc
      return (lambda_, acc), (acc)
  
    (_, _), (vs_minus_v_xs) = jax.lax.scan(compute_vs_minus_v_xs, (lambda_, acc),
                                           (truncation_mask, deltas),
                                           length=int(truncation_mask.shape[0]),
                                           reverse=True)
    # Add V(x_s) to get v_s.
    vs = jnp.add(vs_minus_v_xs, values)
  
    vs_t_plus_1 = jnp.concatenate([vs[1:], jnp.expand_dims(bootstrap_value, 0)],
                                  axis=0)
    advantages = (rewards + discount * vs_t_plus_1 - values) * truncation_mask
    return jax.lax.stop_gradient(vs), jax.lax.stop_gradient(advantages)
  
  
  @flax.struct.dataclass
  class StepData:
    """Contains data for one environment step."""
    obs: jnp.ndarray
    rewards: jnp.ndarray
    dones: jnp.ndarray
    actions: jnp.ndarray
    logits: jnp.ndarray
  
  
  @flax.struct.dataclass
  class TrainingState:
    """Contains training state for the learner."""
    optimizer: flax.optim.Optimizer
    key: jnp.ndarray
    normalizer_params: Any
  
  
  def compute_ppo_loss(
      models: Dict[str, Any],
      data: StepData,
      rng: jnp.ndarray,
      num_agents: int,
      parametric_action_distributions: List[distribution.ParametricDistribution],
      policies_apply: List[Any],
      value_apply: Any,
      entropy_cost: float = 1e-4,
      discounting: float = 0.9,
      reward_scaling: float = 1.0,
      lambda_: float = 0.95,
      ppo_epsilon: float = 0.3):
    """Computes PPO loss."""
    policies_params = [models[f'policy{i}'] for i in range(num_agents)]
    value_params = models['value']
  
    flattened_data_obs = jnp.concatenate([data.obs[..., i, :] for i in range(num_agents)], axis=-1)
    baseline = value_apply(value_params, flattened_data_obs)
  
    # Use last baseline value (from the value function) to bootstrap.
    bootstrap_value = baseline[-1]
    baseline = baseline[:-1]
  
    # At this point, we have unroll length + 1 steps. The last step is only used
    # as bootstrap value, so it's removed.
  
    # already removed at data generation time
    # actions = actions[:-1]
    # logits = logits[:-1]
  
    rewards = data.rewards[1:] * reward_scaling
    # dones is scalar
    truncation = jnp.repeat(jnp.expand_dims(data.dones[1:], axis=-1), 
                            repeats=num_agents, 
                            axis=-1)
  
    vs, advantages = compute_gae(
        truncation=truncation,
        rewards=rewards,
        values=baseline,
        bootstrap_value=bootstrap_value,
        lambda_=lambda_,
        discount=discounting)
  
    # Value function loss
    v_error = vs - baseline
    v_loss = jnp.mean(v_error * v_error) * 0.5 * 0.5
  
    obs_list = [jnp.take(data.obs[:-1], i, axis=-2) for i in range (num_agents)] # 1 is agent axis, 0 is batch axis
    actions_list = [jnp.take(data.actions, i, axis=-2) for i in range (num_agents)] # 1 is agent axis, 0 is batch axis
    logits_list = [jnp.take(data.logits, i, axis=-2) for i in range (num_agents)] # 1 is agent axis, 0 is batch axis
  
    policies_logits = jax.tree_multimap(
      lambda policy_apply, policy_params, obs: policy_apply(policy_params, obs), # :-1 is it like a squeeze??
      policies_apply, policies_params, obs_list
      )
  
    def compute_rho_entropy(key, parametric_action_distr, 
                            policy_logits, 
                            data_actions, data_logits,
                            ):
      target_action_log_probs = parametric_action_distr.log_prob(
          policy_logits, data_actions)
      behaviour_action_log_probs = parametric_action_distr.log_prob(
          data_logits, data_actions)
  
      rho = jnp.exp(target_action_log_probs - behaviour_action_log_probs)
      entropy = parametric_action_distr.entropy(policy_logits, key)
      return rho, entropy
  
    # should we pass back RNG?
    rng_list = list(jax.random.split(rng, num_agents+1))
    rho_s, entropies = jax.tree_multimap(
      compute_rho_entropy, 
      rng_list[1:], parametric_action_distributions,
      policies_logits, actions_list, logits_list
      )
    rho_s = jnp.concatenate(rho_s, axis=-1)
    entropies = jnp.concatenate(entropies, axis=-1)
  
    advantages = jnp.concatenate([advantages[..., i] for i in range(num_agents)], axis=-1)
  
    # Policy loss
    # surrogate_loss_concat = jnp.dot(rho_s_concat, advantages_concat)
    surrogate_loss1 = rho_s * advantages
    surrogate_loss2 = jnp.clip(rho_s, 1 - ppo_epsilon, 
                               1 + ppo_epsilon) * advantages
    policy_loss = -jnp.mean(jnp.minimum(surrogate_loss1, surrogate_loss2))
  
    # Entropy reward
    entropy = jnp.mean(entropies)
    entropy_loss = entropy_cost * -entropy
  
    return policy_loss + v_loss + entropy_loss, {
        'total_loss': policy_loss + v_loss + entropy_loss,
        'policy_loss': policy_loss,
        'v_loss': v_loss,
        'entropy_loss': entropy_loss
    }
  
  
  def multiagent_forward_pass_eval(key, policy_model, policy_params, 
                                   parametric_action_distr, agent_obs):
    agent_logits = policy_model.apply(policy_params, agent_obs)
    agent_actions = parametric_action_distr.sample(agent_logits, key)
    return agent_actions
  
  
  def multiagent_forward_pass(key, policy_model, policy_params, 
                              parametric_action_distr, agent_obs):
    agent_logits = policy_model.apply(policy_params, agent_obs)
    agent_actions = parametric_action_distr.sample_no_postprocessing(agent_logits, key)
    agent_actions = parametric_action_distr.postprocess(agent_actions)
    return agent_actions, agent_logits
  
  
  def train(
      environment_fn: Callable[..., envs.Env],
      num_agents: int,
      num_timesteps,
      episode_length: int,
      action_repeat: int = 1,
      num_envs: int = 1,
      max_devices_per_host: Optional[int] = None,
      num_eval_envs: int = 128,
      learning_rate=1e-4,
      entropy_cost=1e-4,
      discounting=0.9,
      seed=0,
      unroll_length=10,
      batch_size=32,
      num_minibatches=16,
      num_update_epochs=2,
      log_frequency=10,
      normalize_observations=False,
      reward_scaling=1.,
      progress_fn: Optional[Callable[[int, Dict[str, Any]], None]] = None,
  ):
    """PPO training."""
    assert batch_size * num_minibatches % num_envs == 0
    xt = time.time()
  
    process_count = jax.process_count()
    process_id = jax.process_index()
    local_device_count = jax.local_device_count()
    local_devices_to_use = local_device_count
    if max_devices_per_host:
      local_devices_to_use = min(local_devices_to_use, max_devices_per_host)
    logging.info(
        'Device count: %d, process count: %d (id %d), local device count: %d, '
        'devices to be used count: %d',
        jax.device_count(), process_count, process_id, local_device_count,
        local_devices_to_use)
  
    key = jax.random.PRNGKey(seed)
    key, key_models, key_env, key_eval = jax.random.split(key, 4)
    # Make sure every process gets a different random key, otherwise they will be
    # doing identical work.
    key_env = jax.random.split(key_env, process_count)[process_id]
    key = jax.random.split(key, process_count)[process_id]
    # key_models should be the same, so that models are initialized the same way
    # for different processes
  
    core_env = environment_fn(
        action_repeat=action_repeat,
        batch_size=num_envs // local_devices_to_use // process_count,
        episode_length=episode_length)
    key_envs = jax.random.split(key_env, local_devices_to_use)
    tmp_env_states = []
    for key in key_envs:
      first_state, step_fn = env.wrap(core_env, key)
      tmp_env_states.append(first_state)
    first_state = jax.tree_multimap(lambda *args: jnp.stack(args),
                                    *tmp_env_states)
  
    core_eval_env = environment_fn(
        action_repeat=action_repeat,
        batch_size=num_eval_envs,
        episode_length=episode_length)
    eval_first_state, eval_step_fn = env.wrap(core_eval_env, key_eval)
  
    parametric_action_distributions = [distribution.NormalTanhDistribution(
          event_size=core_env.agent_action_size) for i in range(num_agents)]
  
    policy_models, value_model = networks.make_multiagent_models(
        [act_distr.param_size for act_distr in parametric_action_distributions],
        num_agents,
        core_env.observation_size)
    key_policy, key_value = jax.random.split(key_models)
  
    optimizer_def = flax.optim.Adam(learning_rate=learning_rate)
    optimizer = optimizer_def.create({**{f'policy{i}': policy_models[i].init(key_policy) for i in range(num_agents)},
                                      'value': value_model.init(key_value)
                                      }
                                      )
    optimizer = normalization.bcast_local_devices(
        optimizer, local_devices_to_use)
  
    normalizer_params, obs_normalizer_update_fn, obs_normalizer_apply_fn = (
        normalization.create_observation_normalizer(
            (num_agents, core_env.agent_observation_size), normalize_observations,
            num_leading_batch_dims=2, pmap_to_devices=local_devices_to_use)) # why is # leading batch dims 2?
  
    # TODO: replace all core_env.observation_size, core_env.action_size
    key_debug = jax.random.PRNGKey(seed + 666)
  
    loss_fn = functools.partial(
        compute_ppo_loss,
        num_agents=num_agents,
        parametric_action_distributions=parametric_action_distributions,
        policies_apply=[policy_model.apply for policy_model in policy_models],
        value_apply=value_model.apply,
        entropy_cost=entropy_cost,
        discounting=discounting,
        reward_scaling=reward_scaling)
  
    grad_loss = jax.grad(loss_fn, has_aux=True)
  
    def do_one_step_eval(carry, unused_target_t):
      state, policies_params, normalizer_params, key = carry
      # TODO: Make this nicer ([0] comes from pmapping).
      obs = obs_normalizer_apply_fn(
          jax.tree_map(lambda x: x[0], normalizer_params), state.core.obs)
      obs_list = [jnp.take(obs, i, axis=1) for i in range (num_agents)] # 1 is agent axis, 0 is batch axis
      keys = jax.random.split(key, num_agents+1)
      keys = [k for k in keys]
  
      actions = jax.tree_multimap(
        multiagent_forward_pass_eval, keys[1:], policy_models, policies_params, parametric_action_distributions, obs_list
        )
      actions = jnp.stack(actions, axis=1) # shape is (256, 17); want shape (128, 2, 17)
      nstate = eval_step_fn(state, actions)
      return (nstate, policies_params, normalizer_params, keys[0]), ()
  
    @jax.jit
    def run_eval(state, key, policies_params, normalizer_params):
      # TODO: does this tree_map need to go down one more level??
      policies_params = jax.tree_map(lambda x: x[0], policies_params)
      (state, _, _, key), _ = jax.lax.scan(
          do_one_step_eval, (state, policies_params, normalizer_params, key), (),
          length=episode_length // action_repeat)
      return state, key
  
    def do_one_step(carry, unused_target_t):
      state, normalizer_params, policies_params, key = carry
      normalized_obs = obs_normalizer_apply_fn(normalizer_params, state.core.obs)
      obs_list = [jnp.take(normalized_obs, i, axis=1) for i in range(num_agents)] # 1 is agent axis, 0 is batch axis
      keys = jax.random.split(key, num_agents+1)
      keys = [k for k in keys]
  
      ret = jax.tree_multimap(
        multiagent_forward_pass, keys[1:], policy_models, policies_params, parametric_action_distributions, obs_list
        )
      postprocessed_actions, logits = list(zip(*ret))  # unzip list of results
      postprocessed_actions = jnp.stack(postprocessed_actions, axis=1) # shape is (256, 17); want shape (128, 2, 17)
      logits = jnp.stack(logits, axis=1)
      nstate = step_fn(state, postprocessed_actions)
      return (nstate, normalizer_params, policies_params, keys[0]), StepData(
          obs=state.core.obs,
          rewards=state.core.reward,
          dones=state.core.done,
          actions=postprocessed_actions,
          logits=logits)
  
    def generate_unroll(carry, unused_target_t):
      state, normalizer_params, policies_params, key = carry
      (state, _, _, key), data = jax.lax.scan(
          do_one_step, (state, normalizer_params, policies_params, key), (),
          length=unroll_length)
      data = data.replace(
          obs=jnp.concatenate(
              [data.obs, jnp.expand_dims(state.core.obs, axis=0)]),
          rewards=jnp.concatenate(
              [data.rewards, jnp.expand_dims(state.core.reward, axis=0)]),
          dones=jnp.concatenate(
              [data.dones, jnp.expand_dims(state.core.done, axis=0)]))
      return (state, normalizer_params, policies_params, key), data
  
    def update_model(carry, data):
      optimizer, key = carry
      key, key_loss = jax.random.split(key)
      loss_grad, metrics = grad_loss(optimizer.target, data, key_loss)
      loss_grad = jax.lax.pmean(loss_grad, axis_name='i')
      optimizer = optimizer.apply_gradient(loss_grad)
      return (optimizer, key), metrics
  
    def minimize_epoch(carry, unused_t):
      optimizer, data, key = carry
      key, key_perm, key_grad = jax.random.split(key, 3)
      permutation = jax.random.permutation(key_perm, data.obs.shape[1])
      def convert_data(data, permutation):
        data = jnp.take(data, permutation, axis=1, mode='clip')
        data = jnp.reshape(
            data, [data.shape[0], num_minibatches, -1] + list(data.shape[2:]))
        data = jnp.swapaxes(data, 0, 1)
        return data
      ndata = jax.tree_map(lambda x: convert_data(x, permutation), data)
      (optimizer, _), metrics = jax.lax.scan(
          update_model, (optimizer, key_grad), ndata, length=num_minibatches)
      return (optimizer, data, key), metrics
  
    def run_epoch(carry, unused_t):
      training_state, state = carry
      key_minimize, key_generate_unroll, new_key = jax.random.split(
          training_state.key, 3)
      (state, _, _, _), data = jax.lax.scan(
          generate_unroll, (state, 
                            training_state.normalizer_params,
                            [training_state.optimizer.target[f'policy{i}'] for i in range(num_agents)],
                            key_generate_unroll), (),
          length=batch_size * num_minibatches // num_envs)
      # make unroll first
      data = jax.tree_map(lambda x: jnp.swapaxes(x, 0, 1), data)
      data = jax.tree_map(
          lambda x: jnp.reshape(x, [x.shape[0], -1] + list(x.shape[3:])), data)
  
      # Update normalization params and normalize observations.
      normalizer_params = obs_normalizer_update_fn(
          training_state.normalizer_params, data.obs[:-1])
      data = data.replace(
          obs=obs_normalizer_apply_fn(normalizer_params, data.obs))
  
      (optimizer, _, _), metrics = jax.lax.scan(
          minimize_epoch, (training_state.optimizer, data, key_minimize),
          (), length=num_update_epochs)
  
      new_training_state = TrainingState(optimizer=optimizer,
                                         normalizer_params=normalizer_params,
                                         key=new_key)
      return (new_training_state, state), metrics
  
    num_epochs = num_timesteps // (
        batch_size * unroll_length * num_minibatches * action_repeat)
  
    def _minimize_loop(training_state, state):
      (training_state, state), losses = jax.lax.scan(
          run_epoch, (training_state, state), (),
          length=num_epochs // log_frequency)
      losses = jax.tree_map(jnp.mean, losses)
      return (training_state, state), losses
  
    minimize_loop = jax.pmap(_minimize_loop, axis_name='i')
  
    training_state = TrainingState(
        optimizer=optimizer,
        key=jnp.stack(jax.random.split(key, local_devices_to_use)),
        normalizer_params=normalizer_params)
    training_walltime = 0
    eval_walltime = 0
    sps = 0
    eval_sps = 0
    losses = {}
    state = first_state
    metrics = {}
  
    for it in range(log_frequency + 1):
      logging.info('starting iteration %s %s', it, time.time() - xt)
      t = time.time()
  
      if process_id == 0:
        eval_state, key_debug = (
            run_eval(eval_first_state, key_debug,
                     [training_state.optimizer.target[f'policy{i}'] for i in range(num_agents)],
                     training_state.normalizer_params))
        eval_state.total_episodes.block_until_ready()
        eval_walltime += time.time() - t
        eval_sps = (episode_length * eval_first_state.core.reward.shape[0] /
                    (time.time() - t))
        metrics = dict(
            dict({f'eval/episode_{name}': value / eval_state.total_episodes
                  for name, value in eval_state.total_metrics.items()}),
            **dict({
                'eval/total_episodes': eval_state.total_episodes,
                'speed/sps': sps,
                'speed/eval_sps': eval_sps,
                'speed/training_walltime': training_walltime,
                'speed/eval_walltime': eval_walltime,
                'speed/timestamp': training_walltime,
                'losses/total_loss': jnp.mean(losses.get('total_loss', 0)),
                'losses/policy_loss': jnp.mean(losses.get('policy_loss', 0)),
                'losses/value_loss': jnp.mean(losses.get('v_loss', 0)),
                'losses/entropy_loss': jnp.mean(losses.get('entropy_loss',
                                                           0))}))
        logging.info(metrics)
        if progress_fn:
          progress_fn(int(training_state.normalizer_params[0][0]) * action_repeat,
                      metrics)
  
      if it == log_frequency:
        break
  
      t = time.time()
      previous_step = training_state.normalizer_params[0][0]
      # optimization
      (training_state, state), losses = minimize_loop(training_state, state)
      jax.tree_map(lambda x: x.block_until_ready(), losses)
      sps = ((training_state.normalizer_params[0][0] - previous_step) /
             (time.time() - t)) * action_repeat
      training_walltime += time.time() - t
  
    # To undo the pmap.
    normalizer_params = jax.tree_map(lambda x: x[0],
                                     training_state.normalizer_params)
    # Do we need to undo the pmap some other way for n policies?
    policies_params = jax.tree_map(lambda x: x[0],
                                 [training_state.optimizer.target[f'policy{i}'] for i in range(num_agents)]
                                 )
  
    logging.info('total steps: %s', normalizer_params[0] * action_repeat)
  
    _, inference = make_params_and_inference_fn(core_env.agent_observation_size,
                                                core_env.agent_action_size, 
                                                num_agents,
                                                normalize_observations)
    params = normalizer_params, policies_params
  
    if process_count > 1:
      # Make sure all processes stay up until the end of main.
      x = jnp.ones([jax.local_device_count()])
      x = jax.device_get(jax.pmap(lambda x: jax.lax.psum(x, 'i'), 'i')(x))
      assert x[0] == jax.device_count()
  
    return (inference, params, metrics)
  
  
  def make_params_and_inference_fn(agent_obs_size, agent_action_size, num_agents,
                                   normalize_observations):
    """Creates params and inference function for the PPO agent."""
    obs_normalizer_params, obs_normalizer_apply_fn = normalization.make_data_and_apply_fn(
        normalize_observations)
    parametric_action_distributions = [distribution.NormalTanhDistribution(
          event_size=agent_action_size) for i in range(num_agents)]
    policy_models, _ = networks.make_multiagent_models(
        [act_distr.param_size for act_distr in parametric_action_distributions],
        num_agents,
        agent_obs_size)
  
  
    def inference_fn(params, obs, key):
      normalizer_params, policies_params = params
      obs = obs_normalizer_apply_fn(normalizer_params, obs)
      obs_list = [jnp.take(obs, i, axis=0) for i in range (num_agents)] # no batch axis
      keys = list(jax.random.split(key, num_agents))
  
      agent_actions = jax.tree_multimap(
        multiagent_forward_pass_eval, keys, policy_models, policies_params, parametric_action_distributions, obs_list
        )
      agent_actions = jnp.stack(agent_actions, axis=1)
      return agent_actions
  
    key = jax.random.PRNGKey(0)
    policy_keys = jax.random.split(key, num_agents) # TODO: pass key out of inference fn??
    policies_params = [policy_models[i].init(policy_keys[i]) for i in range(num_agents)]
  
    params = (obs_normalizer_params, policies_params)
    return params, inference_fn
  