Exploration via Planning for Information about the Optimal Trajectory
Authors: Viraj Mehta
Ian Char
Joseph Abbate
Rory Conlin
Mark D Boyer
Stefano Ermon
Jeff Schneider
Willie Neiswanger
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
information gain
exploration
planning

Abstract: Many potential applications of reinforcement learning (RL) are stymied by the large numbers of samples required to learn an effective policy. This is especially true when applying RL to real-world control tasks, e.g. in the sciences or robotics, where executing a policy in the environment is costly. In popular RL algorithms, agents typically explore either by adding stochasticity to a reward-maximizing policy or by attempting to gather maximal information about environment dynamics without taking the given task into account. In this work, we develop a method that allows us to plan for exploration while taking both the task and the current knowledge about the dynamics into account.  The key insight to our approach is to plan an action sequence that maximizes the expected information gain about the optimal trajectory for the task at hand. We demonstrate that our method learns strong policies with 2x fewer samples than strong exploration baselines and 200x fewer samples than model free methods on a diverse set of low-to-medium dimensional control tasks in both the open-loop and closed-loop control settings.

Exploration via Planning for Information about the Optimal Trajectory
Viraj Mehta', Ian Char?, Joseph Abbate', Rory Conlin', Mark D. Boyer", Stefano Ermon Jeff Schneider', Willie Neiswanger
Robotics Institute, Machine Learning Department, Carnegie Mellon University
SPrinceton Plasma Physics Laboratory, *Princeton University
ă Computer Science Department, Stanford University
(virajm, ichar, schneide)@cs. cmu. edu, [jabbate, wconl Spriniinttprnittton.i mboyer@pppl. gov, (ermon, meiswanger) Hecs. stanford. edu
Abstract
Many potential applications of reinforcement learning (RL) are stymied by the large numbers of samples required to learn an effective policy. This is especially true when applying RL to real-world control tasks, e.g. the sciences robotics, where executing a policy in the environment is costly. In popular RL algorithms, agents typically explore either by adding stochasticity to reward-maximizing policy or by attempting gather maximal information about environment dynamics without taking the given task into account. In this work, we develop a method that allows us to plan for exploration while taking both the task and the current knowledge about the dynamics into account. The key insight to our approach is to plan an action sequence that maximizes the expected information gain about the optimal trajectory for the task at hand. We demonstrate that our method learns strong policies with 2x fewer samples than strong exploration baselines and 200x fewer samples than model free methods on a diverse set of now-to-medium dimensional control tasks in both the open-loop and closed-loop control settings.
Introduction
The potential of reinforcement learning (RL) as general-purpose method of learning solutions to sequential decision making problems is difficult to overstate. Ideally, RL could allow for agents that learn to accomplish all manner of tasks solely through a given reward function and the agent's s experience; however, RL has so far broadly fallen short of this. Chief among the difficulties in realizing this potential is the fact that typical RL methods in continuous problems require very large numbers of samples to achieve a near-optimal policy. For many interesting applications RL this problem is exacerbated by the fact that collecting samples from the true environment costs. For example, expensive scientific experiments are needed for tokamak control [31, 13, 20] and design of molecules [70], and collecting experience the road is both costly and runs the risk of accidents for autonomous vehicles [66].
incur huge
Though model-based methods like Deisenroth and Rasmussen [21], Chua et al. [15], and Curi et al. [16] are much more efficient than typical model-free methods, they do not explicitly reason about the prospective task-relevant information content of future observations. Moreover, these methods typically require thousands of timesteps of environment dynamics data to solve reasonably simple Markov decision processes (MDPs). Though progress has been made in using Thompson sampling [46], entropy bonuses [27], and upper confidence bounds (UCB) [16. 5] to more intelligently explore
Code available at: attps://githhbl com/ wuwom//uuioonooeeo ml/trajectorr.eeee inf onmaatiionn
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Posterior of Optimall Trajcciory
Tra (oinn Information Gain
(b)
(c)
(d)
Trojyy Information Planning
1
(a)
(c)
Figure 1: schematic depiction Trajectory Information Planning (TIP). Suppose the agent in (a) aims to determine where to explore next from its current state. To do so, in (b) the agent samples dynamics models T' ~ P(T D) from its current posterior and finds approximately optimal trajectories T ~ P(r" T') for each sample. Then in (c) it pools these samples of posterior optimal trajectories r". In (d) it constructs function that gives the joint expected information gain about the optimal trajectory r' given a planned exploration trajectory (i.e. EIG,, over the set points visited). Finally, in (e) the agent can plan an action sequence which maximizes this joint expected information gain.
the state-action space of an MDP, these methods still do not explicitly reason about how information that the agent gathers will affect the estimated MDP solution. Furthermore, in continuous state- action settings these methods must make coarse approximations to be computationally tractable (c.g. bootstrapped Q networks to approxiate a posterior or one-step perturbations for approximate UCB).
Many methods in the vein of Pathak et al. [47] and Shyam et al. [61] explore directly to gather new information based on current uncertainty about environment dynamics, but they do not in general specialize the information that they aim to acquire for particular task specified by an initial state distribution and reward function. We believe that an ideal exploration strategy for sample efficiency should take into account this task, as well as uncertainty about the environment dynamics.
In this work, we start by showing how many methods can be cast as Bayesian planning problem over a specific cost function. This framework helps illuminate the importance of the cost function and what impact it has on exploration. Viewed in this light, it becomes clear that many previous state-of-the ahetart methods rely on cost functions that either result in behavior that is too greedy-i.c. i.c. the policy tries maximize returns during exploration to explore the environment dynamics and does not consider the task at hand. We therefore present a acquisition
too exploratory, i.c. the policy is incentivized
cost funetion that balances out these function introduced in Mehta et al. [40] to apply to a set of future hypothetically acquired data. In particular, our cost function captures the amount of information that would be gained about the optimal trajectory, if the agent were to explore by following a particular planned trajectory. As depicted in Figure this involves the agent sampling what it would hypothetically do given different realizations of the dynamics, and then planning actions that are informative about those possibilities.
extremes, by generalizing an information
In summary, the contributions of this work are as follows: exploration that explicitly accounts for both the specific task and uncertainty about environment dynamics, method for planning which applies the cost function to explore in MDPs with continuous states and actions, and a thorough empirical evaluation of our method across 5 closed-loop and open-loop environments (with a focus expensive RL tasks in plasma physics) compared against 14 baselines. We find that our proposed method is able to learn policies that perform as well as an agent with access to the ground truth dynamics using half or fewer samples than comparison methods.
develop novel cost function for
2 Related Work
Exploration in Reinforcement Learning The most common strategy for exploration in RL is to execute a greedy policy with some form of added stochasticity. The simplest approach, e-greedy exploration as used in Mnih et al. [41], takes the current action thought to be best with probability 1 - and a random action with probability Other methods use added Ornstein-Uhlenber noise [37] to the greedy policy, or entropy bonuses [27] to the policy or value function objectives to add noise to policy which is otherwise optimizing the RL objective.
action

---

Tabular RL is often solved by choosing actions based on upper confidence bounds on the value function| [14, 36]. but explicitly computing and optimizing these bounds in the continuous setting is substantially more challenging. Recent work [16] approximates this method by computing one-step confidence bounds on the dynamies and training 'hallucinated' policy which chooses perturbations within these bound to maximize expected policy performance. Another recent work [5] uses anti- concentration inequalities to approximate upper confidence bounds in MDPs with discrete actions.
Thompson sampling (TS) [55], which samples a realization of the MDP from the posterior and acts optimally as if the realization was the true model, can be applied for exploration in model-free manner as in [45] or in a model-based manner as in [63]. As the posterior over MDP dynamics or value functions can be high-dimensional and difficult to represent, the performance of TS can be hindered by approximation errors using both Gaussian processes and ensembles of neural networks. Curi et al. [16] recently investigated this and found that this was potentially due to an insufficiently expressive posterior over entire transition functions, implying that it may be quite difficult to solve tasks using sampled models. Similarly, the posterior over action-value functions in Oshand et al. [45] is only roughly approximated by training a bootstrapped ensemble of neural networks.
There is also rich literature of Bayesian methods for exploration, which are typically computationally expensive and hard to use, though they have attractive theoretical properties. These methods build upon the fundamental idea of the Bayes-adaptive MDP [53], which we detail in Section E. alongside a discussion of this literature.
Additionally, a broad set of methods explore to learn about the environment without addressing specified task. This line of work is characterized by Pathak et al. [47], which synthesizes a task- agnostic reward function from model errors. Other techniques include MAX [61], which optimizes the information gain about the environment dynamics, Random Network Distillation [11]. which forces the agent learn about a random neural network across the state space, and Plan2Explore [60], which prospectively plans to find areas of novelty where the dynamics are uncertain.
Bayesian Experimental Design: BOED, BO, BAX, and BARL There is a large literature on Bayesian optimal experiment design (BOED) [12] which focuses on efficiently querying a process or function to get maximal information about some quantity of interest. When the quantity of interest is the location of a function optimum, related strategies have been proposed as the entropy search family of Bayesian optimization (BO) algorithms [29, 30J. Recently, a flexible framework known as Bayesian algorithm execution (BAX) [43] has been proposed to efficiently estimate properties of expensive black-box functions; this framework gives a general procedure for sampling points which are informative about the future execution of given algorithm that computes the property of interest, thereby allowing the function property to be estimated with far less data.
A subsequent related work [40], known as Bayesian Active Reinforcement Learning (BARL), uses ideas from BOED and BAX to sample points that are maximally informative about the optimal trajectory in an MDP. However, BARI. relies on setting the authors call Transition Query Reinforce- ment Learning (TQRL), which assumes that the environment dynamics can be iteratively queried at an arbitrary sequence of state-action pairs chosen by the agent. TQRL is thus a highly restrictive setting which is not suitable when data can only be accessed via a trajectory (rollout) of environment dynamics; it typically relies on an accurate environment simulator of sufficient expense to warrant its use. Even then, there will likely be differences between simulators and ground truth dynamics for complex systems. Thus, one would ideally like to collect data in real environments. However, this often requires leaving the TQRL setting, and instcad collecting data via trajectories only.
In this paper, we aim to apply the information-thhoreee ideas from BARL but generalize them to the general MDP setting as well as learn open loop model-based controllers. The typical method for learning to solve open-loop control problems was demonstrated successfully in Tesch et al. [65], where a value function was learned from action sequences to task success. Our method takes a model-based approach to this problem, using similar exploration strategies as Bayesian optimization but benefitting from the more substantial supervision that is typical in dynamics model learning.
3 Problem Setting
In this work we deal with finite-horizon discrete-time Markov decision processes (MDPs) which consist of a sextuple AS, ro, H) where S is the state space, A is the action space, T is the transition function T S A + P(S) (using the convention that P(x)) is the set of probability
measures over X), 7: S x A S + R is reward function, po(s) is distribution over S of start
3

