old_code/train_atari_agent.py

import os
import math
import time
import gym
import random
import utils
import keras
import numpy as np

from collections import deque
from matplotlib import pyplot as plt
from sklearn.preprocessing import OneHotEncoder

class ReplayBuffer():
    """
        Thank you: https://github.com/BY571/
    """

    def __init__(self, max_size):
        self.max_size = max_size
        self.buffer = []
     
    def add_sample(self, states, actions, rewards):
        episode = {"states": states, "actions":actions, "rewards": rewards, "summed_rewards":sum(rewards)}
        self.buffer.append(episode)
    
    def sort(self):
        #sort buffer
        self.buffer = sorted(self.buffer, key = lambda i: i["summed_rewards"],reverse=True)
        # keep the max buffer size
        self.buffer = self.buffer[:self.max_size]

    def get_random_samples(self, batch_size):
        self.sort()
        idxs = np.random.randint(0, len(self.buffer), batch_size)
        batch = [self.buffer[idx] for idx in idxs]
        return batch
    
    def get_n_best(self, n):
        self.sort()
        return self.buffer[:n]
    
    def __len__(self):
        return len(self.buffer)

class UpsideDownAgent():
    def __init__(self, environment, approximator):
        self.environment = gym.make(environment)
        self.approximator = approximator
        self.state_size = (84, 84, 4)
        self.action_size = 3
        self.warm_up_episodes = 1 #50 
        self.render = False
        self.memory = ReplayBuffer(700)
        self.last_few = 50
        self.batch_size = 256
        self.command_size = 2 # desired return + desired horizon
        self.desired_return = 1 
        self.desired_horizon = 1
        self.horizon_scale = 0.02
        self.return_scale = 0.02
     
        self.behaviour_function = utils.get_atari_behaviour_function(self.action_size)
       
        self.testing_rewards = []
        self.warm_up_buffer()

    def warm_up_buffer(self):
        print('Warming up')

        for i in range(self.warm_up_episodes):
            
            states = []
            rewards = []
            actions = []
            
            dead = False
            done = False
            desired_return = 1
            desired_horizon = 1

            step, score, start_life = 0, 0, 5
            observe = self.environment.reset()

            for _ in range(random.randint(1, 30)):
                observe, _, _, _ = self.environment.step(1)

            state = utils.pre_processing(observe)
            history = np.stack((state, state, state, state), axis=2)
            history = np.reshape([history], (1, 84, 84, 4))


            while not done:
                
                states.append(history)
                command = np.asarray([desired_return * self.return_scale, desired_horizon * self.horizon_scale])
                command = np.reshape(command, [1, len(command)])

                action = self.get_action(history, command)
                actions.append(action)
               
                if action == 0:
                    real_action = 1
                elif action == 1:
                    real_action = 2
                else:
                    real_action = 3 

                next_state, reward, done, info = self.environment.step(real_action)
                next_state = utils.pre_processing(observe)
                next_state = np.reshape([next_state], (1, 84, 84, 1))
                next_history = np.append(next_state, history[:, :, :, :3], axis = 3)
 
                rewards.append(reward)

                state = next_state
        
                if start_life > info['ale.lives']:
                    dead = True 
                    start_lide = info['ale.lives']

                if dead:
                    dead = False
                else:
                    history = next_history

                desired_return -= reward  # Line 8 Algorithm 2
                desired_horizon -= 1 # Line 9 Algorithm 2
                desired_horizon = np.maximum(desired_horizon, 1)
                          
            self.memory.add_sample(states, actions, rewards)


    def get_action(self, observation, command):
        """
            We will sample from the action distribution modeled by the Behavior Function 
        """
        
        observation = np.float32(observation / 255.0)

        action_probs = self.behaviour_function.predict([observation, command])
        action = np.random.choice(np.arange(0, self.action_size), p=action_probs[0])

        return action
 
    def get_greedy_action(self, observation, command):

        action_probs = self.behaviour_function.predict([observation, command])
        action = np.argmax(action_probs)

        return action
 
    def train_behaviour_function(self):

        random_episodes = self.memory.get_random_samples(self.batch_size)
       
        training_observations = np.zeros((self.batch_size, self.state_size[0], self.state_size[1], self.state_size[2]))
        training_commands = np.zeros((self.batch_size, 2))

        y = []
        
        for idx, episode in enumerate(random_episodes):
            T = len(episode['states'])
            t1 = np.random.randint(0, T-1)
            t2 = np.random.randint(t1+1, T)
            
            state = np.float32(episode['states'][t1] / 255.)
            desired_return = sum(episode["rewards"][t1:t2]) 
            desired_horizon = t2 -t1
 
            target = episode['actions'][t1]
    
            training_observations[idx] = state[0]
            training_commands[idx] = np.asarray([desired_return*self.return_scale, desired_horizon*self.horizon_scale])
            y.append(target) 
         
        _y = keras.utils.to_categorical(y, num_classes=self.action_size) 
      
        self.behaviour_function.fit([training_observations, training_commands], _y, verbose=0)
       

    def sample_exploratory_commands(self):
        best_episodes = self.memory.get_n_best(self.last_few)
        exploratory_desired_horizon = np.mean([len(i["states"]) for i in best_episodes])
        
        returns = [i["summed_rewards"] for i in best_episodes]
        exploratory_desired_returns = np.random.uniform(np.mean(returns), np.mean(returns)+np.std(returns))

        return [exploratory_desired_returns, exploratory_desired_horizon]

    def generate_episode(self, environment, e, desired_return, desired_horizon, testing):
       
        env = gym.make(environment)

        tot_rewards = []
        
        done = False
        dead = False
  
        scores = []
        states = []
        actions = []
        rewards = []

        step, score, start_life = 0, 0, 5 

        observe = env.reset()
        for _ in range(random.randint(1, 30)):
            observe, _, _, _ = env.step(1)

        state = utils.pre_processing(observe)
        history = np.stack((state, state, state, state), axis=2)
        history = np.reshape([history], (1, 84, 84, 4))

        while not done:            
            states.append(history)
            
            command = np.asarray([desired_return * self.return_scale, desired_horizon * self.horizon_scale])
            command = np.reshape(command, [1, len(command)])

            if not testing:
                action = self.get_action(history, command)
                actions.append(action)
            else:
                action = self.get_greedy_action(history, command)

            if action == 0:
                real_action = 1
            elif action == 1:
                real_action = 2
            else:
                real_action = 3

            next_state, reward, done, info = env.step(real_action)
            next_state = utils.pre_processing(observe)
            next_state = np.reshape([next_state], (1, 84, 84, 1))
            next_history = np.append(next_state, history[:, :, :, :3], axis = 3)

            clipped_reward = np.clip(reward, -1, 1)
            rewards.append(clipped_reward)
            
            score += reward

            if start_life > info['ale.lives']:
                dead = True 
                start_life = info['ale.lives']

            if dead:
                dead = False
            else:
                history = next_history
           
            desired_return -= reward  # Line 8 Algorithm 2
            desired_horizon -= 1 # Line 9 Algorithm 2
            desired_horizon = np.maximum(desired_horizon, 1)
            
        self.memory.add_sample(states, actions, rewards)
        
        self.testing_rewards.append(score)

        if testing:
            print('Querying the model ...')
            print('Testing score: {}'.format(score))

        return score

