159 lines
4.6 KiB
Python
159 lines
4.6 KiB
Python
# coding:utf-8
|
|
import logging
|
|
import random
|
|
|
|
import gym
|
|
import numpy as np
|
|
from gym import wrappers
|
|
|
|
np.random.seed(9999)
|
|
|
|
logger = logging.getLogger()
|
|
logger.setLevel(logging.INFO)
|
|
|
|
"""
|
|
References:
|
|
Sutton, Barto (2017). Reinforcement Learning: An Introduction. MIT Press, Cambridge, MA.
|
|
"""
|
|
|
|
|
|
class DQN(object):
|
|
def __init__(
|
|
self,
|
|
n_episodes=500,
|
|
gamma=0.99,
|
|
batch_size=32,
|
|
epsilon=1.0,
|
|
decay=0.005,
|
|
min_epsilon=0.1,
|
|
memory_limit=500,
|
|
):
|
|
"""Deep Q learning implementation.
|
|
|
|
Parameters
|
|
----------
|
|
|
|
min_epsilon : float
|
|
Minimal value for epsilon.
|
|
epsilon : float
|
|
ε-greedy value.
|
|
decay : float
|
|
Epsilon decay rate.
|
|
memory_limit : int
|
|
Limit of experience replay memory.
|
|
|
|
"""
|
|
|
|
self.memory_limit = memory_limit
|
|
self.min_epsilon = min_epsilon
|
|
self.gamma = gamma
|
|
self.epsilon = epsilon
|
|
self.n_episodes = n_episodes
|
|
self.batch_size = batch_size
|
|
self.decay = decay
|
|
|
|
def init_environment(self, name="CartPole-v0", monitor=False):
|
|
self.env = gym.make(name)
|
|
if monitor:
|
|
self.env = wrappers.Monitor(
|
|
self.env, name, force=True, video_callable=False
|
|
)
|
|
|
|
self.n_states = self.env.observation_space.shape[0]
|
|
self.n_actions = self.env.action_space.n
|
|
|
|
# Experience replay
|
|
self.replay = []
|
|
|
|
def init_model(self, model):
|
|
self.model = model(self.n_actions, self.batch_size)
|
|
|
|
def train(self, render=False):
|
|
max_reward = 0
|
|
|
|
for ep in range(self.n_episodes):
|
|
state = self.env.reset()
|
|
|
|
total_reward = 0
|
|
|
|
while True:
|
|
if render:
|
|
self.env.render()
|
|
|
|
if np.random.rand() <= self.epsilon:
|
|
# Exploration
|
|
action = np.random.randint(self.n_actions)
|
|
else:
|
|
# Exploitation
|
|
action = np.argmax(self.model.predict(state[np.newaxis, :])[0])
|
|
|
|
# Run one timestep of the environment
|
|
new_state, reward, done, _ = self.env.step(action)
|
|
self.replay.append([state, action, reward, new_state, done])
|
|
|
|
# Sample batch from experience replay
|
|
batch_size = min(len(self.replay), self.batch_size)
|
|
batch = random.sample(self.replay, batch_size)
|
|
|
|
X = np.zeros((batch_size, self.n_states))
|
|
y = np.zeros((batch_size, self.n_actions))
|
|
|
|
states = np.array([b[0] for b in batch])
|
|
new_states = np.array([b[3] for b in batch])
|
|
|
|
Q = self.model.predict(states)
|
|
new_Q = self.model.predict(new_states)
|
|
|
|
# Construct training data
|
|
for i in range(batch_size):
|
|
state_r, action_r, reward_r, new_state_r, done_r = batch[i]
|
|
target = Q[i]
|
|
|
|
if done_r:
|
|
target[action_r] = reward_r
|
|
else:
|
|
target[action_r] = reward_r + self.gamma * np.amax(new_Q[i])
|
|
|
|
X[i, :] = state_r
|
|
y[i, :] = target
|
|
|
|
# Train deep learning model
|
|
self.model.fit(X, y)
|
|
|
|
total_reward += reward
|
|
state = new_state
|
|
|
|
if done:
|
|
# Exit from current episode
|
|
break
|
|
|
|
# Remove old entries from replay memory
|
|
while len(self.replay) > self.memory_limit:
|
|
self.replay.pop(0)
|
|
|
|
self.epsilon = self.min_epsilon + (1.0 - self.min_epsilon) * np.exp(
|
|
-self.decay * ep
|
|
)
|
|
|
|
max_reward = max(max_reward, total_reward)
|
|
logger.info(
|
|
"Episode: %s, reward %s, epsilon %s, max reward %s"
|
|
% (ep, total_reward, self.epsilon, max_reward)
|
|
)
|
|
logging.info("Training finished.")
|
|
|
|
def play(self, episodes):
|
|
for i in range(episodes):
|
|
state = self.env.reset()
|
|
total_reward = 0
|
|
|
|
while True:
|
|
self.env.render()
|
|
action = np.argmax(self.model.predict(state[np.newaxis, :])[0])
|
|
state, reward, done, _ = self.env.step(action)
|
|
total_reward += reward
|
|
if done:
|
|
break
|
|
logger.info("Episode: %s, reward %s" % (i, total_reward))
|
|
self.env.close()
|