---
title: "Q-Learning: Exploration vs Exploitation"
format:
  html:
    self-contained: TRUE
---

Groupmates Present: ______

In the last unit 2 classwork, you used trial-and-error learning to estimate action values using the total payoff from an episode.

Today you will learn Q-Learning, which updates $Q(s,a)$ one step at a time, using the idea:

"If I take action $a$ in state $s$ and then behave optimally from the next state onward, what payoff should I expect?"

The key ideas:

1. Action values $Q(s,a)$ store how good it is to take action $a$ in state $s$.
2. Exploration vs exploitation: sometimes try random actions (explore) so you can learn, but usually choose the best-known action (exploit).
3. Temporal-difference (TD) learning: update $Q(s,a)$ from a single transition $(s, a, r, s')$ using a bootstrapped target $r + \gamma \max_{a'} Q(s',a')$.

## GridWorld

We'll use the same GridWorld as before, with no shocks yet:

```{}
┌───┬───┬───┐
│ 0 │ 0 │ 0 │
├───┼───┼───┤
│ 1 │ 2 │ 1 │
├───┼───┼───┤
│ 0 │ 0 │ 0 │
└───┴───┴───┘
```

Use the same `move()` function to handle transitions:

```{r, message = F}
library(tidyverse)

actions <- c("north", "south", "east", "west", "stay")

move <- function(cell, action) {
  if (action == "stay") {
    return(cell)
  } else if (action == "south") {
    if (cell <= 6) return(cell + 3) 
    else return(cell)
  } else if (action == "north") {
    if (cell >= 4) return(cell - 3) 
    else return(cell)
  } else if (action == "east") {
    if (cell %in% c(1, 2, 4, 5, 7, 8)) 
      return(cell + 1) 
    else return(cell)
  } else if (action == "west") {
    if (cell %in% c(2, 3, 5, 6, 8, 9)) 
      return(cell - 1) 
    else return(cell)
  }
}

payoffs <- function(position) {
  c(0, 0, 0, 1, 2, 1, 0, 0, 0)[position]
}
```

## Part 1: The Q-Learning Update

Q-learning uses this update rule: **New estimate = Old estimate + Learning rate x Prediction Error**