def run_experiment():

    import argparse

    parser = argparse.ArgumentParser()
    
    parser.add_argument('--approximator', type=str, default='neural_network')
    parser.add_argument('--environment', type=str, default='PongDeterministic-v4')
    parser.add_argument('--seed', type=int, default=1)

    args = parser.parse_args()

    approximator = args.approximator
    environment = args.environment
    seed = args.seed

    episodes =  1500  
    returns = []

    agent = UpsideDownAgent(environment, approximator)

    for e in range(episodes):

        print("Episode {}".format(e))

        for i in range(100):
            agent.train_behaviour_function()

        print("Finished training B!")
        
        for i in range(15):            
            tmp_r = []
            exploratory_commands = agent.sample_exploratory_commands() # Line 5 Algorithm 1
            desired_return = exploratory_commands[0]
            desired_horizon = exploratory_commands[1]
            r = agent.generate_episode(environment, e, desired_return, desired_horizon, False)
            tmp_r.append(r)

        print(np.mean(tmp_r))
        returns.append(np.mean(tmp_r))

        exploratory_commands = agent.sample_exploratory_commands()
        
    #agent.generate_episode(environment, 1, 200, 200, True)

        utils.save_results(environment, approximator, seed, returns)
    
    if approximator == 'neural_network':
        utils.save_trained_model(environment, seed, agent.behaviour_function)

    plt.plot(returns)
    plt.show()

if __name__ == "__main__":
    run_experiment()