Optimistic Active Exploration of Dynamical Systems
Authors: Bhavya Sukhija
Lenart Treven
Cansu Sancaktar
Sebastian Blaes
Stelian Coros
Andreas Krause
Venue: NeurIPS 2023 poster (2023)
Keywords: Active Exploration
Reinforcement Learning
Dynamical Systems
TL;DR: We study how to actively explore a dynamical system in an efficient manner.

Abstract: Reinforcement learning algorithms commonly seek to optimize policies for solving one particular task. How should we explore an unknown dynamical system such that the estimated model allows us to solve multiple downstream tasks in a zero-shot manner? 
In this paper, we address this challenge, by developing an algorithm -- OPAX -- for active exploration. OPAX uses well-calibrated probabilistic models to quantify the epistemic uncertainty about the unknown dynamics. It optimistically---w.r.t. to plausible dynamics---maximizes the information gain between the unknown dynamics and state observations. We show how the resulting optimization problem can be reduced to an optimal control problem that can be solved at each episode using standard approaches.  We analyze our algorithm for general models, and, in the case of Gaussian process dynamics, we give a sample complexity bound and
show that the epistemic uncertainty converges to zero. In our experiments, we compare OPAX with other heuristic active exploration approaches on several environments. Our experiments show that OPAX is not only theoretically sound but also performs well for zero-shot planning on novel downstream tasks.

