Title: Optimistic Active Exploration of Dynamical Systems
Venue: NeurIPS 2023
Accepted: yes
Keywords: Active Exploration

Reinforcement Learning

Dynamical Systems
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

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 globally approximates the dynamics and 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 first-of-its-kind 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.

1 Introduction

Most reinforcement learning (RL) algorithms are designed to maximize cumulative rewards for a single task at 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 if only one control task is considered, it does not scale to the setting where the system is used to perform several tasks, i.e., under the same dynamics optimized for different reward functions. As a result, when presented with a new reward function, they often need to relearn a policy from scratch, requiring many interactions with the system, or 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-predictive control (García et al., 1989) assume knowledge of the system's dynamics. They leverage the dynamics model to solve an optimal control problem for each task. Moreover, in the presence of an 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 a model using just physics' first principles is very challenging. A promising approach is to leverage data for learning the dynamics, i.e., system identification or 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?

While active learning for regression and classification tasks is well-studied, active learning in RL is much less understood. In particular, active learning methods that yield strong theoretical and practical results, generally query data points based on information-theoretic criteria (Krause et al., 2008; Settles, 2009; Balcan et al., 2010; Hanneke et al., 2014; Chen et al., 2015). In the context of

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 is unrealistic . Rather, we can only execute policies on the real system and observe the resulting trajectories. Accordingly, an active learning algorithm for RL 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 dynamics 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 a 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 a 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 a 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 1 of OPAX in JAX (Bradbury et al., 2018).

2 Problem Setting

We study an unknown discrete-time dynamical system f ∗ , with state x ∈ X ⊂ R d x and control inputs u ∈ U ⊂ R d u .

Here, w k represents the stochasticity of the system for which we assume w k i.i.d. ∼ N ( 0 , σ 2 I ) (Assumption 3). Most common approaches in control, such as trajectory optimization and model-predictive control (MPC), assume that the dynamics model f ∗ is known and leverages the model to control the system state. Given a cost function c : X × U → R , such approaches formulate and solve an optimal control problem to obtain a sequence of control inputs that drive the system's state

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 (Åström and Eykhoff, 1971). Specifically, we are interested in devising a cost-agnostic algorithm that focuses solely on learning the dynamics model. Once a 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 n = 1 , . . . , N . At the beginning of the episode n , we deploy an exploratory policy π n , chosen from a policy space Π for a horizon of T on the system. Next, we obtain trajectory τ n = ( x n, 0 , . . . , x n,T ) , which we save to a dataset of transitions D n = { ( z n,i = ( x n,i , π n ( x n,i )) , y n,i = x n,i +1 ) 0 ≤ i<T } . We use the collected data to learn an estimate µ n of f ∗ . 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.e., D 1: n -1 to determine a policy π n ∈ Π for the next step of data collection, that is, Alg ( D 1: n -1 , n ) → π n . The proposed algorithm should be consistent, i.e., µ n ( z ) → f ∗ ( z ) for n →∞ for all z ∈ R , where R is the reachability set defined as

and efficient w.r.t. rate of convergence of µ n 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 ∗ , a natural objective for active exploration is the mutual information (Lindley, 1956) between f ∗ and observations y D n .

Definition 1 (Mutual Information, Cover and Thomas (2006)) . The mutual information between f ∗ and its noisy measurements y D n for points in D n , where y D n is the concatenation of ( y D n ,i ) i<T is defined as,

where H is the Shannon differential entropy.

The mutual information quantifies the reduction in entropy of f ∗ conditioned on the observations. Hence, maximizing the mutual information w.r.t. the dataset D n 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., I ( f ∗ ; y D 1: N ) . This requires maximizing the mutual information over a sequence of policies, which is 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:

Here f ∗ τ π = ( f ∗ ( z n, 0 ) , . . . , f ∗ ( z n,T -1 )) , y τ π = ( y n, 0 , . . . , y n,T -1 ) , τ π 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 it to the epistemic uncertainty of the learned model. Accordingly, while this notion of information gain stems from Bayesian literature, we can use it to motivate our objective in 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 µ n ( x , u ) and the epistemic uncertainty σ n ( x , 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-calibrated :

Definition 2 (All-time calibrated statistical model of f ∗ , Rothfuss et al. (2023)) . Let, z = ( x , u ) and Z := X × U . An all-time calibrated statistical model of the function f ∗ is a sequence ( µ n , σ n , β n ( δ )) n ≥ 0 , such that

Here µ n,l and σ n,l are the l-th element in the vector valued functions µ n and σ n respectively. The scalar function, β n ( δ ) ∈ R ≥ 0 quantifies the width of the 1 -δ confidence intervals. We assume w.l.o.g. that β n monotonically increases with n , and that σ n,l ( z ) ≤ σ max for all z ∈ Z , n ≥ 0 , and l ∈ { 1 , . . . , d x } .

Assumption 1 (Well calibration assumption) . Our learned model is an all-time-calibrated statistical model of f ∗ , i.e., there exists a sequence of ( β n ( δ )) n ≥ 0 such that our model satisfies the wellcalibration 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 σ n are L f and L σ Lipschitz continuous, respectively. Moreover, we define Π to be the policy class of L π Lipschitz continuous functions.

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 σ n also holds for GPs with common kernels such as the linear or radial basis function (RBF) kernel (Rothfuss et al., 2023).

Finally, we reiterate the assumption of the system's stochasticity.

Assumption 3 (Process noise distribution) . The process noise is i.i.d. Gaussian with variance σ 2 , i.e., w i.i.d ( 0 , σ 2 I ) .

We focus on 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).

