<commit_before>from matplotlib import pyplot as plt
import tensorflow as tf
import numpy as np
import gym
import sys

# TODO: It could be interesting to work on a version of this problem with no time limit.

N_EPISODES = 10000
# It sort of makes sense that this wants to be low for this problem. There isn't
# really any long-term planning involved. In fact, could probably even set this to 0.
GAMMA = 0.8
RENDER_EPISODES = 0
L2 = 1
L2_WEIGHT = 0.001
COMPLETION_BONUS = 0

INPUT_LAST_CHAR = 1
INPUT_LAST_ACTION = 1
INPUT_LAST_REWARD = 0

def bias_var(shape):
  initial = tf.constant(.1, shape=shape)
  return tf.Variable(initial)

def weight_var(shape):
  initial = tf.truncated_normal(shape, stddev=0.1)
  return tf.Variable(initial)

def episode(sess, env, pactions, nchars, actions_size, x, render=False):
  obs = env.reset()
  # Last position represents no previous action taken
  last_action = [0] * actions_size + [1]
  last_char = None
  last_reward = 0
  rewards = []
  xs = []
  action_history = []
  total_reward = 0
  def pack_x():
    xarr = [0]*nchars
    xarr[obs] = 1
    if INPUT_LAST_CHAR:
      last_char_arr = [0]*nchars
      if last_char is not None:
        last_char_arr[last_char] = 1
      xarr += last_char_arr
    if INPUT_LAST_ACTION:
      xarr += last_action
    if INPUT_LAST_REWARD:
      xarr += [last_reward]
    return xarr
    #return curr_char_arr + last_char_arr + last_action + [last_reward]

  counter = 0
  for i in range(100):
    counter += 1
    if render:
      env.render()
    curr_x = pack_x()
    xs.append(curr_x)
    action_probs = sess.run(pactions, feed_dict={x: [curr_x]})
    # choose next action
    actions = []
    actions_onehot = [0] * actions_size
    onehot_offset = 0
    for i, probs in enumerate(action_probs):
      assert probs.shape[0] == 1
      probs = probs[0]
      sample = np.sum(np.cumsum(probs) < np.random.rand())
      actions_onehot[onehot_offset+sample] = 1
      onehot_offset += probs.shape[0]
      actions.append(sample)

    action_history.append(actions_onehot)
    assert env.action_space.contains(actions)
    last_char = obs
    obs, last_reward, done, _ = env.step(actions)
    total_reward += last_reward
    rewards.append(last_reward)
    last_action = actions_onehot + [0]
    if done:
      if render:
        env.render()
      break
      
  return xs, action_history, rewards, total_reward, counter

def policy_gradient(paction_cat, actions_size):
  rewards = tf.placeholder(tf.float32, shape=[None])
  actions = tf.placeholder(tf.float32, shape=[None, actions_size])
  # For each transition, isolate the probabilities of the actions that
  # were actually performed.
  focused_probs = tf.reduce_sum(tf.mul(actions, paction_cat), reduction_indices=[1])
  # Starting from gradients that encourage the performed actions, multiply by
  # the 'advantage'. If the performed actions were bad, this should be a negative
  # multiplier, thus discouraging that set of actions.
  # (The particular loss function I'm using here is in imitation of this example:
  #     https://github.com/kvfrans/openai-cartpole/blob/master/cartpole-policygradient.py
  #  not sure if there are other reasonable ways to do it.)
  
  #tempered = tf.mul(tf.log(focused_probs), rewards)
  tempered = tf.mul(focused_probs, rewards)
  loss = -tf.reduce_sum(tempered)
  if L2:
    l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in tf.trainable_variables()])
    loss = loss + L2_WEIGHT * l2_loss
  optimizer = tf.train.AdamOptimizer(0.001).minimize(loss)
  return (
    rewards, actions, optimizer,
      focused_probs, tempered, loss,
  )

