SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathcal {O} \left(\beta_ {N ^ {*} (\epsilon)} (\delta) T ^ {3 / 2} \sqrt {\frac {\gamma N ^ {*} (\epsilon)}{N ^ {*} (\epsilon)}}\right) <   \epsilon .
$$
$$
\begin{array}{l} \sum_ {n = 0} ^ {\infty} \Pr (S _ {n} > \epsilon) = \sum_ {n = 0} ^ {N ^ {*} (\epsilon) - 1} \Pr (S _ {n} > \epsilon) + \sum_ {n = N ^ {*} (\epsilon)} ^ {\infty} \Pr (S _ {n} > \epsilon) \\ = \sum_ {n = 0} ^ {N ^ {*} (\epsilon) - 1} \Pr (S _ {n} > \epsilon) + \sum_ {n = N ^ {*} (\epsilon)} ^ {\infty} \delta_ {n} \\ \leq N ^ {*} (\epsilon) + \sum_ {n = N ^ {*} (\epsilon)} ^ {\infty} \delta_ {n}. \\ \end{array}
$$
Note, since  $\lim_{n\to \infty}\delta_n = 0$  , we have
$$
\sum_ {n = N ^ {*} (\epsilon)} ^ {\infty} \delta_ {n} <   \infty .
$$
In particular,  $\sum_{n=0}^{\infty} \operatorname{Pr}(S_n > \epsilon) < \infty$  for all  $\epsilon > 0$ . Therefore, we obtain
$$
S _ {n} \xrightarrow {\text {a . s .}} 0 \text {f o r} n \to \infty .
$$
Define the random variable  $V = \lim_{n\to \infty}\left(\sum_{t = 0}^{T - 1}\sum_{j = 1}^{d_x}\frac{1}{2}\log \left(1 + \frac{\sigma_{n,j}^2(\boldsymbol{z}_{n,t}^*)}{\sigma^2}\right)\right)$ , with  $z_{n,t}^{*}\in \tau$  and  $\tau \sim \tau^{\pi_n^*}$ .  $V$  represents the sum of epistemic uncertainties of a random trajectory induced by the sequence of policies  $\{\pi_n\}_{n\geq 0}$ . Note  $V\geq 0$ , therefore we apply Markov's inequality. Moreover, for all  $\epsilon >0$ , we have
$$
\Pr \left(V > \epsilon\right) \leq \frac {\mathbb {E} [ V ]}{\epsilon} = 0.
$$
$$
\Pr (V = 0) = 1 \Longrightarrow V \xrightarrow {\text {a . s .}} 0 \text {f o r} n \to \infty
$$
Accordingly, we get for all  $\pi \in \Pi$
$$
\Pr \left(\lim  _ {n \rightarrow \infty} \sum_ {\boldsymbol {z} _ {t} \in \boldsymbol {\tau} _ {\pi}} \sum_ {j = 1} ^ {d _ {x}} \frac {1}{2} \log \left(1 + \frac {\sigma_ {n , j} ^ {2} (\boldsymbol {z} _ {t})}{\sigma^ {2}}\right)\rightarrow 0\right) = 1. \tag {15}
$$
Assume there exists a  $z \in \mathcal{R}$ , such that for some  $\epsilon$ ,  $\sigma_{n,j}^2(z) > \epsilon$  for all  $n \geq 0$ . Since,  $z \in \mathcal{R}$ , there exists a  $t$  and  $\pi \in \Pi$  such that  $p(\boldsymbol{z}_t = \boldsymbol{z} | \pi, \boldsymbol{f}^*) > 0$ . This implies that  $\operatorname*{Pr}(z \in \tau_{\pi}) > 0$ . However, from Equation (15), we have that  $\sigma_{n,j}^2(z) \to 0$  for  $N \to \infty$  almost surely, which is a contradiction.
Finally, we leverage the results from above to prove Theorem 2.
Proof of Theorem 2. The proof follows directly from Corollary 6 and Lemma 12.
A.5 Zero-shot guarantees
In this section, we give guarantees on the zero-shot performance of OPAX for a bounded cost function. We focus this section on the case of Gaussian noise. However, a similar analysis can be performed for the sub-Gaussian case and Lipschitz continuous costs. Since the analysis for both cases is similar, we only present the Gaussian case with bounded costs here.
Corollary 7. Consider the following optimal control problem
$$
\underset {\boldsymbol {\pi} \in \Pi} {\arg \min } J _ {c} (\boldsymbol {\pi}, \boldsymbol {f} ^ {*}) = \underset {\boldsymbol {\pi} \in \Pi} {\arg \min } \mathbb {E} _ {\boldsymbol {\tau} ^ {\pi}} \left[ \sum_ {t = 0} ^ {T - 1} c \left(\boldsymbol {x} _ {t}, \boldsymbol {\pi} \left(\boldsymbol {x} _ {t}\right)\right) \right], \tag {16}
$$
$$
\boldsymbol {x} _ {t + 1} = \boldsymbol {f} ^ {*} (\boldsymbol {x} _ {t}, \boldsymbol {\pi} (\boldsymbol {x} _ {t})) + \boldsymbol {w} _ {t} \quad \forall 0 \leq t \leq T,
$$
with bounded and positive costs. Then we have for all policies  $\pi$  with probability at least  $1 - \delta$
$$
J _ {c} (\pmb {\pi}, \pmb {\eta} ^ {P}) - J _ {c} (\pmb {\pi}) \leq \mathcal {O} \left(T \mathbb {E} _ {\pmb {\tau} ^ {\pi_ {n}}} \left[ \sum_ {t = 0} ^ {T - 1} \frac {(1 + \sqrt {d _ {x}}) \beta_ {n - 1} (\delta) \| \pmb {\sigma} _ {n - 1} (\pmb {z} _ {t}) \|}{\sigma^ {2}} \right]\right),
$$
where  $J_{c}(\pmb {\pi},\pmb{\eta}^{P}) = \max_{\pmb {\eta}\in \Xi}J_{c}(\pmb {\pi},\pmb {\eta})$
Proof. From Corollary 2 we get
$$
J _ {c} (\boldsymbol {\pi}, \boldsymbol {\eta} ^ {P}) - J _ {c} (\boldsymbol {\pi}) = \mathbb {E} _ {\boldsymbol {\tau} ^ {\pi}} \left[ \sum_ {t = 0} ^ {T - 1} J _ {r, t + 1} (\boldsymbol {\pi}, \boldsymbol {\eta} ^ {P}, \boldsymbol {x} _ {t + 1} ^ {P}) - J _ {r, t + 1} (\boldsymbol {\pi}, \boldsymbol {\eta} ^ {P}, \boldsymbol {x} _ {t + 1}) \right],
$$
with  $\pmb{x}_{t + 1} = \pmb{f}^{*}(\pmb{x}_{t},\pmb{\pi}(\pmb{x}_{t})) + \pmb{w}_{t}$ , and  $\pmb{x}_{t + 1}^{P} = \pmb{\mu}_{n}(\pmb{x}_{t},\pmb{\pi}(\pmb{x}_{t})) + \beta_{n}(\delta)\pmb{\sigma}(\pmb{x}_{t},\pmb{\pi}(\pmb{x}_{t}))\pmb{\eta}^{P}(\pmb{x}_{t}) + \pmb{w}_{t}$ . Furthermore, the cost is positive and bounded. Therefore,
$$
J _ {c} ^ {2} (\boldsymbol {\pi}, \boldsymbol {\eta}, \boldsymbol {x}) \leq T ^ {2} c _ {\max } ^ {2},
$$
for all  $x, \eta, \pi$ . Accordingly, we can now use the same analysis as in Lemma 2 and get
$$
J _ {c} (\boldsymbol {\pi}, \boldsymbol {\eta} ^ {P}) - J _ {c} (\boldsymbol {\pi}) \leq \mathcal {O} \left(T \mathbb {E} _ {\boldsymbol {\tau} ^ {\pi_ {n}}} \left[ \sum_ {t = 0} ^ {T - 1} \frac {(1 + \sqrt {d _ {x}}) \beta_ {n - 1} (\delta) \| \boldsymbol {\sigma} _ {n - 1} (\boldsymbol {z} _ {t}) \|}{\sigma^ {2}} \right]\right),
$$
Lemma 13. Consider the control problem in Equation (16) and let Assumption 3 and 4 hold. Furthermore, assume for every  $\epsilon >0$ , there exists a finite integer  $n^*$  such that
$$
\forall n \geq n ^ {*}; \beta_ {n} ^ {3 / 2} (\delta) T ^ {1 1 / 4} \left(\frac {\gamma_ {n} (k)}{n}\right) ^ {\frac {1}{4}} \leq \epsilon , \tag {17}
$$
and denote with  $\hat{\pi}_n$  the minimax optimal policy, i.e., the solution to  $\min_{\pmb {\pi}\in \Pi}\max_{\pmb {\eta}\in \Xi}J_c(\pmb {\pi},\pmb {\eta})$ . Then for all  $n\geq n^*$ , we have probability at least  $1 - \delta$ ,  $J_{c}(\hat{\pi}_{N}) - J_{c}(\pmb{\pi}^{*})\leq \mathcal{O}(\epsilon)$
Proof of Zero-shot performance. In Theorem 2 we give a rate at which the maximum uncertainty along a trajectory decreases:
$$
\max _ {\boldsymbol {\pi} \in \Pi} \mathbb {E} _ {\boldsymbol {\tau} ^ {\boldsymbol {\pi}}} \left[ \max _ {\boldsymbol {z} \in \boldsymbol {\tau} ^ {\boldsymbol {\pi}}} \frac {1}{2} \| \sigma_ {N, j} (\boldsymbol {z}) \| ^ {2} \right] \leq \mathcal {O} \left(\beta_ {N} T ^ {3 / 2} \sqrt {\frac {\gamma_ {N} (k)}{N}}\right).
$$
Combining this with Corollary 7 we get
$$
\begin{array}{l} J _ {c} (\boldsymbol {\pi}, \boldsymbol {\eta} ^ {P}) - J _ {c} (\boldsymbol {\pi}) \leq \mathcal {O} \left(T \mathbb {E} _ {\boldsymbol {\tau} ^ {\pi_ {n}}} \left[ \sum_ {t = 0} ^ {T - 1} \frac {(1 + \sqrt {d _ {x}}) \beta_ {n - 1} (\delta) \| \boldsymbol {\sigma} _ {n - 1} (\boldsymbol {z} _ {t}) \|}{\sigma^ {2}} \right]\right) \\ \leq \mathcal {O} \left(T ^ {2} \beta_ {n - 1} (\delta) \mathbb {E} _ {\boldsymbol {\tau} ^ {\pi_ {n}}} \left[ \max  _ {\boldsymbol {z} \in \boldsymbol {\tau} ^ {\pi_ {n}}} \| \boldsymbol {\sigma} _ {n - 1} (\boldsymbol {z}) \| \right]\right) \\ \leq \mathcal {O} \left(\sqrt {\mathbb {E} _ {\boldsymbol {\tau} ^ {\pi_ {n}}} \left[ \max  _ {\boldsymbol {z} \in \boldsymbol {\tau} ^ {\pi_ {n}}} T ^ {4} \beta_ {n - 1} ^ {2} (\delta) \| \boldsymbol {\sigma} _ {n - 1} (\boldsymbol {z}) \| ^ {2} \right]}\right) \\ \leq \mathcal {O} \left(\sqrt {\beta_ {n} ^ {3} T ^ {1 1 / 2} \sqrt {\frac {\gamma_ {n} (k)}{n}}}\right) \\ = \mathcal {O} \left(\beta_ {n} ^ {3 / 2} (\delta) T ^ {1 1 / 4} \left(\frac {\gamma_ {n} (k)}{n}\right) ^ {\frac {1}{4}}\right) \\ = \mathcal {O} (\epsilon). \quad (\forall n \geq n ^ {*}) \\ \end{array}
$$
Hence, we have that for each policy  $\pi$ , our upper bound  $\max_{\pmb{\eta}} J_c(\pmb{\pi}, \pmb{\eta}^P)$  is  $\epsilon$  precise, i.e.,
$$
\max  _ {\boldsymbol {\eta} \in \Xi} J _ {c} (\boldsymbol {\pi}, \boldsymbol {\eta}) - J _ {c} (\boldsymbol {\pi}) \leq \mathcal {O} (\epsilon), \forall \boldsymbol {\pi} \in \Pi . \tag {18}
$$
We leverage this to prove optimality for the minimax solution. For the sake of contradiction, assume that
$$
J _ {c} \left(\hat {\pi} _ {n}\right) > J _ {c} \left(\boldsymbol {\pi} ^ {*}\right) + \mathcal {O} (\epsilon). \tag {19}
$$
$$
\begin{array}{l} \max  _ {\boldsymbol {\eta} \in \Xi} J _ {c} (\boldsymbol {\pi} ^ {*}, \boldsymbol {\eta}) \geq \min  _ {\boldsymbol {\pi} \in \Pi} \max  _ {\boldsymbol {\eta} \in \Xi} J _ {c} (\boldsymbol {\pi}, \boldsymbol {\eta}) \\ = J _ {c} \left(\hat {\boldsymbol {\pi}} _ {n}, \hat {\boldsymbol {\eta}} ^ {P}\right) \\ = \max  _ {\boldsymbol {\eta} \in \Xi} J _ {c} (\hat {\boldsymbol {\pi}} _ {n}, \boldsymbol {\eta}) \\ \geq J _ {c} (\hat {\boldsymbol {\pi}} _ {n}) \\ > J _ {c} \left(\boldsymbol {\pi} ^ {*}\right) + \mathcal {O} (\epsilon) (Equation(19)) \\ \geq J _ {c} \left(\boldsymbol {\pi} ^ {*}, \boldsymbol {\eta} ^ {*}, P\right) \quad (Equation(18)) \\ = \max  _ {\boldsymbol {\eta} \in \Xi} J _ {c} \left(\boldsymbol {\pi} ^ {*}, \boldsymbol {\eta}\right). \\ \end{array}
$$
This is a contradiction, which completes the proof.
Lemma 13 shows that OPAX also results in nearly-optimal zero-shot performance. The convergence criteria in Equation (17) is satisfied for kernels  $k$  that induce a very rich class of RKHS (c.f., Table 1).
The environment details and hyperparameters used for our experiments are presented in this section. In Appendix B.2 we discuss the experimental setup of the Fetch Pick & Construction environment (Li et al., 2020) in more detail.
B.1 Environment Details
Table 2 lists the rewards used for the different environments. We train the dynamics model after each episode of data collection. For training, we fix the number of epochs to determine the number of gradient steps. Furthermore, for computational reasons, we upper bound the number of gradient steps by "maximum number of gradient steps". The hyperparameters for the model-based agent are presented in Table 3. Furthermore, we present the iCEM hyperparameters in Table 4. For more detail on the iCEM hyperparameters see Pinneri et al. (2021).
Model-based SAC optimizer For the receiver and MountainCar environment we use scarce rewards (c.f., Table 2), which require long-term planning. Therefore, a receding horizon MPC approach is less suitable for these tasks. Accordingly, we use a policy network which we train using SAC (Haarnoja et al., 2018). Moreover, our model-based SAC uses transitions simulated through the learned model to train a policy. Accordingly, given a dataset of transitions from the true environment, we sample  $P$  initial states from the dataset. For each of the states, we simulate a trajectory of  $H$  steps using our learned model and the SAC policy. We collect the simulated transitions into a simulation data buffer  $\mathcal{D}_{\mathrm{SIM}}$ , which we then use to perform  $G$  gradient steps as suggested by Haarnoja et al. (2018) to train the policy. The algorithm is summarized below, and we provide the SAC hyperparameters in Table 5.
Model-based SAC
**Init:** Statistical Model  $M = (\pmb{\mu}_n, \pmb{\sigma}_n, \beta_n(\delta))$ , Dataset of transitions  $\mathcal{D}_n$ , initial policy  $\pi_0$ , Model Rollout steps  $H$
$\mathcal{D}_{\mathrm{SIM}}\gets \emptyset$
$\succ$  Initialize simulated transitions buffer
for Training steps  $k = 1,\dots ,K$  do
$$
\boldsymbol {x} _ {0, 1: P} \sim \mathcal {D} _ {n}
$$
$\succ$  Sample  $P$  initial states from buffer
for Initial state  $\pmb{x}_0 \in \pmb{x}_{0,1:P}$  do
$$
\mathcal {D} _ {\mathrm {S I M}} \leftarrow \mathcal {D} _ {\mathrm {S I M}} \cup \text {M O D E L R O L L O U T} \left(\boldsymbol {\pi} _ {k}, \boldsymbol {x} _ {0}, H, M\right)
$$
$\triangleright$  Collect simulated transitions
Perform  $G$  gradient updates on  $\pi_k$  as proposed in Haarnoja et al. (2018) using  $\mathcal{D}_{\mathrm{SIM}}$ .
B.2 OPAX in the High-dimensional Fetch Pick & Place Environment
B.2.1 Environment and Model Details
In our experiments, we use an extension of the Fetch Pick & Place environment to multiple objects as proposed in Li et al. (2020) and further modified in Sancaktar et al. (2022) with the addition of a large table. This is a compositional object manipulation environment with an end-effector-controlled robot arm. The robot actions  $\pmb{u} \in \mathbb{R}^4$  correspond to the gripper movement in Cartesian coordinates and the gripper opening/closing. The robot state  $\pmb{x}_{\mathrm{robot}} \in \mathbb{R}^{10}$  contains positions and velocities of the end-effector as well as the gripper-state (open/close) and gripper-velocity. Each object's state  $\pmb{x}_{\mathrm{obj}} \in \mathbb{R}^{12}$  is given by its position, orientation (in Euler angles), and linear and angular velocities.
We follow the free-play paradigm as used in CEE-US (Sancaktar et al., 2022). At the beginning of free play, an ensemble of world models is randomly initialized with an empty replay buffer. At each iteration of free play, a certain number of rollouts (here: 20 rollouts with 100 timesteps each) are collected and then added to the replay buffer. Afterwards, the models in the ensemble are trained for a certain number of epochs on the collected data so far. Note that unlike in the original proposal by Sancaktar et al. (2022), we use Multilayer Perceptrons (MLP) as world models instead of Graph Neural Networks (GNN), for the sake of reducing run-time. This corresponds to the ablation MLP + iCEM presented in Sancaktar et al. (2022), which was shown to outperform all the baselines other than CEE-US with GNNs. As we are interested in exploring the difference in performance via injection of optimism into active exploration, we use the computationally less heavy MLPs in our work. This is reflected in the downstream task performance we report compared to the original CEE-US with GNNs. Details for the environment and models are summarized in Table 6.
After the free-play phase, we use the trained models to solve downstream tasks zero-shot via model-based planning with iCEM. We test for the tasks pick & place, throwing and flipping with 4 objects. The rewards used for these tasks are the same as presented in Sancaktar et al. (2022).
B.2.2 OpAx Heuristic Variant
In the case of Fetch Pick & Place with an high-dimensional observation space, we implement a heuristic of OPAX. Note that as Fetch Pick & Place is a deterministic environment without noise, we only model epistemic uncertainty.
OPAX (Heuristic Variant)
Input: Ensemble  $\{\pmb {f}_i\}_{k = 1}^K$ $\epsilon \ll 1$
for  $n\in 1,\ldots ,N_{\max}$  do
Solve optimal control problem till convergence for the system:  $\pmb{x}_{t + 1} = \pmb{f}_{j_t}(\pmb{x}_t, \pmb{u}_t)$ .
$$
\boldsymbol {u} _ {0: T - 1} ^ {\star}, j _ {0: T - 1} ^ {\star} = \underset {\boldsymbol {u} _ {0: T - 1}, j _ {0: T - 1}} {\operatorname {a r g m a x}} \sum_ {t = 0} ^ {T - 1} \sum_ {i = 1} ^ {d _ {x}} \log \left(\epsilon^ {2} + \sigma_ {n, i} ^ {2} (\boldsymbol {x} _ {t}, \boldsymbol {u} _ {t})\right) \quad > \text {E s t i m a t e}
$$
$\sigma_{n,i}(\pmb{x}_t,\pmb{u}_t)$  with ensemble disagreement.
Rollout  $\pmb{u}_{0:T_1}^*$  on the system and collect data  $\mathcal{D}_n$ .
Update models  $\{\pmb{f}_i\}_{k=1}^K$ .
B.2.3 Controller Parameters
The set of default hyperparameters used for the iCEM controller are presented in Table 7, as they are used during the intrinsic phase for OPAX, CEE-US, and other baselines. The controller settings used for the extrinsic phase are given in Table 8. For more information on the hyperparameters and the iCEM algorithm, we refer the reader to Pinneri et al. (2021).
C Study of exploration intrinsic rewards
The intrinsic reward suggested in Equation (6) takes the log of the model epistemic uncertainty. Another common choice for the intrinsic reward is the epistemic uncertainty or model disagreement without the log (Pathak et al., 2019; Sekar et al., 2020). In the following Lemma, we show that these rewards are closely related.
Lemma 14. Let  $\sigma_{\mathrm{max}} = \sup_{\boldsymbol {z}\in \mathcal{Z};i\geq 0;1\leq j\leq d_x}\sigma_{i,j}(\boldsymbol {z})$  and  $\sigma >0$ . Then for all  $i\geq 0$  and  $j\in \{1,\ldots ,d_x\}$
$$
\sigma_ {i, j} ^ {2} (\pmb {z}) \leq \frac {\sigma_ {\mathrm {m a x}}}{\log (1 + \sigma^ {- 2} \sigma_ {\mathrm {m a x}})} \log (1 + \sigma^ {- 2} \sigma_ {i, j} ^ {2} (\pmb {z})) \leq \frac {\sigma^ {- 2} \sigma_ {\mathrm {m a x}}}{\log (1 + \sigma^ {- 2} \sigma_ {\mathrm {m a x}})} \sigma_ {i, j} ^ {2} (\pmb {z}).
$$
Proof. Curi et al. (2020) derive the first inequality on the left. The second inequality follows since  $\log (1 + x)\leq x$  for all  $x\geq 0$
Due to this close relation between the two objectives, our theoretical findings also apply to the intrinsic reward proposed by (Pathak et al., 2019). Moreover, empirically we notice in Figure 5 that OPAX performs similarly when the sum of epistemic uncertainties is used instead of the objective in Equation (6). However, our objective can naturally be extended to the case of heteroscedastic aleatoric noise, since it trades off the ratio of epistemic and aleatoric uncertainty.