使用model-based RL，可以提高数据利用效率，同时减轻在model-free的off-policy方法训练时，对没见过的数据泛化能力不足的问题。
使用$s_{t+1}, r_t \sim T(s_t, a_t)$表示真实环境的状态转移
使用$s_{t+1}, r_t \sim \hat{T}(s_t, a_t)$表示模型仿真环境的状态转移
使用$M=(S, A, T, r, \rho_0, \gamma)$代表使用真实环境推演的模型交互，S是观测空间，A是策略空间，$r(s, a)$是真实的reward函数或规则，$\rho_0$是初始状态分布。
类似方法可以定义仿真交互模型$\hat{M}$
假设从第j步到j+1步时，使用$M$替换$\hat{M}$，同时策略不变，则沿着轨迹累积的$V(s_0)$就变成了
$W_j = \mathbb{E}_{\tau \sim t<j:\hat{M}; t\ge j:M}\left[ \sum\limits_{t=0}^{\infty} \gamma^t r(s_t, a_t) \right]$
所以$W_0=V_{M}^{\pi}(s_0)$，$W_{\infty}=V_{\hat{M}}^{\pi}(s_0)$
两种交互轨迹的差是：
$V_{\hat{M}}^{\pi}(s_0) – V_{M}^{\pi}(s_0) = \sum\limits_{j=0}^{\infty}(W_{j+1} – W_j)$
稍微改写一下$W$
$W_j = \mathbb{E}_{\hat{M}}[\sum\limits_{t \le j}\gamma^t r_t] + \mathbb{E}_{a_j,s_j \sim \pi, \hat{T}} \left[ \mathbb{E}_{s_{j+1} \sim T}[ \gamma^{j+1} V_{M}^{\pi}(s_{j+1})] \right]$
右边第一项是仿真得到的路径，右边第二项的第一个预期表示$s_j, a_j$是从上一步计算$\pi(a_{j-1} \mid s_{j-1})$，再经过仿真转移$\hat{T}$计算出的。第二个预期表示从j到j+1的转移是使用真实环境。
同样可得
$W_{j+1} = \mathbb{E}_{\hat{M}}[\sum\limits_{t \le j}\gamma^t r_t] + \mathbb{E}_{a_j,s_j \sim \pi, \hat{T}} \left[ \mathbb{E}_{s_{j+1} \sim \hat{T}}[ \gamma^{j+1} V_{M}^{\pi}(s_{j+1})] \right]$
所以
$W_{j+1} – W_j = \gamma^{j+1} \mathbb{E}_{a_j,s_j \sim \pi, \hat{T}} \left[ \mathbb{E}_{s_{j+1} \sim T}[ V_{M}^{\pi}(s_{j+1})] – \mathbb{E}_{s_{j+1} \sim \hat{T}}[ V_{M}^{\pi}(s_{j+1})] \right]$
定义“仅在当前步使用两种不同转移，下一步以及之后都用真实转移，导致的下一步value function的误差”
$G_{\hat{M}}^{\pi}(s, a) := \mathbb{E}_{s’ \sim T(s, a)}[ V_{M}^{\pi}(s’)] – \mathbb{E}_{s’ \sim \hat{T}(s,a)}[ V_{M}^{\pi}(s’)] $
于是
$$V_{\hat{M}}^{\pi}(s_0) – V_{M}^{\pi}(s_0) = \sum\limits_{j=0}^{\infty}\gamma^{j+1} \mathbb{E}_{a_j,s_j \sim \pi, \hat{T}} \left[ G_{\hat{M}}^{\pi}(s_j, a_j) \right] \tag{$1$} $$
使用计算stationary distribution $\rho$的思路来看待$\rho_j=\rho(s=s_j)$的话，无论任何初始状态，推演了多少步的轨迹，只要对齐这些轨迹的$s=s_j$作为原点，就可以把预期
$\mathbb{E}_{a_j,s_j \sim \pi, \hat{T}} \left[ \sum\limits_{j=0}^{\infty} \gamma^j G_{\hat{M}}^{\pi}(s_j, a_j) \right]$
转化成
$\mathbb{E}_{a_t,s_t \sim \tau[s_t \rightarrow s] \mid \pi, \hat{T}} \left[ \sum\limits_{t=0}^{\infty} \gamma^t G_{\hat{M}}^{\pi}(s, a) \right]$
使用geometric sum
$=\frac{1}{1-\gamma} \mathbb{E}_{s \sim \rho(s), a \sim \pi} \left[ G_{\hat{M}}^{\pi}(s, a) \right]$
带入(1)得到
$V_{\hat{M}}^{\pi}(s_0) – V_{M}^{\pi}(s_0) = \frac{\gamma}{1-\gamma} \mathbb{E}_{(s, a) \sim \rho_{\hat{T}}^{\pi}} \left[ G_{\hat{M}}^{\pi}(s, a) \right]$
$\rho_{\hat{T}}^{\pi}$代指在转移$\hat{T}$以及策略$\pi$生成轨迹下，计算出的stationary状态分布$\rho(s)$
注意到其中$V_{\hat{M}}^{\pi}(s) \mid _{s_0=s}=\mathbb{E}_{a_t, s_t \sim \tau[s_t \rightarrow s] \mid \pi, \hat{T}} \left[ \sum\limits_{t=0}^{\infty} \gamma^t r(s, a) \right] = \frac {1} {1-\gamma} \mathbb{E}_{(s, a) \sim \rho_{\hat{T}}^{\pi}} \left[ r(s, a) \right]$
于是可得：
$$V_{M}^{\pi}(s_0) = \frac {1} {1-\gamma} \mathbb{E}_{(s, a) \sim \rho_{\hat{T}}^{\pi}} \left[ r(s, a) – \gamma G_{\hat{M}}^{\pi}(s, a) \right] $$
这个公式描述了如果使用仿真环境而非真实环境推演MDP，则应该对reward做出修正。修正幅度是使用$T$与$\hat{T}$转移到下一步导致的$V(s’)$误差，再乘上衰减系数。
对比一下model-free的优化目标（状态的静态先验预期形式）：$V_{M}^{\pi}(s_0)=\frac {1} {1-\gamma} \mathbb{E}_{(s, a) \sim \rho_{T}^{\pi}} \left[ r(s, a) \right]$
这个结论看似毫无意义，是个明显正确的思路。但能够使用类似于之前计算transition误差上界的方法，同样计算出$ |G_{\hat{M}}^{\pi}(s, a) |$ 的上界。之后我们在仿真环境每一步得到的reward，修正后总会小于真实环境推演下得到的reward，也就是得到了真实环境推演reward的下界。这里有个前提假设是仿真与真实环境的$r(s, a)$规则相同。那么只需要在优化actor时朝着优化真实环境累积reward的下界方向进行policy gradient，就能保证真实环境的$V_{M}^{\pi}(s_0)$也被优化。
由于推演是在仿真环境中进行的，意味着我们的off-policy dataset能够被仿真环境增强，然后通过对reward的修正又能一定程度避免仿真$\hat{T}$带来的误差，于是取得比一般的model-based方法更好的效果。
注意到$ G_{\hat{M}}^{\pi}(s, a)$的定义，在当前步if假设了一条使用$\hat{T}$的轨迹分支，但随后在评估$V(s’)$时，仍然使用的是真实转移$V_{T}^{\pi}$，所以我们无法用$\hat{T}$进行N步TD的采样，而是只推演了1步。不过一般的off-policy也都是使用1步TD来更新的，从bias-variance取舍的角度来看，还是偏向于低variance的。
回到找优化目标下界的问题：
$V_{M}^{\pi}(s_0) = \frac {1} {1-\gamma} \mathbb{E}_{(s, a) \sim \rho_{\hat{T}}^{\pi}} \left[ r(s, a) – \gamma G_{\hat{M}}^{\pi}(s, a) \right] \ge \frac {1} {1-\gamma} \mathbb{E}_{(s, a) \sim \rho_{\hat{T}}^{\pi}} \left[ r(s, a) – \gamma | G_{\hat{M}}^{\pi}(s, a) | \right]$
其中
$| G_{\hat{M}}^{\pi}(s, a) | := | \mathbb{E}_{s’ \sim T(s, a)}[ V_{M}^{\pi}(s’)] – \mathbb{E}_{s’ \sim \hat{T}(s,a)}[ V_{M}^{\pi}(s’)] | $
$ \le \int\limits_{s’} | \left( \hat{p}(s’ , r \mid \hat{T}(s, a)) – p(s’ , r\mid T(s, a)) \right) V_{M}^{\pi}(s’) | ds’$
$= \int\limits_{s’} | \left( \hat{p}(s’, r \mid \hat{T}(s, a)) – p(s’, r \mid T(s, a)) \right) \left[ \frac{1}{1-\gamma} \mathbb{E}_{\bar{s}, \bar{a} \sim \rho_{M}^{\pi}}[r(\bar{s}, \bar{a}) \mid _{\bar{s}=s’, \bar{a}=a}] \right] | ds’ $
当$r_{max} > 0 $ where $ r(s,a) \in \{r_i : Tr(s \rightarrow s’ \mid M, \pi) \}_{i=1…N}$时：
$| G_{\hat{M}}^{\pi}(s, a) | \le \frac{1}{1-\gamma} \int\limits_{s’} |\hat{p}(s’, r \mid \hat{T}(s, a)) – p(s’, r \mid T(s, a))| r_{max} ds’$
注意到$\{r_i : Tr(s \rightarrow s’ \mid M, \pi) \}_{i=1…N}$因为是offline的，所以无法从真实环境中采样获得，而且我们也没有N个真实环境可供采样。所以我们从N个仿真环境采样的集合$\{r_i : Tr(s \rightarrow s’ \mid \hat{M}, \pi) \}_{i=1…N}$采样，并获得这些仿真的r，求出$r_{max}$。只要N足够大且不同仿真环境有一定的差异，就总能估计一个相对靠谱的$r_{max}$，以确保$r_{max} > r_{true}$。再次强调这里N次推演采样并非从一个模拟仿真环境重复N次获得，而是搭建N个模拟仿真环境，各自用maximum likelihood在历史轨迹上独立训练。这样$\hat{T}$作为一套仿真环境组合形成的ensemble model，理论上拥有比单个仿真环境更可靠的$r_{max} > \mathbb{E}_{\bar{s}, \bar{a} \sim \rho_{M}^{\pi}}[r(\bar{s}, \bar{a}) ]$的针对真实reward的较高估算。仿真环境转移里的reward预测，也不可能超过$r_{max}$，于是可以忽略对$r$的预期积分，只保留关于s’的积分。
于是目标的下界
$V_{M}^{\pi}(s_0) \ge \mathbb{E}_{(s, a) \sim \rho_{\hat{T}}^{\pi}} \left[ r(s, a) – \frac{\gamma}{1-\gamma} 2D_{TV}(\hat{p}(\cdot \mid s, a) || p(\cdot \mid s, a)) r_{max} \right]$
之前我们假设了reward拥有上界$r_{max}$作为条件，于是得到$V_{M}^{\pi}(s_0)$的下界里包含$\frac{r_max}{1-\gamma}D_{TV}[]$的一项。但如果更改其他针对reward的约束条件，例如约束reward满足1-Lipschitz连续条件，那么右边项里就会变成一项包含了1-Wasserstein距离的表达式。具体推导见论文原文。总之右边的惩罚项是可能以多种不同形式存在的。
继续回到$r_{max}$的限制方式上来，现在如何计算真实与仿真的状态转移的TV距离就是主要问题了。我们目前有N个样本采样自$\hat{T}(s’^{(i)}, \hat{r} \mid s, a)$，然后我们并不知道真实状态转移$T(s’ , r\mid s, a)$的分布，只有真实转移的一个样本（来自dataset）。
改写一下$D_{TV}[\hat{p} || p] = \mathbb{E}_{\hat{p}}[ |1 – \frac{p}{\hat{p}}| ] \approx \frac{1}{N} \sum\limits_{i=1}^N |1 – \frac {\sigma_i exp \{ \frac{(s’-\mu_i)^2}{\sigma_i^2} \}} { \sigma_{true} exp\{ \frac{(s’ – \mu_{true})^2} {\sigma_{true}^2}\}} |$
我们不知道$\sigma_{true}$与$\mu_{true}$，所以只能暂时假设$\sigma_{true}=\sigma_i$，$\mu_{true} = \frac {1} {N} \sum\limits_{i=1}^N \mu_i$
于是$D_{TV}[\hat{p} || p] = \frac{1}{N} \sum\limits_{i=1}^N | 1 – exp \{ \frac {2s'(\mu_{true}-\mu_i)+ \mu_i^2-\mu_{true}^2} {\sigma_i^2} \}| $
允许$D_{TV}[\hat{p} || p]$被高估，于是可以展开右边的exp，并舍弃二次项以及后面的。
得到$D_{TV}[\hat{p} || p] \le \frac{1}{N} \sum\limits_{i=1}^N | – \frac {2s'(\mu_{true}-\mu_i)+ \mu_i^2-\mu_{true}^2} {\sigma_i^2} |$
接下来的推导不严谨，但可以作为实验参考：
$s’$项在最外层的$\sum$下消失，之后额外加上$\frac{2\mu_i\mu_{true}}{\sigma_i^2}$项允许更进一步地高估，能够拼凑出
$ D_{TV}[\hat{p} || p] \le \frac{1}{N} \sum\limits_{i=1}^{N} |- \frac{(\mu_i – \mu_{true})^2} { \sigma_i^2}|$
$ \le max_i \frac{(\mu_i – \mu_{true})^2} {\sigma_i^2}$
使用$mean_{i}$还是$max_{i}$其实都可以，保守一些就使用max。
reward惩罚的思路从公式中也能看清楚，就是对于最“与众不同”的，而且又“自信满满”的仿真环境施加较大的惩罚，然后让模型即使在最糟糕的惩罚的情况，也能使用策略梯度方法优化$V(s_0)$的下界。多个并行的仿真环境上搜集的轨迹，也会比最糟糕的$\hat{T}$要好一些，所以基本上确保了即使没有优质仿真时的强化学习训练效果。
论文中也有不使用$max_i \frac{(\mu_i – \mu_{true})^2} {\sigma_i^2}$，而是用ensemble penalty $max_i | \mu_i – \mu_{true} |$，以及direct prediction error $| \hat{s’}^{(i)} – s’ |$ 作为代替惩罚项的。
最终真实轨迹数据集$D_{env}$与仿真生成的轨迹数据集$D_{model}$会union在一起，之后使用SAC来训练RL模型。
