Causal ML: Multi-armed bandit

Simulation notebook

Goals:

• Handcode UCB and Thompson sampling in stylized DGP

Bandits

One step

To illustrate bandits, we consider a very simple DGP:

• \(Y(0) \sim N(5,2^2)\)

• \(Y(1) \sim N(5.5,2^2)\)

This means that the optimal policy rule assigns everybody to treatment one. Let’s see how UCB and Thompson sampling figure this out on-the-fly. For a stream of 1000 observations.

if (!require("tidyverse")) install.packages("tidyverse", dependencies = TRUE); library(tidyverse)
if (!require("ggridges")) install.packages("ggridges", dependencies = TRUE); library(ggridges)

set.seed(1234)

# Define parameters
gamma0 = 5
gamma1 = 5.5
sd0 = 2
sd1 = 2
n = 1000

# Draw the potential outcomes of all individuals
Y0 = rnorm(n,gamma0,sd0)
Y1 = rnorm(n,gamma1,sd1)

hist(Y0)
hist(Y1)

As a kick-start, we assign the first four individuals “manually” to control/treatment/control/treatment:

# Assign individual 1 to treatment zero
W = 0
# and observe the potential outcome under non-treatment
Y = Y0[1]
# Assign individual 2 to treatment one
W = c(W,1)
# and observe the potential outcome under treatment
Y = c(Y,Y1[2])
# Assign individual 3 to treatment zero
W = c(W,0)
# and observe the potential outcome under treatment
Y = c(Y,Y0[3])
# Assign individual 4 to treatment one
W = c(W,1)
# and observe the potential outcome under treatment
Y = c(Y,Y1[4])

# Print the realizes values
W

[1] 0 1 0 1

Y

[1] 2.585869 6.102933 7.168882 6.770741

Now calculate the means in each treatment arm:

gamma = c(mean(Y[W==0]), mean(Y[W==1]))
gamma

[1] 4.877375 6.436837

And the standard errors of the means:

se = c(sd(Y[W==0]) / sqrt(sum(W==0)),
       sd(Y[W==1]) / sqrt(sum(W==1)))
se

[1] 2.291507 0.333904

Upper Confidence Bound (UCB)

For \(i=5\), we implement now explicitly UCB \[W_{i+1} = argmax_w (\hat{\gamma}_{i,w} + \alpha \hat{\sigma}_{i,w})\] with alpha = 2:

# Calculate upper confidence bound
alpha = 2
upper_ci = gamma + alpha * se
upper_ci

[1] 9.460389 7.104645

# Choose the treatment with the highest bound
which.max(upper_ci)-1

[1] 0

Let’s visualize the decision.

plot_ucb = function(gamma, se, alpha, unit = FALSE) {
  # Calculate UCB and find the arm with the highest UCB
  ucb <- gamma + alpha * se
  highest_ucb_arm <- which.max(ucb)

  # Data frame for plotting
  data <- data.frame(
    Arm = factor(1:(length(gamma))),
    Estimate = gamma,
    UCB = ucb,
    Highlight = ifelse(1:length(gamma) == highest_ucb_arm, "forestgreen", "orange")
  )

  # Plot
  g = ggplot(data, aes(x = Arm)) +
    geom_errorbar(aes(ymin = Estimate, ymax = UCB, color = Highlight), width = 0.01, linewidth = 1) +
    geom_segment(aes(x = as.numeric(Arm)-0.05, y = UCB, xend = as.numeric(Arm)+0.05, yend = UCB, color = Highlight), linewidth = 1) +
    geom_point(aes(y = Estimate), size = 4, color = "black") +
    scale_color_identity() + geom_hline(yintercept = 0) +
    labs(y = "Estimate / Upper Confidence Bound", x = "Arm",
      title = "UCB Sampling for Multi-Armed Bandits") +
    theme_minimal() + ylim(0,10) + theme(plot.title =  element_text(hjust = 0.5))
  if (is.numeric(unit)) g = g + labs(title = paste("UCB Sampling of treatment for unit",unit+1))
  print(g)
}
plot_ucb(gamma,se,2)

