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
- 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
- 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).
