## State x is number of points
## Each time period, either stop and enjoy utility u(X), or roll a six-sided die
## If roll a 6, lose everything, game ends, get u(0)
## If roll k < 6, move to state x' = x + k

max_x <- 30  # Maximum state to keep things simple -- in "true" problem this is infinite
value <- rep(0, max_x)

utility <- function(x) {
    return(x)  # Linear utility function
}

n_iterations <- 100  # Value function iteration
for(iteration in seq_len(n_iterations)) {
    value_next <- value
    for(x in seq_along(value)) {
        x_next_if_roll <- pmin(x + seq_len(5), max_x)  # Next period's states if roll 1, 2, ... , 5
        value_next[x] <- max(utility(x), (1/6) * (utility(0) + sum(value[x_next_if_roll])))
    }
    distance <- mean(abs(value - value_next))
    value <- value_next
    message("iteration ", iteration, " value function distance ", round(distance, 4))
    if(distance < 10^-8) break
}

max(which(value > seq_along(value)))  # Last index at which it is optimal to keep playing

plot(value, type="b", xlab="state", ylab="value")
abline(a=0, b=1, lty=2, col="red")
abline(v=max(which(value > seq_along(value))) + 1, lty=2, col="grey")  # Stop when reach this state

sum(value[1:5]) / 6  # Value before starting game (i.e. starting from state x=0): around 6.1537