# How are afterstate value functions mathematically defined?

In this answer, afterstate value functions are mentioned, and that temporal-difference (TD) and Monte Carlo (MC) methods can also use these value functions. Mathematically, how are these value functions defined? Yes, they are a function of the next state, but what's the Bellman equation here? Is it simply defined as $$v(s') = \mathbb{E}\left[ R_t \mid S_t = s, A_t = a, S_{t+1} = s' \right]$$? If yes, how can we define it in terms of the state, $$v(s)$$, and state-action, $$q(s, a)$$, value functions, or as a Bellman (recursive) equation?

Sutton & Barto's book (2nd edition) informally describe afterstate value functions in section 6.8, but they don't provide a formal definition (i.e. Bellman equation in terms of reward or other value functions), so that's why I am asking this question.

Based on this and this resources, let me give an answer to my own question, but, essentially, I will just rewrite the contents of the first resource here, for reproducibility, with some minor changes to the notation (to be consistent with Sutton & Barto's book, 2nd edition). Note that I am not fully sure if this formulation is universal (i.e. maybe there are other ways of formulating it), but the contents of the first resource seem to be consistent with the contents in the second resource.

### Setup

Let's assume that we have an infinite-horizon MDP

$$\mathcal{M} = (\mathcal{S}, \mathcal{Y}, \mathcal{A}, \mathcal{T}, \mathcal{R}, \gamma),$$ where

• $$\mathcal{S}$$ is the set of states
• $$\mathcal{Y} \subseteq \mathcal{S}$$ is the set of afterstates (aka post-decision states or "end of period" states [1], which can also be written as after-states)
• $$\mathcal{A}$$ is the set of actions
• $$\mathcal{T}$$ is the transition function
• $$\mathcal{R}$$ is the reward function
• $$\gamma$$ is a discount factor

Let

• $$y \in \mathcal{Y}$$ be an afterstate
• $$f: \mathcal{S} \times \mathcal{A} \rightarrow \mathcal{Y}$$ be a deterministic function (from state-action pairs to afterstates), so we have $$f(s, a) = y$$

The transition function $$\mathcal{T}$$ for $$\mathcal{M}$$ is defined as

\begin{align} \mathcal{T}(s, a, s^{\prime}) &\doteq P ( s^{\prime} \mid f(s, a)) \\ &= P ( s^{\prime} \mid y) \end{align}

A transition is composed of 2 steps

1. a deterministic step, where we apply the deterministic function $$f(s, a) = y$$, which depends on an action $$a$$ taken in the state $$s$$, followed by
2. a stochastic step, where we apply the probability distribution $$P (s^{\prime} \mid y)$$, which does not depend on the action $$a$$ anymore, but only on $$y$$

So, I have denoted afterstates with a different letter, $$y$$, because afterstates are reached with a deterministic function $$f$$, while other states, $$s$$ or $$s'$$, are reached with $$P$$.

After having taken the action $$a$$ in the state $$s$$, we get a reward (i.e. we get a reward in step 1), but we do not get a reward after the stochastic step (given that no action is taken).

So, we can define the reward function $$\mathcal{R}$$ for this MDP as follows

$$\mathcal{R} (s, a, s^{\prime} ) \doteq \mathcal{R}(s, a)$$

The situation is illustrated by the following diagram

So, here, $$P$$ is the stochastic transition function (i.e. a probability distribution) as used above. Note that, here, $$r_t$$ is a specific realization of $$R_t$$ (the random variable) in the formulas below.

### State value function

Let's recall the definition of the state value function $$v_\pi(s)$$ for a given policy $$\pi$$ (as defined in Sutton & Barto, section 3.5)

\begin{align} v_{\pi}(s) &\doteq \mathbb{E}_{\pi}\left[G_{t} \mid S_{t}=s\right] \\ &= \mathbb{E}_{\pi}\left[\sum_{k=0}^{\infty} \gamma^{k} R_{t+k+1} \mid S_{t}=s\right], \end{align} for all $$s \in \mathcal{S}$$ and

\begin{align} G_{t} &\doteq \sum_{k=0}^{\infty} \gamma^{k} R_{t+k+1} \\ &= R_{t+1} + \gamma R_{t+2} + \gamma^{2} R_{t+3}+ \cdots \\ &= \mathcal{R}(s_t, a_t) + \gamma \mathcal{R}(s_{t+1}, a_{t+1})+\gamma^{2} \mathcal{R}(s_{t+2}, a_{t+2}) +\cdots, \end{align} where $$\pi(s_t) = a_t$$ and $$\mathcal{R}(s_t, a_t) = R_{t+1}$$, for $$t=0, 1, 2, \dots$$. (So, note that $$\mathcal{R} \neq R_t$$: the first is the reward function, while the second is a random variable that represents the reward received after having taken action $$a_t$$ in step $$s_t$$)

The optimal state value function is defined as

$$v_{*}(s) \doteq \max _{\pi} v_{\pi}(s)$$

### Afterstate value function

Similarly, we will define the afterstate value function, but we will use the letter $$w$$ just to differentiate it from $$v$$ and $$q$$.

\begin{align} w_{\pi}\left(y\right) &\doteq \mathbb{E}_{\pi}\left[G_{t+1} \mid Y_{t}=y\right] \\ &= \mathbb{E}_{\pi}\left[\sum_{k=0}^{\infty} \gamma^{k} R_{t+k+2} \mid Y_{t}=y\right] \\ &= \mathbb{E}_{\pi}\left[ R_{t+2} + \gamma R_{t+3}+\gamma^{2} R_{t+4} + \cdots \mid Y_{t}=y\right] \\ &= \mathbb{E}_{\pi}\left[ \mathcal{R}(s_{t+1}, a_{t+1})+\gamma \mathcal{R}(s_{t+2}, a_{t+2}) + \gamma^{2} \mathcal{R}(s_{t+3}, a_{t+3}) + \cdots \mid Y_{t}=y\right] , \end{align} where $$\mathcal{R}(s_{t+1}, a_{t+1}) = R_{t+2}$$, for all $$t$$.

In other words, the value of an afterstate $$y$$ (at time step $$t$$, i.e. given $$Y_t = y$$) is defined as the expectation of the return starting from the state that you ended up in after the afterstate $$y$$.

This seems reasonable to me and is similar to my proposal for the definition of the afterstate value function in the question, although I was not considering any deterministic functions in a potential formulation, and I was also not thinking of afterstates as intermediate states, reached by a deterministic step, between the usual states.

Similarly to the optimal state value function, we also define the optimal afterstate value function

$$w_{*}(y) \doteq \max _{\pi} w_{\pi}(y)$$

### Afterstate value function defined in terms of the state value function

We can define the afterstate value function in terms

$$w_{*}(y) = \sum_{s^{\prime}} P (s^{\prime} \mid y ) v_{*} ( s^{\prime} )$$ In other words, $$w_{*}(y)$$ is defined as an expectation over the value of next possible states $$s'$$ from the afterstate $$y$$.

This seems to be correct and consistent with the above definitions.

### More equations

In this and this resources, the state value function is also defined in terms of afterstate value function as follows

$$v_{*}(s)=\max_{a}\left(\mathcal{R}(s, a)+\gamma w_{*}(f(s, a))\right)$$

The Bellman equation for afterstate value function (from which an update rule can be derived) is given by

$$w_{*}(y) = \sum_{s^{\prime}} P(s^{\prime} \mid y ) \max_{a} ( \mathcal{R} (s^{\prime}, a) + \gamma w_{*}(f ( s^{\prime}, a ))),$$ which is really similar to the Bellman equation for the state value function.

Finally, we can also express the state-action value function in terms of the afterstate value function

$$q_\pi(s_t, a_t) = \mathcal{R}\left(s_{t}, a_{t}\right)+\gamma w_{\pi}\left(f\left(s_{t}, a_{t}\right)\right)$$

Given that this answer is already quite long, see the resource for more details (including an algorithm based on the afterstate Bellman equation).

### Implementation

If you are the type of person that understands the concepts by looking at the code, then this Github project, which implements a Monte Carlo method that uses afterstates to play tic-tac-toe, may be useful. Afterstates are useful in tic-tac-toe because it is a 2-player game, where two agents take actions in turn, so we can estimate the action that you should take deterministically (as if it was the $$f$$ above) before the other agent takes an action (probabilistically), at least, this is my current interpretation of the usefulness of afterstates in this game (and similar games/problems).

• In terms of "usefulness", one thing that is not captured so much in the formal treatment is that the mapping function $f(s,a)$ can be a compression (and is in many games). If this applies, it means there are fewer afterstates to estimate values for than action/state pairs, which means less data is required to gain better estimates with lower error bounds. Nov 24 '20 at 17:19
• @NeilSlater interesting ideas, are you working on this idea? Dec 26 '20 at 6:18