SlowGuess/NeurIPS_2023_OCR huggingface dataset:
SDVT, SD, LDM, VariBAD Task-inference-based methods (SDVT, SD, LDM [35], and VariBAD [74]) are adapted to the Meta-World benchmark based on the VariBAD's implementation. We begin with VariBAD's hyperparameter configuration for MuJoCo [58] Ant-goal and modify certain hyperparameters to accommodate the Meta-World benchmark (e.g., batch size, network capacity, etc.). In line with VariBAD and LDM, we train SDVT's policy using PPO [49]. The GMVAE is based on an implementation that employs the Gumbel-Softmax reparameterization trick [23] when sampling the categorical subtask variable. The virtual training processes of SDVT and LDM require agent-environment interaction since they use states from the real environment. Therefore, the virtual training steps are also added when counting the total number of training steps. For a comprehensive list of SDVT hyperparameters, shared among SD, LDM, and VariBAD to ensure a fair comparison, please refer to Appendix C.3 and our source code available at https://github.com/suyoung-lee/SDVT.
Figure 5: Schematic overview of algorithms. SDVT without a Gaussian mixture reduces to LDM, LDM without virtual training reduces to VariBAD, and VariBAD without a VAE decoder reduces to RL $^2$ . MAML and PEARL use fully connected networks.
C.2 Computational Complexity
Table 2: Computational complexity. The total wall-clock time required to generate the results for the ML-10, averaged across all eight random seeds.
SDVT	SDVT-LW	SD	SD-LW	RL2	MAML	PEARL	VariBAD	LDM
Wall-clock time (hours)	142	140	138	135	192	17	258	126	131
Our experiments were conducted using an Nvidia Titan Xp. Due to the considerably larger variety of tasks compared to standard meta-RL benchmarks, non-parametric benchmarks require significantly more computational resources and time. In Table 2, we detail the total wall-clock time expended by our methods and baselines during the training and evaluation of the ML-10 benchmark. This time includes the environmental interaction time for 250M training steps, roughly 100M steps dedicated to
<sup>3</sup>https://github.com/rlworkgroup/garage/pull/2287
<sup>4</sup>https://github.com/lmzintgraf/varibad
<sup>5</sup>https://github.com/jariasf/GMVAE
evaluations, and the time taken to train the neural networks. Despite incorporating the GMVAE and virtual training, our method's computational demand does not substantially surpass that of VariBAD.
Table 3: Hyperparameters of SDVT and SD. Hyperparameters of SDVT used for Meta-World ML-10 and ML-45 along with the notations in the manuscript and the argument names in the source code. SD shares the same hyperparameters except without those of virtual training. Different hyperparameters of our lightweight variant are denoted as (-LW).
Category	Description	Notation	Value (ML-10, 45)	Argument Name
General	Rollout episode horizon	H	500	max_episode_steps
Number of rollout episodes	n_roll	10	max_rollouts_per_task
Discount factor	γ	0.99	policy_gamma
Number of parallel processes	10	num_processes
Total Environment Steps	2.5e8, 4.0e8	num Frames
GMVAE	Optimizer	Adam	optimiser_vae
Learning rate	1e-3	lr_vae
Subsample ELBO	100	vae_subsample_elbos
Subsample decodes	100	vae_subsample_decodes
Buffer size (meta-episodes)	|BVAE|	1000	size_vae_buffer
ELBO categorical coefficient	αc	1.0	cat_loss(coeff
ELBO categorical coefficient (-LW)	αc	0.5	cat_loss(coeff
ELBO Gaussian coefficient	αg	1.0	gauss_loss(coff
ELBO reward coefficient	αR	10	rew_loss(coff
ELBO transition coefficient	αT	1000	state_loss(coff
ELBO occupancy coefficient	αo	1.0	occ_loss(coff
ELBO occupancy coefficient (-LW)	αo	0	occ_loss(coff
Gumbel softmax temperature	1	gumbel_temperature
Number of updates per epoch	10, 20	num_vae Updates
Subtask dimension	K = Ntrain	10, 45	vae_mixture_num
Subtask dimension (-LW)	K	5	vae_mixture_num
Gaussian dimension	dim(z)	5, 10	latent_dim
Dropout rate	pdrop	0.7	dropout_rate
Reward reconstruction objective	MSE	rew_pred_type
State reconstruction objective	MSE	state_pred_type
Policy	Algorithm	PPO	policy
Optimiser	Adam	policy_optimiser
Learning rate	7e-4	lr_policy
Optimizer epsilon	10e-8	policy_EPS
PPO update epochs	5	ppo_num_epochs
Steps per policy update	|Spol|/num_processes	5000	policy_num_steps
Number of minibatches	10	ppo_num_minibatch
PPO clipping parameter	0.1	ppo Clip param
GAE λ	0.9	policy_tau
Initial standard deviation	1.0	policy_init_std
Minimum standard deviation	0.5	policy_min_std
Maximum standard deviation	1.5	policy_max_std
Entropy coefficient	1e-3	policy_entropycoef
Virtual training	Initial virtual ratio	0.0	virtual_ratio
Virtual ratio increment per step	Δpv	5e-8, 2.5e-8	virtual_ratio_increment
ELBO dispersion coefficient	αd	10	ext_loss(coff
Table 4: Hyperparameters of VariBAD and LDM. VariBAD and LDM employ identical hyperparameters for the general and policy categories of SDVT and SD, as shown in Table 3. The sole distinction lies in the structural variation resulting from the utilization of GMVAE and VAE. We choose the Gaussian dimension that yielded the most favorable outcomes among 5, 10, 15, and 20.
Category	Description	Notation	Value (ML-10, 45)	Argument Name
VAE	Optimizer	Adam	optimiser_vae
Learning rate	1e-3	lr_vae
Subsample ELBO	100	vae_subsample_elbos
Subsample decodes	100	vae_subsample_decodes
Buffer size (meta-episodes)	|BVAE|	1000	size_vae_buffer
ELBO KL coefficient	0.1	kl_weight
ELBO reward coefficient	αR	10	rew_loss_coeff
ELBO transition coefficient	αT	1000	state_loss_coeff
Number of updates per epoch	10,20	num_vae Updates
Gaussian dimension	dim(z)	5,10	latent_dim
Dropout rate (VariBAD)	pdrop	0.0	dropout_rate
Dropout rate (LDM)	pdrop	0.7	dropout_rate
Reward reconstruction objective	MSE	rew_pred_type
State reconstruction objective	MSE	state_pred_type
C.4 Network Architecture
The network architectures comprising our method are described in Table 5. Prior to being input into the encoder or decoder, all state, action, and reward inputs pass through embedding networks. The values of  $K = \dim(y)$  and  $\dim(z)$  differ for the ML-10 and ML-45, as indicated in Table 3.
Table 5: Network architecture. Details of the network architecture composing the GMVAE. The numbers in the Layers column represent the dimension of the hidden layers and the output.
Network	Notation	Architecture	Layers	Activations (last layer)
State embedding	MLP	[32]	(Tanh)
Reward embedding	MLP	[16]	(Tanh)
Action embedding	MLP	[16]	(Tanh)
Recurrent encoder	qφh	GRU	[256]	(None)
Categorical parameter	ωφy	MLP	[512, 512, K]	ReLU (Softmax)
Gaussian parameters	μφz, σ2z	MLP	[512, 512, dim(z)]	ReLU (None, SoftPlus)
Gaussian regularization parameters	μθz, σ2z	MLP	[dim(z)]	(None, SoftPlus)
Latent context dispersion	pθh	MLP	[256, 256, 256]	ReLU (None)
Reward decoder	pθR	MLP	[64, 64, 32, 1]	ReLU (None)
Transition decoder	pθT	MLP	[64, 64, 32, 40]	ReLU (None)
Policy	πψ	MLP	[256, 256, 4]	Tanh (Tanh)
C.5 Aggregation method for the Success Rate
The main results in the Meta-World paper (Table 1 and Figure 6 of Yu et al. [71]) present the average of the maximum success rate for each task, diverging from the raw scores in their Figure 17 and 18. For instance, if an agent achieves a  $90\%$  success rate on "Door-close" in one evaluation during training and scores  $10\%$  in all other evaluations at different times, the reported success rate for "Door-close" is  $90\%$ . The mean of these maximum scores across tasks is reported as the aggregated success rate.
This unconventional aggregation method does not accurately represent the meta-RL objective, which aims to train a single agent capable of solving multiple tasks. Meta-World calculates success rates for various tasks at distinct time points during training, even though the agent might specialize in different tasks at various stages. As a result, the agent may excel at specific tasks by chance. Consequently, evaluating the agent more frequently could yield higher maximum scores. We report the conventional final performance, which better reflects the meta-RL objective.
D Detailed Experimental Results
D.1 Performance on Individual Tasks
Table 6: Meta-World V2 ML-10 success rate. We report the final success rates  $(\%)$  of baselines and our methods for training tasks and test tasks of the Meta-World ML-10 benchmark. All results are reported as the mean success rate  $\pm 95\%$  confidence interval of 8 seeds at 250M steps.
Index. Task	SDVT	SDVT-LW	SD	SD-LW	\( RL^2 \)	MAML	PEARL	VariBAD	LDM
1. Reach	50.0±11.0	22.0±3.5	53.8±11.5	56.2±9.1	27.5±19.8	50.0±34.6	70.0±11.9	28.5±5.0	28.5±4.1
2. Push	57.5±18.6	25.2±16.5	61.2±7.3	68.8±10.6	60.0±13.0	0.0±0.0	0.0±0.0	25.8±12.9	25.5±12.8
3. Pick-place	47.5±9.6	39.0±5.4	52.5±7.6	55.0±13.9	42.5±21.0	0.0±0.0	0.0±0.0	24.8±13.3	27.2±13.3
4. Door-open	100.0±0.0	60.5±31.9	100.0±0.0	63.8±32.5	93.8±9.1	79.5±22.1	10.2±11.9	86.8±22.7	62.5±33.5
5. Drawer-close	100.0±0.0	100.0±0.0	100.0±0.0	100.0±0.0	98.8±2.3	99.5±1.1	86.8±5.6	100.0±0.0	100.0±0.0
6. Button-press	98.8±2.3	98.8±1.8	96.2±6.9	100.0±0.0	87.5±8.3	93.0±8.3	46.0±11.2	98.5±0.9	99.5±0.9
7. Peg-insert-side	31.2±13.2	24.5±10.2	50.0±19.9	36.2±11.5	45.0±21.6	0.0±0.0	0.0±0.0	19.8±14.0	21.8±15.3
8. Window-open	100.0±0.0	100.0±0.0	98.8±2.3	100.0±0.0	100.0±0.0	100.0±0.0	19.2±11.5	100.0±0.0	100.0±0.0
9. Sweep	96.2±4.8	72.2±29.0	83.8±22.2	83.8±15.1	75.0±17.3	0.0±0.0	0.0±0.0	60.2±32.4	61.2±32.9
10. Basketball	91.2±6.4	78.8±9.8	73.8±14.3	91.2±8.8	43.8±16.6	0.0±0.0	0.0±0.0	37.2±26.0	40.5±28.9
Train mean	77.2±3.0	62.1±4.1	77.0±5.9	75.5±5.5	67.4±4.4	42.2±4.5	23.2±1.9	58.2±8.9	56.7±12.3
11. Drawer-open	65.0±19.9	30.5±12.9	48.8±23.4	45.0±24.0	2.2±1.9	15.8±19.3	1.5±1.1	12.8±12.8	21.8±12.0
12. Door-close	7.5±9.0	81.2±19.0	33.8±25.4	18.8±24.1	8.2±7.3	3.2±6.0	1.2±1.5	27.0±21.8	30.2±28.2
13. Shelf-place	0.0±0.0	1.0±1.2	0.0±0.0	0.0±0.0	0.2±0.2	0.0±0.0	0.0±0.0	0.0±0.0	0.0±0.0
14. Sweep-into	90.0±8.5	51.2±18.9	71.2±14.9	55.0±21.6	64.5±8.3	0.0±0.0	0.8±1.0	30.5±22.0	46.5±21.8
15. Lever-pull	1.2±2.3	3.2±2.7	0.0±0.0	12.5±10.2	0.5±0.8	0.5±0.9	0.5±0.9	0.2±0.5	0.5±0.9
Test mean	32.8±3.9	33.4±5.0	30.8±7.7	26.2±8.7	15.1±2.7	3.9±3.7	0.8±0.5	14.1±6.1	19.8±6.0
Table 7: Meta-World V2 ML-10 return. We report the final return of baselines and our methods at the last rollout episode on the Meta-World ML-10 analogous to the success rate in Table 6. Results are reported as the mean ± 95% confidence intervals of 8 seeds at 250M steps.
Index. Task	SDVT	SDVT-LW	SD	SD-LW	\( RL^2 \)	MAML	PEARL	VariBAD	LDM
1. Reach	3565±132	3681±115	3708±228	3806±220	2291±276	3899±560	3841±554	3829±153	3740±59
2. Push	3923±339	3493±469	3636±628	3579±870	815±133	17±1	15±2	2609±990	2470±1220
3. Pick-place	2257±239	2241±300	2131±441	2156±593	501±137	6±0	5±1	1395±759	1561±861
4. Door-open	4441±37	3171±825	4447±106	3530±831	1249±115	3561±753	1064±254	4051±780	3249±1134
5. Drawer-close	4813±71	4850±10	4850±13	4854±9	2193±336	4804±55	4081±314	4841±18	4837±18
6. Button-press	3392±59	3437±133	3257±265	3423±209	1189±151	2574±274	956±300	3246±161	3252±140
7. Peg-insert-side	1903±430	1919±348	2634±588	2225±592	682±157	7±0	7±1	1148±761	1237±825
8. Window-open	4486±39	4371±77	4349±165	4398±79	1230±112	3271±732	784±390	4483±32	4458±42
9. Sweep	4235±202	3617±682	3823±702	3557±634	906±223	79±38	51±14	2807±1260	2789±1333
10. Basketball	3549±292	3759±307	3466±511	3726±331	532±116	7±1	6±2	2141±1119	2035±1256
Train mean	3656±62	3454±137	3630±241	3525±297	1159±83	1822±136	1081±77	3055±466	2963±626
11. Driver-open	2558±165	2280±271	2305±298	2310±280	1649±100	1737±270	1130±87	2336±340	2389±386
12. Door-close	566±335	3030±857	1068±740	638±670	568±137	263±278	321±221	1121±689	1343±1214
13. Shelf-place	501±77	550±138	494±102	434±154	578±70	0±0	0±0	211±144	328±204
14. Sweep-into	2221±764	1455±615	1359±632	1532±721	490±46	41±3	48±33	658±358	1509±889
15. Lever-pull	279±38	318±45	335±45	300±48	291±66	156±49	200±58	271±26	260±48
Test mean	1225±160	1527±214	1112±190	1043±234	715±33	439±78	340±54	919±143	1166±264
In Figures 6 through 9, the mean training and test success rates are represented by solid and dashed lines respectively, with the shaded region indicating the  $95\%$  confidence interval. All results are compiled from 8 random seeds. The data corresponding to the final training steps in these plots (250M steps for ML-10 and 400M steps for ML-45) are used to present the main results in Table 1.
Figure 6: Learning curve on ML-10 - success rate.
Figure 7: Learning curve on ML-10 - return.
Figure 8: Learning curve on ML-45 - success rate.
Figure 9: Learning curve on ML-45 - return.
D.3 Subtask Compositions Learned over Training
(b) SDVT  $(K = 10, \alpha_{c} = 1.0, \alpha_{o} = 1.0)$ .
Figure 10: Subtask compositions learned over training. We visualize the development of subtask compositions for (a) SDVT-LW and (b) SDVT during training on ML-10. Similar to Figure 4, each column represents the terminal subtask composition  $(y_{H^{+}})$  averaged across 50 parametric variants. (a) As training progresses, subtask compositions of (4) "Door-open" and (12) "Door-close" merge due to shared transition dynamics, which may cause the "Door-close" policy to keep the door open during testing without virtual training, underscoring the crucial role that virtual training can play in these scenarios. (b) With occupancy regularization, as the training progresses, the decomposition process prioritizes occupying lower indices over higher ones in the subtask compositions.
D.4 Subtask Compositions of all Seeds
(b) SDVT  $(K = 10,\alpha_{c} = 1.0,\alpha_{o} = 1.0)$
Figure 11: Subtask compositions of all seeds. We visualize the subtask compositions of (a) SDVT-LW and (b) on ML-10 after 250M training steps for all seeds. Decomposition processes differ among random seeds due to varying initialization and sample tasks during meta-training. One subtask can be interpreted as a combination of multiple subtasks and vice versa, depending on the seed and the dimension  $K$ . We rarely have a collapse to a single Gaussian as in (a) Seed 5, which lowers the average performance. (b) With occupancy regularization, we find that the 10th element of the composition is rarely occupied, whereas the first element is occupied by many tasks.
E Ablation Results
In order to confirm the source of the gain obtained from the proposed method, we conduct an extensive ablation study based on SDVT and SDVT-LW on ML-10.
E.1 Occupancy Regularization
(a) SDVT  $(\alpha_{o} = 5.0)$ .
(b) SDVT  $(\alpha_{o} = 10.0)$ .
Figure 12: Occupancy ablation. We report the learned subtask compositions of SDVT ( $K = 10, \alpha_{c} = 1.0$ ) for different occupancy coefficients analogous to Figure 4b.
Figures 4b and 12 indicate that our occupancy regularization effectively constrains the use of higher indices as intended. We fix the default choice of  $\alpha_{o} = 1.0$  that we found to work well on ML-10. Our default SDVT with  $\alpha_{o} = 1.0$  scores a test success rate of  $32.8\%$ , which scored the best among  $\alpha_{o} \in \{0.0, 1.0, 5.0, 10.0\}$  that scored  $22.8\%$ ,  $25.5\%$ , and  $21.0\%$ , respectively. Also, our default SD with  $\alpha_{o} = 1.0$  scores a test success rate of  $30.8\%$ , which scored the best among  $\alpha_{o} \in \{0.0, 1.0, 5.0, 10.0\}$  that scored  $22.8\%$ ,  $28.4\%$ , and  $11.2\%$ , respectively.
E.2 Dropout and Dispersion
Table 10: Ablation results. We report the mean success rate (\%) and the difference caused by the changes made from the default SDVT-LW ( $K = 5, \alpha_{c} = 0.5, \alpha_{o} = 0.0$ ) on ML-10 in parenthesis.
SDVT-LW	without	K	αc
Dropout	Dispersion	3	7	10	0.1	1.0	2.0
ML-10 Train	65.6(+3.5)	73.0(+10.9)	60.1(-2.0)	69.5(+7.4)	70.8(+8.7)	59.8(-2.3)	70.3(+8.2)	66.3(+4.2)
ML-10 Test	16.5(-16.9)	21.5(-11.9)	25.0(-8.4)	22.0(-11.4)	21.1(-12.3)	20.5(-12.9)	16.1(-17.3)	17.9(-15.5)
Table 10 demonstrates the significant roles played by both dropout and dispersion in generating extrapolated dynamics during virtual training, thereby enhancing the test success rate. As outlined in Appendix E.1 of the LDM paper [35], virtual training without dropout does not exhibit a significant improvement due to decoder overfitting on the state and action inputs of the training tasks. In fact, virtual training without dropout can even lead to a decrease in test performance.
However, it is important to note that these factors are not the primary contributors to the empirical gains of our method. Dropout is specifically essential for virtual training, but not necessarily for other methods that do not employ virtual training, such as VariBAD and RL2 (as discussed in Appendix E.2 of the LDM paper). We verified this in our ML-10 experiment using VariBAD with dropout. The inclusion of dropout marginally improves the training success rate of vanilla VariBAD from  $58.2\%$  to  $63.0\%$  and the test success rate from  $14.1\%$  to  $16.5\%$ . However, the performance enhancement achieved through dropout is not as substantial for VariBAD as it is for SDVT-LW's virtual training.
E.3 Decomposition Distribution
Table 10 emphasizes the significance of hyperparameters such as subtask dimension  $K$  and categorical coefficient  ${\alpha }_{c}$  in determining subtask compositions. It is crucial to carefully set these hyperparameters based on the number of training tasks and their correlations. In our lightweight (LW) approach, we have selected  $K = 5$  and  ${\alpha }_{c} = {0.5}$  ,which effectively distributes subtasks across the ML-10 tasks. When  ${\alpha }_{c}$  is set to a small value, task classification collapses into a few subtasks, as depicted in Figure 4c. Conversely, with a large value of  ${\alpha }_{c}$  ,the entropy of the composition is maximized,and all tasks exhibit a uniform probability distribution, as illustrated in Figure 4d. However, it is noteworthy that SDVT-LW outperforms VariBAD and LDM in test success rate, regardless of the specific values chosen for  $K$  ,such as  $K = 3,7,{10}$  . To alleviate the burden of hyperparameter tuning in our LW methods, we introduce the occupancy regularization.
E.4 Number of parameters
Table 11: Number of parameters and success rates. We report the number of parameters used by our methods and baselines. We demonstrate that our gain is not from the increased capacity.
Methods	Number of Parameters	ML-10 Success Rate (%)
Encoder	Decoder	Policy	Sum	Train	Test
SDVT-LW	1,047,455	174,821	235,401	1,457,677	62.1	33.4
SD-LW	1,047,455	25,637	235,401	1,308,493	75.5	26.2
LDM	502,580	25,577	202,249	730,406	56.7	19.8
LDM (matched)	2,144,692	175,593	267,401	2,587,686	64.2	22.0
VariBAD	251,290	25,577	202,249	479,116	58.2	14.1
VariBAD (hidden × 2)	894,362	25,577	663,305	1,583,244	62.4	12.0
VariBAD (matched)	1,072,346	175,593	267,401	1,515,340	67.6	17.2
Table 11 indicates that SDVT-LW employs 3 and 2 times more parameters compared to VariBAD and LDM, respectively. The encoder of our method has more parameters than VariBAD due to the added categorical layer. The decoder of SDVT-LW includes the parameters of the dispersion layers. However, we find that the model's improvement is not mainly attributed to the increased capacity. Generally, a larger capacity does not necessarily guarantee better generalization.
VariBAD (hidden  $\times 2$ ) While selecting encoder (GRU) and policy hidden sizes, we experimented with 128, 256, and 512, discovering that 256 works best for all task-inference methods (LDM, SD only, and SDVT). Notably, VariBAD (hidden  $\times 2$ ) with hidden dimensions of 512 possesses more parameters than SDVT but exhibits an even lower test success rate than the vanilla VariBAD.
VariBAD (matched) We add MLP layers with hidden sizes of [1600, 256] into the encoder and [128] into the policy. The decoder's hidden size is increased from [64, 64, 32] to [128, 256, 160] to match the capacities of VariBAD and SDVT-LW. The components of VariBAD possess slightly more parameters than SDVT-LW. LDM employs two VariBAD encoders. While the matched baselines improve, they still lag behind ours.
E.5 Conditioning policy on belief
Table 12: Masking contexts. We present the success rates (\%) on ML-10, achieved by ablating the contexts fed into the policy, to illustrate the effective utilization of learned contexts by the policy.
SDVT-LW	SDVT-LW masked (ωφy)	SDVT-LW masked (μφz, σφz)	SDVT-LW masked (ωφy, μφz, σφz)
ML-10 Train	62.1	58.6	36.9	34.8
ML-10 Test	33.4	33.6	25.8	24.5
We employ VariBAD's stop gradient architecture to avoid the instability that may arise from the concurrent training of policy and ELBO objectives. This may lead to a potential vulnerability where the policy neglects contexts. However, in Meta-World, identical observations could belong to different tasks, such as "Reach" and "Pick-place." Therefore, the agent must rely on the inferred context. To confirm this, we evaluate SDVT by masking either  $(\omega_{\phi_y})$  or  $(\mu_{\phi_z}, \sigma_{\phi_z})$  in the policy input.
Our findings reveal that masking the contexts severely damages the training and test success rates. Especially, masking the Gaussian parameters of  $z$  is more critical than masking the categorical parameter of  $y$ . However, this does not suggest that  $y$  is redundant, given that the continuous context  $z$  is trained to include  $y$ 's information. Interestingly, the test success rate, with all contexts masked (i.e., solely relying on the current state  $s_t$ ), achieves a success rate of  $24.5\%$ , surpassing all baseline scores. Although the policy does not directly employ the context learned by the GMVAE, it can still be effectively trained in the imaginary task generated by the learned GMVAE to prepare for tests with unseen compositions of subtasks.
E.6 Smaller number of rollouts
Table 13: Success rate on ML-10 for smaller rollouts. We report the success rates at three different rollouts,  $n_{\mathrm{roll}} = 1,2,10$ . Note that the results reported in our manuscript are for  $n_{\mathrm{roll}} = 10$ .
Methods	Training Success Rate (%)	Test Success Rate (%)
\(n_{\text{roll}}=1\)	\(n_{\text{roll}}=2\)	\(n_{\text{roll}}=10\)	\(n_{\text{roll}}=1\)	\(n_{\text{roll}}=2\)	\(n_{\text{roll}}=10\)
SDVT-LW	61.4	62.3	62.1	32.8	32.7	33.4
SD-LW	74.9	75.1	75.5	25.8	25.2	26.2
VariBAD	57.0	58.4	58.2	14.5	15.3	14.1
LDM	55.5	56.0	56.7	18.4	18.5	19.8
We use  $n_{\mathrm{roll}} = 10$ , following the original setup of the Meta-World benchmark [69, 71]. Notably, this benchmark utilizes dense rewards, therefore the context converges rather quickly, even within the first rollout (Figure 2b). As a result, the performance across rollouts does not exhibit substantial improvement. To provide further insight, we report the success rates for smaller  $n_{\mathrm{roll}}$  in Table 13.