3 Optimistic Active Exploration

In this section, we 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 a policy under unknown dynamics. Moreover, we show that the optimistic planner is 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.

Lemma1 (Information gain is upper bounded by sum of epistemic uncertainties) . Let y = f ∗ ( z )+ w , with w ∼ N (0 , σ 2 I ) and let σ n -1 be the epistemic uncertainty after episode n -1 . Then the following holds for all n ≥ 1 and dataset D 1: n -1 ,

We prove Lemma 1 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 a tractable objective that also has a frequentist interpretation - collect points with the highest epistemic uncertainty. We can use it to plan a trajectory at each episode n , by solving the following optimal control problem:

The problem in Equation (6) is closely related to previous literature in active exploration for RL. For instance, some works consider different geometries such as 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 µ n -1 in Equation (6) instead of f ∗ for planning (Buisson-Fenet 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 µ n -1 and the epistemic uncertainty σ n -1 , we solve the following optimal control problem

where Ξ is the space of policies η : X → [ -1 , 1] d x . Therefore, we use the policy η to 'hallucinate' (pick) transitions that give us the most information. Overall, the resulting formulation corresponds to a simple optimal control problem with a larger action space, i.e., we increase the action space by another d x dimension. A natural consequence of Assumption 1 is that J n ( π ∗ n ) ≤ J n ( π n , η n ) 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 π n 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 cannot evaluate the optimal exploration policy from eq. (6) and use the optimistic one, i.e., eq. (7) instead, we incur a 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 1 hold. Furthermore, define J n,k ( π n , η n , x ) as

Then, for all n ≥ 1 , with probability at least 1 -δ ,

Lemma 2 gives a bound on the regret of planning optimistically under unknown dynamics. The regret is proportional to the difference in the expected returns for x t and x ′ t . Note, ∥ x t -x ′ t ∥ ∝ β n ( δ ) σ n -1 ( x t -1 , π n ( x t -1 )) . 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 We start by introducing a measure of model complexity as defined by Curi et al. (2020).

This complexity measure captures the difficulty of learning f ∗ given N trajectories. Mainly, the more complicated f ∗ , the larger the epistemic uncertainties σ n , and in turn, the larger corresponding MC N ( f ∗ ) . Moreover, if the model complexity measure is sublinear in N , i.e. MC N ( f ∗ ) /N → 0 for N → ∞ , then the epistemic uncertainties also converge to zero in the limit, which implies

convergence to the true function f ∗ . We present our main theoretical result, in terms of the model complexity measure.

Theorem 1. Let Assumption 1 and 3 hold. Then, for all N ≥ 1 , with probability at least 1 -δ ,

Theorem 1 relates the maximum expected information gain at iteration N to the model complexity of our problem. For deterministic systems, the expectation w.r.t. τ π is redundant. The bound in Equation (9) depends on the Lipschitz constants, planning horizon, and dimensionality of the state space (captured in β N and MC N ( f ∗ ) ). If the right-hand side is monotonically decreasing with N , Theorem 1 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 more data is acquired. In general, deriving a worst-case bound on the model complexity is a challenging and active open research problem. However, in the case of GPs, convergence results can be shown for a 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 kernelized setting, where f ∗ resides in a Reproducing Kernel Hilbert Space (RKHS) of vector-valued functions.

Assumption 4. We assume that the functions f ∗ j , j ∈ { 1 , . . . , d x } lie in a RKHS with kernel k and have a bounded norm B , that is f ∗ ∈ H d x k,B , with H d x k,B = { f | ∥ f j ∥ k ≤ B,j = 1 , . . . , d x } .

In this setting, we model the posterior mean and epistemic uncertainty of the vector-valued function f ∗ with µ n ( z ) = [ µ n,j ( z )] j ≤ d x , and σ n ( z ) = [ σ n,j ( z )] j ≤ d x , where,

Here, y j 1: n corresponds to the noisy measurements of f ∗ j , i.e., the observed next state from the transitions dataset D 1: n , k n = [ k ( z , z i )] i ≤ nT , z i ∈ D 1: n , and K n = [ k ( z i , z l )] i,l ≤ nT , z i , z l ∈ D 1: n is the data kernel matrix. It is known that if f ∗ satisfies Assumption 4, then Equation (10) yields well-calibrated confidence intervals, i.e., that Assumption 1 is satisfied.

Lemma 3 (Well calibrated confidence intervals for RKHS, Rothfuss et al. (2023)) . Let f ∗ ∈ H d x k,B . Suppose µ n and σ n are the posterior mean and variance of a GP with kernel k , c.f., Equation (10) . There exists β n ( δ ) , for which the tuple ( µ n , σ n , β n ( δ )) satisfies Assumption 1 w.r.t. function f ∗ .

Theorem 2 presents our convergence guarantee for the kernelized case to the T -step reachability set R for the policy class π ∈ Π . In particular, R is defined as

There are two key differences from Theorem 1; ( i ) we can derive an upper bound on the epistemic uncertainties σ n , and ( ii ) we can bound the model complexity MC N ( f ∗ ) , with the maximum information gain of kernel k introduced by Srinivas et al. (2012), defined as

Theorem 2. Let Assumption 3 and 4 hold, Then, for all N ≥ 1 , with probability at least 1 -δ ,

If we relax noise Assumption 3 to σ -sub Gaussian. Then, if Assumption 2 holds, we have for all N ≥ 1 , with probability at least 1 -δ ,

Moreover, if γ N ( k ) = O ( polylog ( N )) , then for all z ∈ R , and 1 ≤ j ≤ d x ,

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, we can leverage concentration inequalities to give a bound without the expectation (see Appendix A for more detail).

For the Gaussian noise case, we obtain a tighter bound by leveraging the change of measure inequality from Kakade et al. (2020, Lemma C.2.) (c.f., Lemma 6 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 β T N term in our bound (Equation (12)). Nonetheless, for cases where the RKHS is of a kernel with maximum information gain γ N ( 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 finite-dimensional feature map ϕ ( x ) ) have maximum information gain which grows polylogarithmically 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.

5 Experiments

We evaluate OPAX on the Pendulum-v1 and MountainCar environment from the OpenAI 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 ( 1 ) 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, in our experiments, we compare the optimistic planner with other planning approaches. Moreover, 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 (MEAN-AE), and ( 3 ) trajectory sampling (TS-1) scheme proposed in Chua et al. (2018) (PETS-AE) as our baselines. The mean planner simply uses the mean estimate µ n 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-based RL algorithm. We investigate the following three aspects: ( i ) how fast does active exploration reduce model's epistemic uncertainty σ n with increasing n , ( ii ) can we solve downstream tasks with OPAX, and ( iii ) does OPAX scale to high-dimensional and challenging object manipulation tasks ? For our experiments, we use GPs and probabilistic ensembles (PE, Lakshminarayanan et al. (2017)) for modeling the dynamics. For the planning, we either the soft actor-critic (SAC, Haarnoja et 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-v1 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 2 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 MEAN-AE and PETS-AE. We believe this is due to the additional exploration induced by the optimistic planner.

0

5

10

15

Episodes

10

-

2

10

-

1

10

0

σn

Pendulum-v1 GP - uncertainty

0

10

20

Episodes

10

-

1

10

0

Pendulum-v1 - uncertainty

PETS-AE

Mean-AE

OpAx

Random

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-v1 swing up, ( ii ) Pendulum-v1 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, ( viii ) 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 a 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 ), ( vii ) with PEs. We report the best performance across all episodes.

To make a fair comparison, we use the following evaluation procedure; first, we perform active exploration for each episode on the environment, and then after every few episodes we use the mean estimate µ n to evaluate our learned model on the downstream tasks.

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, we notice that OPAX always performs well and is able to achieve 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 ( 1 ) 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 manipulation tasks? To answer this question, we consider the Fetch Pick & Place Construction environment (Li et al., 2020). We again use the active exploration agents to learn a 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 a 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, a random policy as well as CEE-US (Sancaktar et al., 2022). CEE-US is a model-based active exploration algorithm, for which Sancaktar et al. (2022) reports state-of-the-art performance compared to several other active exploration methods. In all three tasks,

OPAX is at least on par with the best-performing baselines, including CEE-US. We run OPAX and all baselines with the same architecture and hyperparameter settings. See Appendix B for more details.

PETS-AE

Mean-AE

OpAx

Random

H-UCRL

0

10

20

0

.

0

0

.

5

1

.

0

∑

T

t

=0

rt

×

10

4

MountainCar (1)

0

10

20

-

10

2

-

10

1

-

10

0

0

Pendulum-v1 - keep down (4)

0

10

-

2

-

1

∑

T

t

=0

rt

×

10

3

Pendulum-v1 GP - swing up (2)

0

100

-

3

-

2

-

1

×

10

2

Swimmer - go to target (5)

0

50

100

Episodes

0

1

2

∑

T

t

=0

rt

×

10

3

Reacher - go to target (3)

0

200

400

Episodes

0

5

×

10

2

Cheetah - run forward (6)

Downstream tasks where H-UCRL is trained

0

10

20

-

1

.

5

-

1

.

0

-

0

.

5

×

10

3

Pendulum-v1 - swing up (7)

0

50

100

150

3

4

5

×

10

2

Swimmer - go away from target (8)

0

200

400

Episodes

0

5

×

10

2

Cheetah - run backward (9)

New downstream tasks

6 Related Work

System identification is a broadly studied topic (Åström and Eykhoff, 1971; Schoukens and Ljung, 2019; Schön et al., 2011; Ziemann et al., 2022; Ziemann and Tu, 2022). However, system identification 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 are ( i ) model prediction error or 'Curiosity' (Schmidhuber, 1991; Pathak et al., 2017), ( ii ) novelty of transitions (Stadie et al., 2015), and ( iii ) diversity of skills (Eysenbach et al., 2018).

A popular choice for intrinsic rewards is mutual information or entropy (Jain et al., 2018; BuissonFenet et al., 2020; Shyam et al., 2019; Pathak et al., 2019; Sekar et al., 2020). Jain et al. (2018) propose an approach to maximize the information gain greedily wrt the immediate next transition, i.e., one-step greedy, whereas Buisson-Fenet et al. (2020) consider planning full trajectories. Shyam et al. (2019); Pathak et al. (2019); Sekar et al. (2020) and Sancaktar et al. (2022) consider general Bayesian models, such as BNNs, to represent a 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 a Gaussian mixture model (GMM). A similar approach is also proposed in Pathak et al. (2019); Sekar et al. (2020); Sancaktar et al. (2022). The main difference between Shyam et al. (2019) and Pathak et al. (2019) lies in how they represent mutual information. Moreover, Pathak et al. (2019) use the model's epistemic uncertainty, that is the

Success rate

100

50

0

100

200

300

Training Iteration

0

50

100

Pick & Place 4 Objects

0

100

200

300

Training Iteration

0

20

40

Throw 4 Objects

0

100

200

300

Training Iteration

0

15

30

0

100

200

300

Training Iteration

0

Random

Mean-AE

PETS-AE

OpAx

CEE-US

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-the-art 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, e.g., aleatoric noise, this approach is restricted to deterministic systems and susceptible to systems with unstable equilibria. Our approach is the most similar to Pathak et al. (2019); Sekar et al. (2020) since we 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-AE) or does not discuss the planner thoroughly or provide any theoretical results. In general, theoretical guarantees for active exploration algorithms are rather immature (Chakraborty et 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 a rich class of nonlinear systems.

While our work focuses on 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 et 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 et al., 2022) and result in practical algorithms. On the contrary, we focus on solely learning the dynamics. While a 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 to external effects. Furthermore, our proposed method is not only theoretically sound but also practical.

