pg-pong py3

"""
Trains an agent with stochastic policy gradients on Pong
from OpenAI Gym following Karpathy's example
"""

import numpy as np
import pickle
import gym

# hyperparameters
H = 200  # number of hidden layer neurons
batch_size = 2  # update parameters after this many episodes
learning_rate = 1e-4
gamma = 0.99  # discount factor for reward
decay_rate = 0.99  # decay factor for RMS prop leaky sum of grad^2
resume = False  # Resume from previous checkpoint?
render = False

# model initialization
D = 80 * 80  # input dimensionality: 80x80 grid
if resume:
    model = pickle.load(open('save.p', 'rb'))
else:
    model = dict()
    model['W1'] = np.random.randn(H, D) / np.sqrt(D)  # use xavier intialization
    model['W2'] = np.random.randn(H) / np.sqrt(H)

grad_buffer = {k: np.zeros_like(v) for k, v in model.items()}
rms_prop_cache = {k: np.zeros_like(v) for k, v in model.items()}


def sigmoid(x):
  return 1.0 / (1.0 + np.exp(-x))  # sigmoid function to sqash


def prepro(I):
    """ Prepro 210x160x3 uint frame into 6400 (80x80) 1D float vector """
    I = I[35:195]  # crop to game board
    I = I[::2, ::2, 0]  # down sample by a factor of 2
    I[I == 144] = 0  # erase background type 1
    I[I == 109] = 0  # erase background type 2
    I[I != 0] = 1  # everything else is just 1
    return I.astype(np.float).ravel()


def discount_rewards(r):
    """ take 1D float array of rewards and compute discounted reward """
    discounted_r = np.zeros_like(r)
    running_add = 0
    for t in reversed(range(0, r.size)):
        if r[t] != 0:
            running_add = 0
            # reset sum since the ball hit a wall
        discounted_r[t] = running_add
    return discounted_r


def policy_forward(x):
    h = np.dot(model['W1'], x)
    # print('h is {}'.format(h))
    h[h < 0] = 0  # apply ReLU
    logp = np.dot(model['W2'], h)
    p = sigmoid(logp)
    return p, h # probability of action 2 and hidden state


def policy_backward(eph, epdlogp):
    """ backward pass. (epdh is array of intermediate hidden states) """
    dW2 = np.dot(eph.T, epdlogp).ravel()
    dh = np.outer(epdlogp, model['W2'])
    dh[eph <= 0] = 0  # ReLU for back prop
    dW1 = np.dot(dh.T, epx)
    return {'W1': dW1, 'W2': dW2}


env = gym.make("Pong-v0")
observation = env.reset()
prev_x = None  # used to compute difference frame
xs, hs, dlogps, drs = [], [], [], []
running_reward = None
reward_sum = 0
episode_number = 0
while True:
    if render:
        env.render()

    # useful later for scrabble!
    # pre-process observation, set input to network as difference image
    cur_x = prepro(observation)
    x = cur_x - prev_x if prev_x is not None else np.zeros(D)
    #print(x)
    prev_x = cur_x

    # forward the policy network and sample an action from the returned probability
    aprob, h = policy_forward(x)
    action = 2 if np.random.uniform() < aprob else 3

    # record various intermediates necessary later for backprop
    xs.append(x) # observation
    hs.append(h) # hidden state
    y = 1 if action == 2 else 0 # fake label
    dlogps.append(y - aprob) # grad encourages actions taken to be taken
    # http://cs231n.github.io/neural-networks-2/#losses

    # step the environment and get new measurements
    observation, reward, done, info = env.step(action)
    reward_sum += reward

    drs.append(reward) # record reward

    if done:
        print(episode_number)
        episode_number += 1

        # stack together all inputs, hidden states, action gradients, and rewards for this episode
        epx = np.vstack(xs)
        eph = np.vstack(hs)
        epdlogp = np.vstack(dlogps)
        epr = np.vstack(drs)
        xs, hs, dlogps, drs = [], [], [], [] # reset array memory

        # compute discounted reward backwards through time
        discounted_epr = discount_rewards(epr)
        # standardize rewards to unit normal
        discounted_epr -= np.mean(discounted_epr)
        discounted_epr /= np.std(discounted_epr)

        epdlogp *= discounted_epr # modulate gradient with advantage (what...?)
        grad = policy_backward(eph, epdlogp)
        #print(grad)
        for k in model:  grad_buffer[k] += grad[k] # accumulate grad over batch
        #print(grad_buffer)

        # perform rms prop update every batch_size episodes
        if episode_number % batch_size == 0:
          print(model.items())
          for k, v in model.items():
            g = grad_buffer[k] # gradient
            rms_prop_cache[k] = decay_rate * rms_prop_cache[k] + (1 - decay_rate) * g ** 2
            model[k] += learning_rate * g / (np.sqrt(rms_prop_cache[k]) + 1e-5)
            grad_buffer[k] = np.zeros_like(v) # reset gradient buffer

        # bookkeeping
        running_reward = reward_sum if running_reward is None else running_reward * 0.99 + reward_sum * 0.01
        print('Resetting env. Episode reward total was {}. Running mean: {}'.format(reward_sum, running_reward))
        if episode_number % 20 == 0:
          pickle.dump(model, open('save.p', 'wb'))
        reward_sum = 0
        observation = env.reset()
        prev_x = None