Thompson sampling

Thompson sampling draws a value for each arm from an arm specific normal distribution \(\tilde{\gamma}_{i,w} \sim N(\hat{\gamma}_{i,w},\alpha^2\hat{\sigma}^2_{i,w})\) and sets \[W_{i+1} = argmax_w \tilde{\gamma}_{i,w} \]

# Thompson
draws = c(rnorm(1,gamma[1],alpha*se[1]),
          rnorm(1,gamma[2],alpha*se[2]))
draws

[1] 0.4143512 6.3703030

which.max(draws)-1

[1] 1

Let’s visualize the decision.

plot_ts = function(gamma, se, alpha, gamma_tilde=draws, unit = FALSE) {
  # This is probably the most complicated and unelegant plot I ever created, but it makes the point
  # Let me know if you have an alternative to plot it more elegantly

  highest_arm <- which.max(gamma_tilde)

  # Data frame for plotting
  data <- data.frame(
    Arm = 0:1,
    Estimate = gamma,
    Draw = gamma_tilde,
    se = se * alpha
  )

  plot_data <- data %>%
    group_by(Arm) %>%
    slice(rep(1:n(), each = 10000)) %>%
    mutate(
      y = rnorm(n(), mean = Estimate, sd = se),
    ) %>%
    mutate(
      dens = dnorm(y,mean = Estimate, sd = se),
    ) %>%
    ungroup()

  # Plot
  g = ggplot(plot_data, aes(y=Arm)) + 
    geom_point(data=data, aes(y=Arm, x = Draw,color=Draw), size = 5, shape = 4, stroke=1) +
    scale_color_gradient(low = "orange", high = "forestgreen") +
    geom_point(data=data, aes(x=Estimate), size = 3, color = "black") +
    geom_density_ridges(data = plot_data[plot_data$Arm=="0",], aes(y=Arm,x = y, height = dens),
                        alpha=0.7, position = position_nudge(y = 0), stat = "identity") +
    geom_density_ridges(data = plot_data[plot_data$Arm=="1",], aes(y=Arm,x = y, height = dens),
                      alpha=0.7, position = position_nudge(y = 0),  stat = "identity") +
    labs(y = "Arm / Normal distribution", x = "Estimate",
         title = "Thompson Sampling") +
    theme_minimal() + scale_y_continuous(breaks = c(0, 1), labels = c("0", "1")) +
    # # xlim(-10,20) +
    geom_vline(xintercept = 0) +
    theme(legend.position="none", plot.title = element_text(hjust = 0.5)) +
    coord_flip(ylim = c(0, 3),xlim = c(-10, 20))

    if (is.numeric(unit)) g = g + labs(title = paste("Thompson Sampling for unit",unit+1))
  print(g)
}
plot_ts(gamma,se,2,draws)

Note that with this very seed, UCB and Thompson sampling do not take the same choice \(\Rightarrow\) their treatment sequences differ. Note that the decision of UCB is deterministic, while it is random for Thompson sampling.

Dynamic treatment assignment (UCB)

First, proceed with UCB and define a function that takes

• observed outcomes

• assigned treatments

• and the tuning parameter \(\alpha\) as arguments

and returns a treatment assignment for the next individual:

ucb = function(Y, W, alpha, plot = FALSE){
  gamma = c(mean(Y[W==0]), mean(Y[W==1]))
  se = c(sd(Y[W==0]) / sqrt(sum(W==0)), sd(Y[W==1]) / sqrt(sum(W==1)))
  assign = which.max(gamma + alpha * se)
  if (isTRUE(plot)) plot_ucb(gamma,se,2, unit = length(Y))
  return(assign-1)
}

Now we run UCB until \(i=1000\) and plot the evolution:

# Store to use it for Thompson below
Wstart = W
Ystart = Y

