SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Miniworld. While this task is solved from image-based observations, we also note that there are only four distinct tasks (one for each colored box), and the agent does not need to handle new tasks at test time. Hence, the number of in-context datasets required in pretraining is fewer – we use 60K datasets each of length  $H = 50$  steps. So as to reduce computation, the in-context datasets only have only  $(s, a, r)$  tuples. The query states, which consist of image and direction are sampled uniformly from the entire state space, i.e., the agent is place uniformly at random in the environment, pointing in a random direction. The optimal actions are computed as follows: turn towards the correct box if the agent is not yet facing it (within ±15 degrees), otherwise move forward. Of the 60K collections of datasets, query states, and optimal actions, we use 80% for training and the remaining 20% for validation.
B.1 Bandits
This section reports additional experimental results in bandit environments.
Out-of-distribution reward variances. In Figures 2c and 6a, we demonstrate the robustness of the basic pretrained model under shifts in the reward distribution at test time by varying the amount of noise observed in the rewards. DPT maintains robustness to these shifts similar to TS.
Figure 6: (a) Final (after 500 steps) offline suboptimality on out-of-distribution bandits with different Gaussian noise standard deviations. (b) Offline performance on out-of-distribution Bernoulli bandits, given random in-context datasets. (c) Online cumulative regret on Bernoulli bandits. The mean and standard error are computed over 200 test tasks.
(a)
Figure 7: All comparisons in Dark Room evaluated on the tasks that were seen during pretraining, displayed next to their evaluations on test task counterparts from the main text.
Bernoulli rewards. We test the out-of-distribution ability of DPT further by completely changing the reward distribution from Gaussian to Bernoulli bandits. Despite being trained only on Gaussian tasks during pretraining, DPT maintains strong performance both offline and online in Figures 6b and 6c.
B.2 Markov Decision Processes
This section reports additional experimental results in the Dark Room and Miniworld environments.
Performance on training tasks. In Fig. 7, we show the performance of each method on the training tasks in Dark Room. Offline, DPT and AD demonstrate comparable performance as on the training tasks, indicating a minimal generalization gap to new goals. Online, DPT, AD, and  $\mathrm{RL}^2$  also achieve performance on the training tasks similar to that on the test tasks.
Generalization to new dynamics. In this experiment, we study generalization to variations in a different aspect of the MDP, namely the dynamics. We design Dark Room (Permuted), a variant of Dark Room in which the goal is fixed to a corner but the action space is randomly permuted. Hence, the agent must leverage its historical context to infer the effect of each action. On a held-out set of 20 permutations, DPT infers the optimal policy correctly every time offline, given only 100 offline samples, matching the optimal policy at 83 return. Similarly, the online performance immediately snaps to a near optimal policy in one episode once it identifies the novel permutation in Figure 8.
B.3 Sensitivity Analysis
Figure 8: Online evaluation of DPT on Dark Room when tested on novel actions set permutations.
We next seek to understand the sensitivity of DPT to different hyperparameters, including the model size and size of the pretraining
dataset. These experiments are performed in the Dark Room environment. As shown in Fig. 9, the performance of DPT is robust to the model size; it is the same across different embedding sizes,
(b)
Figure 9: Sensitivity analysis of the offline Dark Rook task over the GPT-2 transformer's hyperparameters: (a) layers (b) attention heads (c) embedding dimensions (d) pretraining samples.
number of layers, and number of attention heads. Notably, the performance is slightly worse with 8 attention heads, which may be attributed to slight overfitting. We do see that when the pretraining dataset is reduced to  $10\%$  of its original size (10000 samples) the performance degrades, but otherwise has similar performance with larger pretraining datasets.
C Additional Theory and Omitted Proofs
We start with a well-known concentration inequality for the maximum-likelihood estimate (MLE) to provide some more justification for the approximation made in Assumption 1. We state a version from [89]. Let  $\mathcal{F}$  be a finite function class used to model a conditional distribution  $p_{Y|X}(y|x)$  for  $x\in \mathcal{X}$  and  $y\in \mathcal{V}$ . Assume there is  $f^{\star}\in \mathcal{F}$  such that  $p(y|x) = f^{\star}(y|x)$  (realizable), and  $f(\cdot |x)\in \Delta (\mathcal{V})$  for all  $x\in \mathcal{X}$  and  $f\in \mathcal{F}$  (proper). Let  $D = \{x_{i},y_{i}\}_{i\in [N]}$  denote a dataset of i.i.d samples where  $x_{i}\sim p_{X}$  and  $y_{i}\sim p_{Y|X}(\cdot |x_{i})$ . Let
$$
\hat {f} = \underset {f \in \mathcal {F}} {\operatorname {a r g m a x}} \sum_ {i \in [ N ]} \log f \left(y _ {i} \mid x _ {i}\right) \tag {6}
$$
Proposition C.1 (Theorem 21 of [89]). Let  $D$  and  $\hat{f}$  be given as above under the aforementioned conditions. Then, with probability at least  $1 - \delta$ ,
$$
\mathbb {E} _ {x \sim p _ {X}} \| \hat {f} (\cdot | x) - p _ {Y \mid X} (\cdot | x) \| _ {1} ^ {2} \leq \frac {8 \log (| \mathcal {F} | / \delta)}{N} \tag {7}
$$
The finiteness of  $\mathcal{F}$  is done for simplicity, but we can see that this yields dependence on the log-cardinality, a common measure of complexity. Extensions to infinite  $\mathcal{F}$  of bounded statistical complexity can be readily made to replace this. For our setting, the bound suggests that  $\mathbb{E}_{P_{pre}}\| P_{pre}(\cdot |s_{query},D,\xi_h) - M_\theta (\cdot |s_{query},D,\xi_h)\| _1^2\to 0$  as  $N\to \infty$  with high probability, provided the function class of  $M_{\theta}$  has bounded statistical complexity.
C.1 Posterior Sampling
Posterior sampling is most generally described with the following procedure [13]. Initialize a prior distribution  $\mathcal{T}_1 = \mathcal{T}_{\mathrm{pre}}$  and dataset  $D = \{\}$ . For  $k \in [K]$
1. Sample  $\tau_{k}\sim \mathcal{T}_{k}$  and compute  $\hat{\pi}_{\tau_k}$
2. Execute  $\pi_{\tau_k}^{\star}$  and add interactions to  $D$
3. Update posterior distribution  $\mathcal{T}_{k + 1}(\tau) = P(\tau |D)$ .
The prior and posteriors are typically over models such as reward functions in bandits or transition dynamics in MDPs.
Proof. Without loss of generality, for a task  $\tau$ , we take  $\pi_{\tau}^{\star}(\cdot |s)$  to be deterministic and denote the optimal action in state  $s$  as  $\pi_{\tau}^{\star}(s)$ . Recall that we consider a fixed current task  $\tau_{c}$  and a fixed in-context dataset  $D$ . Define  $\xi_{h} = (s_{1},a_{1},\dots ,s_{h},a_{h})$ .
We now formally state the variant of the full joint distribution from which we sample during pretraining. Let  $\tau$  and  $D^{\prime}$  be an arbitrary task and dataset and let  $a^{\star} \in \mathcal{A}$ ,  $s_{\mathrm{query}} \in \mathcal{S}$ ,  $\xi_{H-1} \in (\mathcal{S} \times \mathcal{A})^{H-1}$ , and  $h \in [0, H-1]$  be arbitrary.
$$
\begin{array}{l} P _ {p r e} (\tau , a ^ {\star}, s _ {\text {q u e r y}}, D ^ {\prime}, \xi_ {H - 1}, h) = \mathcal {T} _ {\text {p r e}} (\tau) \mathcal {D} _ {\text {p r e}} (D ^ {\prime}; \tau) \mathcal {D} _ {\text {q u e r y}} (s _ {\text {q u e r y}}) \mathfrak {S} _ {H} (s _ {1: H}) \pi_ {\tau} ^ {\star} \left(a ^ {\star} | s _ {\text {q u e r y}}\right) (8) \\ \times \operatorname {U n i f} [ 0, H - 1 ] \prod_ {i \in [ H ]} \pi_ {\tau} ^ {\star} \left(a _ {i} \mid s _ {i}\right) (9) \\ \end{array}
$$
The  $\mathrm{Unif}[0,H - 1]$  is due to the fact that we sample  $h\sim \mathrm{Unif}[0,H - 1]$  and then truncate  $\xi_{h}$  from  $\xi_{H - 1}$  (or, equivalently, sample  $\xi_h\sim \mathfrak{S}_h$  directly), marginalizing out the other variables. For  $h^\prime \leq h - 1$ , recall that we also use the notation  $\mathfrak{S}_{h^{\prime}}(s_{1:h^{\prime}})$  to denote the marginalization of the full joint  $\mathfrak{S}_H$ . We will eventually work with the posterior of this distribution given the data  $D$  and history  $\xi_{h}$ :
$$
P _ {\text {p r e}} (\tau | D, \xi_ {h}) \propto \mathcal {T} _ {\text {p r e}} (\tau) \mathcal {D} _ {\text {p r e}} (D; \tau) \prod_ {i \in [ h ]} \pi_ {\tau} ^ {\star} \left(a _ {i} \mid s _ {i}\right) \tag {10}
$$
$$
\propto P _ {p r e} (\tau | D) \prod_ {i \in [ h ]} \pi_ {\tau} ^ {\star} \left(a _ {i} \mid s _ {i}\right) \tag {11}
$$
We define the following random sequences and subsequences:
$$
\Xi_ {p s} (h; D) = \left(S _ {1} ^ {p s}, A _ {1} ^ {p s}, \dots , S _ {h} ^ {p s}, A _ {h} ^ {p s}\right) \tag {12}
$$
where the variables are generated according to the following conditional process:  $\tau_{ps} \sim P(\cdot | D)$ ,  $S_1^{ps} \sim \rho_{\tau_c}$ ,  $A_h^{ps} \sim \pi_{\tau_{ps}}^{\star}(\cdot | S_h^{ps})$ , and  $S_{h+1}^{ps} \sim T_{\tau_c}(\cdot | S_h^{ps}, A_h^{ps})$ . We also define  $\Xi_{ps}(h': h; D)$  to be the last  $h - h'$  elements of  $\Xi_{ps}(h; D)$ . Analogously, we define
$$
\Xi_ {p r e} (h; D) = \left(S _ {1} ^ {p r e}, A _ {1} ^ {p r e}, \dots , S _ {h} ^ {p r e}, A _ {h} ^ {p r e}\right) \tag {13}
$$
where the variables are from the process:  $S_1^{pre} \sim \rho_{\tau_c}$ ,  $A_h^{pre} \sim P_{pre}(\cdot | S_h^{pre}, D, \Xi_{pre}(h - 1; D))$ , and  $S_{h + 1}^{pre} \sim T_{\tau_c}(\cdot | S_h^{pre}, A_h^{pre})$ . Note that  $A_h^{pre}$  is sampled conditioned on the sequence  $\Xi_{pre}(h; D)$  so far.
We will show that  $\Xi_{ps}(h;D)$  and  $\Xi_{pre}(h;D)$  follow the same distribution for all  $h\in [H]$ . For convenience, we will drop notational dependence on  $D$ , except where it resolves ambiguity. Also, because of Assumption 1, we have that  $P_{pre}(\cdot |S_h^{pre},D,\Xi_{pre}(h - 1)) = M_\theta (\cdot |S_h^{pre},D,\Xi_{pre}(h - 1))$  so we will just work with  $P_{pre}$  for the remainder of the proof. We will also make use of the following lemma.
Lemma C.2. If  $\mathcal{D}_{pre}$  is complaint, then  $P_{pre}(\tau | D) = P(\tau_{ps} = \tau | D)$ .
Proof. From the definition of posterior sampling (using the same prior,  $\mathcal{T}_{\mathrm{pre}}$ ), we have that
$$
\begin{array}{l} P \left(\tau_ {p s} = \tau | D\right) \propto P (D | \tau) \mathcal {T} _ {\text {p r e}} (\tau) (14) \\ \propto \mathcal {T} _ {\text {p r e}} (\tau) \prod_ {j \in [ n ]} T _ {\tau} \left(s _ {j} ^ {\prime} \mid s _ {j}, a _ {j}\right) R _ {\tau} \left(r _ {j} \mid s _ {j}, a _ {j}\right) (15) \\ \propto \mathcal {T} _ {\text {p r e}} (\tau) \prod_ {j \in [ n ]} T _ {\tau} \left(s _ {j} ^ {\prime} \mid s _ {j}, a _ {j}\right) R _ {\tau} \left(r _ {j} \mid s _ {j}, a _ {j}\right) \mathcal {D} _ {\text {p r e}} \left(a _ {j} \mid s _ {j}, D _ {j - 1}\right) (16) \\ = \mathcal {T} _ {\text {p r e}} (\tau) \mathcal {D} _ {\text {p r e}} (D; \tau) (17) \\ = P _ {p r e} (\tau | D) (18) \\ \end{array}
$$
where the second line crucially uses the fact that posterior sampling chooses actions based only on the prior and history so far. Similarly, the third line uses the fact that  $\mathcal{D}_{\mathrm{pre}}$  is compliant. Since the two sides are proportional in  $\tau$ , they are equivalent.
We will prove Theorem 1 via induction for each  $h \in [H]$ . First, consider the base case for a sequence of length  $h = 1$ . Recall that  $\rho_{\tau_c}$  denotes the initial state distribution of  $\tau_c$ . We have that the densities can be written as
$$
\begin{array}{l} P \left(\Xi_ {p s} (1) = \xi_ {1}\right) = P \left(S _ {1} ^ {p s} = s _ {1}, A _ {1} ^ {p s} = a _ {1}\right) (19) \\ = \rho_ {\tau_ {c}} (s _ {1}) P \left(A _ {1} ^ {p s} = a _ {1} \mid S _ {1} ^ {p s} = s _ {1}\right) (20) \\ = \rho_ {\tau_ {c}} (s _ {1}) \int_ {\tau} P \left(A _ {1} ^ {p s} = a _ {1}, \tau_ {p s} = \tau \mid S _ {1} ^ {p s} = s _ {1}\right) d \tau (21) \\ = \rho_ {\tau_ {c}} (s _ {1}) \int_ {\tau} \pi_ {\tau} ^ {\star} (a _ {1} | s _ {1}) P _ {p s} (\tau_ {p s} = \tau | D, S _ {1} ^ {p s} = s _ {1}) d \tau (22) \\ = \rho_ {\tau_ {c}} (s _ {1}) \int_ {\tau} \pi_ {\tau} ^ {\star} (a _ {1} | s _ {1}) P _ {p s} (\tau_ {p s} = \tau | D) d \tau (23) \\ = \rho_ {\tau_ {c}} \left(s _ {1}\right) P _ {p r e} \left(A _ {1} ^ {p r e} = a _ {1} \mid s _ {1}, D\right) (24) \\ = P \left(\Xi_ {p r e} (1) = \xi_ {1}\right) (25) \\ \end{array}
$$
where the second line uses the sampling process of  $S_1^{pre}$ ; the third marginalizes over  $\tau_{ps}$ , which is the task that posterior sampling samples to find the optimal policy; the fourth decomposes this into the optimal policy and the posterior over  $\tau_{ps}$  given  $D$  and  $S_1^{ps}$ . Since  $S_1^{ps}$  is independent of sampling of  $\tau_{ps}$  this dependence goes away in the next line. The sixth line applies Lemma C.2 and then, for  $h = 1$ , there is no history to condition on.
Now, we leverage the inductive hypothesis to prove the full statement. Suppose that the hypothesis holds for  $h - 1$ . Then,
$$
P \left(\Xi_ {p s} (h) = \xi_ {h}\right) = P \left(\Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) P \left(S _ {h} ^ {p s} = s _ {h}, A _ {h} ^ {p s} = a _ {h} \mid \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) \tag {27}
$$
By the hypothesis, we have that  $P(\Xi_{ps}(h - 1) = \xi_{h - 1}) = P(\Xi_{pre}(h - 1) = \xi_{h - 1})$ . For the second factor,
$$
\begin{array}{l} P \left(S _ {h} ^ {p s} = s _ {h}, A _ {h} ^ {p s} = a _ {h} \mid \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) (28) \\ = T _ {\tau_ {c}} \left(s _ {h} \mid s _ {h - 1}, a _ {h - 1}\right) \cdot P \left(A _ {h} ^ {p s} = a _ {h} \mid S _ {h} ^ {p s} = s _ {h}, \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) (29) \\ = T _ {\tau_ {c}} \left(s _ {h} \mid s _ {h - 1}, a _ {h - 1}\right) \cdot \int_ {\tau} P \left(A _ {h} ^ {p s} = a _ {h}, \tau_ {p s} = \tau \mid S _ {h} ^ {p s} = s _ {h}, \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) d \tau (30) \\ \end{array}
$$
As before, we can further rewrite the last factor as
$$
\begin{array}{l} P \left(A _ {h} ^ {p s} = a _ {h}, \tau_ {p s} = \tau \mid S _ {h} ^ {p s} = s _ {h}, \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) (31) \\ = \pi_ {\tau} ^ {\star} \left(a _ {h} \mid s _ {h}\right) \cdot P \left(\tau_ {p s} = \tau \mid S _ {h} ^ {p s} = s _ {h}, \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) (32) \\ \end{array}
$$
$$
\begin{array}{l} P \left(\tau_ {p s} = \tau \mid S _ {h} ^ {p s} = s _ {h}, \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right) = \frac {P \left(S _ {h} ^ {p s} = s _ {h} , \Xi_ {p s} (h - 1) = \xi_ {h - 1} \mid \tau_ {p s} = \tau\right) P \left(\tau_ {p s} = \tau \mid D\right)}{P \left(S _ {h} ^ {p s} = s _ {h} , \Xi_ {p s} (h - 1) = \xi_ {h - 1}\right)} (33) \\ \propto P _ {p r e} (\tau | D) \prod_ {i \in [ h - 1 ]} T _ {\tau_ {c}} \left(s _ {i + 1} \mid s _ {i}, a _ {i}\right) \pi_ {\tau} ^ {\star} \left(a _ {i} \mid s _ {i}\right) (34) \\ \propto P _ {p r e} (\tau | D) \prod_ {i \in [ h - 1 ]} \pi_ {\tau} ^ {\star} \left(a _ {i} \mid s _ {i}\right) (35) \\ \propto P _ {p r e} (\tau | D) \mathcal {D} _ {\text {q u e r y}} (s _ {h}) \mathfrak {S} _ {h - 1} \left(s _ {1: h - 1}\right) \prod_ {i \in [ h - 1 ]} \pi_ {\tau} ^ {\star} \left(a _ {i} \mid s _ {i}\right) (36) \\ \propto P _ {p r e} (\tau | s _ {h}, D, \xi_ {h - 1}) (37) \\ \end{array}
$$
where  $\propto$  denotes that the two sides are equal up to multiplicative factors independent of  $\tau$ . In the first line, we used Bayes rule. In the second line, given that  $\tau_{ps} = \tau$  (i.e. posterior sampling selected  $\tau$  to
deploy), we decompose the probability of observing that sequence of states of actions. We also used Lemma C.2. The denominator does not depend on  $\tau$ . Similarly, for the third and fourth lines,  $T_{\tau_c}$  and  $\mathfrak{S}$  do not depend on  $\tau$ . The final line follows from the definition of the joint pretraining distribution in this regime.
Therefore, we conclude that the posterior over the value of  $\tau_{ps}$  is the same as the posterior over the task in the pretraining distribution, given  $s_h, D, \xi_{h-1}$ . Substituting back through all the previous equations, we have
$$
\begin{array}{l} P \left(\Xi_ {p s} (h) = \xi_ {h}\right) (39) \\ = P \left(\Xi_ {p r e} (h - 1) = \xi_ {h - 1}\right) \cdot T _ {\tau_ {c}} \left(s _ {h} \mid s _ {h - 1}, a _ {h - 1}\right) \int_ {\tau} \pi_ {\tau} ^ {\star} \left(a _ {h} \mid s _ {h}\right) P _ {p r e} (\tau \mid s _ {h}, D, \xi_ {h - 1}) d \tau (40) \\ = P \left(\Xi_ {p r e} (h - 1) = \xi_ {h - 1}\right) \cdot T _ {\tau_ {c}} \left(s _ {h} \mid s _ {h - 1}, a _ {h - 1}\right) P _ {p r e} \left(a _ {h} \mid s _ {h}, D, \xi_ {h - 1}\right) (41) \\ = P \left(\Xi_ {p r e} (h) = \xi_ {h}\right) (42) \\ \end{array}
$$
C.3 Proof of Corollary 6.2
Corollary C.3 (Finite MDPs). Suppose that  $\sup_{\tau} \mathcal{T}_{test}(\tau) / \mathcal{T}_{pre}(\tau) \leq \mathcal{C}$  for some  $\mathcal{C} > 0$ . For the above MDP setting, the pretrained model  $M_{\theta}$  satisfies  $\mathbb{E}_{\mathcal{T}_{test}}[Reg_{\tau}(M_{\theta})] \leq \widetilde{\mathcal{O}}(\mathcal{CH}^{3/2}S\sqrt{AK})$ .
Proof. Note that  $\mathcal{D}_{\mathrm{pre}}$  is clearly compliant since it is generated by random sampling. We use the equivalence between  $M_{\theta}$  and posterior sampling established in Theorem 1. The proof then follows immediately from Theorem 1 of [13] to guarantee that
$$
\mathbb {E} _ {\mathcal {T} _ {\text {p r e}}} \left[ \operatorname {R e g} _ {\tau} \left(M _ {\theta}\right) \right] \leq \widetilde {\mathcal {O}} \left(H ^ {3 / 2} S \sqrt {A K}\right) \tag {43}
$$
where the notation  $\widetilde{O}$  omits polylogarithmic dependence. The bound on the test task distribution follows from the assumed bound on the likelihood ratio under the priors:
$$
\int \mathcal {T} _ {\text {t e s t}} (\tau) \operatorname {R e g} _ {\tau} \left(M _ {\theta}\right) d \tau \leq \mathcal {C} \int \mathcal {T} _ {\text {p r e}} (\tau) \operatorname {R e g} _ {\tau} \left(M _ {\theta}\right) d \tau . \tag {44}
$$
C.4 Proof of Corollary 6.3
Corollary C.4 (Latent representation learning in linear bandits). For  $\mathcal{T}_{test} = \mathcal{T}_{pre}$  in the above linear bandit setting,  $M_{\theta}$  satisfies  $\mathbb{E}_{\mathcal{T}_{test}}[Reg_{\tau}(M_{\theta})] \leq \widetilde{\mathcal{O}}(d\sqrt{K})$ .
Proof. The distribution  $\mathcal{D}_{\mathrm{pre}}$  satisfies compliance by definition because it is generated by an adaptive algorithm TS. The proof once again follows by immediately deferring to the established result of [68] (Proposition 3) for linear bandits by the posterior sampling equivalence of Theorem 1. This ensures that posterior sampling achieves regret  $\hat{\mathcal{O}}(d\sqrt{K})$ . It remains, however, to justify that  $P_{pre}(\cdot|D_k)$  will be covered by Gaussian Thompson Sampling for all  $D_k$  with  $k \in [K]$ . This is verified by noting that  $P_{ps}(a|D_k) > 0$  for non-degenerate Gaussian Thompson Sampling (positive variances of the prior and likelihood functions) and finite  $K$ . This guarantees that any  $D_k$  will have support.
C.5 Proof of Proposition 6.4
Proposition C.5. Let  $P_{pre}^{1}$  and  $P_{pre}^{2}$  be pretraining distributions that differ only by their in-context dataset distributions, denoted by  $\mathcal{D}_{pre}^{1}$  and  $\mathcal{D}_{pre}^{2}$ . If  $\mathcal{D}_{pre}^{1}$  and  $\mathcal{D}_{pre}^{2}$  are compliant with the same support, then  $P_{pre}^{1}(a^{\star}|s_{query}, D, \xi_{h}) = P_{pre}^{2}(a^{\star}|s_{query}, D, \xi_{h})$  for all  $a^{\star}, s_{query}, D, \xi_{h}$ .
Proof. The proof follows by direct inspection of the pretraining distributions. For  $P_{pre}^{1}$ , we have
$$
P _ {p r e} ^ {1} \left(a ^ {\star} \mid s _ {\text {q u e r y}}, D, \xi\right) = \int_ {\tau} \pi_ {\tau} ^ {\star} \left(a ^ {\star} \mid s _ {\text {q u e r y}}\right) P _ {p r e} ^ {1} (\tau \mid s _ {\text {q u e r y}}, D, \xi) d \tau \tag {45}
$$
The posterior distribution over tasks is simply
$$
\begin{array}{l} P _ {p r e} ^ {1} (\tau | s _ {\text {q u e r y}}, D, \xi) = \frac {P _ {p r e} ^ {1} (\tau , s _ {\text {q u e r y}} , D , \xi)}{P _ {p r e} ^ {1} (s _ {\text {q u e r y}} , D , \xi)} (46) \\ \propto P _ {p r e} ^ {1} (\tau) P _ {p r e} ^ {1} (\xi | \tau) \mathcal {D} _ {\text {q u e r y}} (s _ {\text {q u e r y}}) \mathcal {D} _ {\text {p r e}} ^ {1} (D; \tau) (47) \\ = P _ {p r e} ^ {2} (\tau) P _ {p r e} ^ {2} (\xi | \tau) \mathcal {D} _ {\text {q u e r y}} \left(s _ {\text {q u e r y}}\right) \mathcal {D} _ {\text {p r e}} ^ {1} (D; \tau) (48) \\ \end{array}
$$
Then, the distribution over the in-context dataset can be decomposed as
$$
\begin{array}{l} \mathcal {D} _ {\text {p r e}} ^ {1} (D; \tau) = \prod_ {i \in [ n ]} R _ {\tau} \left(r _ {i} \mid s _ {i}, a _ {i}\right) T _ {\tau} \left(s _ {i} ^ {\prime} \mid s _ {i}, a _ {i}\right) \mathcal {D} _ {\text {p r e}} ^ {1} \left(a _ {i} \mid s _ {i}, D _ {i - 1}; \tau\right) (49) \\ = \prod_ {i \in [ n ]} R _ {\tau} \left(r _ {i} \mid s _ {i}, a _ {i}\right) T _ {\tau} \left(s _ {i} ^ {\prime} \mid s _ {i}, a _ {i}\right) \mathcal {D} _ {\text {p r e}} ^ {1} \left(a _ {i} \mid s _ {i}, D _ {i - 1}\right) (50) \\ \propto \prod_ {i \in [ n ]} R _ {\tau} \left(r _ {i} \mid s _ {i}, a _ {i}\right) T _ {\tau} \left(s _ {i} ^ {\prime} \mid s _ {i}, a _ {i}\right) \mathcal {D} _ {\text {p r e}} ^ {2} \left(a _ {i} \mid s _ {i}, D _ {i - 1}\right) (51) \\ = \mathcal {D} _ {p r e} ^ {2} (D; \tau), (52) \\ \end{array}
$$
where the second equality holds because  $\mathcal{D}_{\mathrm{pre}}^{1}(a_{j}|s_{j},D_{j};\tau)$  is assumed to be invariant to  $\tau$  by compliance, and the fifth equality holds because  $\mathcal{D}_{\mathrm{pre}}^{2}(a_{j}|s_{j},D_{j};\tau)$  is assumed to be invariant to  $\tau$ .
Therefore, we conclude that, for any  $s,D,\xi$
$$
\begin{array}{l} P _ {p r e} ^ {1} (\tau | s, D, \xi) \propto P _ {p r e} ^ {2} (\tau) P _ {p r e} ^ {2} (\xi | \tau) \mathcal {D} _ {\text {q u e r y}} (s) \mathcal {D} _ {\text {p r e}} ^ {2} (D; \tau) (53) \\ \propto P _ {p r e} ^ {2} (\tau | s, D, \xi). (54) \\ \end{array}
$$
Since also  $\int_{\tau}P_{pre}^{1}(\tau |s,D,\xi) = 1 = \int_{\tau}P_{pre}^{2}(\tau |s,D,\xi)$ , then
$$
P _ {p r e} ^ {1} (\tau | s, D, \xi) = P _ {p r e} ^ {2} (\tau | s, D, \xi). \tag {55}
$$
Substituting this back into Equation45 yields  $P_{pre}^{1}(a^{\star}|s,D,\xi) = P_{pre}^{1}(a^{\star}|s,D,\xi)$ .