Exploration via Planning for Information about the Optimal Trajectory
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Many potential applications of reinforcement learning (RL) are stymied by the largenumbers of samples required to learn an effective policy.
This is especially truewhen applying RL to real-world control tasks, e.g. in the sciences or robotics, whereexecuting a policy in the environment is costly.
In popular RL algorithms, agentstypically explore either by adding stochasticity to a reward-maximizing policy orby attempting to gather maximal information about environment dynamics withouttaking the given task into account.
In this work, we develop a method that allows usto plan for exploration while taking both the task and the current knowledge aboutthe dynamics into account.
The key insight to our approach is to plan an actionsequence that maximizes the expected information gain about the optimal trajectoryfor the task at hand.
We demonstrate that our method learns strong policies with2x fewer samples than strong exploration baselines and 200x fewer samples thanmodel free methods on a diverse set of low-to-medium dimensional control tasksin both the open-loop and closed-loop control settings.
Introduction
The potential of reinforcement learning (RL) as a general-purpose method of learning solutions tosequential decision making problems is difﬁcult to overstate.
Ideally, RL could allow for agentsthat learn to accomplish all manner of tasks solely through a given reward function and the agent’sexperience; however, RL has so far broadly fallen short of this.
Chief among the difﬁculties inrealizing this potential is the fact that typical RL methods in continuous problems require very largenumbers of samples to achieve a near-optimal policy.
For many interesting applications of RL thisproblem is exacerbated by the fact that collecting samples from the true environment can incur hugecosts.
For example, expensive scientiﬁc experiments are needed for tokamak control [30, 13, 20] anddesign of molecules [69] and collecting experience on the road is both costly and runs the risk of caraccidents for autonomous vehicles [65].
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 aboutthe prospective task-relevant information content of future observations.
Moreover, these methodstypically 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[45], entropy bonuses [26], and upper conﬁdence bounds (UCB) [16, 5] to more intelligently explorethe state-action space of an MDP, these methods still do not explicitly reason about how informationthat 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).
[Figure]
Many methods in the vein of Pathak et al.
[46] and Shyam et al.
[60] explore directly to gather newinformation based on current uncertainty about environment dynamics, but they do not in generalspecialize the information that they aim to acquire to a particular task.
We believe that an idealexploration strategy for sample efﬁciency is able to take both a relevant task into account as well asuncertainty about the environment dynamics.
In this work, we start by showing how many methods can be cast as a Bayesian planning problem overa speciﬁc cost function.
This framework helps illuminate the importance of the cost function and whatimpact it has on exploration.
Viewed in this light, it becomes clear that many previous state-of-the-artmethods rely on cost functions that either result in behavior that is too greedy—i.e. the policy tries tomaximize returns during exploration—or too exploratory, i.e. the policy is incentivized to explore theenvironment dynamics and does not consider the task at hand.
We therefore present a cost functionthat balances out these two extremes by substantially generalizing the acquisition function introducedin Mehta et al.
[39] to work for future hypothetically acquired data.
In particular, our cost functioncaptures the amount of information that would be gained about the optimal trajectory , if the agentwere to explore by following a particular planned trajectory.
As depicted in Figure 1, this involvesthe agent planning out what it would hypothetically do given different realizations of the dynamicsand planning actions that are informative about those possibilities.
In summary, the contributions of this work are as follows: we develop a novel cost function forexploration that explicitly accounts for both the speciﬁc task as well as uncertainty about environmentdynamics, a method for planning which applies the cost function to explore in MDPs with continuousstates and actions, and a thorough empirical evaluation of our method across 5 closed-loop and 3open-loop environments (with a focus on expensive RL tasks in plasma physics) compared against 14baselines.
We ﬁnd that our proposed method is able to learn policies that perform as well as an agentwith 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 toexecute a greedy policy with some form of added stochasticity.
The simplest strategy, (cid:15) -greedyexploration as used in Mnih et al.
[40], takes the current action thought to be best with probability1 − (cid:15) and a random action with probability (cid:15) .
Other methods use added Ornstein-Uhlenbeck actionnoise [36] to the greedy policy, or entropy bonuses [26] to the policy or value function objecties toadd 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 valuefunction [14, 35], but explicitly computing and optimizing these bounds in the continuous setting issubstantially more challenging.
Recent work [16] approximates this method by computing one-stepconﬁ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) [54], which samples a realization of the MDP from the posterior and actsoptimally as if the realization was the true model, can be applied for exploration in a model-freemanner as in [44] or in a model-based manner as in [62].
As the posterior over MDP dynamics orvalue functions can be high-dimensional and difﬁcult to represent, the performance of TS can behindered 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ﬁcientlyexpressive posterior over entire transition functions, implying that it may be quite difﬁcult to solvetasks using sampled models.
Similarly, the posterior over action-value functions in Osband et al.
[44]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 computationallyexpensive and hard to use, though they have attractive theoretical properties.
These methods buildupon the fundamental idea of the Bayes-adaptive MDP [52], which we detail in Section E.1 alongsidea discussion of this literature.
Additionally, a broad set of methods explore to learn about the environment without a task inmind.
This line of work is characterized by Pathak et al.
[46], which synthesizes a task-agnosticreward function from model errors.
Other techniques include MAX [60], which optimizes theinformation gain about the environment dynamics, Random Network Distillation [11], which forcesthe agent to learn about a random neural network across the state space, and Plan2Explore [59],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 processor function to get maximal information about some quantity of interest.
When the quantity ofinterest is the location of a function optimum, related strategies have been proposed as the predictiveentropy search family of Bayesian optimization (BO) algorithms [28, 29].
Recently, a ﬂexibleframework known as Bayesian algorithm execution (BAX) [42] has been proposed to efﬁcientlyestimate properties of expensive black-box functions; this framework gives a general procedure forsampling points which are informative about the future execution of a given algorithm that computesthe property of interest, thereby allowing the function property to be estimated with far less data.
A subsequent related work [39], known as Bayesian Active Reinforcement Learning (BARL), usesideas from BOED and BAX to sample points that are maximally informative about the optimaltrajectory 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 queriedat an arbitrary sequence of state-action pairs chosen by the agent.
TQRL is thus a highly restrictivesetting which is not suitable when data can only be accessed via a trajectory (rollout) of environmentdynamics; it typically relies on an accurate environment simulator of sufﬁcient expense to warrant itsuse.
Even then, there will likely be differences between simulators and ground truth dynamics forcomplex systems.
Thus, one would ideally like to collect data in real environments.
However, thisoften 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 tothe general MDP setting as well as learning open loop model-based controllers.
The typical methodfor learning to solve open-loop control problems was demonstrated successfully in Tesch et al.
[64],where a value function was learned from action sequences to task success.
Our method takes amodel-based approach to this problem, using similar exploration strategies to Bayesian optimizationbut 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) whichconsist of a sextuple (cid:104)S , A , T, r, p 0 , H (cid:105) where S is the state space, A is the action space, T is the 119 transition function T :
S ×
A → P ( S ) (using the convention that P ( X ) is the set of probability 120 measures over
X ), r : S ×
A × S → R is a reward function, p 0 ( s ) is a distribution over S of start 121 states, and H ∈ N is the horizon.
We always assume S , A , p 0 , and H are known.
We also assume 122 the reward r is known, though our development of the method can easily be generalized to the casewhere r is unknown.
Our primary function of interest is the transition function T , which we learnfrom 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).
We discuss trajectoriesτ ∼ 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 ( τ ) =
(cid:80) H − 1 i =
r ( s i , a i , s i +1 ) forthe 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 truthtransition 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 [63] but not necessarily unique.
We use τ ∗ to denote optimal trajectories, i.e. τ ∗ ∼ p ( τ | π ∗ , T )
Equivalently for the open-loop setting, we assume a ﬁxed start state s 0 and aim to ﬁnd an actionsequence a 0 , . .
. , a H − 1 that maximizes the sum of rewards in an episode.
We will slightly abusenotation and write τ ∼ p ( τ | a 0: H − 1 ) and J T ( a 0: H − 1 ) with these actions ﬁxed in place of a reactive 136 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 [50] to a dynamics function T willresult in a trajectory that approximates τ ∗ .
We will primarily focus on a Gaussian process (GP) modelof the transition dynamics in order to take advantage of its expressive representation of uncertaintyand grounded methods for sampling, conditioning, and joint inference.
There is substantial prior workusing 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 (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 oftriples D
= { ( s i , a i , s (cid:48) i ) } .
We use unions between the dataset and trajectories τ or next states S (cid:48) to 145 mean a coercion of the other argument into the triples needed to ﬁt into the dataset.
Trajectory Information Planning
Our method consists of a generic framework for Bayesian or approximately Bayesian model-predictivecontrol and a novel cost function for planning that allows us to explicitly plan to ﬁnd the maximalamount of new information relevant to our task.
In Section 4.1, we describe the MPC frameworkand highlight that many prior methods approximate this framework while using the greedy costfunction corresponding to the future negative expected rewards or the pure exploration cost functioncorresponding to future information about the dynamics.
Afterwards in Section 4.2, we derive ournew cost function and describe how it is computed.
The overall method is simply applying theplanning framework with the 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 frommethods such as PILCO [21] and PETS [15].
This formulation highlights these methods’ inherentlygreedy nature and hints at a possible solution.
The objective of Bayesian planning is to ﬁnd the h -stepaction sequence that maximizes the expected future returns under model uncertainty.
That is,
[Equation]
(2) for some cost function C over trajectories and some start state s .
If in the open-loop control setting,the agent executes the sequence of actions found without replanning.
This procedure can also beextended to closed-loop control via model-predictive control (MPC).
This procedure involves replanning (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 tohow the agent will behave.
Prior work has predominantly focused on two types of cost function:
[Equation] [Equation]
Previous works such as Kamthe and Deisenroth [33] and PETS [15] use the greedy exploration costfunction, C g .
This cost function incentivizes trajectories that achieve high rewards over the next htransitions on average.
In works that focus on task-agnostic exploration such as Sekar et al.
[59] and
Shyam et al.
[60], the cost function C e (or similar) is used to encourage the agent to ﬁnd areas of thestate space in which the model is maximally uncertain.
We use π g to refer to the greedy policy givenby using (3) with C g .
The optimization problem in (3) is typically approximately solved in one of 3 ways: Deisenrothand Rasmussen [21] and Curi et al.
[16] directly backpropagate through the estimated dynamics andreward functions to ﬁnd policy parameters that would generate good actions, Janner et al.
[31] usean 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.
[39] use the cross-entropy method [17] to ﬁnd action sequenceswhich directly maximize the reward over the estimated dynamics.
In this work we use a version ofthe last method given in Pinneri et al.
[50] to directly ﬁnd action sequences that optimize the costfunction being used.
We approximate the expectation by playing the actions on multiple samplesfrom the posterior P ( T | D ) .
Algorithm 1 gives a formal description of the method and Section A.5provides further details.
Algorithm
Inputs: transition function episode query budget b , number of posterior function samples k , planning horizon h
Initialize D ← ∅ .
for i ∈ [1 , . . .
, b ] do Sample start state s∼ p 0 .
for t ∈ [0 , . . .
, H − 1] do
Sample posterior functions { T (cid:48) (cid:96) } k(cid:96)
=1 ∼ P ( T (cid:48) | D ) .
Approximately ﬁnd arg min a 0 ,...,a h
[Table]
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.[39], where the EIG τ ∗ acquisition function was introduced for this purpose.
However, in this paperwe generalize this acquisition function in order to allow for sequential information collection thataccounts for the redundant information that could be collected between timesteps.
As discussed atlength in Osband et al.
[45], it is essential to reason about how an action taken at the current timestepwill affect the possibility of learning something useful in future timesteps.
In other words, explorationmust be deep and not greedy.
Explicit examples are given in Osband et al.
[45] where the time toﬁnd an (cid:15) -optimal policy in a tabular MDP is exponential in the state size unless exploration canbe coordinated over large numbers of timesteps rather than being conducted independently at eachaction.
As the EIG τ ∗ acquisition function is only deﬁned over a single state-action pair and mutualinformation is submodular, we cannot naively use the acquisition function as is (or sum it over manydatapoints) to choose actions that lead to good long-term exploration.
As an illustrative example, thisis clear in navigation tasks, where the nearby points visited over trajectories will provide redundantinformation about the local environment.
We therefore give a cost function that generalizes EIG τ ∗ by taking a set of points to query andcomputing the joint expected information gain from observing the set.
Our cost function is nonMarkovian in the state space of the MDP, but it is Markovian in the dataset, which represents a pointin the belief space of the agent about the dynamics.
Let X
= { x : x ⊆ S × A , | x | < ∞} be the 203 set of ﬁnite subsets of the set of all state-action pairs.
Our cost function C τ ∗ :
X → R is deﬁned 204 below to be the negative joint expected information gain about the optimal trajectory τ ∗ for a subset
X ∈ 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 (cid:48) = {
[Equation] [Equation]
This formulation of C τ ∗ forces our method to handle the redundant information among queries—it islikely that there will be some information about τ ∗ that is duplicated between members of the set andour method of computing this cost function must avoid ‘double-counting’ it.
However, as written,this function relies on computing entropies on high-dimensional trajectories where the form of thejoint distribution of the elements is unknown.
To tractably estimate this quantity, we use the fact that
C τ ∗ ( X ) = − I
( S (cid:48) , τ ∗ ) =
− I ( τ ∗ , S (cid:48) ) for the mutual information I .
This allows us to exchange τ ∗ 213 and our set of queries so that τ ∗ is giving information about the posterior predictive distribution ofour 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 experimentsbut derive the procedure in general) and for each start state independently sample n posterior functions
T (cid:48) ij ∼ P ( T (cid:48) | D ) from our posterior over dynamics models.
We then run a planning procedure usingiCEM [50] on each of the posterior functions from s ( i ) 0 using T (cid:48) ij for T (using our assumption thatplanning can generate approximately optimal trajectories given ground-truth dynamics), giving oursampled τ ∗ 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 Gaussianprobability of the next states S (cid:48)
As the entropy of a multivariate Gaussian depends only onthe log-determinant of the covariance, log | Σ | , we can tractably compute the joint entropy of themodel predictions H [ S (cid:48) | D ] and optimize it with a zeroth order optimization algorithm.
Finally,we must calculate the entropy H [ S (cid:48) | D ∪ τ ∗ ij ] .
For this, we follow a similar strategy as Neiswangeret al.
[42]: since τ ∗ i is a set of states output from the transition model, we can treat them as additionalnoiseless datapoints for our dynamics model and condition on them before computing the jointcovariance matrix for S (cid:48) .
Given this newly generalized acquisition function, we can instantiate amethod of planning in order to maximize future information gained.
We give the concrete procedurefor computing our acquisition function in Algorithm 2, noting that trajectories τ ∗ ij do not depend onthe 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: runmodel-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 asthe cost function and greedily attempt to maximize returns over exploration.
We can also formulatean open-loop variant of our method, oTIP, which involves planning once and then executing the entireaction sequence.
4.3 Computational Cost and Implementation Details
Though the TIP algorithm is designed for settings where samples are expensive, it is important tounderstand, both theoretically and practically, the computational cost of this method.
For ease ofnotation, we make the simplifying assumption that the planning algorithm used (in this case, iCEMfrom [50]) evaluates p action sequences consisting of h (the planning horizon) actions and thatour current dataset is of size N .
In order to efﬁciently sample functions from the posterior overdynamics functions, we use the method from Wilson et al.
This reduces the naive complexityof querying these functions from O ( N 3 ) to a one-time O ( N ) cost and then O (1) for additionalqueries.
As we derive in Section A.1, the computational complexity of one TIP planning iterationis O (cid:16) nm (cid:16) ( N + H ) 3 + ph ( N + H ) 2 (cid:17)(cid:17) .
The two asymptotically expensive operations are (1)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 thecovariance matrices Σ S (cid:48) ij | D ∪ τ ∗ ij for each of the p action sequences used by the planning algorithm.
Inputs: dataset D = { ( s k , a k , s (cid:48) 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 ∈ [ m ] do Sample n posterior functions { T (cid:48)
j } nj =1 ∼ P ( T (cid:48) | D ) .
for j ∈ [ n ] do Set π ∗ j ← π MPC using C g and a singleton posterior P ( T | D ) =
δ ( T (cid:48) j ) as in (3).
Compute τ ∗ ij by executing π ∗ j on T (cid:48)
j starting from s ( i ) 0 .
end for end for
Compute the joint posterior covariance Σ S (cid:48)
| D across all points in X .
Compute the joint posterior covariances Σ S (cid:48) ij | D ∪ τ ij i n
, j m all points in X .
return log | Σ S (cid:48) | − 1 nm (cid:80) i ∈ [ n ] ,j ∈ [ m ] log | Σ S (cid:48) ij | .
However, our implementation choices mean that in practice these operations are not the mostexpensive step.
The covariance matrix computations, which are the theoretical bottleneck, areimplemented in JAX [9], allowing them to be compiled to much faster machine code and vectorizedacross large batches of queries.
In fact, the most expensive operation in practice is planning on thesampled transition functions T (cid:48)
i to sample optimal trajectories τ ∗ ij .
This is due to the fact that inpractice p is large and we implemented the planner in NumPy [27] so it cannot be compiled togetherwith the Tensorﬂow [1] code from Wilson et al.
[68], which is used for predicting which states willbe visited for the planner.
We give further information on the implementation in Section A.
Experiments
The aim of our development of the TIP algorithm and the C τ ∗ acquisition function for RL is to reducethe sample complexity of learning effective policies in continuous MDPs given limited access toexpensive dynamics.
In this section we demonstrate the effectiveness of TIP in quickly learning agood policy by comparing against a variety of state-of-the-art reinforcement learning algorithms andstrong baselines (including some that use the TQRL setting from [39], which is also known as RLwith a generative model in Kakade [32] and other works [6, 4]).
[Figure] [Figure]
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 ﬁ
[Figure]
C
[Figure]
i
[Figure]
In particular, we compare the average return across ﬁve evalu- ation episodes across ﬁve random seeds of each algorithm on 5closed-loop control problems.
For sample complexity
we assessthe median amount of data taken by each algorithm to ‘solve’the problem across ﬁve seeds with the threshold performancegiven by an MPC controller using the ground truth dynamics.
We evaluate the open-loop variant of our method, oTIP, againstthree comparison methods on three control problems suitablefor open-loop control.
In particular, to be suitable for open-loopcontrol, the problem cannot be dynamically unstable (as Pendulum and Cartpole famously are) and must have a relatively shortcontrol horizon and ﬁxed start state.
Here too, we assess theaverage return as open-loop trials are conducted as well as thenumber of timesteps required to achieve ‘solved’ performance.
Comparison Methods We use several model-based andmodel-free comparison methods in this work.
We compareto several published model-based methods.
These include
PETS [15], as implemented by Pineda et al.
[49], which usesa probabilistic ensemble of neural networks and CEM overparticle samples to do MPC.
We also compare against threemodel-based techniques from the HUCRL [16] implementation: HUCRL itself, which relies on hallucinating dynamicsperturbations as a way of realizing an upper conﬁdence boundon the policy, model-based Thompson Sampling ( TS ), which
[Table]
samples from the posterior over models and chooses optimal actions for that sample, and a greedymodel-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.
[66], which treatsdirected exploration as a process of actively collecting information for inference on a reward-biasedgenerative model.
A further comparison is with Receding Horizon Curiosity (RHC) [58], whichdoes online Bayesian system identiﬁcation over a linear model in order to quickly ﬁnd a model ofthe environment dynamics.
Our model-free comparison methods, Soft Actor-Critic ( SAC ) [26], anactor-critic method that uses an entropy bonus over the policy to encourage more exploration, andtwo others (TD3 and PPO), are in the appendix.
Finally, we compare against various ablations of the proposed method.
These vary across twoaxes 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 samehyperparameters, 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 ofthe action sequence, the sum of the pointwise entropies of the action sequence, and the individualpointwise entropies in the TQRL setting.
These methods are very similar in spirit to [60, 47] inthat they plan for future information gain about the dynamics, but we chose to compare in a waythat controls for difference in the model and planning algorithm.
MPC uses C g and is very closeto the method in [33].
Like TIP, BARL [39] and sTIP use C τ ∗ .
BARL operates in the TQRLsetting and can therefore use the simpler EIG τ ∗ acquisition function.
sTIP investigates the use of− (cid:80) s i ,a
i ∈ S
EIG τ ∗ ( s i , a i ) as a cost function for planning.
This computes individual informationgains for each future observation without accounting for the information they may have in commonand is therefore an overestimate of the joint information gain.
In the open-loop setting we compareagainst 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 andexecutes the actions found for each open-loop trial, while oMPC does the same thing using C g .
Our closed loop control problems are the standard underactuated Pendulumswing-up task (Pendulum-v0 from Brockman et al.
[10]) with 2D states and 1D actions, a Cartpoleswing-up task with a sparse reward function, 2D s, and 1D actions, a 2-DOF robot arm controlproblem where the end effector is moved to a goal ( Reacher -v2 from Brockman et al.
[10]) with10D states and 2D actions, a simpliﬁed β Tracking problem from plasma control [13, 38] (similarin design but not identical to the one from Mehta et al.
[39]) trained using with 4D states and 1Dactions, and a more complicated problem in plasma control where β + Rotation are tracked with10D states and 2D actions.
Our open loop control problems are a navigation problem with hazards( Lava Path ) from [39] 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 2Dstates and 2D actions.
[Table]
Results As can be seen in Table 1, TIP isable to reach solved performance more quicklyacross the board than the model-based andmodel-free external baselines, often using a fraction or even orders of magnitude less data thanother 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 dynam-
[Figure]
M
[Figure]
l
[Figure]
n C u r r e n
[Figure]
C
[Figure]
e t
[Figure]
I
[Figure]
τ ∗ ( b i t s /a c t i o n )
[Figure]
T
[Figure]
ics. This is especially apparent on the harder plasma control environment, β +Rotation, where TIP isthe only method using our GP which is able to solve the problem.
We believe that this is because thedata acquired through exploration by the ablation methods is less useful for control than the data TIPcollects.
This is underscored by the second column of Figure 3, where it is clear that TIP achievesthe lowest modeling error on the points actually needed during the execution of the policy but noton 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 problemand is optimizing the same quantity with fewer constraints.
We hypothesis that this may be due tothe fact that BARL optimizes the acquisition function EIG τ ∗ by simply uniformly sampling a set ofpoints and choosing the one that evaluates to the largest value.
Our more sophisticated optimizationalgorithm and forced initialization at the start state distribution seems to allow us to collect moreinformation in this case.
This interpretation is bolstered by the fact that on the problems where TIPoutperforms BARL, we see that TIP is actually collecting more information per action than BARL asevidenced by larger EIG values.
We also see clearly that there is value in computing the C τ ∗ functionrather 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 explorationmethods ( 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, oTIPis the most sample efﬁcient.
We believe that this is for much the same reasons as in the closed-loopcase—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 themethod is effective in data reduction, it is computationally expensive and relies on dynamics that arewell-modeled by a GP.
In future work, we aim to scale the method to higher-dimensional and morecomplex control environments.
We also aim to apply this method in the real world.
In particular, weaim to address similar plasma control problems in a small number of trials on a real tokamak.