Integrating a function of a Markov chain on a random length trajectory Suppose that $(X_k)_{k\geq 0}$ is a Markov chain in a finite state space $\mathcal{S}$, with a known transition matrix $\mathbf{P}$, and let $f:\mathcal{S}\to\mathbb{R}$ be a known function that assigns a real value $f(s)$ to each state in $s\in\mathcal{S}$.
If $n$ is a fixed integer and $X_0 = s_0\in\mathcal{S}$ with probability one, then we can easily compute the sum of $f$ along a trajectory of length $n$, given by $\mathbf{E}[\sum_{i=0}^{n}f(X_i)]$. This can be done by computing each term separately and integrating $f$ with respect to the law of $X_i$, which can be expressed using the transition matrix $\mathbf{P}$ and its powers.
I am trying to solve an extension of this problem, where we replace the fixed $n$ by a random integer $N$ that depends on the chain in the following way. Instead of having a trajectory of fixed length $n$ in the sum, I want certain states to be able to increase the length of the trajectory when visited.
To formalize this idea, I define a map $m:\mathcal{S}\to\mathbb{N}\cup\{0\}$ with $m(s_0)\geq 1$, where $m(s)$ represents the increase in trajectory length brought by visiting the state $s \in \mathcal{S}$, and I define the process $(M_k)_{k\geq 0}$ that will keep track of the remaining trajectory length. We set $M_0 = m(s_0) \geq 1$, which is the initial trajectory length, and let $M_k = M_{k-1} + m(X_k) - 1$. With each time step and visited state $s$, the remaining trajectory length decreases by $1$, and increases by $m(s)$.
The summation of $f$ along the trajectory stops when $M_k$ reaches zero, i.e. we define $N = \inf\{k\geq 0 : M_k = 0\}$ and are interested in computing $\mathbf{E}[\sum_{i=0}^{N}f(X_i)]$.
To exclude the possibility that $\mathbf{P}(N=\infty)>0$, I assume that there exists one and only one absorbing state $s_A\in\mathcal{S}$, such that $m(s_A)=0$. This makes $N$ almost surely finite. However, if there exists a state $s\neq s_A$ with $m(s)=0$, it can happen that $X_N \neq s_A$.
Is there a hope to analytically calculate $\mathbb{E}[\sum_{i=0}^N f(X_i)]$ for a general choice of $\mathcal{S}, \mathbf{P}, f$ and $m$ that satisfy the above properties ? The only idea I have is to partition the space according to the value of $N$, but then the difficulty shifts to computing the quantities $\mathbb{P}(N=n)$.
 A: So, you have to solve the following problem. You have a Markov process on some state space $X$ with a transition probability kernel $T$, stopping set $A$ and a reward function $g$. You would like to compute
$$
R(x) := \Bbb E_x\left[\sum_{k=0}^{\tau_A}g(x_k)\right] \tag{1}
$$
where $\tau_A$ is the first time the process visits the set $A$, and expectation is taken with respect to all trajectories that starts at $x \in X$. This function $R$ must satisfy the following Bellman equation
$$
R(x) = 1_{A^c}(x)\left(g(x) + TR(x)\right) \tag{2}
$$
where $1_{A^c}$ is the indicator function of $A^c = X\setminus A$ and
$$
TR(x) = \int_X R(x')T(\mathrm dx'|x) = \sum_{x'\in X} R(x')T(x'|x) \tag{3}
$$
where the last equality holds if the space $X$ is countable (your case). To get an understanding of $(2)$ see that if $x \in A$ then you have already reached the stopping set, so no more terms in $(1)$ will enter the sum. Conversely, if $x \in A^c$ you will get a reward $g(x)$ and an average of the future rewards, depending on where you will end up starting from $x$. The $(3)$ precisely describes this averaging happening.
Now, everything above holds for the most general setting, so what's left is to fit it into your special case. You have a state space $X = \mathcal S \times \Bbb N_0$ (here $\Bbb N_0 := \Bbb N\cup\{0\}$), for reward you get $g(x,n) = f(x)$ and finally
$$
T((x',n')|(x,n)) = \begin{cases}
P(x'|x),&\text{if }n' = n-1+m(x),
\\
0,&\text{otherwise}.
\end{cases}
$$
Ah, and $A = \mathcal S\times \{0\}$. So your expectation satisfies
$$
R(s, n) = 1_{\Bbb N}(n)\cdot\left(f(s) + \sum_{s'\in \mathcal S}R(s', n -1 + m(s))P(s'|s)\right)
$$
which is just $(2)$ written for your case. Looks analytical enought to me. I've spent half of my PhD working out numerical methods for such problems.
