SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$\mathbf{Random}(N,M)$ : We decomposition offline dataset  $\mathcal{D} \coloneqq \{\tau\}$  into  $N$  subsets, each of which contains at most  $M$  trajectories that are randomly sampled from the offline dataset.
Quantization  $(N,M)$ : Leveraging the returns of trajectories in offline data, we first quantize offline trajectories into  $N$  bins, and then randomly sample at most  $M$  trajectories (as a sub-task) from each bin. Specifically, in the  $i$ -th bin, the quantized trajectories  $\{\tau_i\}$  satisfy  $\mathrm{R}_{\mathrm{min}} + \Delta * i < \mathrm{Return}(\tau_i) \leq \mathrm{R}_{\mathrm{min}} + \Delta * (i + 1)$ , where  $\Delta = \frac{(\mathrm{R}_{\mathrm{max}} - \mathrm{R}_{\mathrm{min}})}{N}$ ,  $\mathrm{Return}(\tau_i)$  denotes the return of trajectory  $\tau_i$ , and  $\mathrm{R}_{\mathrm{max}}$  and  $\mathrm{R}_{\mathrm{min}}$  denote the maximum and minimum trajectory returns in the offline dataset respectively.
$\mathbf{Rank}(N,M)$ : We first rank the offline trajectories descendingly based on their returns, and then sequentially sample  $M$  trajectories for each subset. (We adopt this decomposition rule in main paper.)
In Figure 7, we provide the comparison of the above three decomposition rules (see the selected number of sub-tasks and the number of trajectories in each sub-task in Table 9). We can find that across a variety of tasks, the decomposition rule has a fundamental impact on the subsequent model-based optimization. Across different tasks and different embedding inference rules, Random and Quantization decomposition rules tend to exhibit large performance fluctuations, which reveals the importance of choosing a suitable task decomposition rule. In our paper, we adopt the Rank decomposition rule, as it demonstrates a more robust performance shown in Figure 7. In Appendix A.4, we adopt the conditional variational auto-encoder (CVAE) to conduct automatic task decomposition (treating each trajectory in offline dataset as an individual task) and we find such implementation  $(\text{DROP} + \text{CVAE})$  can further improve DROP's performance. In future work, we also encourage better decomposition rules to decompose offline tasks so as to enable more effective model-based optimization for offline RL tasks.
Comparison with filtered behavior cloning. We also note that the Rank decomposition rule leverages more high-quality (high-return) trajectories than the other two decomposition rules (Random and Quantization). Thus, a natural question to ask is, is the performance of Rank better than that of Random and Quantization due to the presence of more high-quality trajectories in the decomposed sub-tasks? That is, whether DROP (using the Rank decomposition rule) only conducts behavioral cloning over those high-quality trajectories, thus leading to better performance.
Figure 7: Comparison of three different decomposition rules on D4RL MuJoCo-Gym suite (*-v0) and AntMaze suite (*-v2), where "Rand", "Quan" and "Rank" denote the Random, Quantization, and Rank decomposition rules respectively. We can find across 18 tasks (AntMaze and MuJoCo-Gym suites) and 3 embedding inference methods (DROP-Grad, DROP-Best-Ada, and DROP-Grad-Ada), Rank is more stable and yields better performance compared with the other two decomposition rules.
Figure 8: Learning curves of DROP, where the x-axis denotes the training steps (k), the y-axis denotes the evaluation return (using DROP-Best embedding inference rule). We only show two seeds for legibility.
To answer the above question, we compare DROP (using the Rank decomposition rule) with filtered behavior cloning (F-BC), where the latter (F-BC) performs behavior cloning after filtering for trajectories with high returns. We provide the comparison results in Table 5. We can find that in AntMaze tasks, the overall performance of DROP is higher than that of F-BC. For the MuJoCoGym suite, DROP-based methods outperform F-BC on these offline tasks that contain plenty of sub-optimal trajectories, including the random, medium, and medium-replay domains. This result indicates that DROP can leverage embedding inference (extrapolation) to find a better policy beyond all the behavior policies in sub-tasks, which is more effective than simply performing imitation learning on a subset of the dataset.
A.3 Online Fine-tuning
Online fine-tuning (checkpoint-level). In Figure 8, we show the learning curves of DROP-Best on four DR4L tasks. We can find that DROP exhibits a high variance (in performance) across training
Table 5: Comparison between our DROP (using the Rank decomposition rule) and filtered behavior cloning (F-BC) on D4RL AntMaze and MuJoCo suites (*-v2). We take the baseline results of BC and F-BC from Emmons et al. [14], where F-BC is trained over the top  $10\%$  trajectories, ordered by the returns. Our DROP results are computed over 5 seeds and 10 episodes for each seed.
Tasks	BC	F-BC	DROP-Grad	DROP-Best-Ada	DROP-Grad-Ada
antmaze-umaze	54.6	60	72 ± 17.2	78 ± 11.7	80 ± 12.6
antmaze-umaze-diverse	45.6	46.5	48 ± 22.3	62 ± 16	66 ± 12
antmaze-medium-play	0	42.1	24 ± 10.2	34 ± 12	30 ± 21
antmaze-medium-diverse	0	37.2	20 ± 19	24 ± 12	30 ± 16.7
antmaze-large-play-	0	28	24 ± 8	36 ± 17.4	42 ± 17.2
antmaze-large-diverse	0	34.3	14 ± 8	20 ± 14.1	26 ± 13.6
halfcheetah	random	2.3	2	2.3 ± 0	2.3 ± 0	2.3 ± 0
hopper	random	4.8	4.1	5.1 ± 0.8	5.4 ± 0.7	5.5 ± 0.6
walker2d	random	1.7	1.7	2.8 ± 1.7	3 ± 1.6	3 ± 1.8
halfcheetah	medium	42.6	42.5	42.4 ± 0.7	42.9 ± 0.4	43.1 ± 0.4
hopper	medium	52.9	56.9	57.5 ± 6.4	60.3 ± 6.1	59.5 ± 5.1
walker2d	medium	75.3	75	76.5 ± 2.4	75.8 ± 3	79.1 ± 1.4
halfcheetah	medium-replay	36.6	40.6	39.5 ± 1	40.4 ± 0.8	40.3 ± 1.2
hopper	medium-replay	18.1	75.9	48 ± 17.7	83.4 ± 6.5	87.4 ± 2.1
walker2d	medium-replay	26	62.5	37.4 ± 13.5	60.9 ± 7.4	61.9 ± 2.3
steps<sup>9</sup>, which means the performance of the agent may be dependent on the specific stopping point chosen for evaluation (such instability also exists in prior offline RL methods, Fujimoto and Gu [19]).
To choose a suitable stopping checkpoint over which we perform the DROP inference (DROP-Grad, DROP-Best-Ada and DROP-Grad-Ada), we propose to conduct checkpoint-level online fine-tuning (see Algorithm 3 in Section B for more details): we evaluate each of the latest  $T$  checkpoint models and choose the best one that leads to the highest episode return.
In Figure 9, we show the total normalized returns across all the tasks in each suite (including Maze2d, AntMaze, and MuJoCo-Gym). We can find that in most tasks, fine-tuning (FT) can guarantee performance improvement. However, we also find such fine-tuning causes negative impacts on performance in AntMaze(*-v0) suite. The main reason is that, in this checkpoint-level fine-tuning, we choose the "suitable" checkpoint model using the DROP-Best embedding inference rule, while we adopt the other three embedding inference rules (DROP-Grad, DROP-Best-Ada and DROP-Grad-Ada) at the test time. Such a finding also implies that the success of DROP's test-time adaptation is not entirely dependent on the best embedding across sub-tasks  $^{10}$  (i.e., the best embedding  $\mathbf{z}_0^* (\mathbf{s}_0)$  in DROP-Best), but requires switching between some "suboptimal" embeddings (using DROP-Best-Ada) or extrapolating new embeddings (using DROP-Grad-Ada).
Online fine-tuning (embedding-level). Beyond the above checkpoint-level fine-tuning procedure, we can also conduct embedding-level online fine-tuning: we aim to choose a suitable gradient update step for the gradient-based embedding inference rules (including DROP-Grad and DROP-Grad-Ada). Similar to the checkpoint-level fine-tuning, we first conduct the test-time adaptation procedure (DROP-Grad and DROP-Grad-Ada) over a set of gradient update steps, and then choose the best step that leads to the highest episode return (see Algorithm 4 in Section B for more details).
In view of such instability, we evaluate our methods over multiple checkpoints for each seed, instead of choosing the final checkpoint models during the training loop (see the detailed evaluation protocol in Appendix B).
Conversely, if the performance of DROP depends on the best embedding across sub-tasks (i.e.,  $\mathbf{z}_0^*(\mathbf{s}_0)$  in DROP-Best), then the checkpoint model we choose by fine-tuning with DROP-Best should enable a consistent performance improvement for rules that perform embedding inference with DROP-Best-Ada and DROP-Grad-Ada. However, we find a performance drop in AntMaze(*-v0) suite, which means there is no explicit dependency between the best embedding  $\mathbf{z}_0^*(\mathbf{s}_0)$  and the inferred embedding using the adaptive inference rules.
Figure 9: Total normalized returns across all the tasks in Maze2d, AntMaze, and MuJoCo-Gym suites.
Table 6: Online fine-tuning results (initial performance  $\rightarrow$  performance after online fine-tuning). The baseline results of AWAC, CQL, and IQL are taken from Kostrikov et al. [31], where they run 1M online steps to fine-tune the learned policy. For our DROP method (DROP-Grad and DROP-Grad-Ada), we run  $0.3\mathrm{M}$  ( $= 6_{\mathrm{checkpoint}} \times 50_{K_{\mathrm{max}}} \times 1000_{\mathrm{steps per episode}}$ ) online steps to fine-tune (embedding-level) the policy, i.e., aiming to find the optimal gradient ascent step that is used to infer the contextual embedding  $\mathbf{z}^{*}(\mathbf{s}_{0})$  or  $\mathbf{z}^{*}(\mathbf{s}_{t})$  for  $\pi^{*}(\mathbf{a}_{t}|\mathbf{s}_{t}) \coloneqq \beta(\mathbf{a}_{t}|\mathbf{s}_{t}, \cdot)$  (see Algorithm 4 for the details). Moreover, for medium- $*$  and large- $*$  tasks, we conduct additional parametric-level fine-tuning, with 0.7M online steps to update the policy's parameters.
Task (*-v0)	CQL	IQL	DROP-Grad	DROP-Grad-Ada
umaze	70.1 → 99.4	86.7 → 96	70 → 96 ± 1.2	76 → 98 ± 0
umaze-diverse	31.1 → 99.4	75 → 84	54 → 88 ± 8	66 → 94 ± 4.9
medium-play	0 23 → 0	72 → 95	20 → 56 ± 8.9	30 → 50 ± 6.3 → 94 ± 2.9
medium-diverse	23 → 32.3	68.3 → 92	12 → 44 ± 4.9	22 → 38 ± 4.9 → 96 ± 0.8
large-play	1 → 0	25.5 → 46	16 → 38 ± 8.9	16 → 40 ± 6.3 → 53 ± 1.3
large-diverse	1 → 0	42.6 → 60.7	20 → 40 ± 13.6	22 → 46 ± 10.2 → 58 ± 4.5
→ 1M	→ 1M	→ 0.3M	→ 0.3M	→ 0.7M
In Table 6, we compare our DROP (DROP-Grad and DROP-Grad-Ada) to three offline RL methods (AWAC [53], CQL [35] and IQL [31]), reporting the initial performance and the performance after online fine-tuning. We can find that the embedding-level fine-tuning (0.3M) enables a significant improvement in performance. The fine-tuned DROP-Grad-Ada (0.3M) outperforms the AWAC and CQL counterparts in most tasks, even though we take fewer rollout steps to conduct the online fine-tuning (basesines take 1M online rollout steps, while DROP-based fine-tuning takes 0.3M steps). However, there is still a big gap between the fine-tuned IQL and the embedding-level fine-tuned DROP (0.3M). Considering that there remain 0.7M online steps in the comparison, we further conduct "parametric-level" fine-tuning (updating the parameters of the policy network) for our DROP-Grad-Ada on medium-* and large-* tasks, we can find which achieves competitive fine-tuning performance even compared with IQL.
A.4 DROP + CVAE Implementation
CVAE-based embedding learning. Similar to LAPO [11] and PLAS [74], we adopt the conditional variational auto-encoder (CVAE) to model offline data. Specifically, we learn the contextual policy and behavior embedding:
$$
\beta (\mathbf {a} | \mathbf {s}, \mathbf {z}), \phi (\mathbf {z} | \mathbf {s}) \leftarrow \arg \max  _ {\beta , \phi} \mathbb {E} _ {(\mathbf {s}, \mathbf {a}) \sim \mathcal {D}} \mathbb {E} _ {(\mathbf {z}) \sim \phi (\mathbf {z} | \mathbf {s})} [ \log \beta (\mathbf {a} | \mathbf {s}, \mathbf {z}) ] - \mathrm {K L} (\phi (\mathbf {z} | \mathbf {s}) \| p (\mathbf {z})). \tag {10}
$$
Then, we learn the score model  $f$  with the TD-error and the conservative regularization:
$$
f \leftarrow \arg \min  _ {f} \mathbb {E} _ {\left(\mathbf {s}, \mathbf {a}, \mathbf {s} ^ {\prime}, \mathbf {a} ^ {\prime}\right) \sim \mathcal {D}} \left[ \left(R (\mathbf {s}, \mathbf {a}) + \gamma \bar {f} \left(\mathbf {s} ^ {\prime}, \mathbf {a} ^ {\prime}, \phi (\mathbf {z} | \mathbf {s})\right) - f (\mathbf {s}, \mathbf {a}, \phi (\mathbf {z} | \mathbf {s}))\right) ^ {2} \right], \tag {11}
$$
$$
\text {s . t .} \mathbb {E} _ {\mathbf {s} \sim \mathcal {D}, \mathbf {z} \sim \mu (\mathbf {z}), \mathbf {a} \sim \beta (\mathbf {a} | \mathbf {s}, \mathbf {z})} [ f (\mathbf {s}, \mathbf {a}, \mathbf {z}) ] - \mathbb {E} _ {\mathbf {s} \sim \mathcal {D}, \mathbf {z} \sim \phi (\mathbf {z} | \mathbf {s}), \mathbf {a} \sim \beta (\mathbf {a} | \mathbf {s}, \mathbf {z})} [ f (\mathbf {s}, \mathbf {a}, \mathbf {z}) ] \leq \eta ,
$$
where  $\bar{f}$  denotes a target network and  $\mu (\mathbf{z})$  denotes the uniform distribution over the  $\mathcal{Z}$ -space.
In testing, we also dynamically adapt the outer-level optimization, setting policy inference with  $\pi^{*}(\mathbf{a}|\mathbf{s}) = \beta (\mathbf{a}|\mathbf{s},\mathbf{z}^{*}(\mathbf{s}))$ , where  $z^{*}(\mathbf{s}) = \arg \max_{z}f\big(\mathbf{s},\beta (\mathbf{a}|\mathbf{s},\mathbf{z}),\mathbf{z}\big)$ .
A.5 Ablation Study
Note that our embedding inference depends on the learned score model  $f$ . Without proper regularization, such inference will lead to out-of-distribution embeddings that are erroneously high-scored (Q2). Here we conduct an ablation study to examine the impact of the conservative regularization used for learning the score model.
In Figure 10, we compare DROP-Grad and DROP-Grad-Ada to their naive implementation (w/o Reg) that ablates the regularization on halfcheetah-medium-expert. We can find that removing the conservative regularization leads to unstable performance when changing the update
Figure 10: Ablation on the conservative regularization. The y-axis represents the normalize return, and the x-axis represents the number of gradient-ascent steps used for embedding inference at deployment. We plot each random seed as a transparent line; the solid line corresponds to the average across 5 seeds.
steps of gradient-based optimization. However, we empirically find that in some tasks such a naive implementation (w/o Reg) does not necessarily bring unstable inference (Appendix C). Although improper gradient update step leads to faraway embeddings, to some extent, embedding-conditioned behavior policy can correct such deviation.
For the practical implementation of DROP, we parameterize the task embedding function  $\phi (\mathbf{z}|n)$ , the contextual behavior policy  $\beta (\mathbf{a}|\mathbf{s},\mathbf{z})$ , and the score model  $f(\mathbf{s},\mathbf{a},\mathbf{z})$  with neural networks. For Equation 8 in the main paper, we construct a Lagrangian and solve the optimization through primal-dual gradient descent. For the choice of  $\mu (\mathbf{z})$ , we simply set  $\mu (\mathbf{z})$  to be the uniform distribution over the  $\mathcal{Z}$ -space and empirically find that such uniform sampling can effectively avoid the out-of-distribution extrapolation at inference.
Lagrangian relaxation. To optimize the constrained objective in Equation 8 in the main paper, we construct a Lagrangian and solve the optimization through primal-dual gradient descent,
$$
\begin{array}{l} \min  _ {f} \max  _ {\lambda > 0} \mathbb {E} _ {\mathcal {D} _ {n} \sim \mathcal {D} _ {[ N ]}} \mathbb {E} _ {(\mathbf {s}, \mathbf {a}, \mathbf {s} ^ {\prime}, \mathbf {a} ^ {\prime}) \sim \mathcal {D} _ {n}} \left[ \left(R (\mathbf {s}, \mathbf {a}) + \gamma \bar {f} (\mathbf {s} ^ {\prime}, \mathbf {a} ^ {\prime}, \phi (\mathbf {z} | n)) - f (\mathbf {s}, \mathbf {a}, \phi (\mathbf {z} | n))\right) ^ {2} \right] + \\ \lambda \left(\mathbb {E} _ {n, \mu (\mathbf {z})} \mathbb {E} _ {\mathbf {s} \sim \mathcal {D} _ {n}, \mathbf {a} \sim \beta (\mathbf {a} | \mathbf {s}, \mathbf {z})} \left[ f (\mathbf {s}, \mathbf {a}, \mathbf {z}) \right] - \mathbb {E} _ {n, \phi (\mathbf {z} | n)} \mathbb {E} _ {\mathbf {s} \sim \mathcal {D} _ {n}, \mathbf {a} \sim \beta (\mathbf {a} | \mathbf {s}, \mathbf {z})} \left[ f (\mathbf {s}, \mathbf {a}, \mathbf {z}) \right] - \eta\right). \\ \end{array}
$$
This unconstrained objective implies that if the expected difference in scores of out-of-distribution embeddings and in-distribution embeddings is less than a threshold  $\eta$ ,  $\lambda$  is going to be adjusted to 0, on the contrary,  $\lambda$  is likely to take a larger value, used to punish the over-estimated value function. This objective encourages that out-of-distribution embeddings score lower than in-distribution embeddings, thus performing embedding inference will not lead to these out-of-distribution embeddings that are falsely and over-optimistically scored by the learned score model.
In our experiments, we tried five different values for the Lagrange threshold  $\eta$  (1.0, 2.0, 3.0, 4.0, and 5.0). We did not observe a significant difference in performance across these values. Therefore, we simply set  $\eta = 2.0$ .
Hyper-parameters. In Table 7, we provide the hyper-parameters of the task embedding  $\phi (\mathbf{z}|\mathbf{s})$ , the contextual behavior policy  $\beta (\mathbf{a}|\mathbf{s},\mathbf{z})$ , and the score function  $f(\mathbf{s},\mathbf{a},\mathbf{z})$ . For the gradient ascent update steps (used for embedding inference), we set  $K = 100$  for all the embedding inference rules in experiments.
In Table 9, we provide the number of sub-tasks, the number of trajectories in each sub-task, and the dimension of the embedding for each sub-task (behavior policy). The selection of hyperparameter N is based on two evaluation metrics: (1) the fitting loss of the decomposed behavioral policies to the offline data, and (2) the testing performance of DROP. Specifically,
Table 7: Hyper-parameters of the task embedding function  $\phi (\mathbf{z}|\mathbf{s})$ , the contextual behavior policy  $\beta (\mathbf{a}|\mathbf{s},\mathbf{z})$ , and the score function  $f(\mathbf{s},\mathbf{a},\mathbf{z})$ . The task embedding function  $\phi (\mathbf{z}|\mathbf{s})$ :  $\mathbf{z}\gets \mathrm{Enc}_{-}0(n)$ . The contextual behavior policy  $\beta (\mathbf{a}|\mathbf{s},\mathbf{z})$ :  $\mathbf{a}\gets \mathrm{Enc}_{-}2(\mathbf{z},\mathrm{Enc}_{-}1(\mathbf{s}))$ . The score function  $f(\mathbf{s},\mathbf{a},\mathbf{z})$ :  $f\gets \mathrm{Enc}_{-}4(\mathbf{z},\mathrm{Enc}_{-}3(\mathbf{s},\mathbf{a}))$ .
Enc_0	Enc_1	Enc_2	Enc_3	Enc_4
Optimizer	Adam	Adam	Adam	Adam	Adam
Hidden layer	2	2	3	2	3
Hidden dim	512	512	512	512	512
Activation function	ReLU	ReLU	ReLU	ReLU	ReLU
Learning rate	1.00E-03	1.00E-03	1.00E-03	1.00E-03	1.00E-03
Mini-batch size	1024	1024	1024	1024	1024
- (Step1) Over a hyperparameter (the number of sub-tasks) set, we conduct the hyperparameter search using the fitting loss of behavior policies, then we choose/filter the four best hyperparameters;
- (Step2) We follow the normal practice of hyperparameter selection and tune the four hypermeters selected in Step1 by interacting with the simulator to estimate the performance of DROP under each hyperparameter setting.
We provide the hyperparameter sets in Table 8. In Step2, we tune the (filtered) hyperparameters using 1 seed, then evaluate the best hyperparameter by training on an additional 4 seeds and finally report the results on the 5 total seeds (see next "evaluation protocol"). In Antmaze domain, a single fixed N works well for many tasks;
Table 8: Hyperparameter (the number of sub-tasks) set.
tasks	the number of sub-tasks
Antmaze	500 (v0), 150 (v2)
Gym-mujoco	10, 20, 50, 100, 200, 500, 800, 1000
Adroit	10, 20, 50, 100, 200, 500, 800, 1000
while in Gym-mujoco and Adroit domains, we did not find a fixed N that provides good results for all tasks in the corresponding domain in D4RL, thus we use the above hyperparameter selection rules (Step1 and Step2) to choose the number.
**Baseline details.** For the comparison of our method to prior iterative offline RL methods, we consider the v0 versions of the datasets in D4RL $^{11}$ . We take the baseline results of BEAR, BCQ, CQL, and BRAC-p from the D4RL paper [18], and take the results of TD3+BC from their origin paper [19]. For the comparison of our method to prior non-iterative offline RL method, we use the v2 versions of the dataset in D4RL. All the baseline results of behavior cloning (BC), Decision Transform (DT), RvS-R, and Onestep are taken from Emmons et al. [14]. In our implementation of COMs, we take the parameters (neural network weights) of behavior policies as the design input for the score model; and during testing, we conduct parameters inference (outer-level optimization) with 200 steps gradient ascent over the learned score function, then the rollout policy is initialized with the inferred parameters. For the specific architecture, we instantiate the policy network with  $\dim(S)$  input units, two layers with 64 hidden units, and a final output layer with  $\dim(\mathcal{A})$ .
Evaluation protocol. We evaluate our results over 5 seeds. For each seed, instead of taking the final checkpoint model produced by a training loop, we take the last  $T$  ( $T = 6$  in our experiments) checkpoint models, and evaluate them over 10 episodes for each checkpoint. That is to say, we report the average of the evaluation scores over  $5_{\mathrm{seed}} \times 6_{\mathrm{checkpoint}} \times 10_{\mathrm{episode}}$  rollouts.
Online fine-tuning (checkpoint-level): Instead of re-training the learned (final) policy with online rollouts, we fine-tune our policy with enumerated trial-and-error over the last  $T$  checkpoint models (Algorithm 3). Specifically, for each seed, we run the last  $T$  checkpoint models in the environment over one episode for each checkpoint. The checkpoint model which achieves the maximum episode return is returned. In essence, this fine-tuning procedure imitates the online RL evaluation protocol: if the current policy is unsatisfactory, we can use checkpoints of previous iterations of the policy.
<sup>11</sup>We noticed that Maze2D-v0 in the D4RL dataset (https://rail.eecs.berkeley.edu/datasets/) is not available, so we used v1 version instead in our experiment. For simplicity, we still use v0 in the paper exposition.
Algorithm 3 DROP: Online fine-tuning (checkpoint-level)
Require: Env, last  $T$  checkpoint models:  $\beta_{t}(\mathbf{a}|\mathbf{s},\mathbf{z})$  and  $f_{t}(\mathbf{s},\mathbf{a},\mathbf{z})$ $(t = 1,\dots ,T)$    
1:  $R_{\mathrm{MAX}} = -\infty$    
2:  $\beta_{\mathrm{best}}\gets$  None.   
3:  $f_{\mathrm{best}}\gets$  None.   
4: while  $t = 1,\dots ,T$  do   
5:  $\mathbf{s}_0 = \mathrm{Env}\cdot \mathrm{Reset}().$    
6:  $\mathbf{z}_0^* (\mathbf{s}_0)\gets$  Conduct embedding inference with DROP-Best.   
7: Return  $\leftarrow$  Evaluate  $\beta_{t}$  and  $f_{t}$  on Env, setting  $\pi^{*}(\mathbf{a}|\mathbf{s}) = \beta (\mathbf{a}|\mathbf{s},\mathbf{z}_{0}^{*}(\mathbf{s}_{0}))$    
8: if  $R_{\mathrm{MAX}} <   \mathrm{Re}$  turn then   
9: Update the best checkpoint models:  $\beta_{\mathrm{best}}\gets \beta_t,f_{\mathrm{best}}\gets f_t$    
10: Update the optimal return:  $R_{\mathrm{MAX}}\gets$  Return.   
11: end if   
12: end while   
Return:  $\beta_{\mathrm{best}}$  and  $f_{\mathrm{best}}$