| import os |
| import sys |
| import gym |
| import random |
| import numpy as np |
| import pickle |
|
|
| from collections import deque |
|
|
| from keras.layers import Dense |
| from keras.optimizers import Adam |
| from keras.models import Sequential |
| from matplotlib import pyplot as plt |
|
|
| WEIGHTS_PATH = './trained_models/CartPole-v0/1/' |
| BUFFER_PATH = './buffers/CartPole-v0/1/' |
|
|
| class Agent: |
| def __init__(self, algorithm, state_size, action_size): |
| self.algorithm = algorithm |
| self.render = False |
| self.state_size = state_size |
| self.action_size = action_size |
| self.memory = deque(maxlen=2000) |
|
|
| if self.algorithm in ['DQN', 'DDQN', 'DQV']: |
| self.model = self.build_model() |
| self.model.load_weights(os.path.join(WEIGHTS_PATH, self.algorithm, 'trained_model.h5')) |
| else: |
| self.model = self.build_actor() |
| self.model.load_weights(os.path.join(WEIGHTS_PATH, self.algorithm, 'trained_model.h5')) |
|
|
|
|
| def build_actor(self): |
| actor = Sequential() |
| actor.add(Dense(24, input_dim=self.state_size, activation='relu', kernel_initializer='he_uniform')) |
| actor.add(Dense(self.action_size, activation='softmax', kernel_initializer='he_uniform')) |
|
|
| return actor |
|
|
| def build_model(self): |
| model = Sequential() |
| model.add(Dense(24, input_dim=self.state_size, activation='relu', |
| kernel_initializer='he_uniform')) |
| model.add(Dense(24, activation='relu', |
| kernel_initializer='he_uniform')) |
| model.add(Dense(self.action_size, activation='linear', |
| kernel_initializer='he_uniform')) |
|
|
| return model |
|
|
| def get_action(self, state): |
| if self.algorithm == 'A2C': |
| policy = self.model.predict(state, batch_size=1).flatten() |
|
|
| return np.random.choice(self.action_size, 1, p=policy)[0] |
|
|
| else: |
| q_value = self.model.predict(state) |
| return np.argmax(q_value[0]) |
|
|
| def append_sample(self, state, action, reward, next_state, done): |
| self.memory.append((state, action, reward, next_state, done)) |
|
|
| def save_buffer(self): |
| if not os.path.exists(os.path.join(BUFFER_PATH, self.algorithm)): |
| os.makedirs(os.path.join(BUFFER_PATH, self.algorithm)) |
|
|
| with open(os.path.join(BUFFER_PATH, self.algorithm, 'memory_buffer.p'), 'wb') as filehandler: |
| pickle.dump(self.memory, filehandler) |
|
|
| def fill_buffer(algorithm): |
| max_len = 10000 |
| results = [] |
| game = 'CartPole-v0' |
|
|
| env = gym.make(game) |
| |
| state_size = env.observation_space.shape[0] |
| action_size = env.action_space.n |
|
|
| agent = Agent(algorithm, state_size, action_size) |
|
|
| while True: |
| done = False |
| score = 0 |
| state = env.reset() |
| state = np.reshape(state, [1, state_size]) |
|
|
| while not done: |
| action = agent.get_action(state) |
| next_state, reward, done, info = env.step(action) |
| next_state = np.reshape(next_state, [1, state_size]) |
|
|
| agent.append_sample(state, action, reward, next_state, done) |
|
|
| score += reward |
| state = next_state |
|
|
| if len(agent.memory) > max_len: |
| agent.save_buffer() |
| break |
|
|
| fill_buffer('DQN') |
|
|