plot_grid = c(5:10,20,50,100,200,500,1000)
for (i in 5:1000) {
  # Assign unit to treatment (and sometimes plot)
  if (i %in% plot_grid) Wi = ucb(Y,W,alpha,plot = TRUE)
  else Wi = ucb(Y,W,alpha,plot = FALSE)
  W = c(W,Wi)
  # and observe the potential outcome under treatment
  Y = c(Y,(1-Wi) * Y0[i] + Wi * Y1[i])
}

additionally plot the share being assigned to the optimal policy

share = cumsum(W) / 1:n
plot(1:n,share)

and the regret

regret = cumsum(gamma1 - (W * gamma1 + (1-W) * gamma0) )  / 1:n
plot(1:n,regret)

We observe that as observations flow in, UCB figures out that treatment one is dominant and starts to solely assign it. Such that at the end of the 1000 observations regret is negligible. Note that the regret of a 50:50 experiment would have been 0.25.

Dynamic treatment assignment (Thompson)

Write a function similar to ucb above:

thompson = function(Y, W, alpha, plot = FALSE){
  gamma = c(mean(Y[W==0]), mean(Y[W==1]))
  se = c(sd(Y[W==0]) / sqrt(sum(W==0)), sd(Y[W==1]) / sqrt(sum(W==1)))
  draws = c(rnorm(1, gamma[1], alpha * se[1]),
          rnorm(1, gamma[2], alpha * se[2]))
  assign = which.max(draws)
  if (isTRUE(plot)) plot_ts(gamma,se,2,draws,unit = length(Y))
  return(assign-1)
}

Now we run Thompson until \(i=1000\) and plot the evolution:

# Restore starting points
W = Wstart
Y = Ystart

plot_grid = c(5:10,20,50,100,200,500,1000)
for (i in 5:1000) {
  # Assign unit to treatment (and sometimes plot)
  if (i %in% plot_grid) Wi = thompson(Y,W,alpha,plot = TRUE)
  else Wi = thompson(Y,W,alpha,plot = FALSE)
  W = c(W,Wi)
  # and observe the potential outcome under treatment
  Y = c(Y,(1-Wi) * Y0[i] + Wi * Y1[i])
}

additionally plot the share being assigned to the optimal policy

share = cumsum(W) / 1:n
plot(1:n,share)

and the regret

regret = cumsum(gamma1 - (W * gamma1 + (1-W) * gamma0) )  / 1:n
plot(1:n,regret)

Again we observe that as observations flow in, Thompson sampling figures out that treatment one is dominant and starts to solely assign it. Such that at the end of the 1000 observations regret is negligible. Note that the regret of a 50:50 experiment would have been 0.25.

Simulation study UCB

We now run this 100 times each for a sequence of values of the tuning parameters alpha = c(0,0.5,1,1.5,2,2.5,3,4,5,10,1000). We track the following quantities:

• Average regret at \(i=10\), \(i=100\) and \(i=1000\)

• Probability to assign the optimal treatment at \(i=10\), \(i=100\) and \(i=1000\)

alpha = c(0,0.5,1,1.5,2,2.5,3,4,5,10,1000)

reps = 1000 # Reduce to 100 if you have no time

temp = matrix(NA,reps,6)
results = matrix(NA,length(alpha),6)
colnames(results) = c("regret10","regret100","regret1000","optimal10","optimal100","optimal1000")

for (a in 1:length(alpha)) {
  for (j in 1:reps) {
    # Draw sample
    Y0 = rnorm(n,gamma0,sd0)
    Y1 = rnorm(n,gamma1,sd1)

    # Get kick-start
    W = c(0,1,0,1)
    Y = (1-W) * Y0[1:4] + W * Y1[1:4]

    # Run UCB
    for (i in 5:n) {
      # Assign individual 4 to treatment one
      Wi = ucb(Y,W,alpha[a])
      W = c(W,Wi)
      # and observe the potential outcome under treatment
      Y = c(Y,(1-Wi) * Y0[i] + Wi * Y1[i])
    }
    regret = cumsum(gamma1 - (W * gamma1 + (1-W) * gamma0) )  / 1:n
    temp[j,1] = regret[10]
    temp[j,2] = regret[100]
    temp[j,3] = regret[1000]
    temp[j,4] = W[10]
    temp[j,5] = W[100]
    temp[j,6] = W[1000]
  }
 results[a,] = colMeans(temp)
}