7 Conclusion

We present OPAX, a novel model-based RL algorithm for the active exploration of unknown dynamical systems. Taking inspiration from Bayesian experiment design, we provide a comprehensive explanation for using model epistemic uncertainty as an intrinsic reward for exploration. By leveraging the optimistic in the face of uncertainty paradigm, we put forth first-of-their-kind theoretical results on 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 to 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.

Acknowledgments and Disclosure of Funding

We would like to thank Jonas Hübotter for the insightful discussions and his feedback on this work. Furthermore, we 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 Automation, grant agreement 51NF40 180545, and the Microsoft Swiss Joint Research Center.

References

Balcan, M.-F., Hanneke, S., and Vaughan, J. W. (2010). The true sample complexity of active learning. Machine learning , 80:111-139.

Berkenkamp, F., Turchetta, M., Schoellig, A., and Krause, A. (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-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 arXiv:1606.01540 .

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, I. (1995). Bayesian experimental design: A review. Statistical science , pages 273-304.

Chen, J., Modi, A., Krishnamurthy, A., Jiang, N., and Agarwal, A. (2022). On the statistical efficiency of reward-free exploration in non-linear rl. Advances in Neural Information Processing Systems , 35:20960-20973.

Chen, Y., Hassani, S. H., Karbasi, A., and Krause, A. (2015). Sequential information maximization: When is greedy near-optimal? In COLT .

Chiuso, A. and Pillonetto, G. (2019). System identification: A machine learning perspective. Annual Review of Control, Robotics, and Autonomous Systems , 2:281-304.

Chowdhury, S. R. 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 a handful of trials using probabilistic dynamics models. In NeurIPS .

Cover, T. M. and Thomas, J. A. (2006). Elements of Information Theory (Wiley Series in Telecommunications 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 , 33:14156-14170.

Eysenbach, B., Gupta, A., Ibarz, J., and Levine, S. (2018). Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.06070 .

García, C. E., Prett, D. M., and Morari, M. (1989). Model predictive control: Theory and practice - a 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 a stochastic actor. In ICML , pages 1861-1870.

Hanneke, S. et al. (2014). Theory of disagreement-based active learning. Foundations and Trends® in Machine Learning , 7(2-3):131-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. In 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. In ICML , pages 2796-2804. PMLR.

Lakshminarayanan, B., Pritzel, A., and Blundell, C. (2017). Simple and scalable predictive uncertainty estimation using deep ensembles. NeurIPS , 30.

Li, R., Jabri, A., Darrell, T., and Agrawal, P. (2020). Towards practical multi-object manipulation using relational reinforcement learning. In ICRA , pages 4051-4058. IEEE.

Lindley, D. V. (1956). On a 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 on model-based reinforcement learning. arXiv preprint arXiv:2112.05244 .

Mutny, M., Janik, T., and Krause, A. (2023). Active exploration via experiment design in markov chains. In AISTATS .

Pasztor, B., Bogunovic, I., and Krause, A. (2021). Efficient model-based multi-agent mean-field reinforcement learning. arXiv preprint arXiv:2107.04050 .

Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. (2017). Curiosity-driven exploration by self-supervised 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. In ICML , pages 8737-8747. PMLR.

Rasmussen, C. 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 - a survey. Automatica , 7(2):123-162.

Rothfuss, J., Sukhija, B., Birchler, T., Kassraie, P., and Krause, A. (2023). Hallucinated adversarial control for conservative offline policy evaluation. UAI .

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, 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 , pages 222-227.

Schön, T. B., Wills, A., and Ninness, B. (2011). System identification of nonlinear state-space models. Automatica , pages 39-49.

Schoukens, J. 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´ skowski, W., and Gomez, F. (2019). Model-based active exploration. In ICML , pages 5779-5788. PMLR.

Simchowitz, M. and Foster, D. (2020). Naive exploration is optimal for online LQR. In ICML , Proceedings of 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-theoretic regret bounds for gaussian process optimization in the bandit setting. IEEE Transactions on Information Theory .

Stadie, B. C., Levine, S., and Abbeel, P. (2015). Incentivizing exploration in reinforcement learning with deep predictive models. arXiv preprint arXiv:1507.00814 .

Sussex, S., Makarova, A., and Krause, A. (2023). Model-based causal bayesian optimization. In ICLR .

Tarbouriech, J., Shekhar, S., Pirotta, M., Ghavamzadeh, 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.00690 .

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 rl in nonlinear systems. arXiv preprint arXiv:2306.09210 .

Wagenmaker, A. J., Chen, Y., Simchowitz, M., Du, S., and Jamieson, K. (2022). Reward-free rl is 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, T. M., and Wang, D. (2016). A survey of transfer learning. Journal of Big data , pages 1-40.

Zhang, Y. and Yang, Q. (2021). A survey on multi-task learning. IEEE Transactions on Knowledge and Data Engineering , pages 5586-5609.

Ziemann, I. and Tu, S. (2022). Learning with little mixing. NeurIPS , 35:4626-4637.

Ziemann, I. M., Sandberg, H., and Matni, N. (2022). Single trajectory nonparametric learning of nonlinear dynamics. In COLT , pages 3333-3364. PMLR.

Contents of Appendix

A Proofs for section 4

We first prove some key properties of our active exploration objective in Equation (6). Then, we prove Theorem 1 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 1 and 2 hold, then the following is true for all n ≥ 0 ,

Proof. The positivity of the reward follows from the positive definiteness of the epistemic uncertainty σ n -1 ,j . For (2), the following holds

From hereon, let J max = 1 / 2 T 2 d 2 x log 2 ( 1 + σ -2 σ 2 max ) . Finally, we show that this reward is Lipschitz continuous.

Where (*) is true because for all x ≥ 0 , log(1 + x ) ≤ x .

Corollary 1 (OPAX gives an optimistic estimate on Equation (6)) . Let Assumption 1 hold and π ∗ n denote the solution to Equation (6) and J n ( π ∗ n ) the resulting objective. Similarly, let π n and η n be the solution to Equation (7) and J n ( π n , η n ) the corresponding value of the objective. Then with probability at least 1 -δ we have for every episode n ∈ { 1 , . . . , N } :

Proof. Follows directly from Assumption 1.

A.1 Proof of Lemma 2

Lemma 5 (Difference in Policy performance) . Let J r,k ( π , x k ) = E τ π [ ∑ T -1 t = k r ( x t , π ( x t )) ] and A r,k ( π , x , a ) = E τ π [ r ( x , a ) + J r,k +1 ( π , x ′ ) -J r,k ( π , x )] with x ′ = f ∗ ( x , a ) + w . For simplicity we refer to J r, 0 ( π , x 0 ) = J r ( π , x 0 ) . The following holds for all x 0 ∈ X :

Proof.

Therefore we obtain

J r ( π ′ , x 0 ) -J r ( π , x 0 ) = E τ π ′ [ A 0 ( π , x 0 , π ′ ( x 0 ))] + E τ π ′ [ J r, 1 ( π ′ , x ′ 1 ) -J r, 1 ( π , x ′ 1 )] . Using the same argument for J r, 1 , J r, 2 , . . . , J r,T -1 and that J r,T ( π , x ) = 0 for all π ∈ Π and x ∈ X completes the proof.

Assume a policy π is fixed and dynamics are of the form:

Here u ∈ [ -1 , 1] d x . Furthermore, assume that the associated running rewards do not depend on u , that is, r ( x t ) , and let η ∈ Ξ denote the policy, i.e., η : X → [ -1 , 1] d x .

Corollary 2. The following holds for all x 0 ∈ X and policy π :

with x t +1 = µ n ( x ′ t , π ( x ′ t )) + β n ( δ ) σ ( x ′ t , π ( x ′ t )) η ( x ′ t ) + w t , and x ′ t +1 = µ n ( x ′ t , π ( x ′ t )) + β n ( δ ) σ ( x ′ t , π ( x ′ t )) η ′ ( x ′ t ) + w t .

Proof. From Lemma 5 we have

Furthermore,

Proof of Lemma 2. From Assumption 1 we know that with probability at least 1 -δ there exists a ¯ η such that f ∗ ( z ) = µ n ( z ) + β n ( δ ) σ ( z ) ¯ η ( x ) for all z ∈ Z .

A.2 Analyzing regret of optimistic planning

In the following, we analyze the regret of optimistic planning for both σ -Gaussian noise and σ -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 ( µ 1 , σ 2 I ) and N ( µ 2 , σ 2 I ) , and for any (appropriately measurable) positive function g , it holds that:

Proof.

Corollary 3 (Regret of optimistic planning for Gaussian noise) . Let π ∗ n , π n denote the solution to Equation (6) and Equation (7) respectively, and z ∗ n,t , z n,t the corresponding state-action pairs visited during their respective trajectories. Furthermore, let Assumption 1 - 3 hold. Then, the following is true for all n ≥ 0 , t ∈ [0 , T -1] , with probability at least 1 -δ

Proof. For simplicity, define g n ( x ) = J n,t +1 ( π n , η n , x ) . Note since w t ∼ N (0 , σ 2 I ) (Assumption 3), we have that x ′ t +1 and x t +1 are also Gaussians. Therefore, we can leverage Lemma 6.

Furthermore,

Lemma 7 (Regret of planning optimistically for sub-Gaussian noise) . Let π ∗ n , π n denote the solution to Equation (6) and Equation (7) respectively, and z ∗ n,t , z n,t the corresponding state-action pairs visited during their respective trajectories. Furthermore, let Assumption 1 and 2 hold, and relax Assumption 3 to σ -sub Gaussian noise. Then, the following is true for all n ≥ 0 with probability at least 1 -δ

Proof. Curi et al. (2020, Lemma 5) bound the regret with the sum of epistemic uncertainties for Lipschitz continuous reward functions, under Assumption 1 and 2 for sub-Gaussian noise (c.f., Rothfuss et al. (2023, Theorem 3.5) for a more rigorous derivation). For the active exploration setting, the reward in episode n +1 is

We show in Lemma 4 that our choice of exploration reward is Lipschitz continuous. Thus, can use the regret bound from Curi et al. (2020).

Compared to the Gaussian case, σ -sub Gaussian noise has the additional exponential dependence on the horizon T , i.e., the β T n term. This follows from the analysis through Lipschitz continuity. Moreover, as we show in Lemma 2, the regret of planning optimistically is proportional to the change in 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.

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 x ∈ R n have covariance K = E [ xx ⊤ ] (i.e., K ij = E [ x i x j ] , 1 ≤ i, j ≤ n ). Then

with equality if and only if x ∼ N ( µ , K ) for µ = E [ x ] .

Lemma 8 (Monotonocity of information gain) . Let τ π denote the trajectory induced by the policy π . Then, the following is true for all n ≥ 0 , policies π

Proof.

(information never hurts)

A direct consequence of Lemma 8 is the following corollary.

Corollary 4 (Information gain at N is less than the average gain till N ) . Let τ π denote the trajectory induced by the policy π . Then, the following is true for all N ≥ 1 , policies π ,

Next, we prove Lemma 1, which is central to our proposed active exploration objective in Equation (6).

Proof of Lemma 1. Let y τ π = { y t } T -1 t =0 = { f ∗ t + w t } T -1 t =0 , where f ∗ t = f ∗ ( z t ) . Furthermore, denote with Σ n ( f ∗ 0: T -1 ) the covariance of f ∗ 0: T -1 .

We can leverage the result from Lemma 1 to bound the average mutual information with the sum of epistemic uncertainties.

Lemma 9 (Average information gain is less than sum of average epistemic uncertainties) . Let Assumption 3 hold and denote with ¯ π N be the solution of Equation (4) . Then, for all N ≥ 1 and dataset D 1: N the following is true

where z ∗ n,t are the state-action tuples visited by the solution of Equation (6) , i.e., π ∗ n .

Proof.

Here (1) follows from the tower property. Note that the second expectation in (1) is wrt τ π conditioned on a realization of D 1: n -1 , where the conditioning is captured in the epistemic uncertainty σ n ( · ) .

We use the results from above, to prove Theorem 1.

Proof of theorem 1. Let ¯ π n denote the solution to Equation (4) at iteration n ≥ 1 . We first relate the information gain from OPAX to the information gain of ¯ π n .

In summary, the maximum expected mutual information at episode N is less than the mutual information of OPAX and the sum of model epistemic uncertainties. Crucial to the proof is the regret

bound for optimistic planning from Corollary 3.

Above, we show that the maximum expected mutual information can be upper bounded with the sum of epistemic uncertainties for the states OPAX visits during learning. Finally, we further upper bound this with the model complexity measure.

Theorem 1 gives a bound on the maximum expected mutual information w.r.t. the model complexity. We can use concentration inequalities such as Markov, to give a high probability bound on the

information gain. In particular, we have for all ϵ > 0

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 no meaning in the frequentist setting, we study the epistemic uncertainty of the GP models.

Corollary 5 (Monotonicity of the variance) . For all n ≥ 0 , and policies π the following is true.

Proof. Follows directly due to the monotonicity of GP posterior variance.

Next, we prove that the trajectory of Equation (6) at iteration n is upper-bounded with the maximum information gain.

Lemma 10. Let Assumption 2 - 4 hold Then, for all N ≥ 1 , with probability at least 1 -δ , we have

Moreover, relax noise Assumption 3 to σ -sub Gaussian. Then, for all N ≥ 1 , with probability at least 1 -δ , we have

Proof. Gaussian noise case : Let z ∗ n,t denote the state-action pair at time t for the trajectory of Equation (6) at iteration n ≥ 1 and π ∗ n the corresponding policy.

Sub-Gaussian noise case : The only difference between the Gaussian and sub-Gaussian case is the regret term (c.f., Corollary 3 and Lemma 7). In particular, the regret for 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 et al. (2020); Rothfuss et al. (2023) for a more detailed derivation.

Lemma 10 gives a sample complexity bound that holds for a richer class of kernels. Moreover, for GP models, β N ∝ √ γ N (Chowdhury and Gopalan, 2017). Therefore, for kernels, where lim N →∞ γ 2 N / N → 0 , we can show convergence (for the Gaussian case). We summarize bounds on γ N from Vakili et al. (2021) in Table 1.

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-Gaussian setting by considering its corresponding bound.

Lemma 11. The following is true for all N ≥ 0 and policies π ∈ Π ,

Proof.

Corollary 6. Let Assumption 2 and 4 hold, and relax noise Assumption 3 to σ -sub Gaussian. Then, for all N ≥ 1 , with probability at least 1 -δ , we have

Lemma 12. Let Assumption 2 and 4 hold, and relax noise Assumption 3 to σ -sub Gaussian. Furthermore, assume lim N →∞ β 2 N ( δ ) γ N ( k ) / N → 0 . Then for all N ≥ 1 , z ∈ R , and 1 ≤ j ≤ d x , with probability at least 1 -δ , we have

Proof. We first show that the expected epistemic uncertainty along a trajectory converges to zero almost surely. Then we leverage this result to show almost sure convergence of all trajectories induced

by π ∈ Π . To this end, let S n = E τ π ∗ n [( ∑ T -1 t =0 ∑ d x j =1 1 2 log ( 1 + σ 2 n,j ( z ∗ n,t ) σ 2 ))] for all n ≥ 0 . So far we have,

Consider a sequence { δ n } n ≥ 0 such that lim n →∞ δ n = 0 , and lim n →∞ β n ( δ n ) T 3 / 2 √ γ n n → 0 . Note, for GP models with lim N →∞ β 2 N ( δ ) γ N ( k ) / N → 0 , such a sequence of δ n exists (Chowdhury and Gopalan, 2017). Consider any ϵ > 0 and let N ∗ ( ϵ ) be the smallest integer such that

Then, we have

Note, since lim n →∞ δ n = 0 , we have

In particular, ∑ ∞ n =0 Pr( S n > ϵ ) < ∞ for all ϵ > 0 . Therefore, we obtain

Define the random variable V = lim n →∞ ( ∑ T -1 t =0 ∑ d x j =1 1 2 log ( 1 + σ 2 n,j ( z ∗ n,t ) σ 2 )) , with z ∗ n,t ∈ τ and τ ∼ τ π ∗ n . V represents the sum of epistemic uncertainties of a random trajectory induced by the sequence of policies { π n } n ≥ 0 . Note V ≥ 0 , therefore we apply Markov's inequality. Moreover, for all ϵ > 0 , we have

Accordingly, we get for all π ∈ Π .

Assume there exists a z ∈ R , such that for some ϵ , σ 2 n,j ( z ) > ϵ for all n ≥ 0 . Since, z ∈ R , there exists a t and π ∈ Π such that p ( z t = z | π, f ∗ ) > 0 . This implies that Pr( z ∈ τ π ) > 0 . However, from Equation (15), we have that σ 2 n,j ( z ) → 0 for N →∞ almost surely, which is a contradiction.

Finally, we leverage the results from above to prove Theorem 2.

Proof of Theorem 2. The proof follows directly from Corollary 6 and Lemma 12.

Hence, we have

A.5 Zero-shot guarantees

In this section, we give guarantees on the zero-shot performance of OPAX for a bounded cost function. We focus this section on the case of Gaussian noise. However, a similar analysis can be performed for the sub-Gaussian case and Lipschitz continuous costs. Since the analysis for both cases is similar, we only present the Gaussian case with bounded costs here.

Corollary 7. Consider the following optimal control problem

with bounded and positive costs. Then we have for all policies π with probability at least 1 -δ

where J c ( π , η P ) = max η ∈ Ξ J c ( π , η ) .

Proof. From Corollary 2 we get

with x t +1 = f ∗ ( x t , π ( x t )) + w t , and x P t +1 = µ n ( x t , π ( x t )) + β n ( δ ) σ ( x t , π ( x t )) η P ( x t ) + w t . Furthermore, the cost is positive and bounded. Therefore,

for all x , η , π . Accordingly, we can now use the same analysis as in Lemma 2 and get

Lemma 13. Consider the control problem in Equation (16) and let Assumption 3 and 4 hold. Furthermore, assume for every ϵ > 0 , there exists a finite integer n ∗ such that

and denote with ˆ π n the minimax optimal policy, i.e., the solution to min π ∈ Π max η ∈ Ξ J c ( π , η ) . Then for all n ≥ n ∗ , we have probability at least 1 -δ , J c (ˆ π N ) -J c ( π ∗ ) ≤ O ( ϵ ) .

Proof of Zero-shot performance. In Theorem 2 we give a rate at which the maximum uncertainty along a trajectory decreases:

Combining this with Corollary 7 we get

Hence, we have that for each policy π , our upper bound max η J c ( π , η P ) is ϵ precise, i.e.,

We leverage this to prove optimality for the minimax solution. For the sake of contradiction, assume that

This is a contradiction, which completes the proof.

Lemma 13 shows that OPAX also results in nearly-optimal zero-shot performance. The convergence criteria in Equation (17) is satisfied for kernels k that induce a very rich class of RKHS (c.f., Table 1).

B Experiment Details

The environment details and hyperparameters used for our experiments are presented in this section. In Appendix B.2 we discuss the experimental setup of the Fetch Pick & Construction environment (Li et al., 2020) in more detail.

B.1 Environment Details

Table 2 lists the rewards used for the different environments. We train the dynamics model after each episode of data collection. For training, we fix the number of epochs to determine the number of gradient steps. Furthermore, for computational reasons, we upper bound the number of gradient steps by 'maximum number of gradient steps'. The hyperparameters for the model-based agent are presented in Table 3. Furthermore, we present the iCEM hyperparameters in Table 4. For more detail on the iCEM hyperparameters see Pinneri et al. (2021).

Then we have,

-

Model-based SAC optimizer For the reacher and MountainCar environment we use scarce rewards (c.f., Table 2), which require long-term planning. Therefore, a receding horizon MPC approach is less suitable for these tasks. Accordingly, we use a policy network which we train using SAC (Haarnoja et al., 2018). Moreover, our model-based SAC uses transitions simulated through the learned model to train a policy. Accordingly, given a dataset of transitions from the true environment, we sample P initial states from the dataset. For each of the states, we simulate a trajectory of H steps using our learned model and the SAC policy. We collect the simulated transitions into a simulation data buffer D SIM, which we then use to perform G gradient steps as suggested by Haarnoja et al. (2018) to train the policy. The algorithm is summarized below, and we provide the SAC hyperparameters in Table 5.

Model-based SAC

Init: Stastistical Model M = ( µ n , σ n , β n ( δ )) , Dataset of transitions D n , initial policy π 0 , Model Rollout steps H

D SIM ←∅ for Training steps k = 1 , . . . , K do n

➤ Initialize simulated transitions buffer

➤ Sample P initial states from buffer for Initial state x 0 ∈ x 0 , 1: P do

D SIM ←D SIM ∪ MODELROLLOUT ( π k , x 0 , H, M

) ➤ Collect simulated transitions

Perform G gradient updates on π k as proposed in Haarnoja et al. (2018) using D SIM.

B.2 OPAX in the High-dimensional Fetch Pick & Place Environment

B.2.1 Environment and Model Details

In our experiments, we use an extension of the Fetch Pick & Place environment to multiple objects as proposed in Li et al. (2020) and further modified in Sancaktar et al. (2022) with the addition of a large table. This is a compositional object manipulation environment with an end-effector-controlled robot arm. The robot actions u ∈ R 4 correspond to the gripper movement in Cartesian coordinates and the gripper opening/closing. The robot state x robot ∈ R 10 contains positions and velocities of the end-effector as well as the gripper-state (open/close) and gripper-velocity. Each object's state x obj ∈ R 12 is given by its position, orientation (in Euler angles), and linear and angular velocities.

We follow the free-play paradigm as used in CEE-US (Sancaktar et al., 2022). At the beginning of free play, an ensemble of world models is randomly initialized with an empty replay buffer. At each iteration of free play, a certain number of rollouts (here: 20 rollouts with 100 timesteps each) are collected and then added to the replay buffer. Afterwards, the models in the ensemble are trained for a certain number of epochs on the collected data so far. Note that unlike in the original proposal by Sancaktar et al. (2022), we use Multilayer Perceptrons (MLP) as world models instead of Graph Neural Networks (GNN), for the sake of reducing run-time. This corresponds to the ablation MLP+iCEM presented in Sancaktar et al. (2022), which was shown to outperform all the baselines other than CEE-US with GNNs. As we are interested in exploring the difference in performance via injection of optimism into active exploration, we use the computationally less heavy MLPs in our work. This is reflected in the downstream task performance we report compared to the original CEE-US with GNNs. Details for the environment and models are summarized in Table 6.

After the free-play phase, we use the trained models to solve downstream tasks zero-shot via modelbased planning with iCEM. We test for the tasks pick & place, throwing and flipping with 4 objects. The rewards used for these tasks are the same as presented in Sancaktar et al. (2022).

(a) Fetch Pick & Place settings.

B.2.2 OPAX Heuristic Variant

In the case of Fetch Pick & Place with an high-dimensional observation space, we implement a heuristic of OPAX. Note that as Fetch Pick & Place is a deterministic environment without noise, we only model epistemic uncertainty.

OPAX (Heuristic Variant)

Input: Ensemble { f i } K k =1 , ϵ ≪ 1

Solve optimal control problem till convergence for the system: x t +1 = f j t ( x t , u t ) .

for n ∈ 1 , . . . , N max do

∑ t =0 ∑ i =1 ( ) σ n,i ( x t , u t ) with ensemble disagreement.

Rollout u ⋆ on the system and collect data

Update models { f i } K k =1 .

0: T 1 D n .

B.2.3 Controller Parameters

The set of default hyperparameters used for the iCEM controller are presented in Table 7, as they are used during the intrinsic phase for OPAX, CEE-US, and other baselines. The controller settings used for the extrinsic phase are given in Table 8. For more information on the hyperparameters and the iCEM algorithm, we refer the reader to Pinneri et al. (2021).

(b) MLP model settings.

C Study of exploration intrinsic rewards

0

5

10

15

10

-

2

10

-

1

10

0

σn

Pendulum-v1 GP - uncertainty

0

5

10

15

-

2

.

0

-

1

.

5

-

1

.

0

-

0

.

5

×

10

3

Pendulum-v1 GP - swing up

0

5

10

15

20

25

Episodes

0

.

0

0

.

5

1

.

0

∑

T

t

=0

rt

×

10

4

MountainCar

0

100

200

300

400

500

Episodes

0

2

4

6

×

10

2

Cheetah - run forward

OpAx

OpAx-sum

The intrinsic reward suggested in Equation (6) takes the log of the model epistemic uncertainty. Another common choice for the intrinsic reward is the epistemic uncertainty or model disagreement without the log (Pathak et al., 2019; Sekar et al., 2020). In the following Lemma, we show that these rewards are closely related.

Lemma 14. Let σ max = sup z ∈Z ; i ≥ 0;1 ≤ j ≤ d x σ i,j ( z ) and σ > 0 . Then for all i ≥ 0 and j ∈ { 1 , . . . , d x }

Proof. Curi et al. (2020) derive the first inequality on the left. The second inequality follows since log(1 + x ) ≤ x for all x ≥ 0 .

Due to this close relation between the two objectives, our theoretical findings also apply to the intrinsic reward proposed by (Pathak et al., 2019). Moreover, empirically we notice in Figure 5 that OPAX performs similarly when the sum of epistemic uncertainties is used instead of the objective in Equation (6). However, our objective can naturally be extended to the case of heteroscedastic aleatoric noise, since it trades off the ratio of epistemic and aleatoric uncertainty.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nThis 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.\n\nLIMITATIONS:\nThe paper does not explicitly discuss the limitations of this work in terms of empirical results or theoretical guarantees.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

comments: 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.

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).

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\"\n\nQUESTIONS:\n(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?\n\nLIMITATIONS:\nThe authors discuss the approach's limitations quite well.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\nSee the weakness section above.\n\nLIMITATIONS:\nMore discussions on related works and discussions on the zero-shot generalizability are needed to better characterize the contribution of this paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nMain 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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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\nWEAKNESSES:\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\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nN/A.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-20"}