9 Exploration in MDPs

9.1Introduction¶

One of the key challenges of reinforcement learning is the exploration-exploitation tradeoff. Should we exploit actions we know will give high reward, or should we explore different actions to discover potentially better strategies? An algorithm that doesn’t explore effectively might easily overfit to certain areas of the state space, and fail to generalize once they enter a region they haven’t yet seen. The algorithms we saw in the chapter on fitted DP 5 Fitted Dynamic Programming Algorithms suffer from this issue.

In 3 Multi-Armed Bandits, where the state never changes so all we care about are the actions, we saw algorithms like Section 3.6 and Thompson sampling that incentivize the learner to explore arms that it is uncertain about. In this chapter, we will see how to generalize these ideas to the MDP setting.

9.1.1Sparse reward¶

Exploration is especially crucial in sparse reward problems where reward doesn’t come until after many steps, and algorithms which do not systematically explore new states may fail to learn anything meaningful (within a reasonable amount of time).

For example, policy gradient algorithms require the gradient to be nonzero in order to learn. If we never observe any reward, the gradient will always be zero, and the policy will never change or improve.

9.1.2Exploration in deterministic MDPs¶

Let us address the exploration problem in a deterministic MDP where taking action $a$ in state $s$ always leads to the state $P(s, a) \in \mathcal{S}$. In this simple setting, there will be no “automatic” exploration due to randomness, so our strategy must actively explore new states. One simple strategy is to visit every possible state-action pair to learn the entire MDP. Then, once the MDP is known, we can use DP to solve for the optimal policy. (This should remind you of the Section 3.4 algorithm.)

9.2Treating an unknown MDP as a MAB¶

We also explored the exploration-exploitation tradeoff in 3 Multi-Armed Bandits. Recall tthat in the MAB setting, we have $K$ arms, each of which has an unknown reward distribution, and we want to learn which of the arms is optimal, i.e. has the highest mean reward.

One algorithm that struck a good balance between exploration and exploitation was the upper confidence bound algorithm Section 3.6: For each arm, we construct a confidence interval for its true mean award, and then choose the arm with the highest upper confidence bound. In summary,

where $N_t^k$ indicates the number of times arm $k$ has been pulled up until time $t$, $R_t^k$ indicates the total reward obtained by pulling arm $k$ up until time $t$, and $\delta > 0$ controls the width of the confidence interval. How might we extend UCB to the MDP case?

Let us formally describe an unknown MDP as an MAB problem. In an unknown MDP, we want to learn which policy is optimal. So if we want to apply MAB techniques to solving an MDP, it makes sense to think of arms as policies. There are $K = (|\mathcal{A}|^{|\mathcal{S}|})^\hor$ deterministic policies in a finite MDP. Then, “pulling” arm π corresponds to using π to act through a trajectory in the MDP, and observing the total reward.

Recall that UCB incurs regret $\tilde{O}(\sqrt{TK})$, where $T$ is the number of pulls and $K$ is the number of arms. So in the MDP-as-MAB problem, using UCB for $T$ episodes would achieve regret

This scales exponentially in $|\mathcal{S}|$ and $\hor$, which quickly becomes intractable. Notably, this method doesn’t consider the information that we gain across different policies. We can illustrate this with the following example:

9.3UCB-VI¶

The approach above is inefficient: We shouldn’t need to consider all $|\mathcal{A}|^{|\mathcal{S}| H}$ deterministic policies to achieve low regret. Rather, all we need to describe the optimal policy is $Q^\star$, which has $H |\mathcal{S}||\mathcal{A}|$ entries to be learned. Can we borrow ideas from UCB to reduce the regret to this order (i.e. polynomial in $|\mathcal{S}|$, $|\mathcal{A}|$, and $H$)?

One way to frame the UCB algorithm is that, when choosing arms, we optimize over a proxy reward that is the sum of the estimated mean reward and an exploration term. In the UCB-VI algorithm, we will extend this idea to the case of an unknown MDP $\mathcal{M}^{?}$ by modelling a proxy MDP $\tilde{\mathcal{M}}$ with a reward function that encourages exploration. Then, we will use DP to solve for the optimal policy in $\tilde{\mathcal{M}}$.