as_tibble(cbind(alpha, results[,1:3])) %>% pivot_longer(!alpha,names_to = "individual", values_to = "regret") %>%
  ggplot(aes(factor(alpha), regret, color=as.factor(individual))) +
  geom_line(aes(group=factor(individual))) +
  labs(x = "alpha", color = "Position") + theme_bw()

as_tibble(cbind(alpha, results[,4:6])) %>% pivot_longer(!alpha,names_to = "individual", values_to = "optimal") %>%
  ggplot(aes(factor(alpha), optimal, color=as.factor(individual))) +
  geom_line(aes(group=factor(individual))) +
  labs(x = "alpha", y = "Prob of optimal choice", color = "Position") + theme_bw()

We observe that a value of \(\alpha = 3\) minimizes regret in the long run and maximizes the probability to assign the best treatment to individual 1000. Very low values of \(\alpha\) result in too little exploration. This means the bandit might commit early on to the wrong treatment. On the other hand, very high values of \(\alpha\) explore too much. \(\alpha = 1000\) is a very extreme case, which is nearly equivalent to running a 50:50 experiment. The bandit ignores basically any knowledge about the better treatment.

Simulation study Thompson sampling

Use the Thompson function in the same loop as before:

alpha = c(0,0.5,1,1.5,2,2.5,3,4,5,10,1000)

reps = 1000 # Reduce to 100 if you have no time

temp = matrix(NA,reps,6)
results = matrix(NA,length(alpha),6)
colnames(results) = c("regret10","regret100","regret1000","optimal10","optimal100","optimal1000")

for (a in 1:length(alpha)) {
  for (j in 1:reps) {
    # Draw sample
    Y0 = rnorm(n,gamma0,sd0)
    Y1 = rnorm(n,gamma1,sd1)

    # Get kick-start
    W = c(0,1,0,1)
    Y = (1-W) * Y0[1:4] + W * Y1[1:4]

    # Run UCB
    for (i in 5:n) {
      # Assign individual 4 to treatment one
      Wi = thompson(Y,W,alpha[a])
      W = c(W,Wi)
      # and observe the potential outcome under treatment
      Y = c(Y,(1-Wi) * Y0[i] + Wi * Y1[i])
    }
    regret = cumsum(gamma1 - (W * gamma1 + (1-W) * gamma0) )  / 1:n
    temp[j,1] = regret[10]
    temp[j,2] = regret[100]
    temp[j,3] = regret[1000]
    temp[j,4] = W[10]
    temp[j,5] = W[100]
    temp[j,6] = W[1000]
  }
 results[a,] = colMeans(temp)
}

as_tibble(cbind(alpha, results[,1:3])) %>% pivot_longer(!alpha,names_to = "individual", values_to = "regret") %>%
  ggplot(aes(factor(alpha), regret, color=as.factor(individual))) +
  geom_line(aes(group=factor(individual))) +
  labs(x = "alpha", color = "Position") + theme_bw()

as_tibble(cbind(alpha, results[,4:6])) %>% pivot_longer(!alpha,names_to = "individual", values_to = "optimal") %>%
  ggplot(aes(factor(alpha), optimal, color=as.factor(individual))) +
  geom_line(aes(group=factor(individual))) +
  labs(x = "alpha", y = "Prob of optimal choice", color = "Position") + theme_bw()

We see the same exploration vs. exploitation trade-off as for UCB. However, the optimal \(\alpha = 1\) and Thompson sampling manages to realize a smaller average regret of around 0.05, while UCB had only 0.07, at least with this very coarse grid.

Potential extensions

• Change the mean and/or variance of the treatment and control group. First think about what you expect and then test your intuition by running the modified notebook.