How is one of the steps in the policy gradient theorem done?

Question

In page 269 of the latest draft of Richard S. Sutton's RL book, he proves the policy gradient theorem for episodic MDPs. I'm trying to follow the proof, but I don't understand one of the steps. How does he go from:

$$ \sum_a \left[\nabla\pi(a|s)q_\pi(s,a)+\pi(a|s)\sum_{s'}p(s'|s,a)\sum_{a'}[\nabla\pi(a'|s')q_\pi(s',a')+\pi(a'|s')\sum_{s''}p(s''|s',a')\nabla v_\pi(s'')] \right] $$

to:

$$ \sum_{r\in \mathcal{S}}\sum_{k=0}^\infty Pr(s\rightarrow x,k,\pi)\sum_a \nabla \pi(a|x)q_\pi(x,a) ~?$$

He mentions that this is obtained after "repeated unrolling", where "unrolling" means (judging from the previous step in the proof) expanding $v_\pi(x)$ into its definition, dependent on $v_\pi(x')$.

Can anyone help me clarify this step?

Answer 1

According to the book, $Pr(s\rightarrow x,k,\pi)$ is the probability of transitioning from state s to state x in k steps. So $$Pr(s\rightarrow x,0,\pi)=\left\{\begin{array}{ll}1 & \mbox{if } x = s \\0 & \mbox{otherwise} \end{array}\right.$$ $$Pr(s\rightarrow x,1,\pi)=\sum_a\pi(a|s)p(x|s,a)$$ $$Pr(s\rightarrow x,2,\pi)=\sum_a\pi(a|s)\sum_{s'}p(s'|s,a)\sum_{a'}\pi(a'|s')p(x|s',a')$$.
Then the above result can be derived by expanding (unrolling) and rewriting. $$ \nabla v_\pi(s) = \sum_a [\nabla\pi(a|s)q_\pi(s,a)+\pi(a|s)\sum_{s'}p(s'|s,a)\nabla v_\pi(s')] $$ $$ = \sum_a \left[\nabla\pi(a|s)q_\pi(s,a)+\pi(a|s)\sum_{s'}p(s'|s,a)\sum_{a'}[\nabla\pi(a'|s')q_\pi(s',a')+\pi(a'|s')\sum_{s''}p(s''|s',a')\nabla v_\pi(s'')] \right] $$ $$ = \sum_a \nabla\pi(a|s)q_\pi(s,a)+\sum_a\pi(a|s)\sum_{s'}p(s'|s,a)\sum_{a'}\nabla\pi(a'|s')q_\pi(s',a')+\sum_a\pi(a|s)\sum_{s'}p(s'|s,a)\sum_{a'}\pi(a'|s')\sum_{s''}p(s''|s',a')\nabla v_\pi(s'')$$ $$ = \sum_{x\in \mathcal{S}}Pr(s\rightarrow x,0,\pi)\sum_a \nabla\pi(a|s)q_\pi(s,a)+\sum_{x\in \mathcal{S}}Pr(s\rightarrow x,1,\pi)\sum_{a'}\nabla\pi(a'|s')q_\pi(s',a')+\sum_{x\in \mathcal{S}}Pr(s\rightarrow x,2,\pi)\nabla v_\pi(s'')$$ repeatedly unrolling $\nabla v_\pi(s'')$ we get $$ \sum_{x\in \mathcal{S}}\sum_{k=0}^\infty Pr(s\rightarrow x,k,\pi)\sum_a \nabla \pi(a|x)q_\pi(x,a)$$