<commit_before>import jax
import jax.random
import jax.numpy as jnp
import numpy as onp
import flax
import time
from functools import partial
from typing import Tuple, List
from queue import Queue
import threading


from models import create_model, create_optimizer
from agent import policy_action
from remote import RemoteSimulator
from test_episodes import test

@partial(jax.vmap, in_axes=(1, 1, 1, None, None), out_axes=1)
@jax.jit
def gae_advantages(rewards, terminal_masks, values, discount, gae_param):
  """Use Generalized Advantage Estimation (GAE) to compute advantages
  Eqs. (11-12) in PPO paper arXiv: 1707.06347"""
  assert rewards.shape[0] + 1 == values.shape[0], ("One more value needed; "
          "Eq. (12) in PPO paper requires V(s_{t+1}) to calculate \delta_t")
  advantages, gae = [], 0
  # Key observation: A_{t} = \delta_t + \gamma*\lambda*A_{t+1}
  for t in reversed(range(len(rewards))):
    # masks 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] 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 = [gae] + advantages
  return jnp.array(advantages)

@jax.jit
def train_step(optimizer, trn_data, clip_param, vf_coeff, entropy_coeff):
  def loss_fn(model, minibatch, clip_param, vf_coeff, entropy_coeff):
    states, actions, old_log_probs, returns, advantages = minibatch
    probs, values = model(states)
    log_probs = jnp.log(probs)
    entropy = jnp.sum(-probs*log_probs, axis=1).mean()
    # we need to choose probs corresponding to actually taken actions
    # log_probs_act_taken = log_probs[jnp.arange(probs.shape[0]), actions])
    # above hits "Indexing mode not yet supported.", hence one hot solution
    act_one_hot = jax.nn.one_hot(actions, num_classes=probs.shape[1])
    log_probs_act_taken = jnp.log(jnp.sum(act_one_hot*probs, axis=1))
    ratios = jnp.exp(log_probs_act_taken - old_log_probs)
    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))
    value_loss = jnp.mean(jnp.square(returns - values))
    return PPO_loss + vf_coeff*value_loss - entropy_coeff*entropy

  batch_size = BATCH_SIZE
  iterations = trn_data[0].shape[0] // batch_size
  trn_data = jax.tree_map(
    lambda x: x.reshape((iterations, batch_size) + x.shape[1:]), trn_data)
  loss = 0.
  for batch in zip(*trn_data):
    grad_fn = jax.value_and_grad(loss_fn)
    l, grad = grad_fn(optimizer.target, batch, clip_param, vf_coeff,
    entropy_coeff)
    loss += l
    optimizer = optimizer.apply_gradient(grad)
  return optimizer, loss


def thread_inference(
  q1 : Queue,
  q2: Queue,<commit_msg>Add advantage normalization<commit_after>import jax
import jax.random
import jax.numpy as jnp
import numpy as onp
import flax
import time
from functools import partial
from typing import Tuple, List
from queue import Queue
import threading


from models import create_model, create_optimizer
from agent import policy_action
from remote import RemoteSimulator
from test_episodes import test

@partial(jax.vmap, in_axes=(1, 1, 1, None, None), out_axes=1)
@jax.jit
def gae_advantages(rewards, terminal_masks, values, discount, gae_param):
  """Use Generalized Advantage Estimation (GAE) to compute advantages
  Eqs. (11-12) in PPO paper arXiv: 1707.06347"""
  assert rewards.shape[0] + 1 == values.shape[0], ("One more value needed; "
          "Eq. (12) in PPO paper requires V(s_{t+1}) to calculate \delta_t")
  advantages, gae = [], 0
  # Key observation: A_{t} = \delta_t + \gamma*\lambda*A_{t+1}
  for t in reversed(range(len(rewards))):
    # masks 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] 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 = [gae] + advantages
  return jnp.array(advantages)

@jax.jit
def train_step(optimizer, trn_data, clip_param, vf_coeff, entropy_coeff):
  def loss_fn(model, minibatch, clip_param, vf_coeff, entropy_coeff):
    states, actions, old_log_probs, returns, advantages = minibatch
    probs, values = model(states)
    log_probs = jnp.log(probs)
    entropy = jnp.sum(-probs*log_probs, axis=1).mean()
    # we need to choose probs corresponding to actually taken actions
    # log_probs_act_taken = log_probs[jnp.arange(probs.shape[0]), actions])
    # above hits "Indexing mode not yet supported.", hence one hot solution
    act_one_hot = jax.nn.one_hot(actions, num_classes=probs.shape[1])
    log_probs_act_taken = jnp.log(jnp.sum(act_one_hot*probs, axis=1))
    ratios = jnp.exp(log_probs_act_taken - old_log_probs)
    # adv. 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))
    value_loss = jnp.mean(jnp.square(returns - values))
    return PPO_loss + vf_coeff*value_loss - entropy_coeff*entropy

  batch_size = BATCH_SIZE
  iterations = trn_data[0].shape[0] // batch_size
  trn_data = jax.tree_map(
    lambda x: x.reshape((iterations, batch_size) + x.shape[1:]), trn_data)
  loss = 0.
  for batch in zip(*trn_data):
    grad_fn = jax.value_and_grad(loss_fn)
    l, grad = grad_fn(optimizer.target, batch, clip_param, vf_coeff,
    entropy_coeff)
    loss += l
    optimizer = optimizer.apply_gradient(grad)
  return optimizer, loss


def thread_inference(
  q1 : Queue,
  q2: Queue,
{'fraction_non_alphanumeric': 0.10438101559907069, 'fraction_numerical': 0.013275804845668768, 'mean_word_length': 3.936117936117936, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '193400', 'n_tokens_mistral': 2166, 'n_tokens_neox': 2098, 'n_words': 562}