def discounted_rewards(r):
  r2 = np.zeros_like(r)
  running_sum = 0
  #assert r.shape[1] == 1
  bonus = COMPLETION_BONUS if (COMPLETION_BONUS and r[-1] == 1) else 1
  for i in reversed(range(len(r))):
    # The min here is basically cheating :/
    #running_sum = min(running_sum * GAMMA,0) + r[i]

    # Also cheating
    #rew = -.2 if r[i] == 0 else r[i]
    rew = r[i]
    running_sum = running_sum * GAMMA + rew
    r2[i] = running_sum * bonus
  return r2


if __name__ == '__main__':
  # TODO: This would all have been easier if I had just used np arrays throughout rather than
  # mixing them with lists. Bleh.
  sess = tf.InteractiveSession()
  #env = gym.make('DuplicatedInput-v0')
  env = gym.make('Copy-v0')

  n_hidden = 30

  nchars = env.observation_space.n 
  n_actions = len(env.action_space.spaces)
  # 2 + 2 + 5 (move head, write or nah, which char to write)
  actions_size = sum(space.n for space in env.action_space.spaces)
  # Currently seen char, last seen char, last action, last reward
  input_size = (nchars + 
    (INPUT_LAST_CHAR * nchars) + 
    # We'll add one more input which represents "no action taken" (i.e. this is the first step)
    INPUT_LAST_ACTION * (actions_size+1) + 
    INPUT_LAST_REWARD * 1
  )
  x = tf.placeholder(tf.float32, shape=[None, input_size])
  
  w1 = weight_var([input_size, n_hidden])
  b1 = bias_var([n_hidden])
  h1 = tf.nn.relu( tf.matmul(x, w1) + b1 )
<commit_msg>some cleanup<commit_after>from matplotlib import pyplot as plt
import tensorflow as tf
import numpy as np
import gym
import sys

# TODO: It could be interesting to work on a version of this problem with no time limit.

N_EPISODES = 10000
# It sort of makes sense that this wants to be low for this problem. There isn't
# really any long-term planning involved. In fact, could probably even set this to 0.
GAMMA = 0.8
RENDER_EPISODES = 0
L2 = 1
L2_WEIGHT = 0.001
COMPLETION_BONUS = 0

INPUT_LAST_CHAR = 1
INPUT_LAST_ACTION = 1
INPUT_LAST_REWARD = 0

# Setting this to -.2 or less seems to help a lot for copy (almost necessary)
# Sweet spot maybe around -.6
# Obviously a high value of this is not good for dedupe, but it can tolerate 
# a penalty <= -.3. A very small penalty like -.01 might even help? Haven't
# collected a ton of data though.
ZERO_PENALTY = -0.01

def bias_var(shape):
  initial = tf.constant(.1, shape=shape)
  return tf.Variable(initial)

def weight_var(shape):
  initial = tf.truncated_normal(shape, stddev=0.1)
  return tf.Variable(initial)

def episode(sess, env, pactions, nchars, actions_size, x, render=False):
  obs = env.reset()
  # Last position represents no previous action taken
  last_action = [0] * actions_size + [1]
  last_char = None
  last_reward = 0
  rewards = []
  xs = []
  action_history = []
  total_reward = 0
  def pack_x():
    xarr = [0]*nchars
    xarr[obs] = 1
    if INPUT_LAST_CHAR:
      last_char_arr = [0]*nchars
      if last_char is not None:
        last_char_arr[last_char] = 1
      xarr += last_char_arr
    if INPUT_LAST_ACTION:
      xarr += last_action
    if INPUT_LAST_REWARD:
      xarr += [last_reward]
    return xarr

  counter = 0
  for i in range(100):
    counter += 1
    if render:
      env.render()
    curr_x = pack_x()
    xs.append(curr_x)
    action_probs = sess.run(pactions, feed_dict={x: [curr_x]})
    # choose next action
    actions = []
    actions_onehot = [0] * actions_size
    onehot_offset = 0
    for i, probs in enumerate(action_probs):
      assert probs.shape[0] == 1
      probs = probs[0]
      sample = np.sum(np.cumsum(probs) < np.random.rand())
      actions_onehot[onehot_offset+sample] = 1
      onehot_offset += probs.shape[0]
      actions.append(sample)

    action_history.append(actions_onehot)
    assert env.action_space.contains(actions)
    last_char = obs
    obs, last_reward, done, _ = env.step(actions)
    total_reward += last_reward
    rewards.append(last_reward)
    last_action = actions_onehot + [0]
    if done:
      if render:
        env.render()
      break
      
  return xs, action_history, rewards, total_reward, counter

