Exploration via Planning for Information about the Optimal Trajectory
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.
Introduction
The potential of reinforcement learning (RL) as a general-purpose method of learning solutions to sequential decision making problems is difﬁcult 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 experience; however, RL has so far broadly fallen short of this.
Chief among the difﬁculties 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 of RL this problem is exacerbated by the fact that collecting samples from the true environment can incur huge costs.
For example, expensive scientiﬁc experiments are needed for tokamak control [31, 13, 20] and design of molecules [70], and collecting experience on the road is both costly and runs the risk of car accidents for autonomous vehicles [66].
Though model-based methods like Deisenroth and Rasmussen [21], Chua et al.
[15], and Curi et al.
[16] are much more efﬁcient 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 conﬁdence bounds (UCB) [16, 5] to more intelligently explore
[Figure]
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 stateaction settings these methods must make coarse approximations to be computationally tractable (e 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 a particular task speciﬁed by an initial state distribution and reward function.
We believe that an ideal exploration strategy for sample efﬁciency 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 a Bayesian planning problem over a speciﬁc 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-art methods rely on cost functions that either result in behavior that is too greedy—i e. the policy tries to maximize returns during exploration—or too exploratory, i.e. the policy is incentivized to explore the environment dynamics and does not consider the task at hand.
We therefore present a cost function that balances out these two extremes, by generalizing an information-gain acquisition 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 1, 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.
In summary, the contributions of this work are as follows: we develop a novel cost function for exploration that explicitly accounts for both the speciﬁc task and uncertainty about environment dynamics, a 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 andopen-loop environments (with a focus on expensive RL tasks in plasma physics) compared againstbaselines.
We ﬁnd 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.
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, ✏ -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-Uhlenbeck action noise [37] to the greedy policy, or entropy bonuses [27] to the policy or value function objectives to add noise to a policy which is otherwise optimizing the RL objective.
Tabular RL is often solved by choosing actions based on upper conﬁdence 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 conﬁdence bounds on the dynamics and training a ‘hallucinated’ policy which chooses perturbations within these bound to maximize expected policy performance.
Another recent work [5] uses anticoncentration inequalities to approximate upper conﬁdence 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 a 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 difﬁcult 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 insufﬁciently expressive posterior over entire transition functions, implying that it may be quite difﬁcult to solve tasks using sampled models.
Similarly, the posterior over action-value functions in Osband et al.
[45] is only roughly approximated by training a bootstrapped ensemble of neural networks.
There is also a 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.1 alongside a discussion of this literature.
Additionally, a broad set of methods explore to learn about the environment without addressing a speciﬁed task.
This line of work is characterized by Pathak et al.
[47], which synthesizes a taskagnostic 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 to learn about a random neural network across the state space, and Plan2Explore [60], which prospectively plans to ﬁnd 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 efﬁciently 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, 30].
Recently, a ﬂexible framework known as Bayesian algorithm execution (BAX) [43] has been proposed to efﬁciently estimate properties of expensive black-box functions; this framework gives a general procedure for sampling points which are informative about the future execution of a 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, BARL relies on a setting the authors call Transition Query Reinforcement 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 sufﬁcient 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 instead collecting data via trajectories only.
In this paper, we aim to apply the information-theoretic 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 beneﬁtting from the more substantial supervision that is typical in dynamics model learning.
Problem Setting
In this work we deal with ﬁnite-horizon discrete-time Markov decision processes (MDPs) which consist of a sextuple h S , A , T, r, p 0 , H i 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 ), r : S ⇥ A ⇥ S !
R is a reward function, p 0 ( s ) is a distribution over S of start states, and H 2 N is the horizon (i.e. the length of an episode).
We always assume S , A , p 0 , and H are known.
We also assume the reward r is known, though our development of the method can easily be generalized to the case where r 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 ﬁnd a policy ⇡ : S !
A that maximizes Objective (1) below.
We will denote trajectories as ⌧ ⇠ p ( ⌧ | ⇡ , T ) where ⌧ =
[( s 0 , a 0 ) , . . .
, ( s H  1 , a H  1 ) , s H ] generated by s 0 ⇠ p 0 , a i = ⇡ ( s i ) , and s i +1 ⇠ T ( s i , a i ) .
We can write the return of a trajectory as
R ( ⌧ ) = P
H  [Equation] i =
r ( s i , a i , s i +1 ) for the states and actions s i , a i that make up ⌧ .
The MDP objective can then be written as
[Equation]
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 ⇡ ⇤ = arg max ⇡ J T ( ⇡ ) , which we can assume to be deterministic [64] but not necessarily unique.
We use ⌧ ⇤ to denote optimal trajectories, i.e. ⌧ ⇤ ⇠ p ( ⌧ | ⇡ ⇤ , T ) .
Similarly, for the open-loop setting, we assume a ﬁxed start state s 0 and aim to ﬁnd an action sequence a 0 , . .
a H  1 that maximizes the sum of rewards in an episode.
We will slightly abuse notation and write ⌧ ⇠ p ( ⌧ | a 0: H  1 ) and J T ( a 0: H  1 ) with these actions ﬁxed in place of a reactive policy, and again use ⌧ ⇤ to refer to the trajectories generated by an optimal action sequence.
We assume in this work that applying planning algorithms like [51] to a dynamics function T will result in a trajectory that approximates ⌧ ⇤ .
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 D
= { ( s i , a i , s 0 i ) } .
In this work, unions D [ ⌧ or D [ S 0 between the dataset D and trajectories ⌧ or next state predictions S 0 coerce ⌧ and S 0 into triples of dynamics data, prior to the union with the dataset.
Trajectory Information Planning
Our method consists of a generic framework for Bayesian (or approximately Bayesian) modelpredictive control and a novel cost function for planning that allows us to explicitly plan to ﬁnd 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 a greedy cost function that corresponds to the future negative expected rewards or a pure exploration cost function that corresponds to future information about the dynamics.
Afterwards, in Section 4.2, we 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-Predictive 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 a possible solution.
The objective of Bayesian planning is to ﬁnd the h -step action sequence that maximizes the expected future returns under model uncertainty.
That is,
[Equation]
for some cost function C over trajectories and some start state s .
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-loop control via model-predictive control (MPC), which involves re-planning (2) at every state the agent visits and playing the ﬁrst action from the optimal sequence.
Concretely, the MPC policy for our Bayesian setting is as follows:
[Equation]
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:
[Equation]
Previous works such as Kamthe and Deisenroth [34] and PETS [15] use the greedy exploration cost function, C g .
This cost function incentivizes trajectories that achieve high rewards over the next h transitions on average.
In works that focus on task-agnostic exploration such as Sekar et al.
[60] and
Shyam et al.
[61], the cost function C e (or similar) is used to encourage the agent to ﬁnd areas of the state space in which the model is maximally uncertain.
Note that we use ⇡ g to refer to the greedy policy given by using (3) with C g .
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 dynamics and reward functions to ﬁnd 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 ﬁnd a policy, and Chua et al.
[15] and Mehta et al.
[40] use the cross-entropy method [17] to ﬁnd 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 ﬁnd 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 1 gives a formal description of the method and Section A.5 provides further details.
Inputs: transition function episode query budget b , number of posterior function samples k , planning horizon h .
Initialize D ; .
for i 2 [1 , . . .
, b ] do Sample start state {+s 
 ⇠+} p 0 .
for t 2 [0 , . . .
, H  1] do
Sample posterior functions { T 0 ` } k ` =1 ⇠ P ( T 0 | D ) .
Approximately ﬁnd arg min a 0 ,...,a h
P k ` =1 E ⌧ ` ⇠ p ( ⌧ | T 0 ` ,a 0 ,...,a h  1 ) [ C ( ⌧ ` )] via iCEM.
Execute action a 0 by sampling s t +1 ⇠ T ( s t , a 0 ) .
Update dataset D D [ { ( s t , a 0 , s t +1 } .
end for end for return ⇡ g for the posterior P ( T 0 | D ) .
4.2 A Task-Speciﬁc Cost Function based on Trajectory Information
In this work, we aim to explore by choosing actions that maximize the conditional expected information gain (EIG) about the optimal trajectory ⌧ ⇤ .
This is 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 in Osband et al.
[46], it is 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 ﬁnd an ✏ -optimal policy in a 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 deﬁned over a single state-action pair and mutual information is submodular, we cannot naively use the acquisition function as is (or sum it over many datapoints) to choose actions that lead to good long-term 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 EIG ⌧ ⇤ by taking a set of points to query and computing the joint expected information gain from observing the set.
Our cost function is non-
Markovian in the state space of the MDP, but it is Markovian in the dataset, which represents a point in the belief space of the agent about the dynamics.
Let X
= { x : x ✓ S ⇥ A , | x | < 1 } be the set of ﬁnite subsets of the set of all state-action pairs.
Our cost function C ⌧ ⇤ : X !
R is deﬁned below to be the negative joint expected information gain about the optimal trajectory ⌧ ⇤ for a subset
X 2 X .
In particular, assuming an existing dataset D , a set of h query points X
= { ( s i , a i ) } i {+
+}
[ h ] , and a random set of next states S {+
+}
= { s 0 i ⇠ T ( s i , a i ) , i 2 [ h ] } ,
[Equation]
This formulation of C ⌧ ⇤ forces our method to handle the redundant information among queries—it is likely that I ( s 0 1 , ⌧ ⇤ )
+ I ( s 0 2 , ⌧ ⇤ ) >
I ( { s 0 1 , s 0 2 } , ⌧ ⇤ ) and our method should avoid this overestimation.
However, as written, this function relies on computing entropies on 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 C ⌧ ⇤ ( X ) =
I ( S 0 , ⌧ ⇤ ) =  I ( ⌧ ⇤ , S 0 ) for the mutual information I .
This allows us to exchange ⌧ ⇤ and our set of queries so that ⌧ ⇤ is giving information about the posterior predictive distribution of our set.
In other words,
[Equation]
In order to compute the right-hand term, we must take samples ⌧ ⇤ ij ⇠ P ( ⌧ ⇤ | D ) , i = 1 , . . .
, m, j = 1 , . . . , n .
To do this, we ﬁrst sample m start states s ( i ) 0 from p
(we always set m = 1 in experiments but derive the procedure in general) and for each start state independently sample n posterior functions
T 0 ij ⇠ P ( T 0 | D ) from our posterior over dynamics models.
We then run a planning procedure using iCEM [51] on each of the posterior functions from s ( i ) 0 using T 0 ij for T (using our assumption that planning can generate approximately optimal trajectories given ground-truth dynamics), giving our sampled ⌧ ⇤ ij .
Formally, we can approximate C ⌧ ⇤ via Monte-Carlo as
[Equation]
Assuming the dynamics are modelled with a Gaussian process, we can compute the joint Gaussian probability of the next states S 0 [52].
As the entropy of a 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 0 | D ] and optimize it with a zeroth order optimization algorithm.
Finally, we must calculate the entropy H [ S 0 | D [ ⌧ ⇤ ij ] .
For this, we follow a similar strategy as Neiswanger et al.
[43]: since ⌧ ⇤ i is a set of states given by the transition model, we can treat them as additional noiseless datapoints for our dynamics model and condition on them before computing the joint covariance matrix for S 0 .
Given this newly generalized acquisition function, we can instantiate a 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 ⌧ ⇤ ij 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 ⌧ ⇤ ( ⌧ ) instead of C g or C e , and compute this cost function using Algorithm 2.
At test time, we return to planning with C g 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, we make the simplifying assumption that the planning algorithm used (in this case, iCEM from [51]) evaluates p action sequences consisting of h (the planning horizon) actions and that our current dataset is of size N .
In order to efﬁciently sample functions from the posterior over dynamics functions, we use the method from Wilson et al.
This reduces the naive complexity of querying these functions from O ( N 3 ) to a one-time O ( N ) cost and then O (1) for additional queries.
As we derive in Section A.1, the computational complexity of one TIP planning iteration is O [Equation] ⌘⌘ .
The two asymptotically expensive operations are (1)
Algorithm Inputs: dataset D = { ( s k , a k , s 0 k ) } , query set X , number of start state samples m , number of posterior function samples n .
Sample m start states { s ( i ) 0 } mi =1 ⇠ p 0 .
for i 2 [ m ] do Sample n posterior functions { T 0 j } nj =1 ⇠ P ( T 0 | D ) .
for j 2 [ n ] do Set ⇡ ⇤ j ⇡ MPC using C g and a singleton posterior P ( T | D ) =
( T 0 j ) as in (3).
Compute ⌧ ⇤ ij by executing ⇡ ⇤ j on T {+
+}
j starting from s ( i ) 0 .
end for end for
Compute the joint posterior covariance ⌃ S 0 | D across all points in X . Compute the joint posterior covariances ⌃ S 0 ij | D [ ⌧ ij 8 i 2 [ n ] , j 2 [ m ] across all points in X .
return log | ⌃ S 0 |  1 nm P i 2 [ n ] ,j 2 [ m ] log | ⌃ S 0 ij | .
computing the Cholesky decompositions of the nm kernel matrices for datasets D [ ⌧ ⇤ ij and (2) solving the triangular systems using the cached Cholesky decompositions in order to compute the covariance matrices ⌃ S 0 ij | D [ ⌧ ⇤ ij for each of 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 {+
+}
i to sample optimal trajectories ⌧ ⇤ ij .
This is due to the fact that in practice p is large and we implemented the planner in NumPy [28] so it cannot be compiled together with the Tensorﬂow [1] code from Wilson et al.
[69], which is used for predicting which states will be visited for the planner.
We give further information on the implementation in Section A.
Experiments
Rollout (Joint)
TQRL
Rollout (Sum)
G r ee d y ( C g ) w / S t o c h a s t
i c
i t y
T a s k - a g
no s t
i c (
C e )
E x p l o r a t i on
T a s k - s p e c
i ﬁ c (
C τ * )
E x p l o r a t i on
N/A
MPC,PETS,
BPTT,PPO,
SAC,TD
EIG T sDIP
DIP
BARL sTIP
TIP,TS,
HUCRL
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 to expensive dynamics.
In this section we demonstrate the effectiveness of TIP in quickly learning a good policy by comparing against a variety of state-of-the-art reinforcement learning algorithms and strong baselines (including some that use the TQRL setting from [40], which is also known as RL with a generative model in Kakade [33] and other works [6, 4]).
In particular, we compare the average return across ﬁve evaluation episodes across ﬁve random seeds of each algorithm on ﬁve closed-loop control problems.
For sample complexity we assess the median amount of data taken by each algorithm to ‘solve’ the problem across ﬁve 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 Pendulum and Cartpole famously are) and must have a relatively short control horizon and ﬁxed start state.
[Table]
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
[Table] 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 a 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 on hallucinating dynamics perturbations as a way of realizing an upper conﬁdence bound on the policy, model-based Thompson Sampling ( TS ), which samples from the posterior over models and chooses optimal actions for that sample, and a 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 a process of actively collecting information for inference on a reward-biased generative model.
A further comparison is with Receding Horizon Curiosity (RHC) [59], which does online Bayesian system identiﬁcation over a linear model in order to quickly ﬁnd a model of the environment dynamics.
Our model-free comparison methods, Soft Actor-Critic ( 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 as 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 C g , C ⌧ ⇤ , and C e .
DIP , oDIP , and EIG T all use C e 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 a way that controls for difference in the model and planning algorithm.
MPC uses C g and is very close to the method in [34].
Like TIP, BARL [40] and sTIP use C ⌧ ⇤ .
BARL operates in the TQRL setting and can therefore use the simpler EIG ⌧ ⇤ acquisition function.
sTIP investigates the use of
P s i ,a i 2 S EIG ⌧
⇤ ( s i , a i ) 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.
oDIP plans an action sequence to minimize the joint C e and executes the actions found for each open-loop trial, while oMPC does the same thing using C g .
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-v0 from Brockman et al.
[10]) with 2D states and 1D actions, a 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 simpliﬁed  Tracking problem from plasma control [13, 39] (similar in design but not identical to the one from Mehta et al.
[40]) trained using with 4D states and 1D actions, and a more complicated problem in plasma control where  + 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.
[Figure]
Environment oTIP
oMPC oDIP
BO
Nonlinear Gain
Nonlinear Gain
[Table]
[Table]
>2k 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 a fraction or even orders of magnitude less data than other methods.
For many of our ablation methods we see failures to solve some of the problems even though the model is demonstrated by
TIP to be able to sufﬁciently predict the dynamics.
This is especially apparent on the harder plasma control environment,  +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 is less useful for control than the data TIP collects.
This is underscored by the second column of Figure 3, where it is clear that TIP achieves the lowest modeling error on the points actually needed during the execution of the policy but not on the uniform test set.
In particular we ﬁnd it interesting that TIP outperforms BARL on the  + Rotation environment, as BARL should in principle have a strictly 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 to 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 is 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 is 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-speciﬁc exploration as the task-agnostic exploration methods ( EIG T , DIP, sDIP) underperform both in returns and model error on the trajectories visited.
For the open-loop experiments (Table 2), we also see strong performance from oTIP.
As the modelbased methods beneﬁt from observing many model transitions for each open-loop trial, it is unsurprising that they are more sample-efﬁcient than the BO method.
Within the model-based techniques, oTIP is the most sample efﬁcient.
We believe that this is for much the same reasons as in the closed-loop case—exciting evidence that the C ⌧ ⇤ cost function can be applied in a variety of settings.
Conclusion
In this work, we presented and evaluated a cost function designed for intelligent, task-aware exploration.
Using this cost function in model-predictive control allows agents to solve continuous
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 on 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 in a small number of trials on a real tokamak.
Acknowledgements
This work was supported in part by US Department of Energy grants under contract numbers
DE-SC0021414 and DE-AC02-09CH1146.
This work is supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE1745016 and DGE2140739.
Any opinions, ﬁndings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reﬂect the views of the National Science Foundation.
This work is additionally supported by NSF (#1651565), AFOSR (FA95501910024), ARO (W911NF21-1-0125), CZ Biohub, and Sloan Fellowship.
[-C ( τ ) = − (cid:80) hi =
-]