Added GPU version

gpuqlearn.moon

@@ -0,0 +1,372 @@
+require 'math'
+require 'nnx'
+require 'os'
+require 'optim'
+require 'cutorch'
+require 'cunn'
+
+math.randomseed os.time!
+torch.setdefaulttensortype 'torch.FloatTensor'
+
+Brain = {}
+
+--  HELPER FUNCTIONS --
+
+export randf = (s, e) ->
+	return (math.random(0, (e - s) * 9999) / 10000) + s
+
+-- new methods for table
+
+table.merge = (t1, t2) ->
+	t = t1
+	for i = 1, #t2
+		t[#t + 1] = t2[i]
+	return t
+
+table.copy = (t) ->
+	u = {k, v for k, v in pairs t}
+	return setmetatable(u, getmetatable t)
+
+table.length = (T) ->
+	count = 0
+	count += 1 for _ in pairs T
+	return count
+
+-- returns experience table for single network decision
+-- contains the state, action chosen, whether a reward was obtained, and the
+-- state that resulted from the action. This is later used to train the network
+-- Remember that the utility of an action is evaluated from the reward gained and
+-- the utility of the state it led to (recursive definition)
+export Experience = (state0, action0, reward0, state1) ->
+	NewExperience =
+		state0: state0
+		action0: action0
+		reward0: reward0
+		state1: state1
+	return NewExperience
+
+-- BRAIN
+
+Brain.init = (num_states, num_actions) ->
+	-- Number of past state/action pairs input to the network. 0 = agent lives in-the-moment :)
+	Brain.temporal_window = 2
+	-- Maximum number of experiences that we will save for training
+	Brain.experience_size = 30000
+	-- experience necessary to start learning
+	Brain.start_learn_threshold = 300
+	-- gamma is a crucial parameter that controls how much plan-ahead the agent does. In [0,1]
+	-- Determines the amount of weight placed on the utility of the state resulting from an action.
+	Brain.gamma = 0.9
+	-- number of steps we will learn for
+	Brain.learning_steps_total = 100000
+	-- how many steps of the above to perform only random actions (in the beginning)?
+	Brain.learning_steps_burnin = 300
+	-- controls exploration exploitation tradeoff. Will decay over time
+	-- a higher epsilon means we are more likely to choose random actions
+	Brain.epsilon = 1.0
+	-- what epsilon value do we bottom out on? 0.0 => purely deterministic policy at end
+	Brain.epsilon_min = 0.05
+	-- what epsilon to use when learning is turned off. This is for testing
+	Brain.epsilon_test_time = 0.01
+
+	[[== states and actions that go into neural net:
+		(state0,action0),(state1,action1), ... , (stateN)
+		this variable controls the size of that temporal window.
+	]]
+	Brain.net_inputs = (num_states + num_actions) * Brain.temporal_window + num_states
+	Brain.hidden_nodes = 16
+	Brain.num_states = num_states
+	Brain.num_actions = num_actions
+	Brain.net_outputs = Brain.num_actions
+
+	[[== Window size dictates the number of states, actions, rewards, and net inputs that we
+		save. The temporal window size is the number of time states/actions that are input
+		to the network and must be smaller than or equal to window_size
+	]]
+	Brain.window_size = math.max Brain.temporal_window, 2
+
+	-- advanced feature. Sometimes a random action should be biased towards some values
+	-- for example in flappy bird, we may want to choose to not flap more often
+	Brain.random_action_distribution = {}
+	if table.length(Brain.random_action_distribution) > 0
+		-- this better sum to 1 by the way, and be of length this.num_actions
+		if table.length(Brain.random_action_distribution) != Brain.num_actions
+			print 'TROUBLE. random_action_distribution should be same length as num_actions.'
+
+		s = 0.0
+
+		for k = 1, table.length Brain.random_action_distribution
+			s += Brain.random_action_distribution[k]
+
+		if math.abs(s - 1.0) > 0.0001
+			 print 'TROUBLE. random_action_distribution should sum to 1!'
+
+
+	-- define architecture
+	Brain.net = nn.Sequential!
+
+	Brain.net\add nn.Linear Brain.net_inputs, Brain.hidden_nodes
+	Brain.net\add nn.Threshold 0, 0
+
+	Brain.net\add nn.Linear Brain.hidden_nodes, Brain.hidden_nodes
+	Brain.net\add nn.Threshold 0, 0
+
+	Brain.net\add nn.Linear Brain.hidden_nodes, Brain.net_outputs
+
+	Brain.net\cuda! -- move network to GPU
+
+	Brain.criterion = nn.MSECriterion!\cuda!
+
+
+	-- other learning parameters
+	Brain.learning_rate = 0.01
+	Brain.learning_rate_decay = 5e-7
+	Brain.batch_size = 16
+	Brain.momentum = 0.9
+
+	-- various housekeeping variables
+	Brain.age = 0 -- incremented every backward!
+
+	-- number of times we've called forward - lets us know when our input temporal
+	-- window is filled up
+	Brain.forward_passes = 0
+	Brain.learning = true
+
+	-- coefficients for regression
+	Brain.coefL1 = 0.001
+	Brain.coefL2 = 0.001
+
+	-- parameters for optim.sgd
+	Brain.parameters, Brain.gradParameters = Brain.net\getParameters!
+
+	-- These windows track old experiences, states, actions, rewards, and net inputs
+	-- over time. They should all start out as empty with a fixed size.
+	-- This is a first in, last out data structure that is shifted along time
+	Brain.experience = {}
+	Brain.state_window = {}
+	Brain.action_window = {}
+	Brain.reward_window = {}
+	Brain.net_window = {}
+	for i = 1, Brain.window_size
+		Brain.state_window[i] = {}
+		Brain.action_window[i] = {}
+		Brain.reward_window[i] = {}
+		Brain.net_window[i] = {}
+
+-- a bit of a helper function. It returns a random action
+-- we are abstracting this away because in future we may want to
+-- do more sophisticated things. For example some actions could be more
+-- or less likely at "rest"/default state.
+Brain.random_action = ->
+	-- if we don't have a random action distribution defined then sample evenly
+	if table.length(Brain.random_action_distribution) == 0
+		return (torch.random! % Brain.net_outputs) + 1
+
+	-- okay, lets do some fancier sampling:
+	else
+		p = randf 0, 1
+		cumprob = 0.0
+
+		for k = 1, Brain.num_actions
+			cumprob += Brain.random_action_distribution[k]
+
+			if p < cumprob
+				return k
+
+-- compute the value of doing any action in this state
+-- and return the argmax action and its value
+Brain.policy = (state) ->
+	tensor_state = torch.Tensor(state)\cuda!
+	action_values = Brain.net\forward tensor_state
+
+	maxval = action_values[1]
+	max_index = 1
+
+	-- find maximum output and note its index and value
+	--max_index = i for i = 2, Brain.net_outputs when action_values[i] > action_values[max_index]
+	for i = 2, Brain.net_outputs
+		if action_values[i] > maxval
+			maxval = action_values[i]
+			max_index = i
+
+	return action: max_index, value: maxval
+
+-- This function assembles the input to the network by concatenating
+-- old (state, chosen_action) pairs along with the current state
+	-- return s = (x,a,x,a,x,a,xt) state vector.
+Brain.getNetInput = (xt) ->
+	w = {}
+	w = table.merge(w, xt) -- start with current state
+
+	-- and now go backwards and append states and actions from history temporal_window times
+	n = Brain.window_size + 1
+	for k = 1, Brain.temporal_window do
+		-- state
+		w = table.merge w, Brain.state_window[n - k]
+		-- action, encoded as 1-of-k indicator vector. We scale it up a bit because
+		-- we don't want weight regularization to undervalue this information, as it only exists once
+		action1ofk = {}
+		action1ofk[i] = 0 for i = 1, Brain.num_actions
+
+		-- assign action taken for current state to be 1, all others are 0
+		action1ofk[Brain.action_window[n - k]] = 1.0 * Brain.num_states
+
+		w = table.merge w, action1ofk
+
+	return w
+
+-- This function computes an action by either:
+-- 1. Giving the current state and past (state, action) pairs to the network
+-- and letting it choose the best acction
+-- 2. Choosing a random action
+Brain.forward = (input_array) ->
+	Brain.forward_passes += 1
+
+	local action, net_input
+
+	-- if we have enough (state, action) pairs in our memory to fill up
+	-- our network input then we'll proceed to let our network choose the action
+	if Brain.forward_passes > Brain.temporal_window
+		net_input = Brain.getNetInput input_array
+
+		-- if learning is turned on then epsilon should be decaying
+		if Brain.learning
+			-- compute (decaying) epsilon for the epsilon-greedy policy
+			new_epsilon = 1.0 - (Brain.age - Brain.learning_steps_burnin)/(Brain.learning_steps_total - Brain.learning_steps_burnin)
+
+			-- don't let epsilon go above 1.0
+			Brain.epsilon = math.min(1.0, math.max(Brain.epsilon_min, new_epsilon))
+		else
+			-- if learning is turned off then use the epsilon we've specified for testing
+			Brain.epsilon = Brain.epsilon_test_time
+
+		-- use epsilon probability to choose whether we use network action or random action
+		if randf(0, 1) < Brain.epsilon
+			action = Brain.random_action!
+		else
+			-- otherwise use our policy to make decision
+			best_action = Brain.policy net_input
+			action = best_action.action -- this is the action number
+	else
+		-- pathological case that happens first few iterations when we can't
+		-- fill up our network inputs. Just default to random action in this case
+		net_input = {}
+		action = Brain.random_action!
+
+	-- shift the network input, state, and action chosen into our windows
+	table.remove Brain.net_window, 1
+	table.insert Brain.net_window, net_input
+
+	table.remove Brain.state_window, 1
+	table.insert Brain.state_window, input_array
+
+	table.remove Brain.action_window, 1
+	table.insert Brain.action_window, action
+
+	return action
+
+-- This function trains the network using the reward resulting from the last action
+-- It will save this past experience which consists of:
+--  the state, action chosen, whether a reward was obtained, and the
+--  state that resulted from the action
+-- After that, it will train the network (using a batch of experiences) using a
+-- random sampling of our entire experience history.
+Brain.backward = (reward) ->
+	-- add reward to our history
+	table.remove Brain.reward_window, 1
+	table.insert Brain.reward_window, reward
+
+	-- if learning is turned off then don't do anything
+	return unless Brain.learning
+
+	Brain.age += 1
+
+	-- if we've had enough states and actions to fill up our net input then add
+	-- this new experience to our history
+	if Brain.forward_passes > Brain.temporal_window + 1
+		-- make experience and fill it up
+		e = Experience nil, nil, nil, nil
+		n = Brain.window_size
+		e.state0 = Brain.net_window[n - 1]
+		e.action0 = Brain.action_window[n - 1]
+		e.reward0 = Brain.reward_window[n - 1]
+		e.state1 = Brain.net_window[n]
+
+		-- if our experience table isn't larger than the max size then expand
+		if table.length(Brain.experience) < Brain.experience_size
+			table.insert Brain.experience, e
+		-- Otherwise replace random experience. finite memory!
+		else
+			ri = torch.random 1, Brain.experience_size
+			Brain.experience[ri] = e
+
+	-- if we have enough experience in memory then start training
+	if table.length(Brain.experience) > Brain.start_learn_threshold
+		inputs = torch.Tensor(Brain.batch_size, Brain.net_inputs)\cuda!
+		targets = torch.Tensor(Brain.batch_size, Brain.net_outputs)\cuda!
+
+		for k = 1, Brain.batch_size
+			-- choose random experience
+			re = math.random 1, table.length Brain.experience
+			e = Brain.experience[re]
+
+			-- copy state from experience
+			x = torch.Tensor(e.state0)\cuda!
+
+			-- compute best action for the new state
+			best_action = Brain.policy e.state1
+
+			-- get current action output values
+			-- we want to make the target outputs the same as the actual outputs
+			-- expect for the action that was chose - we want to replace this with
+			-- the reward that was obtained + the utility of the resulting state
+			all_outputs = Brain.net\forward x
+			inputs[k] = x\clone!
+			targets[k] = all_outputs\clone!
+			targets[k][e.action0] = e.reward0 + Brain.gamma * best_action.value
+
+		-- create training function to give to optim.sgd
+		feval = (x) ->
+			collectgarbage!
+
+			-- get new network parameters
+			Brain.parameters\copy x unless x == Brain.parameters
+
+			-- reset gradients
+			Brain.gradParameters\zero!
+
+			-- evaluate function for complete mini batch
+			outputs = Brain.net\forward inputs
+			f = Brain.criterion\forward outputs, targets
+
+			-- estimate df/dW
+			df_do = Brain.criterion\backward outputs, targets
+			Brain.net\backward inputs, df_do
+
+			-- penalties (L1 and L2):
+			if Brain.coefL1 != 0 or Brain.coefL2 != 0
+				-- locals:
+				norm,sign = torch.norm, torch.sign
+
+				-- Loss:
+				f += Brain.coefL1 * norm Brain.parameters, 1
+				f += Brain.coefL2 * 0.5 * norm(Brain.parameters, 2) ^ 2
+
+				-- Gradients:
+				Brain.gradParameters\add(sign(Brain.parameters)\mul(Brain.coefL1) + Brain.parameters\clone!\mul Brain.coefL2)
+
+			-- return f and df/dX
+			return f, Brain.gradParameters
+
+		-- fire up optim.sgd
+		sgdState =
+			learningRate: Brain.learning_rate
+			momentum: Brain.momentum
+			learningRateDecay: Brain.learning_rate_decay
+
+		optim.sgd feval, Brain.parameters, sgdState
+
+
+
+-- export
+return Brain