---

states, and H E N is the horizon (i.e. the length of an episode). We always assume S, Po, and H are known. We also assume the reward r known, though our development of the method can easily be generalized to the case where is unknown. Our primary object of interest is the transition function T, which we learn from data. We address both open and closed loop control settings. In the more common closed loop setting, our aim is to find a policy x S + AA that maximizes Objective (1) below. We will denote trajectories as T ~ p(T n. T') where T = [(50, ag), ( ag-1), sh] generated by Po. a, (ss), and 8,41 T(s;; ai). We can write the return of trajectory as
R(T) Eio r(s;, ats 8141) for the states and actions 8t, az that make up T. The MDP objective can then be written as
(T) Er- [R(7)].
We aim to maximize this objective while minimizing the number of samples from the ground truth transition function T' that are required to reach good performance. We denote the optimal policy as n² arg max, JT(T),, which we can assume to be deterministic [64] but not necessarily unique. We use r* to denote optimal trajectories, i.e. T ~ p(T TT)
Similarly. for the open-loop setting, we assume a fixed start state so and aim to find an action sequence that maximizes the sum of rewards in an episode. We will slightly abuse notation and write T ~ p(T | 00: and -T( 1) with these actions fixed in place of reactive policy, and again use 7" to refer to the trajectories generated by an optimal action sequence.
ao. ao...., a -0
We assume in this work that applying planning algorithms like [51] to dynamics function T will result in a trajectory that approximates T", We will primarily focus on a Gaussian process (GP) model of the transition dynamics in order to take advantage of its expressive representation of uncertainty and grounded methods for sampling, conditioning, and joint inference. There is substantial prior work using GPs in RL- see Section E.2 for a discussion of this literature. Under this modeling choice, we assume that the dynamics are drawn from a GP prior P(T) (see Section A.3 for further details on our GP model) and use P(T D) for the posterior over transition functions given a dynamics dataset of triples and trajectories t or next state predictions S' coerce and S' into triples of dynamics data, prior to the union with the dataset.
= {(si; as })).. In this work, unions r or D U S' between the dataset D
Trajectory Information Planning
Our method consists of a generic framework for Bayesian (or approximately Bayesian) model- predictive control and novel cost function for planning that allows us to explicitly plan to find the maximal amount of new information relevant to our task. In Section 4.1, we describe the MPC framework and highlight that many prior methods approximate this framework while using greedy cost function that corresponds to the future negative expected rewards or a pure exploration cost function that corresponds future information about the dynamics. Afterwards, in Section 4.2, derive our new cost function and describe how it is computed. The overall method we introduce simply applies this planning framework with our new cost function.
4.1 Model-Prediction Control in Bayesian Model-Based RL
In this section, we give a formulation of Bayesian planning for control that generalizes ideas from methods such as PILCO [21] and PETS [15]. This formulation highlights these methods' inherently greedy nature and hints at possible solution. The objective of Bayesian planning is to find the h-step action sequence that maximizes the expected future returns under model uncertainty. That is,
argmin P(T- 1.T") [C(rr)]
Og
for some cost function C' overr trajectories and some start state If operating in the open-loop control setting, the agent executes the sequence of actions found without replanning. This procedure can also be extended to closed-
control via model peeee--pre control (MPC), which involves re-planning (2) at every state the agent visits and playing the first action from the optimal sequence. Concretely, the MPC policy for our Bayesian setting is as follows:
TMPC (s) arg min min
ET Plo (
[C(Ta))
(3)
01...0h-11
4

---

Whether we do open-loop control or closed-loop control via MPC, the cost function C, is integral to how the agent will behave. Prior work has predominantly focused on two types of cost function:
Ca(T) - Σ H [Ts,s, DJ
Greedy Exploration
Task-Agnostic Exploration
Previous works such as Kamthe and Deisenroth [34] and PETS [15] use the greedy exploration cost function, C.- This cost function incentivizes trajectories that achieve high rewards over the next transitions on average. In works that focus on task-agnostic exploration such Sekar et al. [60] and Shyam et al. [61], the cost function Ce (or similar) is used to encourage the agent to find areas the state space in which the model is maximally uncertain. Note that we use T to refer to the greedy policy given by using (3) with Cg.
The optimization problem in (3) is typically approximately solved in one of three ways: Deisenroth and Rasmussen [21] and Curi et al. [16] directly backpropagate through the estimated dynamies and reward functions to find policy parameters that would generate good actions, Janner et al. [32] use an actor-critic method trained via rollouts in the model alongside the data collected to find policy, and Chua et al. [15] and Mehta et al. [40] use the cross-entropy method [17] to find action sequences which directly maximize the reward over the estimated dynamics. In this work, we use a version of the last method given in Pinneri et al. [51], denoted iCEM, to directly find action sequences that optimize the cost function being used. We approximate the expectation by playing the actions on multiple samples from the posterior P(T) D). Algorithm gives formal description of the method and Section A.5 provides further details.
Algorithm Bayesian Model-Predictive Control with Cost Function C'
Inputs: transition function episode query budget b, number of posterior function samples k. planning horizon h. Initialize D + 0.
for i E (1.... do
Sample start state 80 Po- for E (O...., H - do
Sample posterior functions (T)) ~ P(T' DD.
Approximately find arg min oi........ Li=1 BrepptiTh. [C(Tz)] via iCEM.
Execute action au by sampling 8441 T(s,, ao).
Update dataset + DU [($t, 81-1) end for end for return Fg for the posterior P(T') D).
4.2 A Task-Specific Cost Function based on Trajectory Information
In this work, we aim to explore by choosing actions that maximize the conditional expected informa- tion gain (EIG) about the optimal trajectory r" This the same overall goal as that of Mehta et al. [40], where the EIG,- acquisition function was introduced for this purpose. However, in this paper we generalize this acquisition function in order to allow for sequential information collection, and account for the redundant information that could be collected between timesteps. As discussed at length Osband et al. [46], essential to reason about how an action taken at the current timestep will affect the possibility of learning something useful in future timesteps. In other words, exploration must be deep and not greedy. Explicit examples are given in Osband et al. [46] where the time to find an e-optimal policy in tabular MDP is exponential in the state size unless exploration can be coordinated over large numbers of timesteps rather than being conducted independently at each action. As the EIG,, acquisition function is only defined over single state-action pair and mutual information is submodular, we cannot naively use the acquisition function as is (or sum over many datapoints) to choose actions that lead to good long-ten mengttemmae exploration. This is clear in e.g. navigation tasks, where the nearby points visited over trajectories will provide redundant information about the local environment.
We therefore give a cost function that generalizes FIG,- by taking a set of points to query and computing the joint expected information gain from observing the set. Our cost function is non-
5

---

Markovian in the state space of the MDP, but it is Markovian in the dataset, which represents point
in the belief space of the agent about the dynamics. Let X {a : z C S x A, |æ| < oo) be the set of finite subsets of the set of all state-action pairs. Our cost function C.- X - R is defined below to be the negative joint expected information gain about the optimal trajectory T for a subset X E X. In particular, assuming an existing dataset D, a set of h query points X = {(s,, Q.) (:) [E)) and random set of next states 5" = { ~ T(s az), E [h]],
(5) This formulation of C,- forces method to handle the redundant information among queries- it is likely that I [s1.. T*) + 1 (5/2, r-) > [(154 s2), ,²) and our method should avoid this overestimation. However, as written, this function relies computing entropies
Cx-(X) = Bs
[TH [T' DU S' - [T D).
high-dimensional trajectories where the form of the joint distribution of the elements is unknown. To tractably estimate this quantity, we use the fact that (X)
-I("", -²) - [[5", ,S") for the mutual information J. This allows us to exchange r" and our set of queries so that 7 is giving information about the posterior predictive distribution of our set. In other words,
Cx- = Ey (pp""p(((i))) [S - DJ.
DJ.
(6)
In order to compute the right-hand term, we must take samples Tiyy ~ P(r* D) 1,..., m,j =
1,.... n. To do this, we first sample m start states " from Po (we always set m = 1 in experiments but derive the procedure in general) and for each start state independently sample n posterior functions Ty P(T)) D) from our posterior over dynamies models. We then run a planning procedure using
iCEM [51] on each of the posterior functions from % using Th for (using our assumption that planning can generate approximately optimal trajectories given ground-truth dynamics), giving our sampled T.j. Formally, we can approximate Cy- via Monte-Carlo as
Cx) R
Σ (E(m) JE(n)
Σ DU Tiyl - -II[S' DJ.
Assuming the dynamics are modelled with Gaussian process, we can compute the joint Gaussian probability of the next states S' [52]. As the entropy of multivariate Gaussian depends only on the log-determinant of the covariance, log [Σ], we can tractably compute the joint entropy of the model predictions H [S'] DJ and optimize it with a zeroth order optimization algorithm. Finally, we must calculate the entropy Hi[ H[S'iD U T.y.. For this, we follow a similar strategy as Neiswanger et al. [43]: since T." is set of states given by the transition model, we can treat them as additional noiseless datapoints for our dynamies model and condition on them before computing the joint covariance matrix for S". Given this newly generalized acquisition function, we can instantiate method of planning in order to maximize future information gained. We give the concrete procedure for computing our acquisition function in Algorithm 2, noting that trajectories Tł do not depend on the query set X and can be cached for various values of X as long as the dataset D does not change.
Our ultimate procedure, which we name Trajectory Information Planning (TIP). is quite simple: run model-based RL using MPC as in Algorithm 1, but set the cost function to be C, (r) instead of C. or Ce, and compute this cost function using Algorithm 2. At test time, we return to planning with C2 as the cost function and greedily attempt to maximize returns rather than performing exploration. We can also formulate an open-loop variant of our method, oTIP, which involves planning once and then executing the entire action sequence.
4.3 Computational Cost and Implementation Details
Though the TIP algorithm is designed for settings where samples are expensive, it is important to understand, both theoretically and practically, the computational cost of this method. For ease of notation, make the simplifying assumption that the planning algorithm used (in this case, iCEM from [51]) evaluates action sequences consisting of h (the planning horizon) actions and that our current dataset is of size N. In order efficiently sample functions from the posterior over dynamics functions, we use the method from Wilson et al. [69]. This reduces the naive complexity of querying these functions from O(N3) to one-time ON) cost and then O(1) for additional queries. As we derive in Section A. the computational complexity of one TIP planning iteration
O nm (IN + H)² + ph (N + ) The two asymptotically expensive operations are (1)

---

Algorithm 2 Computation of Cr.
Inputs: dataset D {(5,, ags 5k) query set X, number of start state samples m, number of posterior function samples n.
Sample start states {060 pm Pa-
for i E [m]] do
Sample n posterior functions (T P(T D).
for
[ni] do Set j + TMPC using Cg and singleton posterior P(T) D) = 8(T3) in (3).
Compute Tig by executing Tj on 7 starting from 60
end for
end for
Compute the joint posterior covariance 5
D across all points X..
Compute the joint posterior covariances "
D DUy [n].ú E [m]] across all points in X.
return log ES Eie(n]JE[(] log E
computing the Cholesky decompositions of the nm kernel matrices for datasets D U T. and (2) solving the triangular systems using the cached Cholesky decompositions in order to compute the covariance matrices "
|UU Tij for each the p action sequences used by the planning algorithm.
However, our implementation choices mean that in practice these operations are not the most expensive step. The covariance matrix computations, which are the theoretical bottleneck, are implemented in JAX [9], allowing them to be compiled to much faster machine code and vectorized across large batches of queries. In fact, the most expensive operation in practice is planning on the sampled transition functions T. to sample optimal trajectories T... This is due to the fact that in practice p is large and we implemented the planner in NumPy [28] so cannot be compiled together with the Tensorflow [1] code from Wilson et al. [69], which is used for predicting which states will be visited for the planner. give further information on the implementation in Section A.
5 Experiments
Task aa-asnioticc (Cc)
Task Task specific (Cr:) Exploration
The aim of our development of the TIP algorithm and the C,+ acquisition function for RL is to reduce the sample complexity of learning effective policies in continuous MDPs given limited access expensive dynamics. In this section we demonstrate the effec- tiveness of TIP in quickly learning a good policy by comparing against variety of state-of-the rein- forcement learning algorithms and strong baselines (including some that use the TQRL setting from [40], which is also known as RL with generative model in Kakade [33] and other works [6, 4l).
w/ Stochasticity
Greedy (Cs)
Exploration
TQRL
N/A
EIGr
BARL
Rollout (Sum)
MPC, PETS,
sDIP
sTIP
In particular, we compare the average return across five evaluation episodes across five random seeds of each algorithm on five closed-loop control problems. For sample complexity we assess the median amount of data taken by each algorithm to 'solve" the prob- lem across five seeds with the threshold performance given by an MPC controller using the ground truth dynamics. We evaluate the open-loop variant of our method, oTIP, against three comparison methods on three control problems suitable for open-loop control. In particular, to be suitable for open-loop control, the problem cannot be dynamically unstable (as Pen- dulum and Cartpole famously are) and must have a relatively short control horizon and fixed start state.
BPTT,
PPO.
Rollout (Joint)
SAC,
TIP, TS, HUCRL
TD3
DIP
Figure 2: Our comparison methods can be broken down by the type of cost function used and how methods do do not handle sequential acqui- sition of information. As C, is a sum, it naturally handles future timesteps jointly. For the other in- formation quantities, it is possible to upper-bound information acquired by summing each separate mutual information, or to compute them jointly.

---

Environment
TIP
sTIP
DIP
MPC
PETS
SAC
FEEF
RHC
HUCRL
TS
BARL
EIGr
Pendulum
21
36
36
46
5.6k
7k
800
>40k
>50k
>50k
21
56
Caripole
131
141
161
201
1.63k
32k
>2.5k
>5k
>6k
>6k
111
121
3 Tracking
46
76
276
76
330
12k
300
>3k
480
420
186
>1k
3 + Rotation
201
>500
>500
>500
400
30k
>2k
>2k
>5k
>5k
>500
>1k
Reacher
251
>400
>1k
751
700
233
>5k
1.5k
6.6k
4.5k
251
>1.5k
Table I: Sample Complexity: Median number of samples across five seeds required to reach "solved' perfor- mance, averaged across five trials. We determine 'solved' performance by running an MPC policy (similar to the used for evaluation) on the ground truth dynamics to predict actions. We record unable to solve the problem by the end of training after collecting n datapoints. The methods in the rightmost section operate in the TQRL setting and therefore have more ficxible access to the MDP dynamics for data collection. The full set of methods are shown in Section D as well as boxplots depicting the data in Figure 4. Here too, we assess the average return as open-loop trials are conducted as well as the number of timesteps required to achieve 'solved' performance.
when the median run is
Comparison Methods We use several model-based and model-free comparison methods in this work. We compare to several published model-based methods. These include PETS [15], as implemented by Pineda et al. [50], which uses probabilistic ensemble of neural networks and CEM over particle samples to do MPC. We also compare against three model-based techniques from the HUCRL [16] implementation: HUCRL itself, which relies hallucinating dynamics perturbations as a way of realizing an upper confidence bound on the policy, model-based Thompson Sampling (TS), which samples from the posterior over models and chooses optimal actions for that sample, and greedy model-based neural network method relying on backpropagation through time (BPTT). We also compare against the Free Energy of the Expected Future method from Tschantz et al. [67], which treats directed exploration as process of actively collecting information for inference on reward-biased generative model. A further comparison is with Receding Horizon Curiosity (RHC) [59], which does online Bayesian system identification over linear model in order to quickly find a model of the environment dynamics. Our model-fres -feee comparison methods, Soft Actor-Critie (SAC) [27], an actor-critic method that uses an entropy bonus over the policy to encourage more exploration, and two others (TD3 and PPO), are in the appendix.
Finally, we compare against various ablations of the proposed method. These vary across two axes described in Figure 2: the cost function they use and how they handle sequential queries. Besides these differences, they use the same GP model and iCEM planning algorithm with the same hyperparameters. so they are truly comparable methods. The three cost functions used are Cg. Cr-, and C.. DIP, oDIP, and EIGT all use C. but compute, respectively, the expected joint entropy of the action sequence, the sum of the pointwise entropies of the action sequence, and the individual pointwise entropies in the TQRL setting. These methods are very similar in spirit to [61, 48] in that they plan for future information gain about the dynamics, but we chose to compare in way that controls for difference in the model and planning algorithm. MPC uses Cg and is very close to the method [34]. Like TIP, BARL [40] and sTIP use C,.. BARL operates in the TQRL. setting and can therefore use the simpler FIG,. acquisition function. sTIP investigates the use of -Es EIG,, (s;, a:) as a cost function for planning. This computes individual information gains for each future observation without accounting for the information they may have in common and is therefore an overestimate of the joint information gain. In the open-loop setting we compare against oDIP and oMPC, the open-loop variants of DIP and MPC, and Bayesian optimization (BO) as implemented by Pedregosa et al. [49]. oDIP plans an action sequence to minimize the joint C. and executes the actions found for each open-loop trial, while oMPC does the same thing using Cg- We give additional details on the comparison methods in Section B.
Control Problems Our closed loop control problems are the standard underactuated Pendulum swing-up task (Pendulum-vo from Brockman et al. [10]) with 2D states and 1D actions, Cartpole swing-up task with a sparse reward function, 2D s, and 1D actions, a 2-DOF robot arm control problem where the end effector is moved to a goal (Reacher-v2 from Brockman et al. [10]) with 10D states and 2D actions, a simplified 33 Tracking problem from plasma control [13, 39] (similar in design but not identical to the one from Mehta al. [40]) trained using with 4D states and ID actions, and a more complicated problem in plasma control where 3 Rotation are tracked with 10D states and 2D actions. Our open loop control problems are a navigation problem with hazards (Lava Path) from [40] and two regulation problems with different Nonlinear Gain functions. The Lava Path problem has 4D states and 2D actions and the nonlinear regulation problems have 2D states and 2D actions. Full details on these problems are available in Section C.
8

---

Control Performance ( Rotation
Miale MSE as Cherric MPC
Sool SIN On TUAA
RE Nospookaa
Sunler of Qucss
Namber of Qraviis
Nirdee of Queiis
Ntreber Queike
Control Performance Reacher
Kosta KOU oE Comree SIVC
MATE
Nurber Queiss
Nardoe Performance Pendulum
Nurber
Co 2014
22 Frog
Stnber
Nuxlex
Figure Control Modeling Details for TIP and Ablations. Column I: Learning curves our ablation methods, all of which use the same planner and model. Column 2: Dynamics model accuracy on the points used by the planner to choose actions during MPC. Column 3: Dynamics model accuracy uniformly random test set S. Column EIG,, values normalized by number of actions planned. sTIP was truncated on Reacher as exceeded the wall time budget.
Results As can be seen in Table 1, TIP is able to reach solved performance more quickly across the board than the model-based and model-free external baselines, often using frac- tion or even orders of magnitude less data than other methods. For many our ablation meth- ods we see failures to solve some of the prob- lems even though the model is demonstrated by TIP to be able sufficiently predict the dynam- ics. This is especially apparent on the harder plasma control environment, 3+Rotation, where TIP is the only method using our GP which is able to solve the problem. We believe that this is because the data acquired through exploration by the ablation methods less useful for control than the data TIP collects. This underscored by the second column of Figure where it is clear that TIP achieves the lowest modeling error on the points actually needed during the execution of the policy but not the uniform test set. In particular + Rotation environment, as BARL should in principle have a strietly stronger access to the problem and is optimizing the same quantity with fewer constraints. We hypothesis that this may be due to the fact that BARL optimizes the acquisition function EIG,- by simply uniformly sampling a set of points and choosing the one that evaluates the largest value. Our more sophisticated optimization algorithm and forced initialization at the start state distribution seems to allow us to collect more information in this case. This interpretation bolstered by the fact that on the problems where TIP outperforms BARL, we see that TIP is actually collecting more information per action than BARL as evidenced by larger EIG values. We also see clearly that there value in computing the C,- function rather than summing over EIG,, values, as TIP outperforms sTIP across the board. Additionally, there is clear evidence for the value of task fask-specific exploration as the task-agnostic exploration methods (EIGT, DIP, sDIP) underperform both in returns and model error the trajectories visited.
Environment
oTIP
oMPC
oDIP
BO
Nonlinear Gain 1
41
91
51
210
Nonlincar Gain 2
51
61
>200
60
Lava Path
101
101
>2k
Table 2:
Open Loop Sample Complexity: Median number of samples required reach 'solved" perfor- mance, averaged across five trials. We determine 'solved' performance by running an MPC policy on the ground truth dynamics predict actions. We record the median is unable to solve the prohlem by the end training after collecting v datapoints.
when
find it interesting that TIP outperforms BARL on the

---

For the open-loop experiments (Table 2), we also see strong performance from oTIP. As the model- based methods benefit from observing many model transitions for each open-loop trial, is unsurpris- ing that they are more sample-efficient than the BO method. Within the model-based techniques, oTTP the most sample efficient. We believe that this is for much the same reasons as the closed-loop case exciting evidence that the C,. cost function can be applied in a variety of settings.
6 Conclusion
In this work, we presented and evaluated a cost function designed for intelligent, task-aware ex- ploration. Using this cost function in MDPs with far less data than comparison methods in open- and closed-loop settings. Though the method is effective in data reduction, it is computationally expensive and relies dynamics that are well-modeled by a GP. In future work, we aim to scale the method to higher-dimensional and more complex control environments. We also aim to apply this method in the real world. In particular, we aim to address similar plasma control problems a small number of trials on a real tokamak.
control allows agents to solve continuous
Acknowledgements
This work was supported in part by US Department of Energy grants under contract numbers DE-SC0021414 and DE-AC02-09CHI146
This work is supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE1745016 and DGE2140733 Any opinions, findings, and conclusions or recommendations expressed in this material are those the author(s) and do not necessarily reflect the views of the National Science Foundation.
This work is additionally supported by NSF (#1651565), AFOSR (FA95501910024), ARO (W911NF- 21-1-0125), CZ Biohub, and Sloan Fellowship.
References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffreey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shiens, Benoit Steiner, llya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large- scale machine learning on heterogeneous systems, 2015. URL https: //wp:////wwww./// tensorflow. org/. Software available from tensorflow.org.
[2] Joseph Abbate, R Conlin, and Kolemen. Data-driven profile prediction for diiid. Nuclear Fusion, (14:00427 2021.
[3] Jan Achterhold and Joerg Stueckler. Explore the context: conditional dynamics models. In Arindam Banerjee and Kenji Fukumizu, editors, Proceedings The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 3529-3537. PMLR, 13-15 Apr 2021. URL https://proceeding httprdinssmmm...ees//esee.e http//33/////// htt113//achthhhlalacc.
Optimal data collection for context-
mlr pwwpspeesveee.e.e
html.
[4] Alekh Agarwal, Sham Kakade, and Lin F. Yang. Model-based reinforcement learning with a generative model is minimax optimal. In Jacob Abernethy and Shivani Agarwal, ed- itors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Pro- ceedings of Machine Learning Research, pages 67-83. PMLR, 09-12 Jul 2020. URL html.
mlr
[5] Jordan T. Ash, Cyril Zhang, Surbhi Goel, Akshay Krishnamurthy, and Sham M. Kakade. Anti- concentrated confidence bonuses for scalable exploratio.. In International Conference on Learn- ing Representations, 2022. URL
[6] Mohammad Gheshlaghi Azar, Rémi Munos, and Hilbert J Kappen. Minimax pac bounds on the sample complexity of reinforcement learning with generative model. Machine learning, 91 (3):325-349, 2013.
10

---

[7] Philip Ball, Jack Parker-Holdes Aldo Pacchiano, Krzysztof Choromanski, and Stephen Roberts. Ready policy one: World building through active learning. Hal Daumé III and Aarti Singh. editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 591-601. PMLR, 13-18 Jul 2020. URI. https://proceedingg mlr wwwppresssveeeeee v1 ht//////blllll httpb/a...ala.l......
[8] A. Bondeson and D.J. Ward. Stabilization of external modes in tokamaks by resistive walls and plasma rotation. Physical Review Letters, 72(17):770--2712, 1994.
[9] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://githut com/ google/
[10] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang. and Wojeiech Zaremb.. Openai gym, 2016.
[11] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2019. URL attps://openrevieet
[12] Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: A review. Statistical Science, pages 273-304, 1995.
[13] Ian Char, Youngscog Chung, Willie Neiswanger, Kirthevasan Kandasamy, Andrew O Nelson, Mark Boyer, Egemen Kolemen, and Jeff Schneider. Offline contextual bayesian optimization. Advances in Neural Information Processing Systems, 32:4627 2019.
[14] Richard Y. Chen, Szymon Sidor, Pieter Abbeel, and John Schulman. UCB and infogain exploration via q-ensembles. CoRR, abs/1706.01502, 2017. URL //ax://arxiv. org/ abs/ 1706. 01502.
[15] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep re- inforcement learning in a handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https:///roott httpa//ee/cccc.o...o//oo.. www.ss.eurrssp... www.eeew.u..pppppeeeeeeeee cccppapprr2018/fi22 Pap.es........r........
[16] Sebastian Curi, Felix Berkenkamp, and Andreas Krause.
Efficient model- based reinforcement learning through optimistic policy search and planning. In NeurIPS, 2020. URL pttppe://proceedingg.. neurips cc/paaper/202//aaeee
13bb99aabb3344152422222222222222 wtw.://2270011--ttttttt..... Abstract.html
[17] Pieter-Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Y Rubinstein. A tutorial on the cross-entropy method. Annals of operations research, 134(1):19-67 67, 2005.
[18] Richard Dearden, Nir Friedman, and Stuart Russell. Bayesian q-learning. In Aaaifiaai pages
761-768, 1998.
[19] Richard Dearden, Nir Friedman, and David Andre. Model-based bayesian exploration. CoRR, bbs//130.6660, 1999. URL httt:://ariiv org/ httg//bs/3300 6690.
[20] Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature, 602(7897): 414-419, 2022.
[21] Mare Peter Deisenroth and Carl Edward Rasmussen. Pilco: A model--ased and eafaceefiiceent approach policy search. In Proceedings of the 28th International Conference on Interna- tional Conference on Machine Learning, ICML'11, page 465-472, Madison, WI, USA, 2011. Omnipress. ISBN 9781 450306195.
[22] Max E Fenstermacher, J Abbate, S Abe, T Abrams, M Adams, B Adamson, N Aiba, T Akiyama, Aleynikov, H Allen, et al. Diii-d research advancing the physics basis for optimizing the tokamak approach to fusion energy. Nuclear Fusion, 0424)4 2022.
[23] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine
11

---

Learning Research, pages 1587-1596. PMLR, 10-15 Jul 2018. URL Utthhtpp://proceddins mlr. ppress//800ffuuiooooooBB. html.
[24]
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, and Aviv Tamar. Bayesian reinforce-
ment learning: A survey. CoRR, (bss/1609.04436 2016. URL 04436.
[25] R. J. Groebner, K. H. Burrell, and R. P. Seraydarian. Role of edge electric field and poloidal rotation in the transition. Physical Review Letters, 3-12533, 1990. ISSN 00319007.
doi: 10.1 103//hyssRel.eet..44.. 64.3015.
[26] Arthur Guez, David Silver, and Peter Dayan. Efficient bayes-adaptive reinforcement learning using sample-based search. In Proceedings of the 25th International Conference on Neural Information Processing Systems- - Volume 1, NIPS 12, page 1025-1033, Red Hook, NY, USA,
2012. Curran Associates Inc.
[27] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pages 1861-1870. PMLR, 2018.
[28] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers. Pauli Vir- tanen, David Couraaeeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Hal- dane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre
Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825) September 2020. doi: 10. 00.1038/s41586-02 02002269-22 URL https: //doi. org/ 10. 1038/s41586-00000 -20-5666--00-----22--22
[29] Philipp Hennig and Christian J Schuler. Entropy search for information-efficient global opti-
mization. Journal of Machine Learning Research, 13(6), 2012.
[30] José Miguel Mernández-Lobbrdd entropy search for efficient global optimization of block-box functions. Advances in neural information processing systems, 27, 2014.
Matthew W Hoffman, and Zoubin Ghahramani. Predictive
[31] David Humphreys, G Ambrosino, Peter de Vries, Federico Felici, Sun H Kim, Gary Jackson, A Kallenbach, Egemen Kolemen, J Lister, D Moreau, et al. Novel aspects of plasma control in iter. Physics of Plasmas, (2(2):021806, 2015.
[32] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems, 32, 2019.
[33] Sham Machandranath Kakade. On the sample complexity of reinforcement learning. University of London, University College London (United Kingdom), 2003.
[34] Sanket Kamthe and Mare Deisenroth. Data-efficient reinforcement learning with probabilistic model predictive control. In Amos Storkey and Fernando Perez-Cruz, editors, Proceedings of the Twenty-Firg First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 1701-1710. PMLR, 09-11 Apr 2018. URL https: / /ttpe//ee/ccceec......mm mlr wwwpspree//eepepp hatps//////kmahheeeeeeee html.
[35] J Zico Kolter and Andrew Y Ng. Near-bayesian exploration in polynomial time. In Proceedings of the 26th annual international conference on machine learming. pages 513-520, 2009.
[36] Kimin Lee, Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Sunrise: A simple unified framework for ensemble learning in deep reinforcement learning. In International Conference on Machine Learning, pages 6131-6141. PMLR, 2021.
[37] Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv: w.5...0099997 2015.
[38] David IC MacKay et al. Bayesian nonlinear modeling for the prediction competition. ASHRAE Iransactions, 100(2): 1053-1062, 1994.
[39] Viraj Mehta, Ian Char, Willie Neiswanger, Youngseog Chung, Andrew Nelson, Mark Boyer, Egemen Kolemen, and Jeff Schneider. Neural dynamical systems: Balancing structure and flexibility physical prediction. In 2021 60th IEEE Conference on Decision and Control
(CDC), pages 3735-3742 IEEE, 2021.
12

---

[40] Viraj Mehta, Biswajit Paria, Jeff Schneider, Stefano Ermon, and Willie Neiswanger. An experimental design perspective on model-based reinforcement learning. In International Conference on Learning Representations, 2022.
[41] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv: /3.5.602, 2013.
[42] Radford M Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, 1995.
[43] Willie Neiswanger, Ke Alexander Wang, and Stefano Ermon. Bayesian algorithm execu- tion: Estimating computable properties of black-box functions using mutual information. In International Conference on Machine Learning. PMLR, 2021.
[44] Nikolay Nikolov, Johannes Kirschner, Felix Berkenkamp, and Andreas Krause. Information- directed exploration for deep reinforcement learning. In International Conference on Learning Representations, 2019. URL https://opeereviiee net/ fttt/formmiidd Byyx33300K...
[45] Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep explo- ration via bootstrapped dqn. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Gar- nett, editors. Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016. URL https: http///ooceeenngeeeeeee/eeee www.eeedddinnnnnnns...... httpppcccccpppe/2///11111 14B81866ee1-00644444415555666771 2att://50pPPPppeeeeeee capper.pdf...
[46] Ian Oshand, Benjamin Van Roy, Daniel J. Russo, and Zheng Wen. Deep exploration via randomized value functions. Journal of Machine Learning Research, 20(124):1-62, 2019. URL
100000/813333 339. html.
[47] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference machine learning, pages 2778- 2787. PMLR, 2017.
[48] Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Sulf-super exploration via dis- agreement. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5062-5071. PMLR, 09-15 Jun 2019. URL https: / //proceedings.. mlr. prwss/v97 http///tttttattta..aaa html.
[49] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn Machine learning in Python, Journal of Machine Learning Research, 12:2825-288 22300. 2011.
[50] Luis Pineda, Brandon Amos, Amy Zhang, Nathan O. Lambert, and Roberto Calandra. Mbrl- lib: A modular library for model-based reinforcement learning. Arxiv, 2021. URL https: //arxiv org/ oogg/bs///// /bop/2104..0 10159.
[51] Cristina Pinneri, Shambhuraj Sawant, Sebastian Blaes, Jan Achterhold, Joerg Stueckler, Michal Rolinek, and Georg Martius. Sample-efficient cross-entropy method for real-time planning. arXiv preprint 20Xiv:2008..6389, 2020.
[52] Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning, pages 63-71. Springer, 2003.
[53] Stephane Ross, Brahim Chaib-draa, and Joelle Pineau. Bayes-adaptive pomdps. In NIPS, pages 1225-1232, 2007.
[54] Daniel Russo and Benjamin Van Roy. Learning to optimize via information-diircter sampling. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Wein- berger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL https: htttts://prooeeeeeeeeeeeeeee http///ppoeeen....//////// wtt.oeeeediinn.eu.niispse.eeeepp httpppsccccpappp/0///110 (uia00eluuccbb177722449994422229
Paper. pdf.
[55] Daniel J Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, et al. A tutorial on thompson sampling. Foundations and Trends@ in Machine Learning, 11(1):1-96, 2018.
[56] Ilya Ryzhov and Warren Powell. Information collection on graph. Operations Research, 59(1):188-20 2011.
13

---

[57] Ilya O Ryzhov. Martijn RK Mes, Warren B Powell, and Gerald van den Berg. Bayesian exploration for approximate dynamic programming. Operations research, 67(1):198-22 2019.
[58] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal
policy optimization algorithms. arXiv preprint arXiv:1707.06344 2017.
[59] Matthias Schultheis, Boris Belousov, Hany Abdulsamad, and Jan Peters. Receding horizon curiosity. In Conference on robot learning. pages 1278-1288. PMLR, 2020.
[60] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International Conference on Machine Learning, pages 8583-8592. PMLR, 2020.
[61] Pranav Shyam, Wojciech Jaskowski, and Faustino Gomez. Model-based active exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 5779-5788. PMLR, 09-15 Jun 2019. URL https: htttps://proceediiee mlr salr.press/v9/ shyam19a.. html.
[62] JA Stillerman, TW Fredian, KA Klare, and G Manduchi. Mdsplus data acquisition system. Review of Scientific Instruments, 68(1):939 (111)9)99944) 992, 1997.
[63] Malcolm Strens. A bayesian framework for reinforcement learning. In ICMI,, volume 2000, pages 943-950, 2000.
[64] Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction, MIT Press. Cambridge. MA. USA, 1998. ISBN 0-262-19398-1. URL http: /ttp:://wwww... cs ualberta, ca/
X/TEsuttaa/booozeeool/////// book. html.
[65] Matthew Tesch, Jeff Schneider, and Howie Choset. Expensive function optimization with stochastic binary outcomes. In Sanjoy Dasgupta and David McAllester, editors, Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pages 1283-1291, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR. URL html.
https: https::/pprottt http///proeeeen...eee////o .ml. mlr pwwss/v2//teeseeeeeeeee
[66] Marin Toromanoff, Emilie Wirbel, and Fabien Moutarde.
End-to end model-fire fodel-free reinforce- ment learning for urban driving using implicit affordances. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 7153-7163 2020.
[67] Alexander Tschantz, Beren Millidge, Anil K Seth, and Christopher L Buckley, Reinforcement learning through active inference. arXiv preprini urXiv:2002. /2636, 2020.
[68] Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for regression. 1996.
[69] James Wilson, Viacheslav Borovitskiy, Alexander Terenin, Peter Mostowsky, and Mare Deisen- roth. Efficiently sampling functions from gaussian process posteriors. In International Confer- ence Machine Learning. pages 10292 10302. PMLR, 2020.
[70] Zhenpeng Zhou, Steven Kearnes, Li Li. Richard N Zare, and Patrick Riley. Optimization of molecules via deep reinforcement learning. Scientific reports, 9(1):1-10, 2019.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions of all theoretical results? [Yes] Assumptions were fully stated for all derivations.
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] All code and instructions are included in supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] All training details are specified in the paper or in the code.
(c) Did you report error bars (e-g., with respect the random seed after running experi- ments multiple times)? [Yes] We show error bars computed via the standard error of mean performance between runs.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] We detail computational cost (see Sec A.2).
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] All existing assets are cited.
(b) Did you mention the license of the assets? [Yes] All licenses are properly attributed.
(c) Did you include any assets either in the supplemental material URL? [Yes] New assets are included supplementary material.
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [N/A] No human data was used.
(e) Did you discuss whether the data you are using/curating contains personally identi- fiable information offensive content? [N/A] No data with personally identifiable information was used.
5. If you used crowdsoureing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
15


Reviews:
[{"review_id": "qYZjBE37trX", "reviewer": "Reviewer_ZhTc", "summary": "The paper aims at improving data efficiency for RL, targeting real systems where data acquisition is normally costly.\nThe authors consider model-based control for MDPs with explicitly modelled dynamics.\nControl is relegated to MPC, the procurement of parametric (amortized) policies is left out of scope.\nThe focus is on obtaining useful data for learning the dynamics model as fast as possible, so that MPC can successfully control the agent.\nThe method shows significant improvement in data efficiency in 5 different environments with simple to moderately complex dynamics.\n\nThe main premise is that explorative actions are selected such that maximum information (in the sense of infogain, mutual information, MI) is gained for the optimal future trajectory $\\tau^*$ under the **current** dynamics model, noting that this is better than plain max. entropy RL, as it takes the task into account.\nThis is captured by the *Expected Information Gain* (EIG) objective, which has been proposed for the RL setting in prior work [1, 2].\nIn this paper, the objective is extended to account for the acquisition of data in sequence, i.e. in the form of rollouts in the real system.\nTo me it appears that this is a step forward from the cited prior work [1], where the system could be queried arbitrarily (at any state of the agent, coined the TQRL setting), which is unrealistic on real hardware.\nOne implication of this treatment is that mutual information is evaluated between the optimal future trajectory $\\tau^*$ and joint sets of future states corresponding to candidate future control sequences, i.e. it serves as the objective for evaluating MPC rollouts and accounts for possible information overlaps into the future.\nThe method is called *Trajectory Information Planning* (TIP).\n\n**References**\n\n[1] 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.\n\n[2] Neiswanger, W., Wang, K.A. and Ermon, S., 2021, July. Bayesian algorithm execution: Estimating computable properties of black-box functions using mutual information. In International Conference on Machine Learning (pp. 8005-8015). PMLR.", "questions": "- The **MPC** ablation variant from section 5, i.e. a GP + MPC with no exploration considerations, seems to perform very well in terms of data efficiency. This is expected, as the interpolation properties of GPs are a strong prior and aid generalization when the system dynamics behave nicely in the Lipschitz sense (PILCO is a good example of this). Do you believe the TIP objective would work just as well if this aspect was not there, i.e. without the inherently fast adaptation of a GP due to its interpolation properties?\n- Can you further elaborate on why TIP performs better than BARL for the plasma control tasks? The text mentions that this could be due to the inefficiency of sampling a finite candidate set for BARL, do you believe that ends up being worse than the inherent proximity of sequential trajectory states acquired by TIP?\n- I don't think I agree with the statement that the median amount of data used by algorithms to solve the task is equivalent to an MPC controller with ground-truth dynamics, as stated in sec. 5 -- what is the justification?\n- Have you considered using a critic, to account for infinite horizons?\n\nAlso see my remarks above.\n", "limitations": "I don't see any major negative societal impact. The authors acknowledge the main methodological limitations of the method in the conclusion, although these could be emphasized further.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The posed exploration objective appears sound and widely applicable. I also find it fitting, as it is about reducing the predictive entropy of the optimal trajectories MPC will actually produce, by virtue of training the dynamics model on data that is informative about them.\n- While the core idea of the objective is not novel (see [1]), I believe the extension to realistic data acquisition, i.e. in sequences, is a step forward.\n- The authors make it abundantly clear that information overlap in the acquired data sequences needs to be accounted for, and the method complies with that.\n- Motivation is easy to follow across the manuscript, exposition is clear.\n\n**Weaknesses**\n- The method hinges on the assumption that the posterior over the dynamics model can be updated in closed form (eq. 6 and 7) and that the entropy of predictions is easy to evaluate. This is required for the evaluation of the posed MI objective. Therefore the method limits itself to GP dynamics (just like in [1]), black-box neural dynamics models would encumber the evaluation of the objective.\n- GPs can be limiting both in terms of computation and expressivity for high-dimensional state spaces;\n- The authors use a variation of the cross-entropy method (CEM) for MPC optimization. MC-based optimization techniques are easy to apply, but might scale poorly to high-dimensional control spaces.\n- While I find the selection of experimental tasks reasonable, a lot of the environments are low-dimensional and with nice and smooth dynamics. How would the method fair on something more complex, like the ant locomotion task? Would the GP still be able to capture the transition accurately? For me it remains an open question whether the scheme is applicable to more complex systems (i.e. scalability).\n- The mutual information objective is formulated for the optimal trajectory $\\tau^*$ w.r.t. the current dynamics posterior. I think it is important to highlight that the \"optimal trajectory\" in that sense is not the ideal trajectory, but the best one available based on the current model fit. So overall acquisition success is still subject to the reachability properties of the state-space of the system.\n\nOverall I found this is a good contribution that explores an interesting idea already introduced in prior work [1] further. Thus I lean towards acceptance.\n\n**Minor remarks**\n- The qualifier \"simply\" appears too often in the text (11 times).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "ma7gU0XI80W", "reviewer": "Reviewer_iWun", "summary": "The authors develop Bayesian Model-Predictive Control in which a new cost-function is proposed. The idea of the cost-function goes back to the information gain notion. The aim of the paper is to make the algorithm more sample efficient. ", "questions": "1. There is no description what is $P(T|D)$ in formula (2). Also it is not clear how distributions differ for open-loop and closed-loop problems.\n\n2. In the end of Algorithm 1, it is unclear how to create policy $\\pi_g$ for posteriory distribution $P (T' | D)$.\n\n3. On lines 203 and 204, the use of the notation $S$ for pairs, and $S'$ for one state does not seem quite intuitive to understand. Maybe it's better to use another letter instead of $S$?\n\n4. In formula (5), it is not clear what is the result of $D \\cup S'$  and what symbol $p(S'|S,D)$ stands for. When $D$ – is the set of triples, $S$ – is the set of pairs, $S'$ - is the set of single states. What is the connection between $D$ and $S$?\n\n5. In formula (6), it is unclear what is the result of union $D \\cup \\tau^*$, when $D$ – is the set of triples, $\\tau^*$ is the set of pairs of states and actions and a final state.\n\n6. Perhaps, since line 207 gives an idea of a new cost function, it would be better to give a simple example when it works better than the previous ones.\n\n7. In line 210 the non-obvious equality $C_{τ^*}(S) = −I(S', τ^*) = −I(τ^*, S')$ is given. It would be better to prove it or provide the link to the proof. \n\n8. In formula (7) it is unclear what is the result of union $D \\cup \\tau^∗_{ij}$, when $D$ – is the set of triples, $\\tau^∗_{ij}$ - is the set of pairs of states and actions and a final state.\n\n9. In algorithm 2, it is unclear how to compute the joint posterior covariance matrices. If a gaussian process is used for this, then it would be better to specify what is the kernel function.\n\n10. What hyperparameters were used to run experiments and baseline algorithms (so that they could be repeated)?", "limitations": "The article does not have potentially negative societal impact.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The problem considered in the article seems relevant and the algorithms proposed by the authers seems interesting and original. The experiments show that, apparently, the algorithms are the most effective in terms of using samples. Nevertheless, in my opinion, the Weakness of the article is its descriptive part. To describe the ideas, the authors use mathematical language very carelessly, which can cause the reader serious problems in understanding the article (for more information, see the questions section).  Thus, the article seems to be promising, but requires significant refinement, in particular, in the mathematical presentation – it should be more rigorous.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "giDOo5jADpd", "reviewer": "Reviewer_QRq7", "summary": "The paper extends Bayesian Active Reinforcement Learning (BARL) [38] by lifting its Transition Query RL (TQRL) assumption (TQRL means that the agent can sample individual transitions $(s,a,r,s')$ by querying the simulator at an arbitrary $(s,a)$ pair). Instead, only state-action pairs from collected trajectories are used in this paper.\n\nThe main contribution is Algorithm 2 for computing the cost function $C_{\\tau^*}(S)$, the negative joint expected information gain about the optimal trajectory. The average cost is subsequently optimized via the cross-entropy method (CEM). The computation complexity of one iteration is derived, and experiments on 5 environemnts for closed-loop policies (the algorithm applied in MPC fashion) and 3 environments for open-loop policies are reported.\n\nThe proposed method Trajectory Information Planning (TIP) is shown to perform comparably or even better than BARL in the sense of requiring fewer samples. The fact that TIP is even better than BARL is surprising since BARL can sample the state-action space more freely, but the authors hypothesize that this is because BARL samples uniformly whereas TIP has a more sophisticated scheme and forces initialization at the start distribution.\n\nTo demonstrate that explicit exploration is useful, TIP is shown to require fewer samples than the methods which don't explicitly encourage exploration (MPC, PETS, SAC, TD3, PPO).\n\nFurthermore, TIP is compared to a few baselines — HUCRL and TS — that rely on upper confidence bound (UCB) and Thomson sampling (TS) over the dynamics and shown to provide better exploration by requiring significantly fewer samples.\n\nFinally, a few ablations of the proposed method are considered that either only use the cost function $C_g$ (used in MPC) or the entropy objective $C_e$ (used in DIP, EIG$_T$). A variant of TIP where the information gain is split over time steps called sTIP is also considered. All these methods perform comparably, according to Table 1, however TIP shows the best sample complexity.", "questions": "1) In Fig. 3, \"Model MSE on Random Set\" is bigger than \"Model MSE on Current MPC\" by 4 orders of magnitude. This seems to imply that the model is quite bad outside of the observed data. How should one interpret it? Can one then rely on such model?\n2) Table 1 shows only the median value over 5 runs. It would be usefull to see the whole distributions, at least in the appendix. E.g., display a boxplot for each environment. Otherwise, it is hard to judge how reliable the numbers in Table 1 are and what is the variance in these results.\n3) MPC seems to perform quite comparably to TIP and it requires significanlty fewer samples than SAC/TD3/PPO/HUCRL/TS, despite the fact that it has no notion of exploration and is only optimizing the cost function $C_g$. This seems to indicate that explicit exploration may not be so necessary in the considered tasks or that MPC somehow performs sufficient exploration anyways. It would be worth providing a discussion of this result.\n4) DIP appears very similar to MPC in Table 1 despite the fact that they are using completely different objective functions. DIP is only using the entropy objective $C_e$ whereas MPC is only maximizing the reward  $C_g$. Such similarity seems to indicate that the reward is somehow correlated with the exploration in the considered tasks. Or is there a different explanation?\n5) Results for BARL are better in the BARL paper (Table 1 here https://arxiv.org/pdf/2112.05244.pdf) compared to the results in Table 1 in the present paper. What is the reason for that?\n", "limitations": "Technical limitations are described sufficiently. There is no discussion of the negative societal impact. The authors could indicate potential dangers of active exploration in applications of RL.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- clarity: well-written paper, with good structure and clear presentation\n\nWeaknesses\n1) Significance: medium to low. Too dramatic and misleading claims in the abstract\n    - \"2-200x fewer samples compared to 14 baselines\" — the 200x improvemnt is w.r.t. methods that don't compute information gain (PETS/SAC/TD3/PPO). Such comparison does not seem fair. There has been a lot of algorithms that propose various version of information gain and artifical curiosity, therefore a comparison should be with respect to those algorithms or at least w.r.t. BARL or MPC with entropy objective. Compared to these algorithms, the contribution is still valuable but the statement would be more sober.\n2) Originality: fair\n    - Reading the paper briefly makes an impression that the whole approach of minimizing the information gain about the optimal trajectory is new. In particular, this is also claimed in the abstract. However, all the derivations and ideas are already presented in BARL [1]. The main contribution of this paper compared to BARL is to compute the InfGain using only sampled trajectories. It should be made much more clear what is the contribution of the present paper and what is actually novel here.\n3) Quality: fair. A lot of evaluations but comparison to wrong baselines\n    - As mentioned above, the baselines are either from a different category (SAC/TD3/PPO) or don't use InfGain (MPC/PETS) or seem to be badly tuned or not applicable (HUCRL/TS) or are version of the proposed algorithm (sTIP, DIP, EIG$_T$). On the other hand, a variety of exploration bonuses and surrogates for the InfGain have been proposed, e.g., [2,3]. In particular, the proposed method appears very similar to [4] and [5] and therefore either a comparison would be desirable or a statement why such comparison is not necessary.\n\n\nReferences \\\n[1] Mehta, V., Paria, B., Schneider, J., Ermon, S., & Neiswanger, W. (2021, September). An Experimental Design Perspective on Model-Based Reinforcement Learning. In International Conference on Learning Representations. \\\n[2] Shyam, P., Jaśkowski, W., & Gomez, F. (2019, May). Model-based active exploration. In International conference on machine learning (pp. 5779-5788). PMLR. \\\n[3] Pathak, D., Gandhi, D., & Gupta, A. (2019, May). Self-supervised exploration via disagreement. In International conference on machine learning (pp. 5062-5071). PMLR. \\\n[4] Schultheis, M., Belousov, B., Abdulsamad, H., & Peters, J. (2020, May). Receding horizon curiosity. In Conference on robot learning (pp. 1278-1288). PMLR. \\\n[5] Tschantz, A., Millidge, B., Seth, A. K., & Buckley, C. L. (2020). Reinforcement learning through active inference. arXiv preprint arXiv:2002.12636.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "Q1MYZ6Bbuio", "reviewer": "Reviewer_4S7A", "summary": "This paper presents a method for exploration in MBRL techniques by proposing a method called trajectory information planning. They present an algorithm to maximize the expected information gain for the trajectory during planning. The method have been demonstrated on several control problems and compared with some baselines.", "questions": "1. How would your method work for high dimensional tasks? How would you method compare against a Stochastic MPC approach? like chance constrained MPC? or Tube MPC kind of approaches. In general, they tend to work well on classical control systems.\n2. How would your method scale to high-dimensional systems? Did you try implementing your method on acrobot or humanoid kind of systems?", "limitations": "1. I think authors can try and address some limitations with respect to computations and scaling of the method to high dimensions. That could be useful insight for the readers.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The information gain metric is interesting and seems novel to me.\n2. The authors can show improvement over some chosen baseline methods\nWeaknesses:\n1. I think the evaluation of the paper needs more work.\n2. Would it be possible for authors to perform comparison on some high-dimensional problems? Also, I think it is unfair to compare to model-free methods as they are known to slow compared to model-based method on these classical control tasks. If you would like to compare your method to model-free methods, please select some tasks where the SOTA MBRL techniques dont perform good compared to model-free methods, and try to show that your method succeeds. Otherwise the comparison seems unfair.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept