SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {1}{1 - \gamma} \sum_ {t = 1} ^ {\infty} \gamma^ {t} \mathbb {E} _ {s ^ {\prime}, a \sim p _ {t - 1, M _ {1}} ^ {\pi_ {1}} (s ^ {\prime}, a)} D _ {T V} \left(p _ {M _ {1}} (s | s ^ {\prime}, a) | | p _ {M _ {2}} (s | s ^ {\prime}, a)\right) \\ = \frac {\gamma}{1 - \gamma} \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {s ^ {\prime}, a \sim p _ {t, M _ {1}} ^ {\pi_ {1}} (s ^ {\prime}, a)} D _ {T V} \left(p _ {M _ {1}} (s | s ^ {\prime}, a) \mid \mid p _ {M _ {2}} (s | s ^ {\prime}, a)\right) \\ = \frac {\gamma}{(1 - \gamma) ^ {2}} (1 - \gamma) \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {s ^ {\prime}, a \sim p _ {t, M _ {1}} ^ {\pi_ {1}} \left(s ^ {\prime}, a\right)} D _ {T V} \left(p _ {M _ {1}} \left(s \mid s ^ {\prime}, a\right) \mid \mid p _ {M _ {2}} \left(s \mid s ^ {\prime}, a\right)\right) \tag {27} \\ = \frac {\gamma}{(1 - \gamma) ^ {2}} \mathbb {E} _ {s ^ {\prime}, a \sim d _ {M _ {1}} ^ {\pi_ {1}} (s ^ {\prime}, a)} D _ {T V} \left(p _ {M _ {1}} \left(s \mid s ^ {\prime}, a\right) \mid \mid p _ {M _ {2}} \left(s \mid s ^ {\prime}, a\right)\right) \\ = \frac {\gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {1}} ^ {M _ {2}} \\ \end{array}
$$
Then we bring this result back to Eq.(26) and the proof is complete.
$$
V _ {M _ {2}} ^ {\pi_ {2}} - V _ {M _ {1}} ^ {\pi_ {1}} \geq - 2 R _ {\max } \left(\frac {\epsilon_ {\pi}}{(1 - \gamma) ^ {2}} + \frac {\gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {1}} ^ {M _ {2}}\right) \tag {28}
$$
Theorem 3 (Decomposition TVD Bound). Let  $\epsilon_{M_i}^{\pi_i}$  denote  $\mathbb{E}_{(s,a)}\sim d_{M_i}^{\pi_i}[D_{TV}(p_{M_i}||p_{M^*})]$ . Let  $M_{i}\in \mathcal{M}$  be the evaluated model and  $\pi_i\in \Pi$  be the policy derived from the model. The decomposition terms can be bounded as,
$$
V ^ {\pi_ {2} \mid M _ {2}} - V ^ {\pi_ {1} \mid M _ {1}} \geq \frac {2 R _ {\text {m a x}} \gamma}{(1 - \gamma) ^ {2}} \left(\epsilon_ {M _ {1}} ^ {\pi_ {1}} - \epsilon_ {M _ {2}} ^ {\pi_ {2}} - \epsilon_ {M _ {1}} ^ {M _ {2}}\right) - \frac {2 R _ {\text {m a x}} \epsilon_ {\pi}}{(1 - \gamma) ^ {2}} \tag {29}
$$
Proof: According to CMLO [21] and Eq.(21), the term  $V^{\pi_1|M_1} - V_{M_1}^{\pi_1}$  can be approximated as  $-\frac{2R_{max}\gamma}{(1 - \gamma)^2}\epsilon_{M_1}^{\pi_1}$ , thus we only need to bound the remaining two terms.
For the term  $V^{\pi_2|M_2} - V_{M_2}^{\pi_2}$ , we use Eq.(21) to bound it.
$$
\begin{array}{l} V ^ {\pi_ {2} \mid M _ {2}} - V _ {M _ {2}} ^ {\pi_ {2}} \geq - 2 R _ {\max } \left(\frac {\operatorname* {m a x} _ {s} D _ {T V} \left(\pi_ {2} \mid \mid \pi_ {2}\right)}{(1 - \gamma) ^ {2}} + \frac {\gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {2}} ^ {\pi_ {2}}\right) \tag {30} \\ = - \frac {2 R _ {m a x} \gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {2}} ^ {\pi_ {2}} \\ \end{array}
$$
Similarly, for the term  $V_{M_2}^{\pi_2} - V_{M_1}^{\pi_1}$ , we can get:
$$
V _ {M _ {2}} ^ {\pi_ {2}} - V _ {M _ {1}} ^ {\pi_ {1}} \geq - 2 R _ {\max } \left(\frac {\epsilon_ {\pi}}{(1 - \gamma) ^ {2}} + \frac {\gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {1}} ^ {M _ {2}}\right) \tag {31}
$$
We now combine these three bounds together and complete the proof.
$$
\begin{array}{l} V ^ {\pi_ {2} | M _ {2}} - V ^ {\pi_ {1} | M _ {1}} = \left(V ^ {\pi_ {2} | M _ {2}} - V _ {M _ {2}} ^ {\pi_ {2}}\right) - \left(V ^ {\pi_ {1} | M _ {1}} - V _ {M _ {1}} ^ {\pi_ {1}}\right) + \left(V _ {M _ {2}} ^ {\pi_ {2}} - V _ {M _ {1}} ^ {\pi_ {1}}\right) \\ \geq - \frac {2 R _ {\text {m a x}} \gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {2}} ^ {\pi_ {2}} + \frac {2 R _ {\text {m a x}} \gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {1}} ^ {\pi_ {1}} - 2 R _ {\text {m a x}} \left(\frac {\epsilon_ {\pi}}{(1 - \gamma) ^ {2}} + \frac {\gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {1}} ^ {M _ {2}}\right) \tag {32} \\ = \frac {2 R _ {m a x} \gamma}{(1 - \gamma) ^ {2}} \left(\epsilon_ {M _ {1}} ^ {\pi_ {1}} - \epsilon_ {M _ {2}} ^ {\pi_ {2}} - \epsilon_ {M _ {1}} ^ {M _ {2}}\right) - \frac {2 R _ {m a x} \epsilon_ {\pi}}{(1 - \gamma) ^ {2}} \\ \end{array}
$$
$$
\begin{array}{l} V ^ {\pi_ {2} | M _ {2}} - V ^ {\pi_ {1} | M _ {1}} \\ \geq \kappa \left( \right.\gamma \left( \right.\mathbb {E} _ {(s, a) \sim d _ {M _ {1}} ^ {\pi_ {1}}} \left[ \right. D _ {T V} \left( \right.p _ {M _ {1}} \left. \right\| p _ {M ^ {*}}\left. \right) - D _ {T V} \left( \right.p _ {M _ {1}} \left. \right\| p _ {M _ {2}}\left. \right) - D _ {T V} \left( \right.p _ {M _ {2}} \left. \right\| p _ {M _ {*}}\left. \right) ] + \Delta\left. \right) - \epsilon_ {\pi}) \tag {33} \\ \end{array}
$$
Proof: Based on Eq.(29), we add a new term  $\kappa \mathbb{E}_{(s,a)\sim d_{M_1}^{\pi_1}}[D_{TV}(p_{M_2}||p_{M^*})]$  to reformulate the optimization objective.
Theorem 5 (|Δ| Upper Bound). Let  $M_i \in \mathcal{M}$  be the evaluated model and  $\pi_i \in \Pi$  be the policy derived from the model. The term  $\Delta$  can be upper bounded as:
$$
\left| \Delta \right| \leq \frac {2 \gamma}{1 - \gamma} \mathbb {E} _ {(s, a) \sim d _ {M _ {1}} ^ {\pi_ {1}}} \left[ D _ {T V} \left(p _ {M _ {1}} \mid \mid p _ {M _ {2}}\right) \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \mid \mid p _ {M ^ {*}}\right) \right] + \frac {2 \epsilon_ {\pi}}{1 - \gamma} \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \mid \mid p _ {M ^ {*}}\right) \tag {34}
$$
Proof: First, we combine these two terms.
$$
\begin{array}{l} | \Delta | = | \mathbb {E} _ {(s, a) \sim d _ {M _ {1}} ^ {\pi_ {1}}} [ D _ {T V} (p _ {M _ {2}} | | p _ {M ^ {*}}) ] - \mathbb {E} _ {(s, a) \sim d _ {M _ {2}} ^ {\pi_ {2}}} [ D _ {T V} (p _ {M _ {2}} | | p _ {M ^ {*}}) ] | \\ = (1 - \gamma) \left| \sum_ {t = 0} ^ {\infty} \gamma^ {t} \sum_ {s, a} \left(p _ {t, M _ {1}} ^ {\pi_ {1}} (s, a) - p _ {t, M _ {2}} ^ {\pi_ {2}} (s, a)\right) D _ {T V} \left(p _ {M _ {2}} \left(s ^ {\prime} \mid s, a\right) \right. \right\rvert   \left. \left| p _ {M ^ {*}} \left(s ^ {\prime} \mid s, a\right) \right.\right) \\ \leq (1 - \gamma) \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \left(s ^ {\prime} \mid s, a\right) \mid p _ {M ^ {*}} \left(s ^ {\prime} \mid s, a\right)\right) \sum_ {t = 0} ^ {\infty} \gamma^ {t} \sum_ {s, a} \left| p _ {t, M _ {1}} ^ {\pi_ {1}} (s, a) - p _ {t, M _ {2}} ^ {\pi_ {2}} (s, a) \right| \tag {35} \\ = 2 (1 - \gamma) \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \| p _ {M ^ {*}}\right) \sum_ {t = 0} ^ {\infty} \gamma^ {t} D _ {T V} \left(p _ {t, M _ {1}} ^ {\pi_ {1}} (s, a) \| p _ {t, M _ {2}} ^ {\pi_ {2}} (s, a)\right) \\ \end{array}
$$
Recalling that we get the result of the sum equation above in Eq.(25), and then we have:
$$
\begin{array}{l} \left| \Delta \right| \leq 2 (1 - \gamma) \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \mid \mid p _ {M ^ {*}}\right) \left(\frac {\epsilon_ {\pi}}{(1 - \gamma) ^ {2}} + \frac {\gamma}{(1 - \gamma) ^ {2}} \epsilon_ {M _ {1}} ^ {M _ {2}}\right) \\ = \frac {2 \gamma}{1 - \gamma} \mathbb {E} _ {s, a \sim d _ {M _ {1}} ^ {\pi_ {1}}} \left[ D _ {T V} \left(p _ {M _ {1}} \| p _ {M _ {2}}\right) \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \| p _ {M ^ {*}}\right) \right] + \frac {2 \epsilon_ {\pi}}{1 - \gamma} \max  _ {s, a} D _ {T V} \left(p _ {M _ {2}} \| p _ {M ^ {*}}\right) \tag {36} \\ \end{array}
$$
D.1 Environment Setup
We evaluate the algorithm over a series of MuJoCo [46] continuous control benchmark tasks. To ensure fairness, we use the standard 1000-step version of all the environments. The details of the environment setup are from OpenAI Gym [3], as shown in Table 1.
Table 1: The general outline of the MuJoCo environment.
Environment-Version	State Dim	Action Dim	Termination
Ant-v2	27	8	obs[0] < 0.2 or obs[0] > 1.0
HalfCheetah-v2	17	6	-
Hopper-v2	11	3	obs[1] ≥ 0.2 or obs[0] ≤ 0.7
Humanoid-v2	45	17	obs[0] < 1.0 or obs[0] > 2.0
InvertedPendulum-v2	4	1	obs[1] > 0.2 or obs[1] < -0.2
Walker2d-v2	17	6	obs[0] ≥ 2.0 or obs[0] ≤ 0.8 or obs[1] ≥ 1.0 or obs[1] ≤ -1.0
D.2 Baseline implementation
MFRL Baselines. We use two state-of-the-art model-free algorithms, i.e. SAC [15] and PPO [40], to do baseline comparison. To demonstrate the final performance and sampling efficiency of our method, we train SAC for 3M steps, which is much more than MBRL algorithms. The hyperparameters are consistent with the author's settings.
MBRL Baselines. We use several state-of-the-art model-based algorithms to do baseline comparison, covering CMLO [21], MBPO [20], SLBO [31] and STEVE [4]. The implementation of CMLO is based on the opensource repo published by the author and all of the hyperparameters are set according to the paper [21]. Our algorithm USB-PO is implemented based on the opensource repo published by Janner who is the author of MBPO.
We present the final performance on six continuous benchmark tasks in Table 2. The results demonstrate that our algorithm achieves competitive performance compared to both MBRL and MFRL baselines over these tasks. Each result in the table shows the average and standard deviation on the maximum average returns among different random seeds and we choose 250K for HalfCheetah-v2, 300K for Walker2d-v2, 300K for Humanoid-v2, 250K for Ant-v2, 15K for Inverted-Pendulum-v2, 120K for Hopper-v2.
Table 2: The final performance on six continuous benchmark tasks.
HalfCheetah	Humanoid	Walker2d
MBRL	STEVE	12406.29±458.08	4318.32±853.60	1109.23±1163.74
SLBO	1915.47±1398.73	459.46±34.27	3107.93±1887.09
MBPO	12765.67±594.54	5546.77±221.72	4582.06±67.44
CMLO	10143.55±193.82	5577.01±219.89	4807.60±99.89
USB-PO	15105.91±177.75	5973.75±110.99	5691.62±162.57
MFRL(@3M steps)	SAC	15012	6207	5879
Ant	InvertedPendulum	Hopper
MBRL	STEVE	779.72±45.67	778.54±265.51	1131.61±623.52
SLBO	707.79±218.80	793.24±334.90	898.68±233.21
MBPO	4926.10±818.38	1000.00±0.00	3436.00±120.72
CMLO	5123.71±783.97	1000.00±0.00	3495.41±71.02
USB-PO	6340.84±119.06	1000.00±0.00	3694.22±46.19
MFRL(@3M steps)	SAC	5934	1000	3610
Our algorithm USB-PO is based on MBPO [20] and is implemented according to the opensource repo published by the MBPO author. Except for the learning rate in phase 2 of our USB-PO algorithm, the hyperparameters are completely identical to the MBPO settings for all environments. In all benchmark tasks, we set this learning rate to 1e-4.
D.4 Computing Infrastructure
In Table 3, we list our computing infrastructure and the computational time for training USB-PO on these six continuous benchmark tasks. Note that the time we report is the cost for 4 random seeds simultaneously on one graphics card. For Humanoid, only two random seeds can be run simultaneously because of the limitation of graphics memory.
Table 3: Computing infrastructure and the computational time for each benchmark task compared to MBPO, where the time unit  $\mathrm{d}$  denotes day and  $\mathrm{h}$  denotes hour.
HalfCheetah	Humanoid	Walker2d	Ant	InvertedPendulum	Hopper
CPU	AMD EPYC 7B12 64-Core Processor
GPU	NVIDIA 2080Ti
MBPO times	2.46d	1.64d	1.75d	2.88d	3.43h	17.81h
USB-PO times	2.29d	1.51d	1.65d	2.91d	3.42h	18.28h
E Comparison with Prior Works
In this section, we compare USB-PO with prior theoretical works to emphasize our contribution, as a complementary to the main paper. First, we give a summary and then show the details as follows. MBPO-Style does not consider model shift and CMLO-Style rely on a fixed threshold to constrain model shift. Our algorithm, USB-PO, adaptively adjusts the model updates in a unified manner (unify model shift and model bias) to get the performance improvement guarantee.
MBPO-Style [20, 41, 24, 53]. They use the return discrepancy bound  $V^{\pi |M} \geq V_M^\pi - C(\epsilon_m, \epsilon_\pi)$  to improve the lower bound on the performance under the real environment, i.e. as long as improving  $V_M^\pi$  by more than  $C(\epsilon_m, \epsilon_\pi)$  can guarantee improvement on  $V^{\pi |M}$ . Obviously, this scheme is guaranteed under a fixed model and it does not consider the change in model dynamic during updates nor the performance variation concerning model shift. Even worse, if the model has some excessive updates, it is impractical to find a feasible solution to meet the improvement guarantee.
CMLO-Style [21]. They use the performance difference bound under the model-based setting  $V^{\pi_2|M_2} - V^{\pi_1|M_1} \geq C$  to directly consider model shift and model bias. However, they finally derive a constrained lower-bound optimization problem and use a fixed threshold to constrain model shift, i.e.  $\sup_{s \in S, a \in A} D_{TV}(P_{M_1}(\cdot | s, a) || P_{M_2}(\cdot | s, a)) \leq \sigma_{M_1, M_2}$  and determine when to update the model accordingly. Notably, we find that this fixed threshold plays a key role in the whole algorithm and needs to be carefully adjusted for each environment. If this threshold is set too low, the model bias of the following iteration will be large, which impairs the subsequent optimization process. If this threshold is set too high, the performance improvement can no longer be guaranteed. Additionally, using a fixed threshold during the whole training process makes the algorithm problematic to adjust adaptively.
USB-PO (Ours). Following CMLO-Style [21], we also use the performance difference bound under the model-based setting to directly consider model shift and model bias. Compared to relying on a fixed threshold to constrain model shift, we use a transformation to unify model shift and model bias into one formulation without the constraint (Theorem 4). Due to the intractable property of  $\Delta$ , we further explore the upper bound of  $|\Delta|$ , finding that  $\Delta$  can be ignored with respect to model shift and model bias alone (Theorem 5). Finally, the optimization objective we get can be used to fine-tune  $M_2$  in a unified manner to adaptively adjust the model updates to get a performance improvement guarantee. Notably, our algorithm can use the same learning rate of Phase 2 and our algorithm is robust to this learning rate. To the best of our knowledge, this is the first method that unifies model shift and model bias and adaptively fine-tunes the model updates during the training process.
Figure 10: We choose a specific random seed to show the details of the first 30 training times on all benchmark tasks, covering the difference of optimization objective value, the model shift term and the model bias term before and after the fine-tuning process.
F.1 Working Mechanism Extension
To illustrate that the ability of USB-PO to reduce both model shift and model bias potentially is not a coincidence that exists only in the Walker2d environment, we add experimental results in other MuJoCo [46] environments. As shown in Figure 10, when the fine-tuning actually operates, the difference of the model shift term and the model bias term among all of the benchmark tasks are generally both positive, further validating our superiority.
F.2 Ablation Study Extension
Here, we show the results of the ablation study on all of the MuJoCo benchmark tasks.
As shown in Figure 11, only optimizing the model shift term results in a drop in sample efficiency while only optimizing the model bias term leads to performance deterioration. Only fine-tuning the model updates in a unified manner can achieve excellent performance.
Figure 11: Optimization Objective Variants on all MuJoCo benchmark tasks.
F.3 Ensemble numbers
To illustrate the justification of the ensemble numbers we set, we conduct the ablation experiment on more ensemble numbers containing 3, 5, and 7. As Figure 12 shows, as the number of ensemble models goes up, the performance will be higher and more stable, but it will cost more time. To maintain the balance between performance and time, we finally set the value of this parameter to 7, which is recommended by the MBPO [20] original repo.
Figure 12: Ablation experimental results of the ensemble model numbers.