Assumptions: For simplicity, here we assume the reward function of $\mathcal{M}^{?}$ is known, so we only need to model the state transitions, though the rewards can be modelled similarly. We will also consider the more general case of a time-varying MDP, where the transition and reward functions can change over time. We take the convention that $P_\hi$ is the distribution of $s_{h+1} \mid s_{h}, a_{h}$ and $r_\hi$ is applied to $s_\hi, a_\hi$.

At a high level, the UCB-VI algorithm can be described as follows:

1. Modelling: Use previous data to model the transitions $\hat{P}_0, \dots, \hat{P}_{H-1}$.

2. Reward bonus: Design a reward bonus $b_\hi(s, a) \in \mathbb{R}$ to encourage exploration, analogous to the UCB term.

3. Optimistic planning: Use DP to compute the optimal policy $\hat \pi_\hi(s)$ in the modelled MDP

4. Execution: Use $\hat \pi_\hi(s)$ to collect a new trajectory, and repeat.

We detail each of these steps below. The full definition follows in (9.16).

9.3.1Modelling the transitions¶

We seek to approximate $P_\hi(s_{h+1} \mid s_\hi, a_\hi) = \frac{\pr(s_\hi, a_\hi, s_{h+1})}{\pr(s_\hi, a_\hi)}$. We can estimate these using their sample probabilities from the dataset. That is, define

Then we can model

9.3.2Reward bonus¶

To motivate the reward bonus term $b_\hi^t(s, a)$, recall how we designed the reward bonus term for UCB:

1. We used Hoeffding’s inequality to bound, with high probability, how far the sample mean $\hat \mu_t^k$ deviated from the true mean $\mu^k$.

2. By inverting this inequality, we obtained a $(1-\delta)$-confidence interval for the true mean, centered at our estimate.

3. To make this bound uniform across all timesteps $t \in [T]$, we applied the union bound and multiplied δ by a factor of $T$.

We’d like to do the same for UCB-VI, and construct the bonus term such that $V^\star_\hi(s) \le \hat{V}_\hi^t(s)$ with high probability. However, our construction will be more complex than the MAB case, since $\hat{V}_\hi^t(s)$ depends on the bonus $b_\hi^t(s, a)$ implicitly via DP. We claim that the bonus term that gives the proper bound is

We will only provide a heuristic sketch of the proof; see Agarwal et al. (2022) (Section 7.3) for a full proof.

9.3.3Definition¶

Putting these parts together, we can define the algorithm as follows:

9.3.4Performance of UCB-VI¶

How exactly does UCB-VI strike a good balance between exploration and exploitation? In UCB for MABs, the bonus exploration term is simple to interpret: It encourages the learner to take actions with a high exploration term. Here, the policy depends on the bonus term indirectly: The policy is obtained by planning in an MDP where the bonus term is added to the reward function. Note that the bonuses propagate backwards in DP, effectively enabling the learner to plan to explore unknown states. This effect takes some further interpretation.

Recall we constructed $b^t_\hi$ so that, with high probability, $V^\star_\hi(s) \le \hat{V}_\hi^t(s)$ and so

That is, the l.h.s. measures how suboptimal policy $\pi^t$ is in the true environment, while the r.h.s. is the difference in the policy’s value when acting in the modelled MDP $\tilde{\mathcal{M}}^t$ instead of the true one $\mathcal{M}^{?}$.

If the r.h.s. is small, this implies that the l.h.s. difference is also small, i.e. that $\pi^t$ is exploiting actions that are giving high reward.

If the r.h.s. is large, then we have overestimated the value: $\pi^t$, the optimal policy of $\tilde{\mathcal{M}}^t$, does not perform well in the true environment $\mathcal{M}^{?}$. This indicates that one of the $b_h^t(s, a)$ terms must be large, or some $\hat P^t_\hi(\cdot \mid s, a)$ must be inaccurate, indicating a state-action pair with a low visit count $N^t_\hi(s, a)$ that the learner was encouraged to explore.

It turns out that UCB-VI achieves a per-episode regret of

Comparing this to the UCB regret bound $\tilde{O}(\sqrt{T K})$, where $K$ is the number of arms of the MAB, we see that we’ve reduced the number of effective arms from $|\mathcal{A}|^{|\mathcal{S}|\hor}$ (in (9.4)) to $H^4 |\mathcal{S}||\mathcal{A}|$, which is indeed polynomial in $|\mathcal{S}|$, $|\mathcal{A}|$, and $H$, as desired. This is also roughly the number of episodes it takes to achieve constant-order average regret:

Note that the time-dependent transition matrix has $H |\mathcal{S}|^2 |\mathcal{A}|$ entries. Assuming $H \ll |\mathcal{S}|$, this shows that it’s possible to achieve low regret, and achieve a near-optimal policy, while only understanding a $1/|\mathcal{S}|$ fraction of the world’s dynamics.

9.4Linear MDPs¶

A polynomial dependency on $|\mathcal{S}|$ and $|\mathcal{A}|$ is manageable when the state and action spaces are small. But for large or continuous state and action spaces, even this polynomial factor will become intractable. Can we find algorithms that don’t depend on $|\mathcal{S}|$ or $|\mathcal{A}|$ at all, effectively reducing the dimensionality of the MDP? In this section, we’ll explore linear MDPs: an example of a parameterized MDP where the rewards and state transitions depend only on some parameter space of dimension $d$ that is independent from $|\mathcal{S}|$ or $|\mathcal{A}|$.

9.4.1Planning in a linear MDP¶

It turns out that $Q^\star_\hi$ is also linear with respect to this feature mapping. We can prove this by simply computing it using DP. We initialize $V_{H}^\star(s) = 0 \forall s$. Then we iterate:

9.4.2UCB-VI in a linear MDP¶

9.4.2.1Modelling the transitions¶

This linear assumption on the MDP will also allow us to model the unknown dynamics $P^?_\hi(s' \mid s, a)$ with techniques from supervised learning (SL). Recall that SL is useful for estimating conditional expectations by minimizing mean squared error. We can rephrase the estimation of $P^?_\hi(s' \mid s, a)$ as a least-squares problem as follows: Write $\delta_s$ to denote a one-hot vector in $\mathbb{R}^{|\mathcal{S}|}$, with a 1 in the $s$-th entry and 0 everywhere else. Note that

Furthermore, since the expectation here is linear with respect to $\phi(s, a)$, we can directly apply least-squares multi-target linear regression to construct the estimate

This has a well-known closed-form solution:

where we include a $\lambda I$ term to ensure that the matrix $A^t_h$ is invertible. (This can also be derived by adding a $\lambda \|\mu\|_{\text{F}}^2$ regularization term to the objective.) We can directly plug in this estimate into $\hat{P}^t_h(\cdot \mid s, a) = \hat \mu^t_h \phi(s, a)$.

9.4.2.2Reward bonus¶

Now, to design the reward bonus, we can’t apply Hoeffding anymore, since the terms no longer involve sample means of bounded random variables; Instead, we’re incorporating information across different states and actions. Rather, we can construct an upper bound using Chebyshev’s inequality in the same way we did for the LinUCB algorithm in the MAB setting Section 3.8.1:

Note that this isn’t explicitly inversely proportional to $N_h^t(s, a)$ as in the original UCB-VI bonus term (9.8). Rather, it is inversely proportional to the amount that the direction $\phi(s, a)$ has been explored in the history. That is, if $A_h^t$ has a large component in the direction $\phi(s, a)$, implying that this direction is well explored, then the bonus term will be small, and vice versa.

We can now plug in these transition estimates and reward bonuses into the UCB-VI algorithm (9.16).

9.4.2.3Performance¶

Comparing this to our bound for UCB-VI in an environment without this linear assumption, we see that we go from a sample complexity of $\tilde \Omega(H^4 |\mathcal{S}||\mathcal{A}|)$ to $\tilde \Omega(H^4 d^{3})$. This new sample complexity only depends on the feature dimension and not on the state or action space of the MDP!

9.5Summary¶

In this chapter, we’ve explored how to explore in an unknown MDP.

• We first discussed the explore-then-exploit algorithm Definition 9.2, a simple way to explore a deterministic MDP by visiting all state-action pairs.

• We then discussed how to treat an unknown MDP as a MAB Section 9.2, and how this approach is inefficient since it doesn’t make use of relationships between policies.

• We then introduced the UCB-VI algorithm (9.16), which models the unknown MDP by a proxy MDP with a reward bonus term that encourages exploration.

• Finally, assuming that the transitions and rewards are linear with respect to a feature transformation of the state and action, we introduced the LinUCB-VI algorithm Section 9.4.2, which has a sample complexity independent of the size of the state and action spaces.