def policy_gradient(paction_cat, actions_size):
  rewards = tf.placeholder(tf.float32, shape=[None])
  actions = tf.placeholder(tf.float32, shape=[None, actions_size])
  # For each transition, isolate the probabilities of the actions that
  # were actually performed.
  focused_probs = tf.reduce_sum(tf.mul(actions, paction_cat), reduction_indices=[1])
  # Starting from gradients that encourage the performed actions, multiply by
  # the 'advantage'. If the performed actions were bad, this should be a negative
  # multiplier, thus discouraging that set of actions.
  # (The particular loss function I'm using here is in imitation of this example:
  #     https://github.com/kvfrans/openai-cartpole/blob/master/cartpole-policygradient.py
  #  not sure if there are other reasonable ways to do it.)
  
  #tempered = tf.mul(tf.log(focused_probs), rewards)
  tempered = tf.mul(focused_probs, rewards)
  loss = -tf.reduce_sum(tempered)
  if L2:
    l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in tf.trainable_variables()])
    loss = loss + L2_WEIGHT * l2_loss
  optimizer = tf.train.AdamOptimizer(0.001).minimize(loss)
  return (
    rewards, actions, optimizer,
      focused_probs, tempered, loss,
  )

def discounted_rewards(r):
  r2 = np.zeros_like(r)
  running_sum = 0
  bonus = COMPLETION_BONUS if (COMPLETION_BONUS and r[-1] == 1) else 1
  for i in reversed(range(len(r))):
    # The min here is basically cheating :/
    #running_sum = min(running_sum * GAMMA,0) + r[i]
    rew = r[i]
    # Also cheating
    if ZERO_PENALTY and r[i] == 0:
      rew = ZERO_PENALTY
    running_sum = running_sum * GAMMA + rew
    r2[i] = running_sum * bonus
  return r2


if __name__ == '__main__':
  # TODO: This would all have been easier if I had just used np arrays throughout rather than
  # mixing them with lists. Bleh.
  sess = tf.InteractiveSession()
  env = gym.make('DuplicatedInput-v0')
  #env = gym.make('Copy-v0')

  n_hidden = 30

  nchars = env.observation_space.n 
  n_actions = len(env.action_space.spaces)
  # 2 + 2 + 5 (move head, write or nah, which char to write)
  actions_size = sum(space.n for space in env.action_space.spaces)
  # Currently seen char, last seen char, last action, last reward
  input_size = (nchars + 
    (INPUT_LAST_CHAR * nchars) + 
    # We'll add one more input which represents "no action taken" (i.e. this is the first step)
    INPUT_LAST_ACTION * (actions_size+1) + 
    INPUT_LAST_REWARD * 1
  )
  x = tf.placeholder(tf.float32, shape=[None, input_size])
  
  w1 = weight_var([input_size, n_hidden])
  b1 = bias_var([n_hidden])
  h1 = tf.nn.relu( tf.matmul(x, w1) + b1 )

{'fraction_non_alphanumeric': 0.1013014763014763, 'fraction_numerical': 0.017191142191142192, 'mean_word_length': 3.4174174174174174, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4381923', 'n_tokens_mistral': 3741, 'n_tokens_neox': 3566, 'n_words': 1047}