function q = qlearn

%Q learning in an episodic gridworld task

ntrials = 10000;
worldsize = [5 5]; %2-dimensional world
goal = [3 3]; %state where episode ends
noise = 0;
alpha = .05; %learning rate for Q
eps = .05; %for epsilon-greedy action selection

q=zeros(worldsize(1),worldsize(2),4,ntrials); %initialize Q values
s=randomState(worldsize,goal); %start in random state

for t=1:ntrials
    a = chooseAction(squeeze(q(s(1),s(2),:,t)),eps); %select action based on Q values for current state
    [snew r] = worldstep(s,a,noise,goal,worldsize); %get reward and new state
    TDerr = r + max(q(snew(1),snew(2),:,t)) - q(s(1),s(2),a,t); %TD error
    q(:,:,:,t+1) = q(:,:,:,t); %carry over Q values
    q(s(1),s(2),a,t+1) = q(s(1),s(2),a,t) + alpha*TDerr; %update Q for chosen action
    s = snew; %advance to next trial
    if isequal(s,goal) %end of episode
        s = randomState(worldsize,goal); %restart in random state
    end
end

function s = randomState(worldsize,goal)
%pick a random state other than the goal state
ok = 0;
while ok==0
    s = ceil(rand(1,2).*worldsize);
    if ~isequal(s,goal)
        ok = 1; %have found acceptable state
    end
end

function a = chooseAction(q,eps)
probs = zeros(1,length(q));
probs(q==max(q)) = (1-eps)/sum(q==max(q)); %assign probability 1-eps to optimal action(s)
probs = probs + eps/length(q); %add eps/n to all action probabilities
a = find(rand(1)<cumsum(probs),1); %select action according to probabilities

function [s r] = worldstep(s,a,noise,goal,worldsize)
if rand(1)<noise %random action
    dim = ceil(rand(1)*2); %dimension of movement; 1 or 2
    dir = round(rand(1))*2-1; %direction of movement, +/-1
else %perform chosen action
    dim = ceil(a/2); %dimension of movement
    dir = mod(a,2)*2 - 1; %direction of movement, +/-1
end
s(dim) = min(max(s(dim)+dir,1),worldsize(dim)); %move, while staying in world
if isequal(s,goal)
    r = 0; %end of episode
else
    r = -1; %penalty for not reaching goal
end