Optimistic Active Exploration of Dynamical Systems
Bhavya Sukhija'
Lenart Treven' Cansu Sancaktar
Sebastian Blaes
Stelian Coros' Andreas Krause'
ETH Zürich MPI for Intelligent Systems²
fsukhijab, trevenl, scoros krausea)@ethz. ch
{cansu. sancaktar, sebastian. blaeeeohuebinggnnunp..c.
Abstract
Reinforcement learning algorithms commonly seek to optimize policies for solving one particular task. How should explore unknown dynamical system such that the estimated model globally approximates the dynamics and allows us to solve multiple downstream tasks in zero-shot manner? In this paper, address this challenge, by developing an algorithm OPAX- for active exploration. OPAX uses well-calibrated probabilistic models quantify the epistemic uncertainty about the unknown dynamics. It optimistically-wwa the information gain between the unknown dynamics and state observations. We show how the resulting optimization problem can be reduced to an optimal control problem that can be solved at each episode using standard approaches. We analyze our algorithm for general models, and, in the case of Gaussian process dynamics, we give first-of-its-king sample complexity bound and show that the epistemic uncertainty converges to zero. In our experiments, we compare OPAX with other heuristic active exploration approaches on several environments. Our experiments show that OpAX is not only theoretically sound but also performs well for zero-shot planning on novel downstream tasks.
to plausible dyyamics-maximizee
Introduction
Most reinforcement learning (RL) algorithms are designed to maximize cumulative rewards for single task hand. Particularly, model-based RL algorithms, such as (Chua et al., 2018; Kakade et al., 2020; Curi et al., 2020), excel in efficiently exploring the dynamical system as they direct the exploration in regions with high rewards. However, due to the directional bias, their underlying learned dynamics model fails to generalize in other areas of the state-action space. While this is sufficient only one control task considered, does not scale to setting where the system used to perform several tasks, i.e., under the same dynamics optimized for different reward functions. As result, when presented with a new reward function, they often need to relearn a policy from scratch, requiring many interactions with the system, employ multi-task (Zhang and Yang. 2021) or transfer learning (Weiss et al., 2016) methods. Traditional control approaches such as trajectory optimization (Biagiotti and Melchiorri, 2008) and model-predicte assume knowledge the system's 's dynamics. They leverage the dynamics model to solve an optimal control problem for each task. Moreover, in the presence of accurate model, important system properties such as stability and sensitivity can also be studied. Hence, knowing an accurate dynamics model bears many practical benefits. However, in many real-world settings, obtaining model using just physics' first principles is very challenging. A promising approach is to leverage data for learning the dynamics, i.e., system identification active learning. To this end, the key question we investigate in this work is: how should we interact with the system to learn its dynamics efficiently?
control (García et al., 1989)
While active learning for regression and classification tasks is well-studied, active learning in RL much less understood. particular, active learning methods that yield strong theoretical and practical results, generally query data points based on cnemmation-thooreti criteria (Krause et al., 2008; Settles, 2009; Balcan et al., 2010; Hanneke et al., 2014; Chen et al., 2015). In the context of
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

dynamical systems, this requires querying arbitrary transitions (Berkenkamp et al., 2017; Mehta et al., 2021). However, in most cases, querying a dynamical system at any state-action pair unrealistic. Rather, we can only execute policies on the real system and observe the resulting trajectories. Accordingly, an active learning algorithm for RI. needs to suggest policies that are "informative" for learning the dynamics. This is challenging since it requires planning with unknown dynamics.
Contributions In this paper, we introduce a new algorithm, Optimistic Active eXploration (OPAX), designed to actively learn nonlinear dynamies within continuous state-action spaces. During each episode, OPAX plans an exploration policy to gather the most information possible about the system. It learns a statistical dynamics model that can quantify its epistemic uncertainty and utilizes this uncertainty for planning. The planned trajectory targets state-action pairs where the model's epistemic uncertainty is high, which naturally encourages exploration. In light of unknown dynamics, OPAX uses an optimistic planner that picks policies that optimistically yield maximal information. We show that this optimism paradigm plays crucial role in studying the theoretical properties of OpAX. Moreover, we provide a general convergence analysis for OPAX and prove convergence to the true dynamics for Gaussian process (GP) dynamics models. Theoretical guarantees for active learning in RL exist for limited class of systems (Simchowitz et al., 2018; Wagenmaker and Jamieson, 2020; Mania et al., 2020), but lack for a more general and practical class of dynamics (Chakraborty et al., 2023: Wagenmaker et al., 2023). We are, to the best of our knowledge, the first to give convergence guarantees for rich class of nonlinear dynamical systems.
We evaluate OPAX on several simulated robotic tasks with state dimensions ranging from two to 58. The empirical results provide validation for our theoretical conclusions, showing that OPAx consistently delivers strong performance across all tested environments. Finally, we provide an efficient implementation' of OPAX in JAX (Bradbury al., 2018).
2 Problem Setting
We study an unknown discrete-time dynamical system f", with state x E X C Rdz and control inputs u E U C gd.
Øk-1 = (#*,, ux) + wk-
Here, wa represents the stochasticity of the system for which assume we ' o0, (Assump- tion 3). Most common approaches in control, such as trajectory optimization and preei-predictive control (MPC), assume that the dynamics model f is known and leverages the model to control the system state. Given a cost function e X' x l -> IR, such approaches formulate and solve an optimal control problem to obtain a sequence of control inputs that drive the system's 's state
arg min Empgr- No.r
Σ c(E,, un)
(2)
Lt+1 ( u.) + w. VO 5
Moreover, if the dynamics are known, many important characteristics of the system such as stability, and robustness (Khalil, 2015) can be studied. However, in many real-world scenarios, an accurate dynamics model f" is not available. Accordingly, in this work, we consider the problem of actively learning the dynamics model from data a.k.a. system identification (Aström and Eykhoff, 1971). Specifically, we are interested in devising a cost-agnostic algorithm that focuses solely on learning the dynamics model. Once good model is learned, it can be used for solving different downstream tasks by varying the cost function in Equation (2).
We study an episodic setting, with episodes = 1, 1,...,N. N.. At the beginning of the episode n, we deploy an exploratory policy R., chosen from a policy space II for a horizon of on the system. Next, we obtain trajectory Tn = (mI,,0, ..., ,n,T),,T), which we save to dataset of transitions D, = [(zn.; = (æn,s, Tr (Inz)), Yo.i = In. )o We use the collected data to learn an estimate pnn of ft To this end, the goal of this work is to propose an algorithm Alg, that at each episode n leverages the data acquired thus far, i... Dim-1 to determine a policy n. E II for the next step of data collection, that is, Alg Alg(DD.... (Dim-11,7) n) + Wn. The proposed algorithm should be consistent, ie., p=( + f (z) for n + oo for all E R, where R is the reachability set defined as
R = fz E Z 3( E II, - T), s.t. p(zz = 2/TT. f³) > 0},
'https://github http://tttt..cc/a/ggggr///raaaaa
2

---

and efficient w.r.t. rate of convergence of pn to f".
To devise such an algorithm. we take inspiration from Bayesian experiment design (Chaloner and Verdinelli, 1995). In the Bayesian setting, given a prior over f* natural objective for active exploration is the mutual information (Lindley, 1956) between f and observations yD..
Definition (Mutual Information, Cover and Thomas (2006)). The mutual information between f* and its noisy measurements yD for points in D,, where yD. the concatenation of (yD,,,)eeT defined as,
P( (f = H (yD..) - (YD. f),
(3)
where is the Shannon differential entropy.
The mutual information quantifies the reduction in entropy of f conditioned the observations.
Hence, maximizing the mutual information w.r.t. the dataset 2,, leads to the maximal entropy reduction of our prior. Accordingly, a natural objective for active exploration in RL can be the mutual information between f and the collected transitions over a budget of N episodes, i.e., / (f"; yD..,) ). This requires maximizing the mutual information over a sequence of policies, which a challenging planning problem even in settings where the dynamics are known (Mutny et al., 2023). A common approach is to greedily pick a policy that maximizes the information gain conditioned on the previous observations at each episode:
max E,n [] (frz; yrr D Pi:n-1)] FEll
Here fü- = (fn ...., (*), y/- = (ym,0.-. ... Un.T-1); 7 is the trajectory under the policy π, and the expectation is taken w.r.t. the process noise w.
Interpretation in frequentist setting While information gain is Bayesian in nature (requires a prior over f*), it also has a frequentist interpretation. In particular, later in Section 3 we relate to the epis- temic uncertainty of the learned model. Accordingly, while this notion of information gain stems from Bayesian literature, we can use it to motivate our objective both Bayesian and frequentist settings.
2.1 Assumptions
In this work, we learn a probabilistic model of the function f from data. Moreover, at each episode n, we learn the mean estimator (x, u) and the epistemic uncertainty on( u). which quantifies our uncertainty on the mean prediction. To this end, we use Bayesian models such as Gaussian processes (GPs. Rasmussen and Williams, 2005) or Bayesian neural networks (BNNs, Wang and Yeung. 2020). More generally, we assume our model is well cell-calibrated:
Definition 2 (All-time calibrated statistical model of f*, Rothfuss et al. (2023)). Let, z (a u) and Z == X x U. An all-time calibrated statistical model of the function f is a sequence (Mh, nn (n(0)) nzw such that
Pr (V E 2,41E (1,..., dz), E :: [1n,(2)- Pu 1-8
Here Pn,1 and on, are the I-th element in the vector valued functions pn, and on respectively. The scalar function, (6) E R>o quantifies the width of the 5 confidence intervals. We assume w.l.o.g. that Bn monotonically increases with n, and that oz, S muxx all z E Z, n 2 0, and l (1....,d,)
Assumption 1 (Well calibration assumption). Our learned model is an all-time-calibrane statistical model of f" ie., there exists a sequence of (3x(6)) such that our model satisfies the well- calibrution condition, c.f., Definition 2.
This is a natural assumption on our modeling. It states that we can make a mean prediction and also quantify how far it is off from the true one with high probability. A GP model satisfies this requirement for a very rich class of functions, c.f., Lemma 3. For BNNs, calibration methods (Kuleshov et al., 2018) are often used and perform very well in practice. Next, we make a simple continuity assumption on our function f
Assumption 2 (Lipschitz. Continuity). The dynamics model f and our epistemic uncertainty prediction on are T.g and Log Lipschitz continuous, respectively. Moreover, we define TI be the policy class of Lx Lipschitz continuous functions.
3

---

The Lipschitz continuity assumption on f is quite common in control theory (Khalil, 2015) and learning literature (Curi et al., 2020; Pasztor et al., 2021; Sussex et al., 2023). Furthermore, the Lipschitz continuity of o, also holds for GPs with common kernels such as the linear or radial basis function (RBF) kernel (Rothfuss al., 2023).
Finally, we reiterate the assumption of the system's 's stochasticity.
Assumption 3 (Process noise distribution). The process noise is i..dd Gaussian with variance o²,
Le., we ~ N (0,
We focus the setting where w is homoscedastic for simplicity. However, our framework can also be applied to the more general heteroscedastic and sub-Gaussian case (c.f., Theorem 2).
Optimistic Active Exploration
In this section, propose our optimistic active exploration (OPAX) algorithm. The algorithm consists of two main contributions: (i) First we reformulate the objective in Equation (4) to a simple optimal control problem, which suggests policies that visit states with high epistemic uncertainty, (ii) We leverage the optimistic planner introduced by Curi et al. (2020) to efficiently plan policy under unknown dynamics. Moreover, we show that the optimistic planner crucial in giving theoretical guarantees for the algorithm.
3.1 Optimal Exploration Objective
The objective in Equation (4) is still difficult and expensive to solve in general. However, since in this work, we consider Gaussian noise, c.f., Assumption 3, we can simplify this further.
Lemma 1 (Information gain upper bounded by sum of epistemic uncertainties). Let y (2) uw with w ~ o2 and let dn-1 1 be the epistemic uncertainty after episode n 1. Then the following holds for all n 2 1 and dataset Din-b
d
o (
(ff+i Din--1
log
1+
(5)
j=1
We prove Lemma in Appendix A. The information gain is non-negative (Cover and Thomas, 2006). Therefore, if the right-hand side of Equation (5) goes to zero, the left-hand side goes to zero as well. Lemma 1 relates the information gain to the model epistemic uncertainty. Therefore, it gives tractable objective that also has frequentist interpretation collect points with the highest epistemic uncertainty. We can control problem:
plan a trajectory at each episode n, by solving the following optimal
)
T-1 d=
0 Pn π (z:)) ((z))) a²
T. arginax Jn (T) = argmax E,n
Σ
log
1
(6)
0++1 f* IT (x:)) + w.
The problem in Equation (6) closely related to previous literature in active exploration for RL. For instance, some works consider different geometries such the sum of epistemic uncertainties (Pathak et al. (2019); Sekar et al. (2020), c.f., appendix C for more detail).
3.2 Optimistic Planner
The optimal control problem in Equation (6) requires knowledge of the dynamics f for planning, however, f is unknown. A common choice is to use the mean estimator f4n-1 in Equation (6) instead f for planning "Buisson-Fened et al., 2020). However, in general, using the mean estimator is susceptible to model biases (Chua et al., 2018) and is provably optimal only in the case of linear systems (Simchowitz and Foster, 2020). To this end, we propose using an optimistic planner, as suggested in Curi et al. (2020), instead. Accordingly, given the mean estimator Mn-1 and the epistemic uncertainty o Fx-1, we solve the following optimal control problem
o ( T (2x)) (()2))
] 17)
nu, 7n argmax Ja(n,, argmax Eyn. *EIL,EE
ΣΣ
log 1
rEll.nEE
2
+ ( T (2t)) + (5)- n- (#,, n (2x))) + wt:

---

OPAX: OPTIMISTIC ACTIVE EXPLORATION
Init: Aleatoric uncertainty o, Probability 8, Statistical model (Ho, oo. Bo(o)) for episode n 1, , N do
T
dy
1n (Bt: π (a:t))
T. argmax max E MEI
log
1 +
Prepare policy
o²
1=0 1=1
DR + ROLLOUT(TT))
Collect measurements
Update (He, nn Bn())) Din
Update model
where E is the space of policies 17 X' -> [-1, Therefore, we use the policy 77 to "hallucinate" (pick) transitions that give the most information. Overall, the resulting formulation corresponds to a simple optimal control problem with larger action space, i.e., we increase the action space by another d, dimension. A natural consequence of Assumption 1 is that Jn (TT) S Jn.(TT 7n) with high probability (c.f., Corollary 1 in Appendix A). That is by solving Equation (7), we get an optimistic estimate on Equation (6). Intuitively, the policy TT, that OPAX suggests, behaves optimistically with respect to the information gain at each episode.
4 Theoretical Results
We theoretically analyze the convergence properties of OPAX. We first study the regret of planning under unknown dynamics. Specifically, since we camnot evaluate the optimal exploration policy from eq. (6) and use the optimistic one, i.e., eq. (7) instead, we incur regret. We show that due to the optimism in the face of uncertainty paradigm, we can give sample complexity bounds for the Bayesian and frequentist settings. All the proofs are presented in Appendix A.
Lemma 2 (Regret of optimistic planning under unknown dynamics). Let Assumption I hold. Fur- thermore, define Ja,k (T Th, x) as
de log 1=k
1 of 1n-15 (2,, n (x:))
Tnn())))
Jn,k( Fas Tn x) = E, Eyn.o
0²
s.t. 3++1 ( Tr (ā:)) + 0n ( we and to I.
Then, for all n 2 with probability at least 1-
(T( (Jn)-J((m (T( Σ Eymo [Jn,t+ nu; 2/411) J (Ta #141)]
with 01+11 f" (x)e)) + Wt:
and x/4,1 Mn- (I,, n. (x+)) + S
TI, (z;)) (xx)) wi
Lemma 2 gives bound the regret of planning optimistically under unknown dynamics. The regret is proportional to the difference in the expected returns for æ; and æt. Note, ||æ; - x'1 x (0) ( (x;-))) Hence, when we have low uncertainty in our predictions, planning optimistically suffers smaller regret. Next, we leverage Lemma 2 to give a sample complexity bound for the Bayesian and frequentist setting.
Bayesian Setting (2020).
We start by introducing a measure of model complexity as defined by Curi et al.
MCx
max D......Dx22.. x=1=ED.
1on-n
(
This complexity measure captures the difficulty of learning f given N trajectories. Mainly, the more complicated f", the larger the epistemic uncertainties o,,, and in turn, the larger corresponding MCN (. Moreover, if the model complexity measure sublinear N, i.e. MCx(f*)/N for N + then the epistemic uncertainties also converge zero in the limit, which implies
5

---

convergence to the true function f- We present our main theoretical result, in terms of the model complexity measure.
Theorem 1. Let Assumption and 3 hold. Then, for all 2 1, with probability at least 1 - 6,
max Ey [1 (f7=i PLN- So 8xT%
(f
E
(9)
N
Theorem I relates the maximum expected information gain iteration N to the model complexity of our problem. For deterministic systems, the expectation w.r.t. T" is redundant. The bound in Equation (9) depends the Lipschitz constants, planning horizon, and dimensionality of the state space (captured in By and MCx(f*)) If the right-hand side is monotonically decreasing with N, Theorem guarantees that the information gain at episode N is also shrinking with N, and the algorithm is converging. Empirically, Pathak et al. (2019) show that the epistemic uncertainties go to zero as data is acquired. In general, deriving worst-case bound on the model complexity is challenging and active open research problem. However, in the case of GPs, convergence results can be shown for very rich class of functions. We show this in the following for the frequentist setting.
Frequentist Setting with Gaussian Process Models We extend our analysis to the frequentist kemnelized setting, where f resides in Reproducing Kernel Hilbert Space (RKHS) of vector-valued functions.
Assumption 4. We assume that the functions fj. j (l....,dy lie in a RKHS with kernel k and have a bounded norm B, that is f E HH with Hy, = {f IIfsll S B,j = 1,...., dz
In this setting, we model the posterior mean and epistemie uncertainty of the vector-valued function f with (z) = [unr())lisd_ and (= = (on.(=)lisad, where,
Pnu (2) = k (z) K, +021)= - ) (2) k(æ w) kij (=)(K, +
(10)
Here, vi.n corresponds to the noisy measurements of fj. i.e., the observed next state from the transitions dataset Di... kn, = 2:))isTT,
Din, and K, = 21)l.ssTT Zi, 2) E Dio is the data kernel matrix. It is known that if f satisfies Assumption 4, then Equation (10) yields well-calibrated confidence intervals, i.c., that Assumption 1 is satisfied.
Lemma (Well calibrated confidence intervals for RKHS, Rothfuss al. (2023)). Let f E H Suppose μa, and nn are the posterior mean and variance of GP with kernel cf., Equation (10). There exists (8)(8). for which the tuple (pn. on (,)) satisfies Assumption I w.r.. function f*.
Theorem presents convergence guarantec for the kernelized case to the T-step reachability set R. for the policy class π E IT. In particular, R is defined as
R= = fz E 2 3( E II, < 7), s.t., p.t., 2| > 0}
There two key differences from Theorem 1; (i) we derive upper bound the epistemic uncertainties ou, and (ii) we can bound the model complexity MCN (f*), with the maximum information gain of kernel i introduced by Srinivas et al. (2012), defined
(k)
log det +o -==²kx)
Di..... D.....DD: D I ST
Theorem 2. Let Assumption 3 and 4 hold, Then, for all N 2 1, with probability at least
de
< s0 (BNTT
max
Ey WEll
max ZET
N.j (z)
(11)
N
1=1
If we relax noise Assumption 3 to o-sub Gaussian. Then, if Asssupption holds, we have for all N 1, with probability least 8,
d,
5 (8hr21 ( (w)(k) (z)
max E,- FENI
max SET*
2 j
N.j
N
(12)

---

Moreover, if (k) = O (polylog(N)), then for all = E R, and 1 S S dz.
(N)( -اك for N +
00. co.
(13)
We only state Theorem 2 for the expected epistemic uncertainty along the trajectory at iteration N. For deterministic systems, the expectation is redundant and for stochastic systems, concentration inequalities to give a bound without the expectation (see Appendix A for more detail).
leverage
For the Gaussian noise case, obtain tighter bound by leveraging the change of measure inequality from Kakade et al. (2020, Lemma C.2.) (c.f., Lemma in Appendix A for more detail). In the more general case of sub-Gaussian noise, we cannot use the same analysis. To this end, we use the Lipschitz continuity assumptions (Assumption 2) similar to Curi et al. (2020). This results in comparing the deviation between two trajectories under the same policy and dynamics but different initial states (see Lemma 2). For many systems (even linear) this can grow exponentially in the horizon T. Accordingly, we obtain a 3% term in our bound (Equation (12)). Nonetheless, for cases where the RKHS is of a kernel with maximum information gain 7xx (k) = O (polylog(N)), we can give sample complexity bounds and an almost sure convergence result in the reachable set R (Equation (13)). Kernels such as the RBF kernel or the linear kernel (kernel with a mnite-dimensional feature map ((z)) have maximum information gain which grows polylogarihhmicaly with n (Vakili et al. (2021)) )) Therefore, our convergence guarantees hold for a very rich class of functions. The exponential dependence of our bound on T imposes the restriction on the kernel class. For the case of Gaussian noise, we can include a richer class of kernels, such as Matèrn.
In addition to the convergence results above, we also give guarantees on the zero-shot performance of OPAX in Appendix A.5.
Experiments
We evaluate OPAX on the Pendulum-vl and MountainCar environment from the OpenAl gym benchmark suite (Brockman et al., 2016), on the Reacher, Swimmer, and Cheetah from the deep mind control suite (Tassa et al., 2018), and a high-dimensional simulated robotic manipulation task introduced by Li et al. (2020). See Appendix B for more details on the experimental setup.
Baselines We implement four baselines for comparisons. To show the benefit of our intrinsic reward, we compare OPAX to (/) a random exploration policy (RANDOM) which randomly samples actions from the action space. As we discuss in Section 3 our choice of objective in Equation (6) is in essence similar to the one proposed by Pathak et al. (2019) and Sekar et al. (2020). Therefore, compare the optimistic planner with other planning approaches. Moreover,
in our experiments, most work on active exploration either uses the mean planner or does not specify the planner (c.f., Section 6). We use the most common planners: (2) mean AEE) and (3) trajectory sampling (TS-1) scheme proposed in Chua et al. (2018) (PETS- AAE) as our baselines. The mean planner simply uses the mean estimate μt, of the well-calibrated model. This is also used in Buisson- Fenet et al. (2020). Finally, we compare OPAX to (4) H-UCRL (Curi et al., 2020), a single-task model- RL algorithm. We investigate the following three aspects: (i) how fast does active exploration reduce model's 's epistemic uncertainty T, with increasing n. (i) can we solve downstream tasks with OPAX, and (iii) does OPAX scale to dineeniimnsion and challenging object manipulation tasks? For our experiments, we use GPs and probabilistic ensembles (PE, Lakshminarayanaa et al. (2017)) for modeling the dynamics. For the planning, we either the soft actor-critic (SAC, Haarnoja al. (2018)) policy optimizer, which takes simulated trajectories from our learned model to train a policy, or MPC with the iCEM optimizer (Pinneri et al., 2021).
How fast does active exploration reduce the epistemic uncertainty? For this experiment, we consider the Pendulum-vi environment. We sample transitions at random from the pendulum's reachable state-action space and evaluate our model's epistemic uncertainty for varying episodes and baselines. We model the dynamics with both GPs and PE. We depict the result in Figure 1. We conclude that the RANDOM agent is slower in reducing the uncertainty compared to other active exploration methods for both GP and PE models. In particular, from the experiment, we empirically validate Theorem for the GP case and also conclude that empirically even when using PE models, we find convergence of epistemic uncertainty. Moreover, we notice for the PE case that OPAX reaches smaller uncertainties slightly faster than MFAN-AE and PETS-AT We believe this due to the additional exploration induced by the optimistic planner.
7

---

Pendulum-vl GP unneraainty
Penculuu-vl uncertainty
10%
10"
20
10--3
10
10
Epieodes
Episodes
PETS-AE
MEAN-AE
OPAX
RANDOM
Figure 1: Reduction in maximum epistemic uncertainty in reachable state-action space for the Pendulum-vI environment over 10 different random seeds. We evaluate OPAX with both GPs and PE and plot the mean performance with two standard error confidence intervals. For both, active exploration reduces epistemic uncertainty faster compared random exploration. All active exploration baselines perform well for the GP case, whereas for the PE case OPAX gives slightly lower uncertainties.
Can the model learnt through OPAX solve downstream tasks? We use OPAX and other active exploration baselines to actively learn a dynamics model and then evaluate the learned model on downstream tasks. We consider several tasks. (i) Pendulum-vl swing up, (it) Pendulum-vl keep down (keep the pendulum at the stable equilibria), (iii) MountainCar, (iv) Reacher go to target, (v) Swimmer go to target, (vi) Swimmer- go away from target (quickly go away from the target position), (vii) Cheetah run forward, (viiz) Cheetah run backward. For all tasks, we consider PEs, except for (i) where we also use GPs. Furthermore, for the MountainCar and Reacher, we give reward once the goal is reached. Since this requires long-term planning, we use a SAC policy for these tasks. We use MPC with iCEM for the remaining tasks. We also train H-UCRL on tasks (i) with GPs, and (ii, (iii), (iv), (v), (vil) with PEs. We report the best performance across all episodes.
To make fair comparison, exploration for each episode on the environment, and then after every few episodes we use the mean estimate nn to evaluate our learned model on the downstream tasks.
the following evaluation procedure; first, we perform active
Figure 2 shows that all active exploration variants perform considerably better than the RANDOM agent. In particular, for the MountainCar, the RANDOM agent is not able to solve the task. Moreover, PETS-AE performs slightly worse than the other exploration baselines in this environment. In general, notice that OPAX always performs well and is able to achieve H H-UCRL's performance on all the tasks for which H-UCRL is trained. However, on tasks that are new/unseen for H-UCRL, active exploration algorithms outperform H-UCRL. From this experiment, we conclude two things (/) apart from providing theoretical guarantees, the model learned through OPAX also performs well in downstream tasks, and (2) active exploration agents generalize well to downstream tasks, whereas H-UCRL performs considerably worse on new/unseen tasks. We believe this is because, unlike active exploration agents, task-specific model-based RL agents only explore the regions of the state-action space that are relevant to the task at hand.
Does OPAX scale to high-dimensional and challenging object manip- ulation tasks? To answer this question, consider the Fetch Pick & Place Construction environment (Li et al., 2020). We again use the active exploration agents learn model and then evaluate the success rate of the learned model in three challenging downstream tasks: (i) Pick & Place, (ii) Throw, and (iii) Flip (see Figure 4). The environment contains 7-DoF robot arm and four 6-DoF blocks that can be manipulated. In total, the state space is 58-dimensional. The 4-dimensional actions control the end-effector of the robot in Cartesian space as well as the opening/closing of the gripper. We compare OPAX to PETS-AE, MEAN-AE, random policy as well as CEE-US (Sancaktar et al., 2022). CEE-US is model-based active exploration algorithm, for which Sancaktar et al. (2022) reports mate-of-the-ar perfor- mance compared to several other active exploration methods. In all three tasks, OPAX least on par with the best-performing baselines, including CHE-US. We run OPAx and all basclines with the same architecture and
Figure 3: Fetch Pick & Place Construc- tion environment.
hy hyperparameter settings. See Appendix B for more details.
8

---

Downstream tasks where H-UCRL is trained
New downstream tasks
MountainCar (1)
Pendulum-vl swing up (7)
x10
Pendulum-yl keep down (4)
x10
10
10
-10-
0.55 0.5
10F
1.5
10
20
10
20
Pendulum-vl GP swing (2)
Swimmer go to target (5) x10F
Swimmer go away from target (8
-2
2
31
10
100
100
150
Reacher go to target (3)
Cheetah run forward (6)
Cheetah run backward (9) x10
x10
x10²
3-
50
100
0
200 Episodes
400
0
200
400
Episedes
Episodes
PETS-AE
MEAN-AE
OPAx
RANDOM
H-UCRL
We evaluate the downstream performance of our agents over 10 different random seeds and plot the mean performance with two standard eтror confidence intervals. For all the environments we use PE as models, except plot (1), for which we use a GP model (see plot (2) in the figure above). For tasks (1)-(6), we also train H-UCRL, model-based RL algorithm. Tasks (7)-(9) are new/unseen for H-UCRL. From the Figure, we conclude that (?) compared to other active exploration baselines, OPAX constantly performs well and is on par with H-UCRI, and (ii) on the new/unseen tasks the active exploration baselines and OPAX outperform H-UCRL. by a large margin.
Figure 2:
Related Work
System identification is broadly studied topic (Aström and Eykhoff, 1971; Schoukens and Ljung, 2019; Schön et al., 2011; Ziemann et al., 2022; Ziemann and Tu, 2022). However, system identifica- tion from the perspective of experiment design for nonlinear systems is much less understood (Chiuso and Pillonetto, 2019). Most methods formulate the identification task through the maximization of intrinsic rewards. Common choices of intrinsic rewards (/) model prediction error "Curios- ity" (Schmidhuber, 1991; Pathak et al., 2017), (ii) novelty of transitions (Stadie et al., 2015), and (iii) diversity of skills (Eysenbach al., 2018).
popular choice for intrinsic rewards is mutual information or entropy (Jain al., 2018; Buisson- Fenet et al., 2020; Shyam et al., 2019: Pathak et al., 2019; Sekar et al., 2020). Jain et al. (2018) propose an approach maximize the information gain greedily wrt the immediate next transition. i... one-step greedy, whereas Buisson-Fenet et al. (2020) consider planning full trajectories. Shyam al. (2019); Pathak et al. (2019); Sckar et (2020) and Sancaktar al. (2022) consider general Bayesian models, such BNNs, to represent probabilistic distribution for the learned model. Shyam et al. (2019) propose using the information gain of the model with respect to observed transition as the intrinsic reward. To this end, they learn an ensemble of Gaussian neural networks and represent the distribution over models with Gaussian mixture model (GMM), similar approach is also proposed in Pathak et al. (2019); Sekar et al. (2020); Sancaktar et al. (2022). The main difference between Shyam al. (2019) and Pathak al. (2019) lies how they represent mutual information. Moreover, Pathak et al. (2019) use the model's epistemic uncertainty, that is the

---

Pick & Place Objects
Throw 4 Objects
Flip Objects
100
40
50
20-
15-
100
200
300
100
2000
300
100
200
300
Training Iteration
Training Ireration
Training Iteeation
RANDOM
MEAN-AE
PETS-AE
OPAX
CEE-US
Figure 4: Success rates for pick & place, throwing and flipping tasks with four objects in the Fetch Pick & Place Construction environment for OPAX and baselines. We evaluate task performance via planning zero-shot with models learned using different exploration strategies. We report performance on three independent seeds. OPAX is on par with the best-performing baselines in all tasks.
disagreement between the ensemble models as an intrinsic reward. Sekar et al. (2020) link the model disagreement (epistemic uncertainty) reward to maximizing mutual information and demonstrate state-of- of-tthe-rttttttte performance on several high-dimensional tasks. Similarly, Sancaktar et al. (2022), use the disagreement in predicted trajectories of an ensemble of neural networks to direct exploration. Since trajectories can diverge due to many factors beyond just the model epistemic uncertainty, c.g., aleatoric noise, this approach is restricted deterministic systems and susceptible systems with unstable equilibria. Our approach is the most similar to Pathak et al. (2019); Sekar et al. (2020) since also propose the model epistemic uncertainty as the intrinsic reward for planning. However, we thoroughly and theoretically motivate this choice of reward from a Bayesian experiment design perspective. Furthermore, we induce additional exploration in OPAX through our optimistic planner and rigorously study the theoretical properties of the proposed methods. On the contrary, most of the prior work discussed above either uses the mean planner (MEAN-AB. or does not discuss the planner thoroughly or provide any theoretical results. In general. theoretical guarantees for active exploration algorithms are rather immature (Chakraborty al., 2023; Wagenmaker et al., 2023) and mostly restrictive to a small class of systems (Simchowitz et al., 2018; Tarbouriech et al., 2020; Wagenmaker and Jamieson, 2020; Mania et al., 2020). To the best of our knowledge, we are the first to give convergence guarantees for rich class of nonlinear systems.
While work focuses the active learning of dynamics, there are numerous works that study exploration in the context of reward-free RL (Jin et al., 2020; Kaufmann et al., 2021; Wagenmaker al., 2022; Chen et al., 2022). However, most methods in this setting give guarantees for special classes of MDPs (Jin et al., 2020; Kaufmann et al., 2021: Wagenmaker et al., 2022: Qiu et al., 2021: Chen al., 2022) and result in practical algorithms. On the contrary, we focus on solely learning the dynamics. While good dynamics model may be used for zero-shot planning, it also exhibits more relevant knowledge about the system such as its stability or sensitivity external effects. Furthermore, our proposed method is not only theoretically sound but also practical.
7 Conclusion
We present OpAX, novel model-based RL. algorithm for the active exploration of unknown dynamical systems. Taking inspiration from Bayesian experiment design, we provide a compre- hensive explanation for using model epistemic uncertainty an intrinsic reward for exploration. By leveraging the optimistic in the face of uncertainty paradigm, we put forth first-of-their-king theoretical results the convergence of active exploration agents in reinforcement learning. Specifically, we study convergence properties of general Bayesian models, such as BNNs. For the frequentist case of RKHS dynamics, we established sample complexity bounds and convergence guarantees for OPAX for a rich class of functions. We evaluate the efficacy of OPAX across various RL environments with state space dimensions from two 58. The empirical results corroborate our theoretical findings, as OPAX displays systematic and effective exploration across all tested environments and exhibits strong performance in zero-shot planning for new downstream tasks.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Jonas Hübotter for the insightful discussions and his feedback on this work. Furthermore, also thank Alex Hägele, Parnian Kassraie, Scott Sussex, and Dominik Baumann for their feedback.
This project has received funding from the Swiss National Science Foundation under NCCR Automa- tion, grant agreement 51NF40 180545, and the Microsoft Swiss Joint Research Center.
References
Balcan, M.-F., Hanneke, S., and Vaughan, W. (2010). The true sample complexity of active learning. Machine learning, 80:111-133
Berkenkamp. F., Turchetta, M., Schoellig, A., and Krause, (2017). Safe model-based reinforcement learning with stability guarantees. NeurIPS, 30.
Biagiotti, L.. and Melchiorri, C. (2008). Trajectory Planning for Automatic Machines and Robots. Springer Publishing Company, Incorporated, 1st edition.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke,
A., VanderPlas, J., Wanderman-Minn Milne, S., and Zhang, Q (2018). JAX: composable transformations of Python+NumPy programs.
Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W. (2016). Openai gym. arXiv preprint 166.0116000011400..
Buisson-Fenet, M., Solowjow, F., and Trimpe, S. (2020). Actively learning gaussian process dynamics. In Bayen, A. M., Jadbabaie, A., Pappas, G., Parrilo, P. A., Recht, B., Tomlin, C., and Zeilinger, M., editors, L4DC.
Chakraborty, S., Bedi, A., Koppel, A., Wang, M., Huang, F., and Manocha, D. (2023). STEERING Stein information directed exploration for model-based reinforcement learning. pages 3949-3978.
Chaloner, K. and Verdinelli, (1995). Bayesian experimental design: review. Statistical science, pages 273-304.
Chen, J., Modi, A., Krishnamurthy, A., Jiang, N., and Agarwal, A. (2022). On the statistical efficiency of reward--fre -free exploration in non-linear rl. Advances in Neural Information Processing Systems, 0-20960-209773
Chen, Y., Hassani, S. H., Karbasi, A., and Krause, A. (2015). Sequential information maximization: When is greedy near-optimal? In COLT.
Chiuso, and Pillonetto, G. (2019). System identification: machine learning perspective. Annual Review of Control, Robotics, and Autonomous Systems, 2-24.
Chowdhury, S. and Gopalan, A. (2017). On kernelized multi-armed bandits. In ICML, pages 844 853. PMLR.
Chua, K., Calandra, R., McAllister, R., and Levine, S. (2018). Deep reinforcement learning in handful of trials using probabilistic dynamics models. NeurIPS.
Cover, T. M. and Thomas, J. A. (2006). Elements of Information Theory (Wiley Series in Telecommu- nications and Signal Processing). Wiley-Interscience
Curi, S., Berkenkamp, F., and Krause, A. (2020). Efficient model-based reinforcement learning through optimistic policy search and planning. NeurIPS, 14117006-4170
Eysenbach, B., Gupta, A., Iharz, J., and Levine, S. (2018). Diversity is all you need: Learning skills without reward function. arXiv preprint /
García, C. E., Prett, M., and Morari, M. (1989). Model predictive control: Theory and practice survey, Automatica, pages 335-348.

---

Haarnoja, T.. Zhou, A., Abbeel, P., and Levine, S. (2018). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICML, pages 1861-1870.
Hanneke, S. et al. (2014). Theory of disagreement- active learning. Foundations and Trends@ in Machine Learning, 7(2-3):111-309
Jain, A., Nghiem, T., Morari, M., and Mangharam, R. (2018). Learning and control using gaussian processes. In ICCPS, pages 140-149.
Jin, C., Krishnamurthy, A., Simchowitz, M., and Yu, T. (2020). Reward-free exploration for reinforcement learning. In ICML, pages 4870-4879, PMLR.
Kakade, S., Krishnamurthy, A., Lowrey, K., Ohnishi, M., and Sun, W. (2020). Information theoretic regret bounds for online nonlinear control. NeurIPS, 33:15312-15325.
Kaufmann, E., Ménard, P., Domingues, O, D., Jonsson, A., Leurent, E., and Valko, M. (2021). Adaptive reward-free exploration. ALT, pages 865-891. PMLR.
Khalil, H. K. (2015). Nonlinear control, volume 406. Pearson New York.
Krause, A., Singh, A., and Guestrin, C. (2008). Near-optimal sensor placements in gaussian processes: Theory, efficient algorithms and empirical studies. Journal of Machine Learning Research, 9(2).
Kuleshov, V., Fenner, N., and Ermon, S. (2018). Accurate uncertainties for deep learning using calibrated regression. ICML, pages 2796-2804. PMLR.
Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles. NeurlPS, 30.
Li, R., Jabri, A., Darrell, T., and Agrawal, (2020). Towards practical multi-object manipulation using relational reinforcement learning. In ICRA, pages 4051-4058. IEEE.
Lindley, D. V. (1956). On Measure of the Information Provided by an Experiment. The Annals of Mathematical Statistics, pages 986 1005.
Mania, H., Jordan, M. I., and Recht, B. (2020). Active learning for nonlinear system identification with guarantees. arXiv preprint arXiv:2006. 10277.
Mehta, V., Paria, B., Schneider, J., Ermon, S., and Neiswanger, W. (2021). An experimental design
perspective model-based reinforcement learning. arXiv preprint arXiv:21
Mutny, M., Janik, T., and Krause, A. (2023). Active exploration via experiment design markov chains. In AISTATS.
Pasztor, B., Bogunovic, I., and Krause, (2021). Efficient model-based multi-agent mean-fiell reinforcement learning. arXiv preprint arXiv:2107.040500
Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. (2017). Curiosity-driven exploration by sule-super prediction. In ICML, pages 2778-2787. PMLR.
Pathak, D., Gandhi, D., and Gupta, A. (2019). Self-supervised exploration via disagreement. In ICML, pages 5062-5071 PMLR.
Pinneri, C., Sawant, S., Blaes, S., Achterhold, J., Stueckler, J., Rolinek, M., and Martius, G. (2021). Sample-efficient cross-entropy method for real-time planning. In CORL, Proceedings of Machine Learning Research, pages 1049-1065.
Qiu, S.. Ye, J.. Wang, Z., and Yang, Z. (2021). On reward-free rl with kernel and neural function approximations: Single-agent mdp and markov game. ICML, pages 8737-8747. PMLR.
Rasmussen, E. and Williams, C. K. I. (2005). Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press.
Åström, K. and Eykhoff, P. (1971). System identification survey, Automatica, 122):123-
12

---

Rothfuss. J., Sukhija, B., Birchler, T., Kassraie, P., and Krause, A. (2023). Hallucinated adversarial control for conservative offline policy evaluation. UAL.
Sancaktar, C., Blaes, S., and Martius, G. (2022). Curious exploration via structured world models yields zero-shot object manipulation. In NeurIPS 35 (NeurIPS 2022).
Schmidhuber, (1991). possibility for implementing curiosity and boredom in model-building neural controllers. In Proc. of the international conference simulation of adaptive behavior: From animals to animats, pages 222-227.
Schön, T. B., Wills, A., and Ninness, B. (201 System identification of nonlinear state-space models. Automatica, pages 39-49.
Schoukens, and Ljung. L. (2019). Nonlinear system identification: A user-oriented road map. IEEE Control Systems Magazine, 39(6):28-99.
Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, D. (2020). Planning to explore via self-supervised world models, In ICML, pages 8583-8592. PMLR.
Settles, B. (2009). Active learning literature survey.
Shyam, P., Jaškowski, W., and Gomez, (2019). Model-based active exploration. ICML, pages 5779-5788. PMLR.
Simehowitz, M. and Foster, D. (2020). Naive exploration is optimal for online LQR. In ICML, Proceedings Machine Learning Research, pages 8937-8948. PMLR.
Simchowitz, M., Mania, H., Tu, S., Jordan, M. I., and Recht, B. (2018). Learning without mixing: Towards a sharp analysis of linear system identification. In COLT, pages 439-473. PMLR.
Srinivas, N., Krause, A., Kakade, S. M., and Seeger, M. W. (2012). Information-theerrtic regret bounds for gaussian process optimization in the bandit setting. IEEE Transactions Information Theory.
Stadie, B. C., Levine, S., and Abbeel, P. (2015). Incentivizing exploration in reinforcement learning with deep predictive models. arXiv preprint 4rXiv:1507.00814.
Sussex, S., Makarova, A., and Krause, A. (2023). Model-based causal bayesian optimization. In ICLR.
Tarbouriech, J., Shekhar, S., Pirotta, M., Ghavamzzdeh, M., and Lazaric, A. (2020). Active model estimation in markov decision processes. In UAI, pages 1019-1028. PMLR.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., Casas, D. d. L., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., et al. (2018). Deepmind control suite. arXiv preprint arXiv:1801.0069
Vakili, S., Khezeli, K., and Picheny, V. (2021). On information gain and regret bounds in gaussian process bandits. In AISTATS.
Wagenmaker, A. and Jamieson, K. (2020). Active learning for identification of linear dynamical systems. In COLT, pages 3487-3582. PMLR.
Wagenmaker, A., Shi, G., and Jamieson, K. (2023). Optimal exploration for model-based in nonlinear systems. arXiv preprint arXiv:2306.092. 092/0.
Wagenmaker, A. J., Chen, Y., Simchowitz, M., Du, S., and Jamieson, K. (2022). Reward-free no harder than reward-aware rl in linear markov decision processes. In ICML, pages 22430-22456 PMLR.
Wang, H. and Yeung, D.-Y. (2020). A survey on bayesian deep learning. ACM computing surveys (csur), pages 1-37.
Weiss, K., Khoshgoftaar, M., and Wang, D. (2016). A survey of transfer learning. Journal of Big data, pages 1-40.
13

---

Zhang, Y. and Yang, Q. (2021). A survey on multi-task learning, IEEE Transactions on Knowledge and Data Engineering. pages 5586-5609.
Ziemann, and Tu, S. (2022). Learning with little mixing. NeurIPS, 35:4626-4637.
Ziemann, M., Sandberg, H., and Matni, N. (2022). Single trajectory nonparametric learning of nonlinear dynamics. In COLT, pages 3333-3364 PMLR.
14

---

Contents of Appendix
Proofs for section 4
16
A.1 Proof of Lemma 2
17
A.2
Analyzing regret of optimistic planning
18
A.3 Proof for general Bayesian models
20
A.4 Proof of GP results
23
A.5
Zero-shot guarantees
26
B Experiment Details
27
B.I Environment Details
27
B.2
OPAX in the High-dimensional Fetch Pick & Place Environment
29
Study of exploration intrinsic rewards
32
15

---

A Proofs for section 4
We first prove some key properties of our active exploration objective in Equation (6). Then, we prove Theorem which holds for general Bayesian models, and finally we prove Theorem 2, which guarantees convergence for the frequentist setting where the dynamics are modeled using a GP.
Lemma 4 (Properties of OPAX's objective). Let Assumption I and 2 hold, then the following is true for all n 2
1+ o - 1, (oe, (x²))
log
20
g²
27
T-1 de
1 o on-1,(2t, π (Tx)))
max
sup IEEL,
E
Σ log
S T de log?
1
(2)
0²
x
0²
j=1
di
(2)
d maxloo
log
1+
log
1+
|| - z'1l-
2
x2
g²
j=1
where Fmax = 2.2. (z).
Proof. The positivity of the reward follows from the positive definiteness of the epistemic uncertainty On-1.j- For (2), the following holds
27
T du log
1 02 1n (I T (xx)))
E
a2
t=0
j=1
T-1 1 d2
1+ o -n π (x;))
s E
Σ
Td log?
0²
j=1
de
22 1
< E
Σ
Th log²
max o²
j=1
Fuax o2
< T de log?
1+
From hereon, let Jmax 1/2 log? (1 + a 20 max)
Finally, we show that this reward is Lipschitz continuous.
[ong(()) oz')(²)) = (2..(2))nn - ons (2)(2))n)(2)) + on (=n()')onn
< Log ||z z'll (z))))(()) + Lo 1|z - 2'll (2n)
< 20 1onnet Loo 1|2 - -211..
(2) o²
log (1t a²e (2') 1+
1
d
02 (2)
log
1+
log
1
o²
o²
2
1 + 2x'
1
d
2 n.j
(z) o (21)
<
log
1+
2
j=l
onn (2) o (2')l
2
20
d d Luo
||2 - 1z'1
Where (*) is true because for all 2 0, log(1 + x) s
16

---

Corollary (OPAX gives an optimistic estimate on Equation (6)). Let Assumption hold and n denote the solution to Equation (6) and (T)) the resulting objective. Similarly, let n and nh be the solution to Equation (7) and [n(Tn: 7m) the corresponding value of the objective. Then with probability least 6 we have for every episode (1,...,N]. N]:
Jn Jn (T {m))
Proof. Follows directly from Assumption 1.
A.1 Proof of Lemma
Lemma 5 (Difference in Policy performance). Let Jr (T, = E,- ETS r(æ, ((x)) and
A (T, x. a) E, (r( a) + Jr.x-I(T, ( (T, x)] with x' fxx a) + w. For simplic-
ity we refer to Jn, 22) = (T, The following holds for all xn E X::
T-1
( - Jr( = E,
Σ
(T, a'y क (x())
Proof.
Jr(T'' = =E,=' Σ r( x (a\)) = Ey [r( x (xo)) + 1 (x11
-.,,' (r(wn, x' (an)) + Jr,1 (TT, a²4) + J,,1 a1)
Jr,1 (T,, x1)
B, (r(zo, N' (zo)) + ( x11) - (T,, 2o)]
[T 20) En |Jr,1 (T', x²1) - (
-, (A,,OTT, 2o N' (2r) (TT,, mn) B [Jr, (T', Jr,1 (TT,
Therefore we obtain
J,(T',, 2n) (, xn) = Epr [Ao(T, Tu ा (xx))) + Epr ( (T',,')) - Jr.1 (TT, x1))
Using the same argument for Jr,1, Jr.2 .... JrTT-1 and that (n,((T, x) 0 for all E and E X' completes the proof.
Assume policy is fixed and dynamics are of the form:
x' = (2. (x)) + (ā)or π r<)) w.
(14)
Here u E {]=,, Furthermore, assume that the associated running rewards do not depend on u, that is, r(z;), and let v E E denote the policy, i.e., n X' + [-1,1]
Corollary 2. The following holds for all E X and policy ग::
T-1
Σ Jr (T,, n. x'41) - Jr.tii(T n7 27, 1=
Jr zu)) Jr,(T m7 (T, xo) = E
with x(+1 = u,, ग (x()) + Bn (8) (a()))(()) + wi and 2/41 Mn(22 T (xx))) + (a)o (z/. π. (z())r (a() + wt.
Proof. From Lemma 5 we have
J, xb) J, m. (00) = Eyn'
Σ
And An. (n. z((,)((2()))
Furthermore,
Eyn' (n, ,/,,,' (x(x)))) = Epar [r(x() Jn n. Jr.( x7
-Ep [r(x() Jr.t1 (7, r(x)) - (
n.
Epp [Jr,41 (T, n. - 01+1))]
17

---

Proof of Lemma 2. From Assumption we know that with probability at least 1 there exists such that f (z) (=) 3, (z) (z))(x) for z E Z.
IT(TT - (TT) S Jn (T mm) - Jn (T
(Corollary 1)
Jn n. mm) Jn (T 77)
Ey
Inat+ TT, nn l((11) - Un,t2(TT Ths L++1)
(Corollary 2)
= E, Σ
Jasti (Tn Thi x/41) Jn (T,, Ta Tt+1))
(Expectation wrt n, under true dynamics f*)
with x++1 = (2 Tr Tn (xx)) + wt:
and r//1 pn (aee; nn (a=+) + 3n-1 (2%, nn (m:) + wn.
A.2 Analyzing regret of optimistic planning
In the following, we analyze the regret of optimistic planning for both o-Gaussian noise and o-sub Gaussian noise case. We start with the Gaussian case.
Lemma 6 (Absolute expectation Difference Under Two Gaussians (Lemma C.2. Kakade et al. (2020))). For Gaussian distribution N(u1, and (122, and for any (appropriately mea- surable) positive function g. it holds that:
BENAN [g(2] PewN)lo(=)) & min {
Ill1 - -12ll .1 v EEENN [g²(2)]
o²
Proof.
BgrMi g(z) (1 (1- .22
Ni }
E [9(=)]
Be [9(2)l1
(1-
.25 27
<
EznNi [g((2]]
Be
N1
(1
N2 N
BEENs [o((
EenAis
uzl σ² (Lemma C.2. Kakade et al. (2020)
||
S v BEmnN (2) min
Corollary 3 (Regret of optimistic planning for Gaussian noise). Let πI n. denote the solution Equation (6) and Equation (7) respectively, and zh.r Zn,t the corresponding state state-action pairs visited during their respective trajectories. Furthermore, let Assumption 1 hold. Then, the following is true for all n 2 0, t E [o.T - 1). with probability at least 1
(1+ Vd,) (8) I 11 (2n+))) 1) g²
JT(TT;)) - (T( 50 TEy.n
Σ
18

---

Proof. For simplicity, define In (x) = Jn,411 Th. x). Note since w, ~ (V o (Assump- tion 3). we have that 2/41 and x+41 are also Gaussians. Therefore, we can leverage Lemma
E Eyan
[Jn,t+1 (Ty Th 2/41)
Jn,t-1 ( 7h x++1)] E (gx(2<(11) In (2(+1)
x'441 1t+1
El Elgh (z4 min
(Lemma 6)
|0/441 xc+1
S v J Jmax min
(Lemma 4)
0²
Furthermore,
|4/41 Et+1
Mn- De, n n (2+)) )} Bn- 1 16 0 1 n (n()))))(( f TT, (x:))))
Mn-1 (Ze, n (xx)) f (x,, n. (xx)))l
+ Bn (6) || - (2) nr, (xx ))] 7m (a:))]
(1 + v d2) 1 (x:)
(Assumption 1)
Next, we use Lemma
(TT - (TT) v nn
Jn,ti Mns 2l(41) Jn,.t1 Tn Pett))
[ v1 3n |1 1 Tr (20)))]
E n
Jmnx min
(1+ Vd Bn (6)
o (2t, n (2x
S V Jmax Bn
r²
(1+ v da B, (0)
I (Ii, Tn (n(())))]
O Teprn
g²
Lemma 7 (Regret of planning optimistically for sub-Gaussian noise). Let Tē, nu denote the solution to Equation (6) and Equation (7) respectively, and z Zn, the corresponding state-action pairs visited during their respective trajectories. Furthermore, let Assumption I and hold, and relax Assumption 3 to o-sub Gaussian noise. Then, the following is true for all n 2 0 with probability at least 1 8
(n(T))) - Jn (n())) S
LT 18
Σ (om-1w((m:)))
Proof. Curi et al. (2020, Lemma 5) bound the regret with the sum of epistemic uncertainties for Lipschitz continuous reward functions, under Assumption and 2 for uub-Gaussian noise (c.f., Rothfuss et al. (2023, Theorem 3.5) for more rigorous derivation). For the active exploration setting, the reward in episode
1 de
of ) 1
r(z)
Σ log
2
o²
We show in Lemma 4 that choice of exploration reward is Lipschitz continuous. Thus, can the regret bound from Curi et al. (2020).
Compared to the Gaussian case, o-bub Gaussian noise has the additional exponential dependence on the horizon T, i.e. the 8T term. This follows from the analysis through Lipschitz continuity. Moreover, we show in Lemma 2, the regret of planning optimistically is proportional to the change value under the same optimistic dynamics and policy, but different initial states. The Lipschitz continuity property of our objective allows us to relate the difference in values to the discrepancy in the trajectories. Even for linear systems, trajectories under the same dynamics and policy but different initial states can deviate exponentially in the horizon.
19

---

A.3 Proof for general Bayesian models
In this section, we analyze the information gain for general Bayesian models and prove Theorem 1.
Theorem 3 (Entropy of a RV with finite second moment is upper bounded by Gaussian entropy (Theorem 8.6.5 Cover and Thomas (2006))). Let the random vector In have covariance K E vx (i.e., Kig- = [,, S i,j S n). Then
1 (og(((2T))) K]) 2
((X:
with equality if and only if x ~ NN, K) for = [z)].
Lemma (Monotonocity of information gain). Let T" denote the trajectory induced by the policy æ.
Then, the following is true for all 2 0, policies π
ED [T if DI:n)
ED (T (1:n-1)]-
Proof.
Ep,,, [T (ffx: : Vr 11m-- - [fū V
Di:n)
ED H (U+* DIm-1) H (Ur"
f D Dim
(H (yr D 1.m) H (gr= f D 11m))]
Ep,.. [H (yr= Dım-1) H (Ur- DI.n]]
+ ED... [H (yr"
fzx) - H (yr" f
2
(information never hurts)
direct consequence of Lemma 8 the following corollary.
Corollary 4 (Information gain at N is less than the average gain till N). Let T denote the trajectory induced by the policy π. Then, the following is true for all N 2 1. policies π,
EDIx-1 // (f=+i yr- DINN-
ED 17 (f=+;; yrn Dinn-
Next, we prove Lemma 1, which is central to our proposed active exploration objective in Equation (6).
Proof of Lemma I. Let yr- denote with E, (för-1) the covariance of fū.r-1-
( (fr w.) where f (zx). Furthermore,
Tff+: Vrn (ism) (fór Bo.r (Dim))
H1 (3o.T-1 Disn) - 11 (yo:T fö.r-1 Dim))
log (1o21+ + En (f6TT - log ((o21l) 2
(Theorem 3)
1
log (Idiag (I + a 25,
(Hadamard's inequality)
2
T-1
o p 1
log
2
o²
We can leverage the result from Lemma to bound the average mutual information with the sum of epistemic uncertainties.
20

---

Lemma 9 (Average information gain is less than sum of average epistemic uncertainties). Let Assumption hold and denote with AN be the solution of Equation (4). Then, for all N 1 and dataset Da the sollowing is true
1 N
D 1
[1 [f "
D Dimn-1)
1
E
T-1
Σ
log 1+ (
where z are the state-action tuples visited by the solution of Equation (6), iee. n,
Proof.
N
N
D Din-1
xx
[I
D Pin-1)
12
Σ
d2
0²
m
(ze)
log
1+
(Lemma 1)
im 1,7
2
ENE
j=l
du
and (zt)
EDi.-
max E *Ell
Σ
log
1
(I)
EIETH j=1
dp
o ong
log
1
2
Here (1) follows from the tower property. Note that the second expectation in (I) is wrt conditioned on a realization of D 11 where the conditioning is captured in the epistemic uncertainty o., G
We use the results from above, to prove Theorem
Proof of theorem 1. Let n. denote the solution to Equation (4) at iteration n 2 1. We first relate the information gain from OPAx to the information gain of R..
Ep 7 UUiNiiUry U (D)
Ep, E
(ff#... : 17 D1mn-1)]
(Corollary 4)
de
02 1 (z
D
log
1+
(Lemma 9)
2
1 2
)
Ep EDi--IA
log
Jn (=())) Jn (TT)
g²
d
o² 1+ on-Le(Zne) ]
E D D
Eran
log
o²
+0 Br- (O
In- (Zn,))l
(Corollary 3)
In summary, the maximum expected mutual information episode N less than the mutual information of OPAX and the sum of model epistemic uncertainties. Crucial to the proof is the regret
21

---

bound for optimistic planning from Corollary 3.
N
T-1
1
dx
1 T )
E D
ETTh
log
N
0²
n=1
1=0
==1
)
| 1 (Smx)lla)
+0
du
on ) (1++
Erap
log
N
d2
n=1
1=
1Ia- (Emx))ll)
+0
(Bn=1(8)Em
de
1 on-lz(Znee)
ED
BAnn
log
N
lIo (2m,)]12 )
+0
T3 18)E
On-1s(Znrt)
E E
Erar
+0
18)-1
Σ
|Ion (22
(log(1 +x) S x for x 2 0.)
N
| (za, 11
5
ED ED
18)
Σ 1=0
N
n=1
Above, we show that the maximum expected mutual information can be upper bounded with the sum of epistemic uncertainties for the states OPAX x visits during learning. Finally, further upper bound this with the model complexity measure.
N
Σ | Mel 11)
1
o
E
Bn (8)TEEE
N
mm
N Σ
Σ o (Zn,)))22 1))
EDY,N
(T (8))
N
m=1
Σ | (z. Ellon(mm))11 |1)
50
(B)(
TNNDII
Σ E BTan
N
MMC(((
50
(N(O)T
N
Theorem gives a bound on the maximum expected mumual information w.r.t. the model complexity. We can concentration inequalities such Markov, give high probability bound on the
22

---

information gain. In particular, we have for all E > 0
ED EDINNI [I (Sūxx
i ip
Di:N-1)]
Pr (f((xiii (D:N-1) 26)
5 (TAS
Mfx Ne
A.4 Proof of GP results
This section presents our results for the frequentist setting where the dynamics are modeled using GPs. Since the information gain has meaning in the frequentist setting, study the epistemic uncertainty of the GP models.
Corollary (Monotonicity of the variance). For all 2 0, and policies the following is true.
o (211
T-1 de
0 (7-1,(z))
log 1
log 1
Proof. Follows directly due the monotonicity of GP posterior variance.
Next, we prove that the trajectory of Equation (6) at iteration upper-bounded with the maximum information gain.
Lemma 10. Let Assumption 2-4 hold Then, for all N 2 1, with probability least 1 ō, have
))
max E WEll
<ΣΣ
da ΣΣ
log 1
(2)(21)
O BN (8)T
IN
σ²
N
1=
Moreover, relax noise Assumption a-sub Gaussian. Then, for all N 2 1, with probability least 1 - 8, we have
T-1 de
a RN)(2)) 1
1)
max E well
Σ
log
< 1H
IN
o²
N
Proof. Gaussian noise case: Let n denote the state-action pair at time for the trajectory of Equation (6) iteration n 2 and n, the corresponding policy.
T-1 d=
1+ oh (#N.1)
ETag
log
2
0²
nn (2
E, B
log
1
(Corollary 5)
2
T-1
(2
Bre;
log
1+
(By definition of );
2
r²
n=1
MCN(F*)
50
Bx(o)T
(See proof of Theorem 1)
N
50 (N(O)TTh
IN
(Curi al., 2020, Lemma 17)
N
Lub-Gaussian noise case: The only difference between the Gaussian and sub- uub-Gaussian case is the regret term (c.f., Corollary and Lemma 7). particular, the regret the sub-Gaussian case leverages the Lipschitz continuity properties of the system (Assumption 2). This results in an exponential dependence on the horizon for our bound. We refer the reader to Curi al. (2020); Rothfuss et al. (2023) for more detailed derivation.
23

---

Lemma 10 gives sample complexity bound that holds for richer class of kernels. Moreover, for GP models, By x v'N (Chowdhury and Gopalan, 2017). Therefore, for kernels, where lim 1 - O, we can show convergence (for the Gaussian case). We summarize bounds on ON from Vakili et (2021) Table 1.
Table 1: Maximum information gain bounds for common choice of kernels.
Kernel
k<x; x²)
mn
Linear
x xTx' x²'
O (dlog(n))
||===|||
RBF
223
O (og
V20
22v a
20 O n autd log 2r (n)
Matèrn
Be
F(uy2u-
From hereon, we focus on deriving the results for the case of Gaussian noise case. All our results can be easily extended for the sub-Gal Lub-Gaussian setting by considering its corresponding bound.
Lemma 11. The following is true for all N 2 0 and policies π E II.
d
T-1 d
2 1
(2)
E
max 2ET
o
(2)
< C.o
log
2
1=0
with C.
log(1+o smass
Proof.
T-1 d2
Co
log
1
o²
t= j=1
T- dz
> ET,
Σ
(2+)
(Curi et al., 2020, Lemma. 15)
2
7
t= j=l
d=
2 Er,
max ZET"
o
N.)
2
j=1
Corollary 6. Let Assumption 2 and 4 hold, and relax noise Assumption to o-sub Gaussian. Then, for all N 2 with probability at least 1 5, we have
0o (B V
IN N
max E ET HEIl
max #ET
2 NoN
Let Assumption 2 and 4 hold, and relax noise Assumption 3 to o-bu Gaussian. Fur- thermore, assume lim 4-20 8 (K)) O. Then for all N 2 1, R, and $ dz, with probability at least - 8, we have
Lemma 12.
On)(2)
as. 0 -
Proof. We first show that the expected epistemic uncertainty along a trajectory converges zero almost surely. Then we leverage this result to show almost sure convergence of all trajectories induced

---

[Eto Ej 2 log (1 (1+ sلسساالe)
o² ))
by π E IT. To this end, let Sn = B far have,
for all n 2 0. So
(x(6))T 50
V In )
Pr Sn so
2 21-
Consider sequence {6n (n20 such that limn-roe 8n = 0, and lim,- Bn )T V + 0. Note, for GP models with N-m (8)) - 0, such a sequence of &n exists (Chowdhury and Gopalan, 2017). Consider any e and let (e) be the smallest integer such that
o ( (0)
(e)
<c
N-(c)
Then, we have
DC
(6)-11
2
Σ
Pr(S, >6) =
Pr(Sp 2))+
Σ
Pr(Sn >)
n=0
n= N' (6)
(0)
Σ
Pr(S,, > +)+
n=0
n=N
S (c +
Σ
En
n= (r)
Note, since limm-+0o 8n = 0, we have
6n < oc.
N= (c)
In particular, En Pr(Sn, 2) < oo for all Therefore, we obtain
Sn ah 0 for n oc.
E Z log 1 -..
Define the random variable V = lim
with 2 E
and = V represents the sum of epistemic uncertainties of random trajectory induced by the sequence of policies [Tn)220. Note V 2 0, therefore we apply Markov's inequality. Moreover, for all > 0, we have
(v > 6)
Hence, we have
Pr Pr = 0) =
0 for - 00
Accordingly, we get for all π e II.
2 (zt) 1+ 0²
Pr
lim
log
0 =1
(15)
2 j=1
Assume there exists a z E R, such that for some e, (z) > e for all 2 0. Since, 2 E R, there ex-
ists a / and E II such that p( ž 2 I|, f*) > 0. This implies that Pr(z E Tx) > 0. However, from Equation (15), have that o (=)
for N almost surely, which is contradiction. O
Finally, we leverage the results from above to prove Theorem
Proof of Theorem 2. The proof follows directly from Corollary 6 and Lemma 12.
25

---

A.5 Zero-shot guarantees
In this section, give guarantees the zero-shot performance of OPAX for a bounded cost function. We focus this section on the case Gaussian noise. However, a similar analysis can performed for the Gaussia case and Lipschitz continuous costs. Since the analysis for both cases similar, we only present the Gaussian case with bounded costs here.
Corollary 7. Consider the following optimal control problem
arg min Jo(TT, f*) = arg min E, mETl #(+1 = (2 T (xx)) + wi
Σ da m(x=))
(16)
*Ell
VO
with bounded and positive costs. Then we have for all policies with probability at least 1
T-1
(14 v dz
1n |1o (2:)1 ))
Je(T,n)) - Je(TT) 5O
TE
Σ 1=0
where Je(TT, maxncz Jo(TT, n))
Proof. From Corollary we get
Je( (0) - Je(TT) = =E E
Σ Jrati a Jr n² m2-1))
with Øz+1 f (x,, π (x:)) Wt, and xf+1 pn (2,, (xx))) Br (0) (())))) (2+) wc Furthermore, the cost is positive and bounded. Therefore,
fT < 2 2
for all a, n, Accordingly, we can now use the same analysis as in Lemma 2 and get
(1 + Bm (0) |0 (z))1 (23)11 ))
Je(T,nn) Ja(Tt) 50 TE
Lemma 13. Consider the control problem in Equation (16) and let Assumption 3 and 4 hold. Furthermore, assume for every >> 0, there exists funite integer n² such that
Vn n²: (7 (k:
Sc,
(17)
and denote with ii, the minimax optimal policy, i.e., the solution min max Jo(T,, m)) Then for all n 2 n", we have probability at least 1 - 8, J(( - Je(T') S O(e).
Proof of Zero-shot performance. In Theorem 2 we give rate at which the maximum uncertainty along a trajectory decreases:
[oNN)())11 5o
IN (R)
max E
max #E**
REIT
N
26

---

Combining this with Corollary 7 we get
(1+ / 1(8) (
Jd(T, (1) Jc(T)) 00
TE
x²
0
T 10)
max
|on-1
SET
0
Bynn max T
1 (0) o (o-1(2))1 (2)1
VATT (k)
00
n
(hn(A)) $1
=
(1)("))"
n
O(e).
(Vn 2 n²)
Hence, have that for each policy π,, our upper bound max, Jo(TT, (T,, is precise, i.e., Jo( n) Jo(T)= s (e),, VT E el.
max
(18)
ME
We leverage this to prove optimality for the minimax solution. For the sake contradiction, assume that
Je(#tn) +
(19)
Then we have,
max Je( ,n) NEE
min max Je(TT, 17) xell nEE
Je(Atn ,i
max Jol Jo( 77)
nEE
2 J((#n)
> Ja(TT) + O(E))
(Equation (19))
" ,n",.))
(Equation (18))
max Je(TT" .n). MEB
This is contradiction, which completes the proof.
Lemma 13 shows that OPAX also results in nearly-optimal zero zero-shot performance. The convergence criteria in Equation (17) is satisfied for kernels & that induce a very rich class of RKHS (c.f., Table 1).
B Experiment Details
The environment details and hyperparameters used for our experiments are presented 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 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, steps by "maximum number of gradient steps". The hyperparameters for the model-based agent are presented Table Furthermore, we present the iCEM hyperparameters in Table For more detail on the iCEM hyperparameters see Pinneri et al. (2021).
computational reasons, we upper bound the number of gradient
27

---

Table Downstream task rewards for the environments presented Section 5.
Environment
Reward r1
Pendulum-v/ swing up
0 - 0.10, 0.001
Pendulum-vi v1 keep down
(0, ④ + 0.10, + 0.001
MountainCar
-0.1u, + 100 goaal
Reacher go to target
100 { De Baarges 12 1 0.05
Swimmer go to target
a uargeet 2
Swimmer go away from target
xe Burger l
Cheetah forward
U
Cheetah run backward
Uorward
Table 3: Hyperparameters for results in Section 5.
Hyperparameters
Pendulum-vI
Pendulum-yl
MountainCar
Swimmer
Cheetah
Action repeat
N/A
N/A
N/A
NIA
4
Exploration horizon
100
200
200
50
1000
1000
Downstream task horizon
200
200
1000
Hidden layers
N/A
2
2
4
Neurons per layens
N/A
256
128
256
256
256
Number of ensembles
N/A
5
Butch size
N/A
64
64
64
64
64
Learning rate
0.1
10
x 10
10-3
10
10
Number epochs
50
50
50
50
Maximum number of gradient steps
N/A
5000
5000
6000
7500
7500
2.0
Model-based SAC optimizer For the reacher and MountainCar environment use scarce rewards (c.f., Table 2), which require long-term planning. Therefore, receding horizon MPC approach is less suitable for these tasks. Accordingly, we use 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, simulate trajectory of // steps using learned model and the SAC policy. We collect the simulated transitions into a simulation data buffer DSIM, which we then the policy. The algorithm summarized below, and we provide the SAC hyperparameters Table
perform G gradient steps suggested by Haarnoja et al. (2018) train
Table 4: Parameters of iCEM optimizer for experiments in Section
Hyperparameters
Pendulum-vi GP
Pendulum-v1
Swimmer
Cheetah
Number of samples P
500
500
250
200
Horizon !!
20
15
10
Size of elite-set K
50
50
25
20
Colored-noise exponent 3
0.25
0.25
0.25
0.25
Number of particles
10
10
10
10
CEM-iter iterations
10
5
5
Fraction of clites reused &
0.3
0.3
0.3
0.3
28

---

Model-based SAC
Init: Stastistical Model M (Jis on (6))),, Dataset of transitions Dn, initial policy no Model Rollout steps H
DSIM t for Training steps l = , do
Initialize simulated transitions buffer
E0,1:P ~ Dn
Sample initial states from buffer
for Initial state To E E0,1:P do
Collect simulated transitions Perform G gradient updates on as proposed in Haarnoja et al. (2018) using DSIM-
DSIM +- DSIM U MODELROLLOUTTE Wo, H, M)
Table 5: Parameters of model-based SAC optimizer for experiments in Section 5.
Parameters
MountainCar
Reacher
Discount factor
0.99
0.99
Learning rate actor
x 10
5 X 10-4
Learning rate critic
5 x 10
5 x 10-4
Learning rate entropy coefficient
x 10
5 x 10-4
Actor architecture
[64, 64]
[250, 250]
Critic architecture
[256, 256]
2250, 250]
Batch size
32
64
Gradient steps G
350
500
Simulation horizon H
4
10
Initial state sample size P
500
2000
Total SAC training steps K
35
350
B.2
OPAX in the High-diimensional 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 al. (2022) with the addition of large table. This is compositional object manipulation environment with an end-effecttor-conrollee robot arm. The robot actions u E RRd correspond to the gripper movement in Cartesian coordinates and the gripper opening/closing. The robot state Frobor E R110 contains positions and velocities of the end-effector as well as the gripper-state (open/close) and gripper-velocity, Each object's state Lobj IR is given by its position, orientation (in Euler angles), and linear and angular velocities.
We follow the free-play paradigm 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, 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 certain number of epochs on the collected data 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 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 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 objects. The rewards used for these tasks are the same as presented in Sancaktar et al. (2022).
29

---

Table 6: Environment and model settings used for the experiment results shown in Figure 4.
(a) Fetch Pick & Place settings.
(b) MLP model settings.
Parameter
Value
Parameter
Value
Episode Length
100
Network Size
256
Train Model Every
20 Ep.
Activation function
SiLU
Action Dim.
Ensemble Size
Robot/Agent State Dim.
10
Optimizer
ADAM
Object Dynamic State Dim.
12
Batch Size
256
Number of Objects
4
Epochs
50
Learning Rate
0.0001
Weight decay
5 10-3
Weight Initialization
Truncated Normal
Normalize Input
Yes
Normalize Output
Yes
Predict Delta
Yes
B.2.2 OPAX Heuristic Variant
In the case of Fetch Pick & Place with an high-dimensional observation space, implement 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 {f:) <
for I,..., Nmax do Solve optimal control problem till convergence for the system: x1+1 = (10, u.).
T-1 de
10:T--1111
argmax Ma. - Jo:7-1
ΣΣ
log (e + (a ue))
Estimate
(n ue) with ensemble disagreement.
Rollout uo.T, on the system and collect data Dn- Update models {f
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 intrinsie 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 Pinneri et al. (2021).
30

---

Table 7: Base settings for iCEM as they are used in the intrinsic phase. Same settings are used for all methods.
Parameter
Value
Number of samples P
128
Horizon H
30
Size elite-set K
10
Colored-noise exponent
3.5
CEM-iterations
3
Noise strength ainit Momentum a
0.5
0.1
use lse_mean_aaee actions ons chift_elited
Yes
Yes
keep_elites
Yes
Fraction of elites reused
0.3
Cost along trajectory
sum
Table 8: iCEM hyperparameters used for zero-shot generalization in the extrinsic phase. Any settings not specified here are the same as the general settings given in Table 7.
Task
Controller Parameters use_mean___ actioon aus_meen_actions
Horizon h
Colored-nood -oose exponent
Noise strength
Cost Along
B
Cinit
Trajectory
Pick & Place
30
3.5
Yes
0.5
best
Throwing
35
2.0
Yes
0.5
Flipping
30
3.5
0.5
sum
31

---

C Study of exploration intrinsic rewards
Pendnhum-vl GP -moeraiinty
x10 Pendulum-vl GP swing
1//
-1.0
-1.-
10
-2.0
10
16
10
X.00
MountainCar
x10F
Choctah run forward
0.3
0.0
3
25
200
400
500
Episodes
Episodes
OpAx
OPAX-SUM
Figure 5: Comparison OPAx between intrinsic reward proposed in Equation (6) and the sum of epistemic uncertainties proposed by (Pathak et al.. 2019), i.e., OPAX-SUM. For both choices of intrinsic rewards, OPAX reduces the epistemic uncertainty and performs well on downstream tasks.
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, show that these rewards are closely related.
Lemma 14. Let Tmax (1....,dz) d
suP (z)(z) and a > 0. Then for all i 0 and j
0 max
(2) log(1
mmax
log(1
1 t -0-30) S
(z).
- 20 max)
log(1 20 Fmax))
Proof. Curi et al. (2020) derive the first inequality on the left. The second inequality follows since log(1 (5(1+ + x) < for all x 2 O.
Due to this close relation between the two objectives, 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.
32


Reviews:
[{"review_id": "JyzocIyxxo", "reviewer": "Reviewer_kKQP", "summary": "This paper proposes and studies a rather intuitive algorithm for active learning in nonlinear dynamical systems in the episodic setting. They establish consistency (in terms of mutual information) and provide supporting numerical experiments. ", "questions": "* Could the authors comment on the necessity of the exponential dependence on the horizon in their main result? It would be interesting to see if this indeed necessary or whether it can be overcome by an improved proof method/algo.\n\n* Is the dependence on $\\delta$ hidden in the main results (and if so what is it)? Seems a little strange to me as the result is claimed to hold wp at least $1-\\delta$.\n\n* Are there any interesting cases where $ L \\times \\max \\beta \\leq 1$ or is the exponential dependence always present?", "limitations": "N/A.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The proposed algorithm is intuitive and it is satisfying that it \"works\" (in terms of consistency).\n\n* The paper is generally quite well-written and I did not have (m)any issues in terms of clarity and level of writing.\n\n* The exact setting is relatively novel and well-motivated (some caveats below) and the question is interesting and definitely deserves further study.\n\n* The experiments look relatively thorough (but I am not the best person to judge their significance).\n", "weaknesses": "\n* While I am sympathetic to the fact it can be hard to keep track of everything appearing in ML conferences,  unfortunately the manuscript does miss some of the most closely related recent references especially when it comes to learning in nonlinear dynamical systems or general mixing processes. My concern here is that missing these makes the contribution of the present paper appear larger than it actually is, since the below references [A,B,C,D] also treat learning in rather general dynamical systems/time-series.  In this light, the last part of the stated contributions (cf. line 52) are  somewhat misleading as it states that related work is not general enough to treat the present system dynamics. However, [A,B,C,D] are actually general enough (or at least very close to general enough) and the authors might want to complement their references  listed starting from line 50 (which certainly aren't the most general setups considered in the recent literature). In particular [B,C] also explicitly study RKHS-like dynamical systems.\n\n\n[A] Roy, Abhishek, Krishnakumar Balasubramanian, and Murat A. Erdogdu. \"On empirical risk minimization with dependent and heavy-tailed data.\" Advances in Neural Information Processing Systems 34 (2021): 8913-8926.\n\n[B] Ziemann, Ingvar M., Henrik Sandberg, and Nikolai Matni. \"Single trajectory nonparametric learning of nonlinear dynamics.\" conference on Learning Theory. PMLR, 2022.\n\n[C] Ziemann, Ingvar, and Stephen Tu. \"Learning with little mixing.\" Advances in Neural Information Processing Systems, 2022.\n\n[D] Li, Yingcong, et al. \"Transformers as algorithms: Generalization and stability in in-context learning.\" International Conference on Machine Learning. 2023.\n\n* As the above references do not treat the active learning setting, this limitation can easily be overcome by including the above references and making the qualifying distinction that the contributions are not \"first in system identification/supervised learning of nonlinear systems\" but rather first in terms of experiment design/active learning.\n\n* The exponential dependence on the horizon in the main results appears overly pessimistic to me and I question whether the the bound is informative at all in any meaningful setting. While convergence is not guaranteed in terms of MI the above references achieve convergence even from a single trajectory, at least indicating at a glance that such a dependence ought to be removable. This exponential dependence on the horizon in the provided bounds is a major caveat here---I have therefore estimated the theoretical results to be asymptotic consistency results and not bona fide finite sample guarantees.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and for providing additional references.\nAs the reviewer highlighted, the shared references do not consider the challenging active learning/unsupervised learning setting and focus more on the supervised learning problem. We have included the references in the revised paper as suggested\n\n## Weaknesses\n*W1. Exponential dependency on horizon T in regret bound*: The exponential dependency may be removed, however, in accordance to prior work, we believe it will require further assumptions (cf. [1] last paragraph, page 7). For instance, restricting the noise to be sampled from a Gaussian distribution (as opposed to sub-Gaussian) gives a $\\mathcal{O}(\\beta_N\\sqrt{T\\frac{\\gamma_N}{N}})$ bound in the RKHS setting. This analysis was presented in Appendix B, Theorem 4. In the revised paper, we moved this result to the main text.\n\n*W2. Comparison to the single trajectory setting*: The single trajectory works mostly deal with linear systems that are stable or assume knowledge of a stabilizing controller [2]. We make no such assumptions. Works that first learn a stabilizing controller from the data, also suffer from an exponential growth in the cost until a controller is learned [3]. Our general setting will require stronger assumptions to avoid the exponential dependence (see response W1.). We also emphasize that we give a worst-case theoretical upper bound. In practice, the epistemic uncertainty decreases much faster as depicted in Figure 1 in the paper. \n\n## Additional Questions\n*Q1 Dependence on $\\delta$*: The calibration factor $\\beta_n$ depends on $\\delta$ (see definition 2). For the RKHS case, the dependence of $\\beta_n$ on $\\delta$ is well studied ( $\\beta_n \\propto \\sqrt{\\log(1/\\delta)}$) [e.g. in 4]. \n\n\nHaving addressed your concerns we kindly ask you to consider revising our score. For any remaining questions, we are happy to provide further clarification. \n\n## References\n[1] Curi, S., Berkenkamp, F., and Krause, A. (2020). Efficient model-based reinforcement learning through optimistic policy search and planning. Advances in Neural Information Processing Systems.\n\n[2] Simchowitz, M. and Foster, D. (2020). Naive exploration is optimal for online lqr. In International Conference on Machine Learning, pages 8937–8948. PMLR.\n\n[3] Chen, Xinyi, and Elad Hazan. \"Black-box control for linear dynamical systems.\" In Conference on Learning Theory, pp. 1114-1143. PMLR, 2021.\n\n[4] Chowdhury, S.R., et al (2017). On kernelized multi-armed bandits. International Conference on Machine Learning.\n\n"}, {"review_id": "ueJ5jx7dH9", "reviewer": "Reviewer_meUw", "summary": "This paper proposes a task-agnostic active exploration algorithm for non-linear dynamic systems as long as it can be well calibrated. By combining the optimistic exploration principle and some standard baysian techinques, they give a general convergence bound as well as the more specfic bound in gaussian process case. Besides the theoretical proofs, they also provide several downstream experiments, showing their algorihthm can achieve better performance.", "questions": "Main questions:\n1. Can you give more dicussion on the computational complexity. Or maybe give me some examples of exsiting oracles?\n2. Currently paper aims to upper bound the maximum expected information gain (or the epistemic uncertainty). While I believe it makes sense, it is hard for me to connect that with previous paper. For example, [Wagenmaker and Jamieson, 2020] gives the upper bound ob $||\\hat{A} - A^*||$, which is very easy to understand. Can you show what this maximum expected information gain implies in those simple models so reader can have more intuition on your results?\n3. Can you give some discussions on the lower bound?\n4. Can you give more dicussions on Assumpton 2 (line 114-115) and Line 160 \"Incorporating constraints Equation (7) can impose input constraints by considering a restrictive class of policies $\\Pi$. \"  I understand this is one standard assumption. But theoretically, it is still unclear to me how difficult it is to lift this assumption in order to get the generalization result, and how it will affect the computational complexity.", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper gives the first theoretical gaurantees on general non-linear dynamic models, which is significant.\n2. The bayesian-based framework is somewhat novel and can be extended in many cases. \n3. Their theoretical analysis are solid. Their results on guassian process model helps reader to further understand this problem.\n4. I am not familiar with control experiments, but according to what stated in their paper, their approach give nontrivial improvements.", "weaknesses": "1. There is no discussion on the computational complexity. For eqn.(7), it is unclear to me how to efficiently solve $\\pi$, $\\eta$ when the policy space is large or the $T$ is large.\n2. The techinque itself is not very surprising to me. Exploration by optimism is a widely used in all the RL related paper, and the estimation on confidence bound seems to me just a standard derivation from exsiting techiniques.\n3. There is no dicussion on the lower-bound, therefore it is hard for me to understand how good this result is.  \n\nBut I am willing to raise my scores if my questions are addressed, or there are some techinical difficulty I am missing, or someone else want to advocate contribution of the experiment results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback, our response follows.\n\n## Weaknesses\n*W1 Computational complexity*:  Solving an optimization problem for general nonlinear systems is challenging, and out of scope for this work. For our problem formulation, standard trajectory optimizers such as iLQR[1], iCEM[2] or policy optimizers such as BPTT [3], and SAC [4] can be and have been commonly used. We also provide a heuristic variant of our algorithm in Appendix C.2.2., which, unlike the theoretical variant, does not scale with the dimension of the state space. This heuristic variant was used for the experiments in the “Fetch, Pick & Place Construction” environment. We have added this explanation in the updated version of the paper.\n\n*W2 Novelty of technique*: As we highlight in the general summary above, the novelty of our method lies in leveraging ideas from Bayesian experiment design, and the OFU paradigm, to provide novel and first-of-its-kind analysis for active learning. We relate the regret of planning under unknown dynamics and the information gained over a whole trajectory, to give a novel sample complexity bound for active/curious exploration. To our knowledge, we are the first to do so. This lack of theoretical results is also acknowledged by concurrent works [6, 7].  \n\n*W3 There is no discussion on the lower bound*: Providing lower bounds is a challenging open research problem. However, in more restrictive classes of problems [7, 8, 9] there exist nearly minimax optimal lower bounds, which compared to the upper bound, only vary in constants and logarithmic terms (wrt to N). In our setting, we would expect a similar outcome, but we acknowledge that this is yet to be solved as future work.\n\n## Additional Questions:\n\n*Q1 Intuition of the upper bound*: We are happy to give some intuition. Mainly, for the general RKHS setting we give a bound on how fast the epistemic uncertainty decays (Theorem 2). \nIn Lemma 2, we show that with high probability for all $j \\in \\{1, \\dots, d_x}$ \n\n$$\n|\\mu_{n, j}(z) - f_j^*(z)|  \\leq \\beta_n(\\delta) \\sigma_{n, j}(z) \\leq \\mathcal{O}\\left(\\sqrt{\\frac{\\gamma^{T+1}_n}{n}}\\right)\n$$\n\nTherefore, the decay rate of the epistemic uncertainty quantifies how fast our mean estimate converges to the true function. Our result holds for general RKHS kernels. Example for the linear kernel (with known matrix B): $f^*(z) = Ax + Bu; \\mu_n(z) = \\hat{A}x + Bu$.\nThen our bound is of the following sort:\n\n$$\n||\\mu_{n, j}(z) - f^*(z)||_1 = ||(\\hat{A} - A)x||_1 \\leq \\mathcal{O}\\left(\\sqrt{\\frac{d^{T} \\log^T(n)}{n}}\\right)\n$$\n\nfor all $x$. Note, this bound holds for a more general setting than the one considered in [Wagenmaker and Jamieson, 2020]. Furthermore, under stronger assumptions on the noise (Gaussian, instead of sub-Gaussian), we can avoid the exponential dependence on $T$ (see Theorem 4 Appendix B for more detail).\n\n*Q2 Constraints on inputs*: For input constraints like $||a|| \\leq a_{\\max}$, we can restrict our search to policies that only map to actions between $[-a_{\\max}, a_{\\max}]$. Most RL optimizers, such as SAC, squash actions to a fixed range. \n\nHaving addressed all of the questions provided by the reviewer, and given the contributions of this paper, we would appreciate it if the reviewer would increase their score for our paper. We would be happy to answer any remaining questions or concerns.\n\n## References\n[1] Li, Weiwei, and Emanuel Todorov (2004). Iterative linear quadratic regulator design for nonlinear biological movement systems. First International Conference on Informatics in Control, Automation and Robotics.\n\n[2] Pinneri, C., et al (2021). Sample-efficient cross-entropy method for real-time planning. Conference on Robot Learning.\n\n[3] Clavera, I. et al (2020). Model-augmented actor-critic: Backpropagating through paths. arXiv preprint arXiv:2005.08068 .\n\n[4] Haarnoja, T., et al (2018). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. International conference on machine learning.\n\n[5] Kakade, S., et al (2020). Information theoretic regret bounds for online nonlinear control. Advances in Neural Information Processing Systems.\n\n[6] Chakraborty, S. et al (2023). STEERING: Stein Information Directed Exploration for Model-Based Reinforcement Learning. arXiv preprint arXiv:2301.12038.\n\n[7] Wagenmaker, A. et al (2023). Optimal Exploration for Model-Based RL in Nonlinear Systems. arXiv preprint arXiv:2306.09210.\n\n[8] Wagenmaker, A. et al (2020). Active learning for identification of linear dynamical systems. In Conference on Learning Theory.\n\n[9] Scarlett, J., et al (2017). Lower bounds on regret for noisy gaussian process bandit optimization. Conference on Learning Theory.\n\n"}, {"review_id": "JzYtVsI7Qq", "reviewer": "Reviewer_V3Yd", "summary": "This paper studies provable exploration in model-based reinforcement learning and proposes an algorithm with optimistic active exploration based on information gain. Theoretical results and experimental results are provided to support their method.", "questions": "See the weakness section above.", "limitations": "More discussions on related works and discussions on the zero-shot generalizability are needed to better characterize the contribution of this paper.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper presents a practical algorithmic implementation of active exploration with optimism and information gain, by incorporating the techniques from Curi et al. [1] to introduce a hallucinate policy.\n2. The authors provide thorough theoretical and experimental analysis.\n3. The paper is well written and easy to follow.\n\n[1] Sebastian Curi et al. Efficient Model-Based Reinforcement Learning through Optimistic Policy Search and Planning.", "weaknesses": "1. My biggest concern is the novelty of this paper. Provable exploration based on optimism (mutual information or uncertainty) is not novel. Although the proposed algorithm provides an efficient way for practical implementation, which is good, the novelty is still limited compared to H-UCRL, which is also an optimism-based MBRL algorithm with the same hallucinate policy technique.\n2. The authors claim that their method can achieve better zero-shot performance compared to baselines. But it is not clear why is this. Can the authors explain in more detail? The current theory does not indicate this result.\n3. What makes the algorithm better in terms of generalization ability compared to H-UCRL? They seem to be developed from very similar optimism perspectives with the same hallucinate technique. \n4. From the experimental results, it seems that H-UCRL outperforms the proposed algorithm in most training tasks, despite that the proposed algorithm generalizes better (again, more explanation needed). Can the authors comment on this?\n5. Can the authors provide the training curves of H-UCRL and CEE-US instead of asymptotic performance?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. Below we address the reviewer’s concerns. \n\n### Weaknesses\n\n*W1 Novelty of our work*: We refer the reviewer to the author rebuttal section. If the reviewer is still concerned with the novelty of our work, we’d be happy to receive more detailed feedback on this concern. Additionally, a key component of our algorithm is active exploration which H-UCRL does not perform, i.e., our algorithm takes a generalist/uniform approach towards exploration of dynamical systems, whereas H-UCRL takes a specialist one (see W3).\n\n*W2 Baselines*: We consider four baselines PETS-AE, MEAN-AE, random, and H-UCRL. Random uniformly samples actions from the action space and does not use any intrinsic rewards to guide exploration. Accordingly, it underperforms. PETS-AE and MEAN-AE perform well in all environments. They use our proposed intrinsic rewards but use a greedy (not optimistic) planning approach that has no theoretical guarantees, except for linear systems [1]. Our method comes with strong theoretical guarantees, while empirically performing on par with SOTA baselines. \n\n*W3 H-UCRL Baseline and zero-shot performance*: H-UCRL is a task-specific model-based RL algorithm therefore its exploration is guided for the specific task it is trained for. Accordingly, its learned model does not generalize to novel unseen tasks. This is particularly observable in Figure 2, where we show that H-UCRL does not perform well on new tasks.\nOpAx is a task-agnostic algorithm that performs undirected exploration. Therefore, OpAx is a generalist. Accordingly, OpAx achieves on-par performance with H-UCRL on the tasks H-UCRL is trained for, however, it outperforms H-UCRL on novel unseen tasks.\n\n*W4 Training curves for H-UCRL and CEE-US*: We provide the training curves for H-UCRL for the swimmer environment in order to explain this trade-off between the specialists H-UCRL and generalists OpAx (cf., Figure 2 in the attached document). From the figure, it is noticeable that H-UCRL achieves higher rewards for the tasks it is trained for faster but fails to solve unseen/novel downstream tasks.\nWe have added these plots in the appendix of the revised version of the paper as well for further clarification. We also provide training curves for CEE-US (cf., Figure 3).\n\nThank you for your valuable feedback on our paper. We hope we addressed your concerns and if they're resolved, we kindly request you to consider revising our score upwards. We are happy to provide further clarification.\n\n## Reference\n[1] Simchowitz, Max, and Dylan Foster. \"Naive exploration is optimal for online lqr.\" International Conference on Machine Learning. PMLR, 2020.\n"}, {"review_id": "GtLO0fYwhf", "reviewer": "Reviewer_zJUZ", "summary": "The paper presents some insights into active exploration for model-based reinforcement learning. For certain kinds of environments, the authors show a convergence guarantee for model uncertainty. They augment their analysis with an empirical study of an agent using their approach OpAX.", "questions": "(also see \"weaknesses\")\n\nWhy does the presented approach work as well as it does? Are there no trade-offs involved?\n\nWhat does the theoretical argument say about other types of intrinsic reward?\n\nWhere is the novelty in the empirical study?\n\nWhat experiments would be necessary to show the borderline cases of the convergence property?", "limitations": "The authors discuss the approach's limitations quite well.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors point out an interesting research direction. They motivate their approach well and built up almost all of their mathematical framework. The empirical study manages to give a broad picture of the approach's performance.", "weaknesses": "The empirical part of this paper is connected only weakly to the theoretical part. Where are the proven convergence properties in the study? Picking domains where edge cases of the convergence properties might be observed, would help the paper. The authors state that the approach they show in the empirical study is very similar to other approaches found in literature, but they show no direct comparison. Why is the empirical study part of this paper, then?\n\nThe authors' interpretation of the results is somewhat lavish. The performance of the main approach is very similar to other shown approaches and clear advantage is not shown. The comparison on downstream tasks is important to stress the overall benefit of intrinsic reward, but also to be expected. It is not discussed how the competitiveness or even advantage of the said approaches arises for the \"high-dimensional task\" (Fig. 4). Is there no price to pay for the more involved approach? Why are the baselines shown without training times? At some point, the authors argue that their approach is not limited to domains where simple assumption of physicality might help. But why are all tested domains strictly physical, then?\n\nErrors:\n- throughout the paper: \"c.f.\" --> \"cf.\"\n- line 65: \"since\" --> \"when\"\n- line 133: \"since\" --> \"since,\" (add comma)\n- lines 168ff switch to $f^\\star$  from $f^*$\n- line 203: \"Theorem 1;\" --> \"Theorem 1:\" (use colon)\n- line 212: \"RKHS, however\"--> \"RHKS; however,\" (use semicolon and comma)\n- line 222f: \"deep mind\" --> \"DeepMind\" ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and for pointing out our typos. We have addressed them in the paper.\n\n### Weaknesses\n\n*W1 Empirical evaluation of theoretical results*: We designed the pendulum experiment (line 242) on GPs, i.e., RKHS setting, precisely to evaluate our theoretical findings empirically. In our experiments, we show in the example of the pendulum environment how fast our algorithm reduces the model epistemic uncertainty in the full state space of the pendulum. Moreover, in this experiment, we consider both GP models with an RBF kernel and BNNs using probabilistic ensembles. In Figure 1., we show that our algorithm monotonously decreases the epistemic uncertainty with the number of episodes N for the GP case, i.e., validating the findings of theorem 2. Furthermore, for the BNN case, we also show the reduction in the epistemic uncertainty which corresponds to an empirical study of theorem 1. \n\n*W2 Similarity of baselines*: We compare with a similar baseline which is [1]. Our intrinsic reward, the sum of log epistemic uncertainty, is similar to the one used by [1, 2] (they use the sum of epistemic uncertainty instead of the log). However, [1, 2] do not use the optimistic planner. Moreover, the two most commonly used planners in RL are; 1) mean (using the mean model) and 2) PETS [3]. We compare our algorithm to these for the same intrinsic reward (baselines: MEAN-AE, PETS-AE), where we show that our method performs on par with these baselines while also providing theoretical guarantees.\nFurthermore, in Appendix D, we also compare the exact intrinsic reward from [1, 2], i.e., the sum of epistemic uncertainty without the log to ours. Here, we demonstrate that both choices of intrinsic rewards perform similarly. Additionally, we show that the intrinsic reward from [1, 2] is also theoretically sound (Lemma 11, Appendix D) when combined with optimistic planning from OpAx. \n\n*W3 Empirical performance of OpAx and baselines*: The strength of our work lies in the theoretical guarantees. In the empirical results, we show that our algorithm performs at least on par with our baselines that do not yield such theoretical guarantees. \n\n*W4 High-dimensional task study*: For the high-dimensional task, the state space is considerably larger (58D), which makes optimizing over hallucinated controls $\\eta$ challenging (dimension of $\\eta$ is equal to the dimension of the state). In Appendix C.2.2, we propose a heuristic variant for OpAx, which does not scale with the state space size, and therefore can handle high-dimensional systems. Our results indicate that the heuristic variant also performs well. We have updated the main paper to clarify this further. \n\n*W5 Training Budget of OpAx vs Baselines*: All algorithms are given equal compute budget (cf., Appendix C). Our approach requires policy optimization over a larger domain space that includes the actions and the hallucinated actions $\\eta$. However, as we list in Appendix C, for all our baselines we use the same hyperparameters (training, optimization steps, optimizer samples, etc.) for fairness. Therefore, all methods had equal training budgets. \n\n*W6 Limitation to domains where the simple assumption of physicality hold and evaluation*: We would appreciate it if the reviewer refers us to the line where the argument is made. Then we can address the comment in detail. We study general dynamical systems, which we also consider in our evaluation.\n\n## Additional Questions:\n*Q1 Trade-offs for OpAx*: The main practical trade-off of the algorithm lies in the optimization problem, which is now performed over a larger space that includes the hallucinated controls. In theory, we assume an oracle that solves the problem, in practice this is challenging, for instance, due to compute limitations. To this end, we propose a heuristic variant of our approach in Appendix C.2.2 and show that it works for the high-dimensional setting.\n\n*Q2 Other intrinsic rewards*: Essential to our algorithm and theoretical analysis is the choice of our intrinsic reward. Could the reviewer clarify this question? Is the reviewer asking about theoretical guarantees of another algorithm with another objective?\n\n*Q3 Borderline case of convergence*: Could the reviewer clarify what is meant by a borderline case? Our results give a bound on the model epistemic uncertainty in the worse case (including the borderline case). For kernels such as RBF, this implies convergence for any function in its RKHS.\n\nHaving addressed all of the questions provided by the reviewer, and given the contributions of this paper, we kindly ask the reviewer to reconsider the assessment for our paper. We would be happy to answer any remaining questions or concerns.\n\n## References\n[1] Pathak et al (2019). Self-supervised exploration via disagreement. In International conference on machine learning.\n\n[2] Sekar, R. et al (2020). Planning to explore via self-supervised world models.  In International Conference on Machine Learning.\n\n[3] Chua, K., et al. (2018) Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems.\n"}, {"review_id": "5SSKtI0F1V", "reviewer": "Reviewer_EYiX", "summary": "The paper presents an active exploration method (OPAX) for dynamics model learning. The method seeks to maximize information gain, while being optimistic about unknown dynamics with respect to the achievable information gain. Theoretical results establish a connection between information gain and model complexity, and a convergence guarantee for GP models. Experimental results compare OPAX to baselines in several control and manipulation domains.", "questions": "- Can you identify experimental regimes in which the optimism has a significant impact on the performance of the resulting model? Even if such an experiment might be a little cherry-picked, I think it would still be useful for the reader.", "limitations": "There is no detailed discussion of limitations anywhere in the paper. I think adding this would improve the paper.\n", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "To my knowledge, the paper's principled approach to derive the planning objective in eqs. (6) and (7) from the optimal design perspective is novel. The paper is carefully written and easy to follow. The experiments are relevant and the selection of baselines seems fair. I believe that improving the data efficiency of learning dynamics models for zero-shot task generalization is a relevant research objective.", "weaknesses": "The approach does not really outperform the relatively the baseline methods used in the experiments. I don't think that's a huge problem, as I believe the main contribution of the paper is its theoretical part. Still, I think the paper could be improved by explaining better why the optimism does not lead to a significant difference in model performance (OPAX vs PETS-AE and Mean-AE).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback and acknowledge that the reviewer correctly highlighted the strengths of our work. Below we have addressed the reviewer’s concerns.\n\n## Weaknesses:\n*W1 Outperforming baselines*: We are happy that the reviewer recognized that the main contributions of the paper are theoretical. In [1], the benefits of optimism are discussed in more detail. Particularly, they show that optimism helps in settings with scarce rewards, whereas for non-scarce rewards PETS and optimistic planning perform on par. Since our exploration objective by definition is not scarce (i.e. the reward is the uncertainty which is mostly non-zero everywhere), we observe that OPAX performs similarly to the baselines. We have added this discussion to the revised paper. \n\n## Additional Questions:\n*Q1 Experimental regimes where optimism helps*:  As highlighted in our response to W1, optimism has an impact on settings with scarce rewards. Our proposed objective is not scarce by definition. However, we could penalize large actions in our objective, i.e., augment our reward to:\n$$\nr(s_t, a_t) = \\log\\left(1 + \\frac{\\sigma_n^2(s_t, a_t)}{\\sigma^2}\\right) - \\lambda ||a_t||^2.\n$$\nIn this setting, scarcity in the rewards can be induced by picking large values for $\\lambda$. Practically,  we penalize/avoid exerting large actions on the system. We compare OpAx to PETS-AE on the pendulum environment for a specific choice of $\\lambda=5$ to give the reviewer an intuition, cf., Figure 1 in the attached document. In the figure, it is visible that in this setting, optimism helps in solving the problem faster. We have also added this simple experiment to the appendix of the revised paper.\n\n## Reference\n[1] Curi, S., Berkenkamp, F., and Krause, A. (2020). Efficient model-based reinforcement learning\n"}, {"review_id": "Vjjs5s1sh1", "reviewer": "Reviewer_Y1JX", "summary": "This paper addresses the problem of active exploration of a (Markovian) dynamical system with continuous states and actions. In the absence of any cost function, the proposed objective is the maximization of the one-step information gain on the dynamics model. The paper presents an algorithm, called OpAx, to maximize the introduced objective by alternating optimistic planning with the current estimate of the dynamics model and data collection with the resulting policy to improve the model estimate. Then, the paper provides a convergence analysis of OpAx, which provably converge to the true model asymptotically under various assumptions. Finally, the OpAx algorithm is evaluated against some relevant baselines in continuous control tasks.", "questions": "This paper looks like a sound and neat work in system identification through active exploration and optimism. I do not have major concerns about the quality of the paper per se. However, judging from an RL perspective, this contribution looks two/three years late. Several recent works have demonstrated how to learn a model of the transition dynamics that is just good enough for planning by taking a polynomial number of samples from the process, which goes under the name of reward-free exploration. Instead, this paper takes the hardest route of trying to reduce the epistemic uncertainty \"uniformly\" over the dynamics model, which is arguably far-fetched without strong reachability assumptions.\n\nFor this reason, I struggle to see how this paper contributes to the advancement of RL research, as it does not compare well with reward-free exploration in terms of theoretical guarantees, neither significantly improves over previous practical methods empirically.\nHowever, this might be a useful contribution for the control community, and I am open to raise my score if the authors could show me that.\n\nI report below some detailed comments.\n\n**Dynamical Systems and MDPs**\n\nThe paper motivates the approach as a tool for learning the dynamics model to then solve RL tasks. However, the standard model for RL is the MDP, while the paper focuses on a specific type of Markovian processes in which the transitions are given by a deterministic function plus a Gaussian noise. Can the authors explain how this model relates with MDPs, and the relative expressive power of the two?\n\n**Well Calibration Assumption**\n\nWhereas Ass. 1 might be standard in the control literature, which I am not familiar with, it is kind of black-box from an RL perspective. It looks like it is saying that the problem is learnable, but I would like to know under which conditions on $f^*$ this is the case.\nWhen is the Ass. 1 violated in practice? Does this implicitly assume strong reachability? Can the authors relate the Ass. 1,2,3 to the common structural assumptions on the MDP model?\n\n**Convergence Results**\n\nThe theoretical results focus on showing that the maximum information gain shrinks with $N$, but the rate is actually exponential in $T$. This means that the sample complexity of learning an approximately correct model is also exponential?\n\n**Experiments**\n\nThe OpAx algorithm does not seem to improve significantly over the Mean-AE baseline in the considered experiments. While OpAx arguably comes with a stronger theoretical justification, I would like to see some benefit in the experiments as well, or at least a theoretical result showing that planning with the mean model cannot be provably efficient under the considered assumptions.\n\n**Related Works**\n\nThe paper neglects a recent yet considerable stream of works in reward-free RL, which seems to have overlapping motivation with their active exploration problem. This stream of works started with (Jin et al., Reward-free exploration for reinforcement learning, 2020) and counts dozens of subsequent results for tabular models, linear models, and general function approximation. Can the authors relates their contribution with reward-free RL literature?\n\nAnother important work that seems to be missing is (Tarbouriech et al., Active model estimation in Markov decision processes, 2020), which looks also very close to the problem formulation of this paper, although it focuses on MDPs rather than dynamical systems with Gaussian noise.\n\n**Minor**\n\n- The paper refers to the Appendix A for the proof of Lemma 1, but I cannot find it;\n- I would report confidence intervals instead of standard error in the plots.", "limitations": "The paper does not explicitly discuss the limitations of this work in terms of empirical results or theoretical guarantees.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper tackles an active exploration problem for system identification that is of general interest to both the reinforcement learning and optimal control communities;\n- The algorithm presented in the paper is simple and intuitively sound;\n- The paper combines theoretical justification with empirical validation in challenging domains;\n- The paper is well-presented and looks rigorous in the theoretical statements.", "weaknesses": "- The paper is mainly motivated as a tool for system identification that can be used in preparation to solve RL or control tasks, but this motivation looks somewhat weak for the RL side;\n- The experimental results are not terrible, but I can hardly see improvements over simpler baselines, such as one using planning with the average model instead of optimism;\n- The paper neglects several related works in RL, especially reward-free RL and active model estimation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and for sharing other relevant work. Based on their summary, we believe that the reviewer misunderstood our objective. In particular, our objective is *not the one step but $T$-step information gain on the model*. This is a crucial difference since it allows us to quantify the information gained over a whole rollout instead of just one transition. \n## Weaknesses: \n*W1 Motivation from RL side*: There is a plethora of work on unsupervised RL that studies our problem setting [7-12]. Generally, systems such as robots are used to perform multiple tasks that are often not known apriori. Prior works such as [5] motivate the problem in the context of reward-free RL but with rewards coming from a known reward class. This setting, however, needs a finite covering number of the reward class and is not practical for our general problem.   \nTraditionally, for a given dynamics model, well-established model-based controllers are commonly used to optimize for any task. Therefore, having a generalist approach to efficient exploration and identification of the dynamics *uniformly* is typically preferred.\n\n## Additional Questions: \n*Q1 Dynamical Systems and MDPs*: We consider an MDP model with an infinite number of states and actions, unlike in [13, 14]. We use a classical textbook definition of a dynamical system [6], which includes a large class of systems in RL (MUJOCO physics simulator, etc., ). \n\n*Q2 Well Calibrate Assumption*: Well calibration is a standard assumption in model-based RL with continuous states and action spaces. It has been used for both linear and nonlinear systems [1-4].  For the RKHS settings, this assumption is satisfied (cf., Lemma 2). For BNNs, proving well-calibration is an open problem (cf, line 108-112). The well-calibrated assumption has no connection to strong reachability. This assumption intuitively tells that the algorithm is able to learn a mean as well as the confidence region around the mean function. \n\n*Q3 Intuition of Assumptions 2 and 3*: Assumption 3 is a common assumption on the transition noise, which can further be relaxed to the even more general heteroscedastic noise case (cf., line 123).  Finally, Assumption 2, makes smoothness assumptions on the dynamics that are common in nonlinear systems literature [6].\n\n*Q4 Convergence Results*: The rate of convergence depends on $N$ so the sample complexity (as a rate) for kernels such as RBF, linear, etc is of polynomial order (cf.Theorem 2). Under our assumptions, the dependence on the horizon T is exponential, similar to [3]. However, that doesn’t influence the rate. As also argued in [3], to obtain better dependence on the horizon T, we need stronger assumptions (one such assumption is shown in Appendix B where we obtain polynomial dependence on T).\n\t\n*Q4 Experiments*: The main contribution of our work is theoretical. In [3] the benefits of optimism are discussed in further detail. Here it is shown that optimistic planning performs similarly to mean or PETS planning when the reward landscape is not scarce. Since this is the case for our intrinsic rewards, we empirically observe similar performance. Nonetheless, Mean and PETS, until now, have no theoretical guarantees for the general class of systems we consider. Deriving such guarantees for them is an open research problem.\n\n*Q5 Related Works*: We thank the reviewer for sharing additional related works. We have included them and more works from reward-free RL in the revised paper. In our work, we consider a general class of dynamical systems for which we make very common/practical assumptions.  As we mention in the general comment, theoretical results of our kind, do not exist.\n\n*Q6 Proof of Lemma 1*: Cf., line 516.\n\nHaving addressed all of the questions provided by the reviewer, and given the contributions of this paper, we kindly ask the reviewer to increase their evaluated score for our paper. We would be happy to answer any remaining questions or concerns.\n\n## References:\n[1] Chowdhury, Sayak Ray, and Aditya Gopalan. \"On kernelized multi-armed bandits.\" ICML, 2017.\n\n[2] Yasin Abbasi-Yadkori et al (2011). Regret bounds for the adaptive control of linear quadratic systems. COLT.\n\n[3] Curi, S., et al (2020). \"Efficient model-based reinforcement learning through optimistic policy search and planning.\" NeurIPS.\n\n[4] Abeille, M., et al (2017). \"Thompson sampling for linear-quadratic control problems.\" AISTATS.\n\n[5] Chen, J. et al (2023). On the statistical efficiency of reward-free exploration in non-linear rl. NeurIPS.\n\n[6] Khalil, H. K. (2015). Nonlinear control, volume 406. Pearson New York.\n\n[7] Schmidhuber, J. (1991). A possibility for implementing curiosity and boredom in model-building neural controllers. In Proc. of the international conference on simulation of adaptive behavior: From animals to animats.\n\n[8] Wagenmaker, A. et al (2020). Active learning for identification of linear dynamical\nsystems. COLT.\n\n[9] Mania, H. et al (2022). Active learning for nonlinear system identification with guarantees. JMLR.\n\n[10] Pathak et al (2019). Self-supervised exploration via disagreement. ICML.\n\n[11] Sekar, R. et al (2020). Planning to explore via self-supervised world models.  ICML.\n\n[12] Sancaktar, C., et al (2022). Curious exploration via structured world models yields zero-shot object manipulation. NeurIPS.\n\n[13] Jin, C. et al (2020). Reward-free exploration for reinforcement learning. ICML.\n\n[14] Tarbouriech, J. et al (2020). Active model estimation in markov decision processes. UAI."}]

Author rebuttal:
We thank the reviewers for their feedback. It seems that some key contributions to our work have missed the attention of our reviewers. We have made our contributions more clear in the revised paper and included additional references. We highlight our contributions below for clarification.

1. We consider the problem of learning a dynamical system in a reward-free/task-agnostic manner over *continuous state and action spaces*. We adhere to the general textbook definition of a dynamical system, e.g. [1]. 

2. We introduce a practical exploration objective, based on the concept of information in Bayesian experiment design. Our *novel* derivation technique may be of *independent interest* in active learning and experiment design. 

3. We utilize the principle of optimism in the face of uncertainty (OFU) and give a PAC bound on the epistemic uncertainty of any visited trajectory (Theorem 2). 

4. To our knowledge, this paper is the first to present active learning guarantees on *continuous domains* and for *generic dynamical systems in an RKHS*. Prior work is limited to strictly simpler classes of systems such as finite, linear, and low-rank MDPs [2, 3, 4, 5], linear systems [6], and nonlinear systems with finite-dimensional feature spaces [7]. In fact, concurrent work on model-based RL [8, 9] explicitly highlights the *lack of theoretical guarantees* within the context of curious/active exploration of general dynamical systems. Our paper directly targets this gap.

5. We validate our method extensively over 6 RL tasks, to demonstrate that *in addition to enjoying strong theoretical guarantees*, OPAX performs on par with the state-of-the-art baselines, none of which are supported by theory. 

We have revised our Related Works section and added additional experiments in the Appendix to demonstrate the benefits of active learning and optimism, as asked for by our reviewers. Furthermore, some of our reviewers raised concerns about the exponential dependence of our bound on the horizon $T$. In accordance with prior work [10], we believe stronger assumptions are needed to alleviate this. We had discussed this in Appendix B, where in Theorem 4 we give a tighter bound (polynomial in $T$) under additional assumptions on the noise. We have now moved the theorem to the main paper for more visibility. 

We would be happy to further update the paper, if there are any remaining questions or feedback by the reviewers.

## References
[1] Khalil, H. K. (2015). Nonlinear control, volume 406. Pearson New York.

[2] Jin, C. et al (2020). Reward-free exploration for reinforcement learning. International Conference on Machine Learning.

[3] Tarbouriech, J. et al (2020). Active model estimation in markov decision processes. Conference on Uncertainty in Artificial Intelligence.

[4] Wagenmaker, Andrew J (2022). et al. Reward-free rl is no harder than reward-aware rl in linear markov decision processes. 

[5] Chen, J. et al (2023). On the statistical efficiency of reward-free exploration in non-linear rl. Advances in Neural Information Processing Systems.

[6] Wagenmaker, A. et al (2020). Active learning for identification of linear dynamical
systems. In Conference on Learning Theory.

[7] Mania, H. et al (2022). Active learning for nonlinear system identification with guarantees. The Journal of Machine Learning Research.

[8] Chakraborty, S. et al (2023). STEERING: Stein Information Directed Exploration for Model-Based Reinforcement Learning. arXiv preprint arXiv:2301.12038.

[9] Wagenmaker, A. et al (2023). Optimal Exploration for Model-Based RL in Nonlinear Systems." arXiv preprint arXiv:2306.09210.

[10] Curi, S., Berkenkamp, F., and Krause, A. (2020). Efficient model-based reinforcement learning through optimistic policy search and planning. Advances in Neural Information Processing Systems.


Meta-review: There was quite a discussion about this paper in relation to ‘zero shot’ RL methods. The submitted paper does not explicitly maximize for future zero-shot performance (rather, for learning a model), which has its own strengths and weaknesses. However, the introduction does seem to claim zero-shot performance of the method, without evaluating this aspect against existing methods. The authors should review such (implied) claims to make them in line with the paper's actual results. 

There was also a discussion on novelty. Several reviewers argued the method is not very novel. However, the authors argue that it is the analysis presented that is novel and valuable. On the contribution, there were also questions about the usefulness of the bound and the strength of the connection between theory and practice. The authors pointed at a result from the appendix avoiding exponential dependence, and gave an indication of what a lower bound might look like, alleviating the concerns on this aspect.
Decision: Accept (poster)