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, (cid:15) -greedy exploration as used in Mnih et al.
[41], takes the current action thought to be best with probability 1 − (cid:15) and a random action with probability (cid:15) .
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 (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 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 ∈ 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 ( τ ) =
(cid:80) H − 1 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 (cid:48) i ) } .
In this work, unions D ∪ τ or D ∪ S (cid:48) between the dataset D and trajectories τ or next state predictions S (cid:48) coerce τ and S (cid:48) 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,
argmin a 0 ,...,a h −
[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.
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 − 1 (cid:80) k(cid:96) =
1 E τ (cid:96) ∼ p ( τ | T (cid:48) (cid:96) ,a 0 ,...,a h − 1 )
[ C ( τ (cid:96) )] 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 (cid:48) | 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 (cid:15) -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 | < ∞} 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 ∈ 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) = { s (cid:48) i ∼ T ( s i , a i ) , i ∈ [ h ] } ,
[Equation]
This formulation of C τ ∗ forces our method to handle the redundant information among queries—it is likely that I ( s (cid:48) 1 , τ ∗ ) + I ( s (cid:48) 2 , τ ∗ ) > I ( { s (cid:48) 1 , s (cid:48) 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 (cid:48) , τ ∗ ) = − I ( τ ∗ , S (cid:48) ) 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 (cid:48) ij ∼ P ( T (cid:48) | 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 (cid:48) 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 (cid:48)
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 (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 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 (cid:48) .
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 (cid:16) nm (cid:16) ( N + H ) 3 + ph ( N + H ) 2 (cid:17)(cid:17) .
The two asymptotically expensive operations are (1)
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 Compute return log | Σ S (cid:48) | − 1 nm (cid:80) i ∈ [ n ] ,j ∈ [ m ] log | Σ S (cid:48) 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 (cid:48) 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 (cid:48)
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.
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 − (cid:80) s i ,a
i ∈ 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]
r
[Figure]
τ
[Figure]
i t s /a
c t
i
[Figure]
Environment oTIP oMPC
[Table]
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.
A Implementation Details A.1 Derivation of Computational Cost
In this section, we derive the computational complexity of the TIP algorithm.
For simplicity, we focus on a single TIP planning iteration as might be done at each replanning in closed-loop control or at the start of a trial in open-loop control.
In order to keep the analysis general, we assume that the chosen planning algorithm requires p accesses to the model where h actions are sequentially executed, giving ph total queries per planner execution.
We also assume that the numbers of inducing points and basis functions used in our GP posterior function sampling are constant as are the Monte
Carlo hyperparameters m, n .
The TIP algorithm consists of the following major operations:
• Sample T (cid:48)
i for j ∈ [ m ] : O ( nmN ) total cost from sampling algorithm, where N is the dataset size.
• Sample τ ∗ ij for i ∈ [ n ] , j ∈ [ m ] : phHnm total cost from running the planner ( ph posterior function queries)
H times for each sampled τ ∗ , where H is the MDP horizon.
• Compute Cholesky decomposition for each D ∪ τ ∗ ij .
This takes a total of O ( nm ( N + H ) 3 ) operations as the augmented dataset is of size N + H and Cholesky decompositions are O ( d 3 ) in the matrix size d . • Compute posterior covariance Σ S (cid:48)
| D ∪ τ ij for all τ ij .
This involves several matrix operations but the most computationally intensive is solving h triangular systems of size ( N + H ) × ( N + H ) , which each take O (( N + H ) 2 ) time.
So the total computation here is O ( pnmh ( N + H ) 2 ) .
• Compute determinants of covariance matrices for each of p queries and nm augmented datasets D ∪ τ ij .
Each of these operations is over a matrix of size h × h and therefore costs O ( h 3 ) .
So the total cost is O ( pnmh 3 ) .
Summing these costs gives O ( nmN + phHnm + nm ( N + H ) 3 + pnmh ( N + H ) 2 + pnmh 3 ) .
Clearly the third term dominates the ﬁrst, the fourth dominates the second, and since
H > h , the fourth dominates the ﬁfth.
So, the computational cost can be summarized as
O (cid:0) nm (cid:0) ( N + H ) 3 + ph ( N + H ) 2 (cid:1)(cid:1) .
A.2 Wall Times
Though TIP and oTIP are designed for applications where samples are expensive and computation is relatively inexpensive, we present in this section data on the running time of these methods.
We ran all experiments on a shared research cluster available to us on large machines with hundreds of GB of memory and between 24 and 88 CPU cores.
In general our implementation did not make use of more than 20 CPU cores concurrently.
In Table 3, we give the running time of the phases of the TIP algorithm.
We note that the bulk of the computation in the planning procedure actually goes towards the just-in-time compilation of the JAX code that computes the cost function C τ ∗ on sampled future trajectories.
In order to allow for this compilation cost, we modiﬁed the iCEM algorithm from [51] to take ﬁxed batch sizes as the compilation (e.g. for the β tracking problem) takes approximately 90% of the time required for planning.
Unfortunately this compilation process must be repeated at every iteration due to the limitations of the JAX compiler.
We believe that a similarly JIT-compiled implementation of the planning algorithm for sampling τ ∗ on posterior samples could lead to a substantial speedup and a more ﬂexible compiler could do more still.
A.3 GP Model Details
For all of our experiments, we use a squared exponential kernel with automatic relevance determination [38, 42].
The parameters of the kernel were estimated by maximizing the likelihood of the parameters after marginalizing over the posterior GP [68].
To optimize the transition function, we simply sampled a set of points from the domain, evaluated the acquisition function, and chose the maximum of the set.
This set was chosed uniformly for every problem but β + Rotation and Reacher, for which we chose a random subset of ∪ i ∪ j τ ∗ ij (the posterior
[Table]
[Table]
samples of the optimal trajectory) since the space of samples is 10-dimensional and uniform random sampling will not get good coverage of interesting regions of the state space.
A.4 Cost Function Details
We set n = 15 and m = 1 for our Monte Carlo estimate of the cost function for each problem.
A.5 Details on Planning Method
As mentioned in the main text, we use the iCEM method from Pinneri et al.
[51] with one major modiﬁcation: a ﬁxed sample batch size.
This is in order to take advantage of the JIT compilation features of JAX and avoid recompiling code for each new batch size.
In Tables 4 and 5, we present the hyperparameters used for the planning algorithm across each problem.
The same hyperparameters were used for the TIP, MPC, EIG T , DIP, sDIP, and sTIP methods.
As recommended by the original paper, we use β = 3 for the scaling exponent of the power spectrum density of sampled noise for action sequences, γ = 1 . 25 for the exponential decay of population size, and ξ = 0 .
3 for the amount of caching.
B Description of Comparison Methods
We compare against 14 different methods across open and closed-loop problems.
Of these, 7 used the same model and planning algorithm (including hyperparameters) as TIP and oTIP.
DIP and oDIP use the cost function C ( τ ) =
− H [ T ( S (cid:48) )
| D ] and sDIP (summed DIP) uses the cost function
{+C ( τ ) = − (cid:80) hi =
+}
H [ T ( s i , a i ) | D ] .
These are all pure exploration methods, but DIP and oDIP are more sophisticated in that they plan for future observations with a large amount of joint information
[Table]
as opposed to sTIP which sums the individual information expected at each timestep.
oDIP is simply the open loop variant of DIP.
EIG T uses the same objective as sDIP but operates in the TQRL setting, querying points that approximately maximize the predictive entropy of the dynamics model.
BARL similarly operates in the TQRL setting but uses the EIG τ ∗ acquisition function from Mehta et al.
We use the authors’ implementation of that work for comparison.
MPC uses C g from (4) and plans to directly maximize expected rewards.
This method can be seen as quite similar to Kamthe and Deisenroth [34] and a close cousin of Deisenroth and Rasmussen [21] in that it optimizes the same objective with a similar model.
oMPC is simply the open loop variant of MPC.
Besides these methods which directly compare cost functions, we include 8 additional baselines from published work.
PETS is a method given in Chua et al.
[15] which uses a similar cross-entropy based planner and a probabilistic ensemble of neural networks for an uncertainty-aware estimate of the dynamics.
PETS also plans to minimize C g .
HUCRL [16] learns a policy via backpropagation through time using a hallucinated perturbation to the dynamics that maximizes discounted rewards subject to the one-step conﬁdence interval of the dynamics.
HUCRL also uses a probabilistic ensemble of neural networks.
Using the same implementation we also tested Thompson Sampling ( TS ), which acts optimally according to a network drawn from the posterior over models, and BPTT which plans to minimize C g using a neural network policy and backpropagation through time.
BPTT can also be viewed as a cousin of PILCO [21] as it attempts to take stochastic gradients of the expected cost.
We also compare against SAC [27], TD3 [23], and PPO [58].
SAC uses entropy bonuses to approximate Boltzmann exploration in an actor-critic framework.
TD3 and PPO include various tricks for stable learning and add Ornstein-Uhlenbeck noise in order to explore.
For our FEEF implementation, we took hyperparameters from the most similar comparison environments in that paper and used them for our results.
We tried several values for ‘expl_weight’ in order to se whether we were inadequately balancing exploration and exploitation.
Ultimately we saw an ‘expl_weight’ of 0.1 was the best value.
We used the author’s implementation of RHC.
RHC makes strong assumptions on the form of the reward function by assuming that all problems are regulation problems where the goal is to drive the system to a given state and keep it there (with some cost for actuation).
We were able to pass the targets for all of our problems (which may change between episodes) to the RHC controller.
We did a light hyperparameter search tuning the number of random Fourier features used in the Bayesian linear model in this method.
Ultimately we were disappointed in the performance of RHC when applied to our problems.
We believe that this might be due to its undirected uncertainty sampling objective and relatively constrained model of environment dynamics.
C Description of Control Problems C.1 Plasma Control Problems
The plasma control problems are based on controlling a tokamak, a toroidally shaped device for conﬁning a thermonuclear plasma using magnetic ﬁelds.
Achieving net positive energy from fusion requires conﬁning a plasma at high enough temperature and density long enough for hydrogen isotopes to collide and fuse.
However, as the temperature and density are increased, a wide variety of instabilities can occur which degrade conﬁnement, leading to a loss of energy.
Full physics simulation of tokamak plasmas requires 10s-1000s of CPU hours to simulate a single trajectory, and often require hand tuning of different parameters to achieve accurate results.
Following the work of Abbate et al.
[2], each of our plasma control problems used neural networks trained on data as the ground truth dynamics models.
We used the MDSPlus tool [62] to fetch historical discharges from the DIII-D tokamak in San Diego [22].
In total, we trained our models on 1,479 historical discharges.
The data was pre-processed following the procedure outlined in Abbate et al.
We describe how each environment was constructed in more detail below.
β
Tracking In this environment the goal is to adjust the total injected power (PINJ) of the neutral beams so that the normalized plasma pressure, β N (deﬁned as the ratio of thermal energy in the plasma to energy in the conﬁning magnetic ﬁelds), reaches a target value of 2%.
Reliably controlling plasmas to sustain high performance is a major goal of research efforts for fusion energy, so even this simple scenario is of interest.
The ground-truth dynamics model takes in the current β N and PINJ, the β N and PINJ at some ∆ t time in the past, and the PINJ at some ∆ t time in the future (we assume that we have complete control over the values of PINJ at all times).
Given these inputs, the model was trained to output what β N will be ∆ t time into the future.
In total, the state space is 4D and the action space is 1D.
For this environment, we set ∆ t = 200 ms , and we specify the reward function to be the negative absolute difference between the next β N and the target β N = 2% .
β
+ Rotation
Tracking This environment is a more complicated version of the β tracking environment in several ways.
First of all, the controller now must simultaneously track both β N and the core toroidal rotation of the plasma.
To do so, the controller is also allowed to set the total torque injected (TINJ) of the neutral beams (DIII-D has eight neutral beam injectors at different positions around the tokamak, so it is generally possible to control both total power and total torque independently).
Controlling both of these quantities simultaneously is of interest since rotation shear often results in better conﬁnement and less chance of instabilities in the plasma [8, 25].
In addition, we assume a multi-task setting where the requested targets for β N and rotation can be set every trajectory.
Speciﬁcally, the β N target is drawn from U (1 . 5% , 2 . 5%) and the rotation target is drawn from U (25 , 125) krad/
s every trajectory.
These targets are appended to the state space.
The learned, ground-truth dynamics model is also more sophisticated here.
In addition to the inputs and outputs used by the β tracking environment model, the inputs for this model also include rotation and TINJ at times t , t − ∆ t , and t + ∆ t for TINJ only.
This model receives additional information about the plasma (e.g. the shape of the plasma); however, we have assumed these inputs are ﬁxed to reasonable values in order to avoid partial observability problems.
In total, the state space of this problem is 10D (targets plus current and past observations for β N , rotation, PINJ, and TINJ) and the action space is 2D (next PINJ and TINJ settings).
C.2 Robotics Problems Pendulum
The pendulum swing-up problem is the standard one found in the OpenAI gym [10].
The state space contains the angle of the pendulum and its ﬁrst derivative and action space simply the scalar torque applied by the motor on the pendulum.
The challenge in this problem is that the motor doesn’t have enough torque to simply rotate the pendulum up from all positions and often requires a back-and-forth swing to achieve a vertically balanced position.
The reward function here penalizes deviation from an upright pole and squared torque.
Cartpole The cartpole swing-up problem has 4-dimensional state (position of the cart and its velocity, angle of the pole and its angular velocity) and a 1-dimensional action (horizontal force applied to the cart).
Here, the difﬁculty lies in translating the horizontal motion of the cart into effective torque on the pole.
The reward function is a negative sigmoid function penalizing the distance betweent the tip of the pole and a centered upright goal position.
Reacher The reacher problem simulates a 2-DOF robot arm aiming to move the end effector to a randomly resampled target provided.
The problem requires joint angles and velocities as well as an indication of the direction of the goal, giving an 8-dimensional state space along with the 2-dimensional control space.
D Additional Results
Due to space constraints in the main paper, we omitted results for the methods sDIP and BPTT.
The are included alongside the rest in Table 6.
They are outperformed across the board by TIP.
E Additional Related Work E.1 Bayesian Exploration Techniques
Given unlimited computation and an accurate prior, solving the Bayes-adaptive MDP [53] gives an optimal tradeoff between exploration and exploitation by explicitly accounting for the updated beliefs that would result from future observations and planning to ﬁnd actions that result in high rewards as quickly as can be managed given the current posterior.
However, this is computationally expensive even in small ﬁnite MDPs and totally intractable in continuous settings.
Kolter and Ng
Environment
TIP sTIP
DIP sDIP
MPC
PETS
SAC
TD
PPO
FEEF
HUCRL
TS
BPTT
BARL
EIG T
Pendulum     5.6k 7k 26k 14k >50k >50k >50k 56 Cartpole     1.63k 32k 18k >1M >2.5k >6k >6k >6k 121 β Tracking      12k 17k 39k     >1k β +Rotation > > > >  30k >50k >50k >2k >5k >5k >5k > >1k Reacher > >1k >1k  23k 13k >100k >5k
6.6k 4.5k 3.7k >1.5k
[Figure]
[35] and Guez et al.
[26] show that even approximating these techniques can result in substantial theoretical reductions in sample complexity compared to frequentist PAC-MDP bounds as in Kakade [33].
Another line of work [18, 19] uses the myopic value of perfect information as a heuristic for similar Bayesian exploration in the tabular MDP setting.
Further techniques for exploration include knowledge gradient policies [57, 56], which approximate the value function of the Bayes-adaptive
MDP and information-directed sampling (IDS) [54], which takes actions based on minimizing the ratio between squared regret and information gain over dynamics.
This was extended to continuousstate ﬁnite-action settings using neural networks in Nikolov et al.
Another very relevant recent paper [7] gives an acquisition strategy in policy space that iteratively trains a data-collection policy in the model that trades off exploration against exploitation using methods from active learning.
Achterhold and Stueckler [3] use techniques from BOED to efﬁciently calibrate a Neural Process representation of a distribution of dynamics to a particular instance, but this calibration doesn’t include information about the task.
A tutorial on Bayesian RL methods can be found in Ghavamzadeh et al.
[24] for further reference.
E.2 Gaussian Processes (GPs) in Reinforcement Learning
There has been substantial prior work using GPs [52] in reinforcement learning.
Most well-known is
PILCO [21], which computes approximate analytic gradients of policy parameters through the GP dynamics model while accounting for uncertainty.
The original work is able to propagate the ﬁrst 2 moments of the occupancy distribution through time using the GP dynamics and backpropagate gradients of the rewards to policy parameters.
In [69], a method is developed for efﬁciently sampling functions from a GP posterior with high accuracy.
One application show in their work is a method of using these samples to backpropagate gradients of rewards through time to policy paramters, which can be interpreted as a different sort of PILCO implementation.
Most related to our eventual
MPC-based method is [34], which gives a principled probabilistic model-predictive control algorithm for GPs.
We combine ideas from this paper, PETS [15], and the ability to sample posterior functions discussed above to give our eventual MPC component as discussed in Section 4.1.