https://github.com/Microsoft/CNTK
Tip revision: f0d63b2c531018ec206a05629ad1a64aeb9c42fb authored by Eldar Akchurin on 30 November 2017, 12:33:49 UTC
Exposing EvaluatorWrapper interface and restructure the code to be able to reuse most of the stuff in Halide inference engine
Exposing EvaluatorWrapper interface and restructure the code to be able to reuse most of the stuff in Halide inference engine
Tip revision: f0d63b2
DeepQNeuralNetwork.py
# Copyright (c) Microsoft. All rights reserved.
# Licensed under the MIT license. See LICENSE.md file in the project root
# for full license information.
# ==============================================================================
from argparse import ArgumentParser
import gym
import numpy as np
from cntk.core import Value
from cntk.initializer import he_uniform
from cntk.layers import Sequential, Convolution2D, Dense, default_options
from cntk.layers.typing import Signature, Tensor
from cntk.learners import adam, learning_parameter_schedule, momentum_schedule
from cntk.logging import TensorBoardProgressWriter
from cntk.ops import abs, argmax, element_select, less, relu, reduce_max, reduce_sum, square
from cntk.ops.functions import CloneMethod, Function
from cntk.train import Trainer
class ReplayMemory(object):
"""
ReplayMemory keeps track of the environment dynamic.
We store all the transitions (s(t), action, s(t+1), reward, done).
The replay memory allows us to efficiently sample minibatches from it, and generate the correct state representation
(w.r.t the number of previous frames needed).
"""
def __init__(self, size, sample_shape, history_length=4):
self._pos = 0
self._count = 0
self._max_size = size
self._history_length = max(1, history_length)
self._state_shape = sample_shape
self._states = np.zeros((size,) + sample_shape, dtype=np.float32)
self._actions = np.zeros(size, dtype=np.uint8)
self._rewards = np.zeros(size, dtype=np.float32)
self._terminals = np.zeros(size, dtype=np.float32)
def __len__(self):
""" Returns the number of items currently present in the memory
Returns: Int >= 0
"""
return self._count
def append(self, state, action, reward, done):
""" Appends the specified transition to the memory.
Attributes:
state (Tensor[sample_shape]): The state to append
action (int): An integer representing the action done
reward (float): An integer representing the reward received for doing this action
done (bool): A boolean specifying if this state is a terminal (episode has finished)
"""
assert state.shape == self._state_shape, \
'Invalid state shape (required: %s, got: %s)' % (self._state_shape, state.shape)
self._states[self._pos] = state
self._actions[self._pos] = action
self._rewards[self._pos] = reward
self._terminals[self._pos] = done
self._count = max(self._count, self._pos + 1)
self._pos = (self._pos + 1) % self._max_size
def sample(self, size):
""" Generate size random integers mapping indices in the memory.
The returned indices can be retrieved using #get_state().
See the method #minibatch() if you want to retrieve samples directly.
Attributes:
size (int): The minibatch size
Returns:
Indexes of the sampled states ([int])
"""
# Local variable access is faster in loops
count, pos, history_len, terminals = self._count - 1, self._pos, \
self._history_length, self._terminals
indexes = []
while len(indexes) < size:
index = np.random.randint(history_len, count)
if index not in indexes:
# if not wrapping over current pointer,
# then check if there is terminal state wrapped inside
if not (index >= pos > index - history_len):
if not terminals[(index - history_len):index].any():
indexes.append(index)
return indexes
def minibatch(self, size):
""" Generate a minibatch with the number of samples specified by the size parameter.
Attributes:
size (int): Minibatch size
Returns:
tuple: Tensor[minibatch_size, input_shape...], [int], [float], [bool]
"""
indexes = self.sample(size)
pre_states = np.array([self.get_state(index) for index in indexes], dtype=np.float32)
post_states = np.array([self.get_state(index + 1) for index in indexes], dtype=np.float32)
actions = self._actions[indexes]
rewards = self._rewards[indexes]
dones = self._terminals[indexes]
return pre_states, actions, post_states, rewards, dones
def get_state(self, index):
"""
Return the specified state with the replay memory. A state consists of
the last `history_length` perceptions.
Attributes:
index (int): State's index
Returns:
State at specified index (Tensor[history_length, input_shape...])
"""
if self._count == 0:
raise IndexError('Empty Memory')
index %= self._count
history_length = self._history_length
# If index > history_length, take from a slice
if index >= history_length:
return self._states[(index - (history_length - 1)):index + 1, ...]
else:
indexes = np.arange(index - history_length + 1, index + 1)
return self._states.take(indexes, mode='wrap', axis=0)
class History(object):
"""
Accumulator keeping track of the N previous frames to be used by the agent
for evaluation
"""
def __init__(self, shape):
self._buffer = np.zeros(shape, dtype=np.float32)
@property
def value(self):
""" Underlying buffer with N previous states stacked along first axis
Returns:
Tensor[shape]
"""
return self._buffer
def append(self, state):
""" Append state to the history
Attributes:
state (Tensor) : The state to append to the memory
"""
self._buffer[:-1] = self._buffer[1:]
self._buffer[-1] = state
def reset(self):
""" Reset the memory. Underlying buffer set all indexes to 0
"""
self._buffer.fill(0)
class LinearEpsilonAnnealingExplorer(object):
"""
Exploration policy using Linear Epsilon Greedy
Attributes:
start (float): start value
end (float): end value
steps (int): number of steps between start and end
"""
def __init__(self, start, end, steps):
self._start = start
self._stop = end
self._steps = steps
self._step_size = (end - start) / steps
def __call__(self, num_actions):
"""
Select a random action out of `num_actions` possibilities.
Attributes:
num_actions (int): Number of actions available
"""
return np.random.choice(num_actions)
def _epsilon(self, step):
""" Compute the epsilon parameter according to the specified step
Attributes:
step (int)
"""
if step < 0:
return self._start
elif step > self._steps:
return self._stop
else:
return self._step_size * step + self._start
def is_exploring(self, step):
""" Commodity method indicating if the agent should explore
Attributes:
step (int) : Current step
Returns:
bool : True if exploring, False otherwise
"""
return np.random.rand() < self._epsilon(step)
def huber_loss(y, y_hat, delta):
""" Compute the Huber Loss as part of the model graph
Huber Loss is more robust to outliers. It is defined as:
if |y - y_hat| < delta :
0.5 * (y - y_hat)**2
else :
delta * |y - y_hat| - 0.5 * delta**2
Attributes:
y (Tensor[-1, 1]): Target value
y_hat(Tensor[-1, 1]): Estimated value
delta (float): Outliers threshold
Returns:
CNTK Graph Node
"""
half_delta_squared = 0.5 * delta * delta
error = y - y_hat
abs_error = abs(error)
less_than = 0.5 * square(error)
more_than = (delta * abs_error) - half_delta_squared
loss_per_sample = element_select(less(abs_error, delta), less_than, more_than)
return reduce_sum(loss_per_sample, name='loss')
class DeepQAgent(object):
"""
Implementation of Deep Q Neural Network agent like in:
Nature 518. "Human-level control through deep reinforcement learning" (Mnih & al. 2015)
"""
def __init__(self, input_shape, nb_actions,
gamma=0.99, explorer=LinearEpsilonAnnealingExplorer(1, 0.1, 1000000),
learning_rate=0.00025, momentum=0.95, minibatch_size=32,
memory_size=500000, train_after=200000, train_interval=4, target_update_interval=10000,
monitor=True):
self.input_shape = input_shape
self.nb_actions = nb_actions
self.gamma = gamma
self._train_after = train_after
self._train_interval = train_interval
self._target_update_interval = target_update_interval
self._explorer = explorer
self._minibatch_size = minibatch_size
self._history = History(input_shape)
self._memory = ReplayMemory(memory_size, input_shape[1:], 4)
self._num_actions_taken = 0
# Metrics accumulator
self._episode_rewards, self._episode_q_means, self._episode_q_stddev = [], [], []
# Action Value model (used by agent to interact with the environment)
with default_options(activation=relu, init=he_uniform()):
self._action_value_net = Sequential([
Convolution2D((8, 8), 16, strides=4),
Convolution2D((4, 4), 32, strides=2),
Convolution2D((3, 3), 32, strides=1),
Dense(256, init=he_uniform(scale=0.01)),
Dense(nb_actions, activation=None, init=he_uniform(scale=0.01))
])
self._action_value_net.update_signature(Tensor[input_shape])
# Target model used to compute the target Q-values in training, updated
# less frequently for increased stability.
self._target_net = self._action_value_net.clone(CloneMethod.freeze)
# Function computing Q-values targets as part of the computation graph
@Function
@Signature(post_states=Tensor[input_shape], rewards=Tensor[()], terminals=Tensor[()])
def compute_q_targets(post_states, rewards, terminals):
return element_select(
terminals,
rewards,
gamma * reduce_max(self._target_net(post_states), axis=0) + rewards,
)
# Define the loss, using Huber Loss (more robust to outliers)
@Function
@Signature(pre_states=Tensor[input_shape], actions=Tensor[nb_actions],
post_states=Tensor[input_shape], rewards=Tensor[()], terminals=Tensor[()])
def criterion(pre_states, actions, post_states, rewards, terminals):
# Compute the q_targets
q_targets = compute_q_targets(post_states, rewards, terminals)
# actions is a 1-hot encoding of the action done by the agent
q_acted = reduce_sum(self._action_value_net(pre_states) * actions, axis=0)
# Define training criterion as the Huber Loss function
return huber_loss(q_targets, q_acted, 1.0)
# Adam based SGD
lr_schedule = learning_parameter_schedule(learning_rate)
m_schedule = momentum_schedule(momentum)
vm_schedule = momentum_schedule(0.999)
l_sgd = adam(self._action_value_net.parameters, lr_schedule,
momentum=m_schedule, variance_momentum=vm_schedule)
self._metrics_writer = TensorBoardProgressWriter(freq=1, log_dir='metrics', model=criterion) if monitor else None
self._learner = l_sgd
self._trainer = Trainer(criterion, (criterion, None), l_sgd, self._metrics_writer)
def act(self, state):
""" This allows the agent to select the next action to perform in regard of the current state of the environment.
It follows the terminology used in the Nature paper.
Attributes:
state (Tensor[input_shape]): The current environment state
Returns: Int >= 0 : Next action to do
"""
# Append the state to the short term memory (ie. History)
self._history.append(state)
# If policy requires agent to explore, sample random action
if self._explorer.is_exploring(self._num_actions_taken):
action = self._explorer(self.nb_actions)
else:
# Use the network to output the best action
env_with_history = self._history.value
q_values = self._action_value_net.eval(
# Append batch axis with only one sample to evaluate
env_with_history.reshape((1,) + env_with_history.shape)
)
self._episode_q_means.append(np.mean(q_values))
self._episode_q_stddev.append(np.std(q_values))
# Return the value maximizing the expected reward
action = q_values.argmax()
# Keep track of interval action counter
self._num_actions_taken += 1
return action
def observe(self, old_state, action, reward, done):
""" This allows the agent to observe the output of doing the action it selected through act() on the old_state
Attributes:
old_state (Tensor[input_shape]): Previous environment state
action (int): Action done by the agent
reward (float): Reward for doing this action in the old_state environment
done (bool): Indicate if the action has terminated the environment
"""
self._episode_rewards.append(reward)
# If done, reset short term memory (ie. History)
if done:
# Plot the metrics through Tensorboard and reset buffers
if self._metrics_writer is not None:
self._plot_metrics()
self._episode_rewards, self._episode_q_means, self._episode_q_stddev = [], [], []
# Reset the short term memory
self._history.reset()
# Append to long term memory
self._memory.append(old_state, action, reward, done)
def train(self):
""" This allows the agent to train itself to better understand the environment dynamics.
The agent will compute the expected reward for the state(t+1)
and update the expected reward at step t according to this.
The target expectation is computed through the Target Network, which is a more stable version
of the Action Value Network for increasing training stability.
The Target Network is a frozen copy of the Action Value Network updated as regular intervals.
"""
agent_step = self._num_actions_taken
if agent_step >= self._train_after:
if (agent_step % self._train_interval) == 0:
pre_states, actions, post_states, rewards, terminals = self._memory.minibatch(self._minibatch_size)
self._trainer.train_minibatch(
self._trainer.loss_function.argument_map(
pre_states=pre_states,
actions=Value.one_hot(actions.reshape(-1, 1).tolist(), self.nb_actions),
post_states=post_states,
rewards=rewards,
terminals=terminals
)
)
# Update the Target Network if needed
if (agent_step % self._target_update_interval) == 0:
self._target_net = self._action_value_net.clone(CloneMethod.freeze)
def _plot_metrics(self):
"""Plot current buffers accumulated values to visualize agent learning
"""
if len(self._episode_q_means) > 0:
mean_q = np.asscalar(np.mean(self._episode_q_means))
self._metrics_writer.write_value('Mean Q per ep.', mean_q, self._num_actions_taken)
if len(self._episode_q_stddev) > 0:
std_q = np.asscalar(np.mean(self._episode_q_stddev))
self._metrics_writer.write_value('Mean Std Q per ep.', std_q, self._num_actions_taken)
self._metrics_writer.write_value('Sum rewards per ep.', sum(self._episode_rewards), self._num_actions_taken)
def as_ale_input(environment):
"""Convert the Atari environment RGB output (210, 160, 3) to an ALE one (84, 84).
We first convert the image to a gray scale image, and resize it.
Attributes:
environment (Tensor[input_shape]): Environment to be converted
Returns:
Tensor[84, 84] : Environment converted
"""
from PIL import Image
return np.array(Image.fromarray(environment).convert('L').resize((84, 84)))
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument('-e', '--epoch', default=100, type=int, help='Number of epochs to run (epoch = 250k actions')
parser.add_argument('-p', '--plot', action='store_true', default=False, help='Flag for enabling Tensorboard')
parser.add_argument('env', default='Pong-v3', type=str, metavar='N', nargs='?', help='Gym Atari environment to run')
args = parser.parse_args()
# 1. Make environment:
env = gym.make(args.env)
# 2. Make agent
agent = DeepQAgent((4, 84, 84), env.action_space.n, monitor=args.plot)
# Train
current_step = 0
max_steps = args.epoch * 250000
current_state = as_ale_input(env.reset())
while current_step < max_steps:
action = agent.act(current_state)
new_state, reward, done, _ = env.step(action)
new_state = as_ale_input(new_state)
# Clipping reward for training stability
reward = np.clip(reward, -1, 1)
agent.observe(current_state, action, reward, done)
agent.train()
current_state = new_state
if done:
current_state = as_ale_input(env.reset())
current_step += 1