QLearning.qmd

---
title: "Reinforcement Learning: A Q-Learning Agent"
author: "Michael Hahsler"
format: 
  html: 
    theme: default
    toc: true
    number-sections: true
    code-line-numbers: true
    embed-resources: true
---

This code is provided under [Creative Commons Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) License.](https://creativecommons.org/licenses/by-sa/4.0/)

![CC BY-SA 4.0](https://licensebuttons.net/l/by-sa/3.0/88x31.png)
```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
knitr::opts_chunk$set(tidy = TRUE)
options(digits = 2)
```

# Introduction

AIMA chapter 22 is about _reinforcement learning_ where the agent learns from reward signals. We will implement the 
key concepts using R for the AIMA 3x4 grid world example. 
A more detailed treatment of the algorithm and its properties can be found in 
[Reinforcement Learning: An Introduction (RL)](http://incompleteideas.net/book/the-book-2nd.html) by Sutton and Barto (2020)
in Chapter 6: Temporal-Difference Learning.

The used environment is an MDP, but instead of trying to solve the MDP directly
and estimating the value function estimates $U(s)$, we will try to learn 
a policy from interactions with the environment.
The MDP's transition model will only be used to simulate the response of the environment
to actions by the agent.

The code in this notebook defines explicit functions 
matching the textbook definitions and is for demonstration purposes only. Efficient implementations for larger problems use fast vector multiplications
instead. 


{{< include _AIMA-4x3-gridworld.qmd >}}

# Q-Learning

Q-Learning is an off-policy TD control algorithm to estimate $Q \approx q_*.$
The pseudo code of a Q-Learning algorithm is given in AIMA Figure 22.8 as:

![AIMA Figure 22.8: An exploratory Q-learning agent. It is an active learner that learns the value
$Q(s, a)$ of each action in each situation. It uses the exploration function $f$.](figures/AIMA_Figure_22_8.png)

The algorithm uses a __temporal-difference (TD) learning__ since it updates
using the TD error given by $Q[s',a'] - Q[s, a]$.

The algorithm performs __off-policy learning__ since it learns a different policy
than the one used to perform actions. 

* **Behavior policy:** is defined by the exploration function $f$. It needs
  to guarantee exploration. That is, every available action needs to have a 
  chance to be tried. A popular policy is an $\epsilon$-greedy policy, but any 
  $\epsilon$-soft policy that is increasing action probability with Q and 
  decreases it with N can be used.
  
* **Target policy:** This is the learned policy. Q-learning learns a 
  deterministic greedy policy represented by the $max_{a'} Q[s', a']$ in
  the update of Q.

## Agent

The Q-learning agent parameters and helper functions.

```{r}
# helper: argmax breaking ties randomly.
argmax <- function (x) 
{
    y <- seq_along(x)[x == max(x)]
    if (length(y) > 1L) 
        sample(y, 1L)
    else y
}
```

We use a class to store persistent information: Agent state, $Q$ and $N$ tables.
Reference classes use `<<-` to assign values to member variables (fields). 
```{r}
QL_Agent <- setRefClass(
  "QL_Agent",
  fields = c("S", "A", "Q", "N", "s", "a", "gamma"),
  methods = list(
    initialize = function(S, A, gamma = 1) {
      S <<- S
      A <<- A
      gamma <<- gamma
      reset()
    },
    
    next_episode = function() {
      "Start a new episode."
      
      s <<- NULL
      a <<- 'None'
    },
    
    reset = function() {
      "Reset the agent's state and erase all persistent information."
      
      # set q-values to 0 except use NA for unavailable actions
      Q <<-
        matrix(NA_real_,
               nrow = length(S),
               ncol = length(A),
               dimnames = list(S, A))
      
      for (s in S) Q[s, actions(s)] <<- 0
      
      N <<-
        matrix(0L,
               nrow = length(S),
               ncol = length(A),
               dimnames = list(S, A))
      
      next_episode()
    }
    
  )
)
```



The agent function is called `run` with the precepts including the new state 
and the associated reward.

```{r}
QL_Agent$methods(
  # Learning rate:
  # Fixed learning rate
  #alpha = function(n, rate = 0.1) return(rate)

  # A decreasing learning rate converges faster. n is the number a 
  # state was visited. The numbers depend on the problem.
  alpha = function(n) {
    return (10 / (9 + n))
  },

  # Exploration function (behavior policy): We use a simple 
  # epsilon-greedy policy here.
  maxf = function(s_prime, a_best, epsilon = .8) {
    if (runif(1) < epsilon)
      action <- a_best
    else
      action <- sample(actions(s_prime), size = 1)
  
    return(action)
  },
 
  # Q-learning algorithm 
  run = function(s_prime, r) {
    "Agent function. The pecepts are the current state and the last reward."
    
    # how often did the agent try this s/a combination
    N[s, a] <<- N[s, a] + 1L
    
    # Greedy target policy: find best available action for s' (breaking ties randomly).
    available_actions <- actions(s_prime)
    a_best <- available_actions[argmax(Q[s_prime,available_actions])]
    
    # No update for start of new episode
    if (!is.null(s)) {
      Q[s, a] <<-
        Q[s, a] + alpha(N[s, a]) * (r + gamma * Q[s_prime, a_best] - Q[s, a])
    }
    
    if (is_terminal(s_prime))
      a <<- 'None'
    else
      a <<- maxf(s_prime, a_best)
    
    s <<- s_prime
    
    return(a)
  }
  
)
```

Create an agent.

```{r}
agent <- QL_Agent$new(S, A)
agent
```

## Environment

This is an environment that can simulate $n$ episodes and returns the 
utility the agent received for each episode.

```{r}
simulate_environment <-
  function(n = 1000,
           agent,
           reset_agent = TRUE,
           verbose = FALSE) {
    
    reward_history <- numeric(n)
    state <- start
    
    if (reset_agent)
      agent$reset()
    agent$next_episode()
    
    i <- 1L # episode counter
    j <- 1L # step in episode counter
    reward <- 0
    reward_episode <- 0
    while (TRUE) {
      reward_episode <- reward_episode + GAMMA ^ j * reward
      
      # call agent function with percepts (current state and reward)
      action <- agent$run(state, reward)
      
      if (verbose) {
        if (j == 1L)
          cat("\n*** Episode", i , "***\n")
        cat("Step",
            j,
            paste0("(s = ", state, ", r = ",
                   reward, ") -> ", action),
            "\n")
      }
      
      j <- j + 1L
      
      if (is_terminal(state)) {
        reward_history[i] <- reward_episode
        if (verbose) {
          cat("Episode reward: ", reward_episode, "\n")
          
          cat("Q:\n")
          print(agent$Q)
          
          cat("U:\n")
          print(show_layout(apply(
            agent$Q, MARGIN = 1, max, na.rm = TRUE
          )))
        }
        
        j <- 1L
        
        # finished n episodes?
        i <- i + 1L
        if (i > n)
          break
        
        state <- start
        reward <- 0
        reward_episode <- 0
        
        agent$next_episode()
      }
      else
        if (action %in% actions(state)) {
          state_prime <- sample_transition(state, action)
          reward <- R(state, action, state_prime)
          state <- state_prime
        } else {
          stop("The agent chose an illegal action!")
          #reward <- -1e9
        }
    }
    reward_history
  }
```

## Analyzing the First Few Episodes

Run a few episodes. 

```{r}
simulate_environment(n = 10, agent, verbose = TRUE)
```

Note that the agent has initial no idea about the optimal policy and
runs around randomly. After it randomly finds the goal, the rewards starts to 
spread one square every episode where the agent reaches the goal. 
Also, the agent starts to find the goal faster.

Calculate the value function $U$ from the learned Q-function as the largest 
Q value of any action in a state.
```{r}
U <- apply(agent$Q, MARGIN = 1, max, na.rm = TRUE)
show_layout(U)
```

Note that not all states are reached the same amount of times and some estimates 
may be worse than others. Since Q-Learning uses a $\epsilon$-greedy behavior policy,
the best actions are used most often leading to better estimates around the optimal policy.

Extract the learned policy from $Q$ as the action with the highest Q-value 
for each state.
```{r}
pi <- A[apply(agent$Q, MARGIN = 1, which.max)]
show_layout(pi)
```

Even after simulating just 10 episodes, we get a reasonable policy that will lead 
the agent mostly to the 
goal state. We can estimate the expected utility of the policy using simulation.

```{r}
mean(simulate_utilities(pi, s0 = 1, N = 100))
```

## Leaning a Good Policy

Simulating more episodes results in better estimates (also for the rarely 
visited states).

```{r}
reward_history <- simulate_environment(n = 10000, agent, 
    verbose = FALSE, reset_agent = TRUE)

plot(cumsum(reward_history[1:500]), type = "l", 
     ylab = "Total reward collected", xlab = "Episode")
```


The total reward approaches a straight line after just a few epochs, meaning that
a reasonable policy was already found.

```{r}
show_layout(LAYOUT)
agent$Q
```

Learned state-value function.

```{r}
U <- apply(agent$Q, MARGIN = 1, max, na.rm = TRUE)
show_layout(U)
```

Extract the greedy policy.

```{r}
pi <- A[apply(agent$Q, MARGIN = 1, which.max)]
show_layout(pi)
```

We can estimate the expected utility of the policy using simulation.

```{r}
mean(simulate_utilities(pi, s0 = 1, N = 100))
```