$$Q(s, a) \leftarrow Q(s, a) + \alpha(r + 0.9 \max_{a'} Q(s', a') - Q(s, a))$$

- $\alpha \in (0, 1]$ is the learning rate (we'll use 0.2)
- $r$ is the payoff (reward) you got this step
- $s'$ is the next state after taking action $a$

**Question 1**: Numeric Warm-up

Suppose current $Q(s, a) = 5$, $\alpha = 0.2$, you take action $a$ in state $s$ and get a payoff of $r = 1$, the next state is $s'$ and $\max_{a'} Q(s', a') = 8$. Compute the updated value of $Q(s, a)$.

```{r}
Q_old <- 5
alpha <- 0.2
r <- 1
max_next <- 8

Q_next <- ___
Q_next == 5.64
```

## Part 2: Q-table helpers

We'll store Q-values in a tibble with one row per (state, action).

```{r}
init_Q <- function() {
  expand.grid(
    state = 1:9,
    action = actions,
    stringsAsFactors = FALSE
  ) %>%
    mutate(Q = 0)
}

init_Q()
```

Helper functions: `get_Q()` and `set_Q`:

- `get_Q(Qtab, s, a)` returns the current Q value for a state, action pair
- `set_Q(Qtab, s, a, value)` returns an updated Qtab with that value written.

**Question 2**: use filter and pull to complete `get_Q`. `pull()` takes a variable in a tibble returns the column vector, pulling a vector out of a tibble.

```{r}
get_Q <- function(Qtab, s, a) {
  Qtab %>%
    ___ %>%
    ___
}

set_Q <- function(Qtab, s, a, value) {
  idx <- which(Qtab$state == s & Qtab$action == a)
  Qtab$Q[idx] <- value
  return(Qtab)
}

library(testthat)
test_that("get_Q and set_Q work", {
  Q <- init_Q()
  Q <- set_Q(Q, 3, "north", 7)
  expect_equal(get_Q(Q, 3, "north"), 7)
})
```

## Part 3: Greedy vs epsilon-greedy actions

A **greedy** policy always chooses the action with the highest current $Q(s,a)$ (exploitation). An **epsilon-greedy** policy chooses a random action with some probability like $\epsilon = 0.3$ (exploration), and with probability $1 - \epsilon$, chooses a greedy action.

**Question 3**: finish writing the functions `greedy_action` and `epsilon_greedy_action` with epsilon = 0.3.

```{r}
Qtab <- init_Q()

greedy_action <- function(Qtab, s) {
  # Find the action with the highest Q value.
  # If there is a tie, randomly break ties.
  # Return the action.
  Qtab %>%
    filter(___) %>%
    slice_max(___) %>%
    pull(___) %>%
    sample(size = 1)
}

epsilon_greedy_action <- function(Qtab, s, eps) {
  if (sample(1:10, size = 1) <= (eps * 10)) {
    # with probability 0.3, sample randomly from actions 
    # (explore)
    sample(actions, size = 1)
  } else {
    # with probability 0.7, take the greedy action
    # (exploit)
    greedy_action(Qtab, s)
  }
}

# Verify your function works:
replicate(10, epsilon_greedy_action(Qtab, 5, eps = .3))
```

## Part 4: One episode of Q-Learning

Next we'll generate an episode, updating Q every step.

**Question 4**: Write `run_episode_qlearn()` that:

- Has the episode begin in a random cell (1 to 9)
- After each step, the episode ends with probability 0.1
- If the episode continues, pick an action using `epsilon_greedy_action`
- Observe the payoff r = payoffs(state)
- Compute the next state `s_next = move(state, action)`
- Perform the Q-learning update
- Store the transition in a tibble

```{r}
q_update <- function(Qtab, s, a, r, s_next) {
  Q_old <- get_Q(Qtab, s, a)
  
  max_next <- Qtab %>% 
    filter(state == s_next) %>% 
    pull(Q) %>%
    max()
  
  # Q learning rule:
  Q_next <- ___ + 0.2 * (r + 0.9 * ___ - ___)
  
  set_Q(Qtab, s, a, Q_next)
}

run_episode_qlearn <- function(Qtab, eps) {
  state <- sample(1:9, size = 1)
  game_continues <- TRUE
  
  states <- c()
  actions_taken <- c()
  rewards <- c()
  next_states <- c()
  
  while (game_continues) {
    a <- epsilon_greedy_action(Qtab, state, eps)
    r <- payoffs(state)
    s_next <- move(state, a)
    
    Qtab <- q_update(Qtab, state, a, r, s_next)
    
    states <- c(states, state)
    actions_taken <- c(actions_taken, a)
    rewards <- c(rewards, r)
    next_states <- c(next_states, s_next)
    
    game_continues <- sample(c(TRUE, FALSE), size = 1, prob = c(0.9, 0.1))
    state <- s_next
  }
  
  list(
    Qtab = Qtab,
    episode = tibble(
      t = 1:length(states),
      state = states,
      action = actions_taken,
      payoff = rewards,
      next_state = next_states
    )
  )
}

# example use
set.seed(123)
run_episode_qlearn(Qtab, eps = .3)
```

Here's a function `animate_episode` to visualize an episode.

```{r}
library(gganimate)

animate_episode <- function(episode, payoff_vec = c(0, 0, 0, 1, 2, 1, 0, 0, 0)) {

  xy3 <- function(s) tibble(
    x = ((s - 1) %% 3) + 1,
    y = 3 - ((s - 1) %/% 3)
  )

  grid_df <- tibble(state = 1:9, payoff = payoff_vec) %>%
    bind_cols(xy3(1:9))

  states_path <- c(episode$state, tail(episode$next_state, 1))

  traj <- tibble(
    step = seq_along(states_path),
    state = states_path
  ) %>%
    bind_cols(xy3(states_path)) %>%
    mutate(
      x = x + rnorm(n(), sd = 0.05),
      y = y + rnorm(n(), sd = 0.05)
    )

  # Dynamic color map for whatever values are in payoff_vec
  payoff_levels <- sort(unique(payoff_vec))
  payoff_levels_chr <- as.character(payoff_levels)

  # Make a palette (base R) and name it to match factor(payoff) levels
  pal <- grDevices::hcl.colors(length(payoff_levels), palette = "Zissou 1")
  color_map <- stats::setNames(pal, payoff_levels_chr)

  p <- ggplot(grid_df, aes(x, y)) +
    geom_tile(aes(fill = factor(payoff, levels = payoff_levels)), color = NA) +
    geom_text(aes(label = payoff), color = "white", size = 7) +
    scale_fill_manual(values = color_map, guide = "none") +
    geom_path(
      data = traj,
      aes(x = x, y = y, group = 1),
      linewidth = 1,
      color = "black"
    ) +
    coord_equal() +
    theme_void() +
    gganimate::transition_reveal(step)

  gganimate::animate(
    p,
    nframes = nrow(traj),
    fps = 5,
    width = 200,
    height = 200,
    renderer = gganimate::gifski_renderer(loop = FALSE)
  )
}
```

Run this several times to see different episodes:

```{r, eval = F}
animate_episode(run_episode_qlearn(Qtab, eps = .3)$episode)
```

## Part 5: Train Q with 1000 episodes

**Question 5**: Run 1000 episodes, updating Qtab the whole time. Afterward, inspect Qtab for the highest Q value in each state: what stands out? Animate one more episode with eps set to 0 so that you always choose the greedy action.

```{r}
Qtab <- init_Q()

for (___) {
  out <- run_episode_qlearn(Qtab, eps = .3)
  Qtab <- out$Qtab
}

Qtab %>%
  ___ %>%
  ___
```

Run this several times to see different episodes:

```{r, eval = F}
animate_episode(run_episode_qlearn(Qtab, eps = 0)$episode)
```

## Part 6: Cliff GridWorld

**Question 6**: Consider a new GridWorld, which I'll call the cliff GridWorld. Learn Q with 1000 episodes, inspect Q for the highest Q value in each state, and animate a final episode with epsilon set to 0.

```{}
┌─────┬─────┬──────┐
│  2  │ -10 │ -10  │
├─────┼─────┼──────┤
│  0  │  1  │ -10  │
├─────┼─────┼──────┤
│  0  │  0  │  1   │
└─────┴─────┴──────┘
```

```{r}
cliff_payoffs <- function(position) {
  c(___)[position]
}

run_episode_qlearn <- function(Qtab, eps) {
  ___
}

Qtab <- init_Q()

for (___) {
  out <- run_episode_qlearn(Qtab, eps = .3)
  Qtab <- out$Qtab
}

Qtab %>%
  ___ %>%
  ___
```

```{r, eval = F}
animate_episode(
  run_episode_qlearn(Qtab, eps = 0)$episode,
  payoff_vec = c(2, -10, -10, 0, 1, -10, 0, 0, 1)
  )
```
