Original: A N E XPERIMENTAL D ESIGN P ERSPECTIVE ON M ODEL -B ASED R EINFORCEMENT L EARNING
Revised: A N E XPERIMENTAL D ESIGN P ERSPECTIVE ON M ODEL -B ASED R EINFORCEMENT L EARNING

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In many practical applications of RL, it is expensive to observe state transitions from the environment.
Revised: In many practical applications of RL, it is expensive to observe state transitions from the environment.

Original: For example, in the problem of plasma control for nuclear fusion, computing the next state for a given state-action pair requires querying an expensive transition function which can lead to many hours of computer simulation or dollars of scientiﬁc research.
Revised: For example, in the problem of plasma control for nuclear fusion, computing the next state for a given state-action pair requires querying an expensive transition function which can lead to many hours of computer simulation or dollars of scientiﬁc research.

Original: Such expensive data collection prohibits application of standard RL algorithms which usually require a large number of observations to learn.
Revised: Such expensive data collection prohibits application of standard RL algorithms which usually require a large number of observations to learn.

Original: In this work, we address the problem of efﬁciently learning a policy while making a minimal number of state-action queries to the transition function.
Revised: In this work, we address the problem of efﬁciently learning a policy while making a minimal number of state-action queries to the transition function.

Original: In particular, we leverage ideas from Bayesian optimal experimental design to guide the selection of state-action queries for efﬁcient learning.
Revised: In particular, we leverage ideas from Bayesian optimal experimental design to guide the selection of state-action queries for efﬁcient learning.

Original: We propose an acquisition function that quantiﬁes how much information a state-action pair would provide about the optimal solution to a Markov decision process.
Revised: We propose an acquisition function that quantiﬁes how much information a state-action pair would provide about the optimal solution to a Markov decision process.

Original: At each iteration, our algorithm maximizes this acquisition function, to choose the most informative state-action pair to be queried, thus yielding a data-efﬁcient RL approach.
Revised: At each iteration, our algorithm maximizes this acquisition function, to choose the most informative state-action pair to be queried, thus yielding a data-efﬁcient RL approach.

Original: We experiment with a variety of simulated continuous control problems and show that our approach learns an optimal policy with up to 5 – 1 , 000 × less data than modelbased RL baselines and 10 3 – 10 5 × less data than model-free RL baselines.
Revised: We experiment with a variety of simulated continuous control problems and show that our approach learns an optimal policy with up to 5 – 1 , 000 × less data than modelbased RL baselines and 10 3 – 10 5 × less data than model-free RL baselines.

Original: We also provide several ablated comparisons which point to substantial improvements arising from the principled method of obtaining data.
Revised: We also provide several ablated comparisons which point to substantial improvements arising from the principled method of obtaining data.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Reinforcement learning (RL) has suffered for years from a curse of poor sample complexity.
Revised: Reinforcement learning (RL) has suffered for years from a curse of poor sample complexity.

Original: State-ofthe-art model-free reinforcement learning algorithms routinely take tens of thousands of sampled transitions to solve very simple tasks and millions to solve moderately complex ones (Haarnoja et al., 2018; Lillicrap et al., 2015).
Revised: State-ofthe-art model-free reinforcement learning algorithms routinely take tens of thousands of sampled transitions to solve very simple tasks and millions to solve moderately complex ones (Haarnoja et al., 2018; Lillicrap et al., 2015).

Original: The current best model-based reinforcement learning (MBRL) algorithms are better, requiring thousands of samples for simple problems and hundreds of thousands of samples for harder ones (Chua et al., 2018).
Revised: The current best model-based reinforcement learning (MBRL) algorithms are better, requiring thousands of samples for simple problems and hundreds of thousands of samples for harder ones (Chua et al., 2018).

Original: In settings where each sample is expensive, even this smaller cost can be prohibitive for the practical application of RL.
Revised: In settings where each sample is expensive, even this smaller cost can be prohibitive for the practical application of RL.

Original: For example, in the physical sciences many simulators require the solution of computationally demanding spatial PDEs in plasma control (Breslau et al., 2018) or aerodynamics applications (Jameson & Fatica, 2006).
Revised: For example, in the physical sciences, many simulators require the solution of computationally demanding spatial PDEs in plasma control (Breslau et al., 2018; Char et al., 2019) or aerodynamics applications (Jameson & Fatica, 2006).

Original: In robotics, due to the difﬁculty of simulating more complicated objects (Heiden et al., 2021),typical RL methods operate on fast but limited rigid-body simulators (Todorov et al., 2012).
Revised: In robotics, due to the cost of simulating more complicated objects (Heiden et al., 2021), RL methods are typically constrained to fast but limited rigid-body simulators (Todorov et al., 2012).

Original: These costly transition functions prompt the question: “If we were to collect one additional datapoint from anywhere in the state-action space to best improve our solution to the task, which one would it be?”
Revised: These costly transition functions prompt the question: “If we were to collect one additional datapoint from anywhere in the state-action space to best improve our solution to the task, which one would it be?”

Original: An answer to this question can be used to guide data collection in RL.
Revised: An answer to this question can be used to guide data collection in RL.

Original: Across the ﬁelds of black-box optimization and experimental design, techniques have been developed which choose data to collect that are particularly useful in improving the value of the objective function of the problem.
Revised: Across the ﬁelds of black-box optimization and experimental design, techniques have been developed which choose data to collect that are particularly useful in improving the value of the objective function of the problem.

Original: For example, Bayesian optimization (BO) focuses on maximizing a function for which only a black-box model is available and queries are expensive (Frazier, 2018; Shahriari et al., 2015).
Revised: For example, Bayesian optimization (BO) focuses on maximizing an unknown (black-box) function where queries are expensive (Frazier, 2018; Shahriari et al., 2015).

Original: More generally, Bayesian optimal experimental design (BOED) aims to choose data to collect which are maximally informative about the value of some derived quantity (Chaloner & Verdinelli, 1995).
Revised: More generally, Bayesian optimal experimental design (BOED) aims to choose data to collect which are maximally informative about the value of some derived quantity (Chaloner & Verdinelli, 1995).

Original: We aim to leverage these ideas for data-efﬁciency in reinforcement learning.
Revised: We aim to leverage these ideas for data-efﬁciency in reinforcement learning.

Original: Along these lines, several works in the realm of Bayesian RL address this problem in the sequential setting.
Revised: Along these lines, several works in the realm of Bayesian RL address this problem in the sequential setting.

Original:  Bayes-adaptive MDP (Ross et al., 2007) constructs a modiﬁed MDP by augmenting it using the posterior  of the MDP, leading to a policy that can optimally trade off between acquiring more information and exploiting the knowledge it already has.
Revised: A Bayes-adaptive MDP (Ross et al., 2007) constructs a modiﬁed MDP by augmenting the state space with the posterior belief of the MDP, leading to a policy that can optimally trade off between acquiring more information and exploiting the knowledge it already has.

Original: However such an MDP is intractable to exactly solve in large spaces so approximations and heuristics have been developed for the solution (Smith, 2007; Guez et al., 2012).
Revised: However such an MDP is intractable to exactly solve in large spaces so approximations and heuristics have been developed for the solution (Smith, 2007; Guez et al., 2012).

Original: A particularly relevant heuristic is the value of perfect information (VPI) from Dearden et al.
Revised: A particularly relevant heuristic is the value of perfect information (VPI) from Dearden et al.

Original: (1998), which attempts to capture the potential change in value of a state if the value of a particular action at that state was perfectly known, speciﬁcally in a tabular setting.
Revised: (1998), which attempts to capture the potential change in value of a state if the value of a particular action at that state was perfectly known, speciﬁcally in a tabular setting.

Original: However, VPI doesn’t attempt to distinguish between states that are visited during the execution of the optimal policy and those that aren’t.
Revised: However, VPI doesn’t attempt to distinguish between states that are visited during the execution of the optimal policy and those that aren’t.

Original: This is a critical distinction when evaluating datato collect in continuous spaces asotherwise queries would be wasted learning the optimal policy in irrelevant parts of the state space.
Revised: This is a critical distinction when collecting data in continuous spaces, as queries may be wasted learning an optimal policy in irrelevant parts of the state space.

Original: Motivated by our opening question, in this paper we study the setting where the agent collects data by sequentially making queries to the transition function with free choice of both the initial state and the action.
Revised: Motivated by our opening question, in this paper we study the setting where the agent collects data by sequentially making queries to the transition function with free choice of both the initial state and the action.

Original: We refer to this setting as transition-query reinforcement learning (TQRL) and formally deﬁne it in Section 3.1.
Revised: We refer to this setting as transition-query reinforcement learning (TQRL) and formally deﬁne it in Section 3.1.

Original: An effective algorithm for this setting would need to choose informative initial states as well as actions, which to the best of our knowledge  has not been studied in the continuous MDP literature.
Revised: Although this setting has been studied in the tabular case, to the best of our knowledge it has not been studied in the continuous MDP literature.

Original: In this work, we draw a connection between MBRL and the world of BOED by deriving an acquisition function that quantiﬁes how much information a state-action pair would provide about the optimal solution to a MDP.
Revised: In this work, we draw a connection between MBRL and the world of BOED by deriving an acquisition function that quantiﬁes how much information a state-action pair would provide about the optimal solution to a MDP.

Original: Like the techniques in Bayesian RL, our acquisition function is able to determine which state-action pairs are worth acquiring in a way which takes into account the reward function and the uncertainty in the dynamics.
Revised: Like the techniques in Bayesian RL, our acquisition function is able to determine which state-action pairs are worth acquiring in a way which takes into account the reward function and the uncertainty in the dynamics.

Original: However, like thefull Bayes-Adaptive MDP and unlike the VPI heuristic, this function takes into account the current estimates of which states the optimal policy willactually visit and values potential queries accordingly.
Revised: However, like the Bayes-Adaptive MDP and unlike the VPI heuristic, this function takes into account the current estimates of which states the optimal policy will visit and values potential queries accordingly.

Original: It isalso scalable enough to apply to multidimensional continuous control problems.
Revised: Furthermore, our acquisition function is scalable enough to apply to multidimensional continuous control problems.

Original: In particular, our acquisition function is the expected information gain (EIG) about the trajectory taken by an optimal policy in the MDP that would be achieved if we were to query the transition function at a given state-action pair.
Revised: In particular, our acquisition function is the expected information gain (EIG) about the trajectory taken by an optimal policy in the MDP that would be achieved if we were to query the transition function at a given state-action pair.

Original: Finally, we assess the performance of our acquisition function as a data selection strategy in the
Revised: Finally, we assess the performance of our acquisition function as a data selection strategy in the

Original: TQRL setting.
Revised: TQRL setting.

Original: Using this method we are able to solve several continuous reinforcement learning tasks (including a nuclear fusion example) using orders of magnitude less data than a variety of competitor methods.
Revised: Using this method we are able to solve several continuous reinforcement learning tasks (including a nuclear fusion example) using orders of magnitude less data than a variety of competitor methods.

Original: In summary, the contributions of our paper are:
Revised: In summary, the contributions of our paper are:

Original: • We construct a novel acquisition function that quantiﬁes how much information a stateaction pair would provide about the optimal solution to a continuous MDP if the next state were observed from the ground truth transition function.
Revised: • We construct a novel acquisition function that quantiﬁes how much information a state-action pair would provide about the optimal solution to a continuous MDP if the next state were observed from the ground truth transition function.

Original: Our  function is able to select relevant datapoints for control purposes leading to much-improved data efﬁciency.
Revised: Our acquisition function is able to select relevant datapoints for control purposes leading to improved data efﬁciency.

Original: • We propose a practical algorithm for computing this acquisition function and use it to solve continuous MDPs in the TQRL setting.
Revised: • We propose a practical algorithm for computing this acquisition function and use it to solve continuous MDPs in the TQRL setting.

Original: • We evaluate the algorithm on ﬁve widely varying control tasks, where it is often orders of magnitude more sample-efﬁcient than competitor methods and reaches similar asymptotic performance.
Revised: • We evaluate the algorithm on ﬁve diverse control tasks, where it is often orders of magnitude more sample-efﬁcient than competitor methods and reaches similar asymptotic performance.

Original:  R ELATED W ORK Transition Query Reinforcement Learning
Revised:  R ELATED W ORK Transition Query Reinforcement Learning

Original: In many RL algorithms, data is collected by initializing a policy at a start state and executing actions in the environment in an episodic manner.
Revised: In many RL algorithms, data is collected by initializing a policy at a start state and executing actions in the environment in an episodic manner.

Original: Kearns et al.
Revised: Kearns et al.

Original: (2002) introduced the setting where the agent collects data by  sampling transitionsin a sequential manner from the ground truth transition model by querying at a state and action of its choice, which they refer to as RL with access to a generative model .
Revised: (2002) introduced the setting where the agent collects data by sequentially sampling transitions from the ground truth transition model by querying at a state and action of its choice, which they refer to as RL with access to a generative model .

Original: We refer to this setting for brevity as TQRL.
Revised: We refer to this setting for brevity as TQRL.

Original: This setting is relevant in a variety of real-world applications where there is a simulator of the transition model available.
Revised: This setting is relevant in a variety of real-world applications where there is a simulator of the transition model available.

Original: We in particular see the setting in nuclear fusion research, where plasma dynamics are modeled by solving large partial differential equations where 200ms of plasma time can take up to an hour in simulation (Breslau et al., 2018).
Revised: In particular, we see the setting in nuclear fusion research, where plasma dynamics are modeled by solving large partial differential equations where 200ms of plasma time can take up to an hour in simulation (Breslau et al., 2018).

Original: There is substantial theoretical work on TQRL for ﬁnite MDPs.
Revised: There is substantial theoretical work on TQRL for ﬁnite MDPs.

Original: In particular, Azar et al.
Revised: In particular, Azar et al.

Original: (2013) give matching log-linear upper and lower PAC sample complexity bounds, a substantial speedup to the upper bound for the standard problem which is quadratic in state size (Kakade, 2003).
Revised: (2013) give matching log-linear upper and lower PAC sample complexity bounds, a substantial speedup to the upper bound for the standard problem which is quadratic in state size (Kakade, 2003).

Original: This is achieved simply by the naive algorithm of learning a transition model by uniformly sampling the space and then performing value iteration on the estimate of the MDP for an optimal policy.
Revised: This is achieved simply by the naive algorithm of learning a transition model by uniformly sampling the space and then performing value iteration on the estimate of the MDP for an optimal policy.

Original: More recently, the bound for this setting was tightened to hold for smaller numbers of samples by Li et al.
Revised: More recently, the bound for this setting was tightened to hold for smaller numbers of samples by Li et al.

Original: (2020), meaning that for any dataset size in a continuous problem, the PAC performance can be quantiﬁed.
Revised: (2020), meaning that for any dataset size in a continuous problem, the PAC performance can be quantiﬁed.

Original: Finally, Agarwal et al.
Revised: Finally, Agarwal et al.

Original: (2020) show that the naive ‘plug-in’ estimator used in the previous works is minimax optimal for this setting.
Revised: (2020) show that the naive ‘plug-in’ estimator used in the previous works is minimax optimal for this setting.

Original: In summary, this setting is thoroughly understood for ﬁnite
Revised: In summary, this setting is thoroughly understood for ﬁnite

Original: MDPs and it gives a sample complexity reduction from quadratic to linear in the state space size.
Revised: MDPs and it gives a sample complexity reduction from quadratic to linear in the state space size.

Original: To our knowledge there do not exist works speciﬁcally solving the TQRLsetting, though every standard RL algorithm can be applied in this setting by continuing to collect data in an online manner.
Revised: To our knowledge there do not exist works speciﬁcally solving the TQRL setting for continuous

Original:  In this work, we give an algorithm speciﬁcally designed for this setting in continuous MDPs. Our work empirically shows sample complexity beneﬁts reminiscent of those theoretically shown in the tabular setting.
Revised: MDPs. In this work, we give an algorithm speciﬁcally designed for this setting, which shows sample complexity beneﬁts reminiscent of those theoretically shown in the tabular setting.

Original: Exploration in Reinforcement Learning To encourage exploration in RL, agents often use an (cid:15) - greedy approach (Mnih et al., 2013), upper conﬁdence bounds (UCB) (Chen et al., 2017), Thompson sampling  (Osband et al., 2016), added Ornstein-Uhlenbeck action noise (Lillicrap et al., 2015), or entropy bonuses (Haarnoja et al., 2018) to add noise to a policy which is otherwise optimizing the
Revised: Exploration in Reinforcement Learning To encourage exploration in RL, agents often use an (cid:15) - greedy approach (Mnih et al., 2013), upper conﬁdence bounds (UCB) (Chen et al., 2017), Thompson sampling (TS) (Osband et al., 2016), added Ornstein-Uhlenbeck action noise (Lillicrap et al., 2015), or entropy bonuses (Haarnoja et al., 2018) to add noise to a policy which is otherwise optimizing the

Original: RL objective.
Revised: RL objective.

Original: Though UCB, Thompson Sampling, andan entropy bonus all try toexplore in a way which can adapt  to the problem, they all tackle which action to take from a predetermined state and don’t explicitly consider which states would be good to acquire data from.
Revised: Although UCB, TS, and entropy bonuses all try to adapt the exploration strategy to the problem, they all tackle which action to take from a predetermined state and don’t explicitly consider which states would be best to acquire data from.

Original: As we mentioned, an ideal method of exploration would be to solve the intractable Bayes-adaptive MDP (Ross et al., 2007), giving an optimal tradeoff between exploration and exploitation.
Revised: An ideal method of exploration would be to solve the intractable Bayes-adaptive MDP (Ross et al., 2007), giving an optimal tradeoff between exploration and exploitation.

Original: Kolter & Ng (2009); Guez et al.
Revised: Kolter & Ng (2009); Guez et al.

Original: (2012) show that even approximating these techniques in the sequential setting can result in substantial theoretical reductions in sample complexity compared to frequentist PACMDP bounds as in Kakade (2003).
Revised: (2012) show that even approximating these techniques in the sequential setting can result in substantial theoretical reductions in sample complexity compared to frequentist PAC-MDP bounds as in Kakade (2003).

Original: Other methods stemming from Dearden et al.
Revised: Other methods stemming from Dearden et al.

Original: (1998; 1999) address this by using the myopic Value of Perfect Information as a heuristic for similar Bayesian exploration.
Revised: (1998; 1999) address this by using the myopic value of perfect information as a heuristic for similar Bayesian exploration.

Original: These methods don’t scale to continuous problems and don’t give any method of choosing at which states to query, only actions given a current state.
Revised: However, these methods don’t scale to continuous problems and don’t provide a way to choose states to query.

Original: These methods were further extended with the development of Knowledge Gradient policies (Ryzhov et al., 2019; Ryzhov & Powell, 2011), which approximate the value function of the Bayes-adaptive MDP, and Information-Directed Sampling (Russo & Van Roy, 2014), which takes actions based on minimizing the ratio between squared regret and information gain over dynamics.
Revised: These methods were further extended with the development of knowledge gradient policies (Ryzhov et al., 2019; Ryzhov & Powell, 2011), which approximate the value function of the Bayes-adaptive MDP, and information-directed sampling (IDS) (Russo & Van Roy, 2014), which takes actions based on minimizing the ratio between squared regret and information gain over dynamics.

Original: This was extended to continuous-state ﬁnite-action settings in Nikolov et al.
Revised: This was extended to continuous-state ﬁnite-action settings in Nikolov et al.

Original: However, this work doesn’t solve fully continuous problems, operates in the rollout setting rather than TQRL, and computes the information gain with respect to the dynamics rather than some notion of the optimal policy.
Revised: However, this work doesn’t solve fully continuous problems, operates in the rollout setting rather than TQRL, and computes the information gain with respect to the dynamics rather than some notion of the optimal policy.

Revised: In a similar spirit, Arumugam & Van Roy (2021) provide a further generalization of IDS which can also be applied to RL.

Revised: One recent work very close to ours is Lindner et al.

Revised: (2021), which actively queries an expensive reward function (instead of dynamics as in this work) to learn a Bayesian model of reward.

Revised: Another very relevant recent paper (Ball et al., 2020) 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.

Revised: Achterhold & Stueckler (2021) 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.

Original: A tutorial on Bayesian  methods can be found in Ghavamzadeh et al.
Revised: A tutorial on Bayesian RL methods can be found in Ghavamzadeh et al.

Original: (2016) for further reference.
Revised: (2016) for further reference.

Original: Separate from the techniques used in RL for a particular task, several methods tackle the problem of unsupervised exploration (Schmidhuber, 1991), where the goal is to learn as much as possible about the transition model without a task or reward function.
Revised: Separate from the techniques used in RL for a particular task, several methods tackle the problem of unsupervised exploration (Schmidhuber, 1991), where the goal is to learn as much as possible about the transition model without a task or reward function.

Original: One approach synthesizes a reward from modeling errors (Pathak et al., 2017).
Revised: One approach synthesizes a reward from modeling errors (Pathak et al., 2017).

Original: Another estimates learning progress by estimating model accuracy (Lopes et al., 2012).
Revised: Another estimates learning progress by estimating model accuracy (Lopes et al., 2012).

Original: Others use an information gain-motivated formulation of model disagreement (Pathak et al., 2019; Shyam et al., 2019) as a reward.
Revised: Others use an information gain-motivated formulation of model disagreement (Pathak et al., 2019; Shyam et al., 2019) as a reward.

Original: Other methods incentivize the policy to explore regions it hasn’t been before using hash-based counts (Tang et al., 2017), predictions mimicking a randomly initialized network (Burda et al., 2019),or a density estimate (Bellemare et al., 2016).
Revised: Other methods incentivize the policy to explore regions it hasn’t been before using hash-based counts (Tang et al., 2017), predictions mimicking a randomly initialized network (Burda et al., 2019), a density estimate (Bellemare et al., 2016), or predictive entropy (Buisson-Fenet et al., 2020).

Original: However, these methods all assume that there is no reward function and would be inefﬁcient if used in that setting  as they spend time exploring areas of state space which can be quickly determined to be bad for maximizing reward on a task.
Revised: However, these methods all assume that there is no reward function and are inefﬁcient for the setting of this paper, as they spend time exploring areas of state space which can be quickly determined to be bad for maximizing reward on a task.

Original: Bayesian Algorithm Execution and BOED
Revised: Bayesian Algorithm Execution and BOED

Original: Recently, a ﬂexible framework known as Bayesian algorithm execution (BAX) (Neiswanger et al., 2021) has been proposed for efﬁciently estimating properties of expensive black-box functions, which builds off of a large literature from Bayesian Optimal Experiment Design (Chaloner & Verdinelli, 1995).
Revised: Recently, a ﬂexible framework known as Bayesian algorithm execution (BAX) (Neiswanger et al., 2021) has been proposed for efﬁciently estimating properties of expensive black-box functions, which builds off of a large literature from Bayesian Optimal Experiment Design (Chaloner & Verdinelli, 1995).

Original: The BAX framework gives a general procedure for sampling points which are informative about the future execution of an algorithm.
Revised: The BAX framework gives a general procedure for sampling points which are informative about the future execution of an algorithm.

Original: In this paper, we extend this framework to the setting of model-predictive control, when we have expensive dynamics (i.e. transition function) which we treat as a black-box function in the BAX framework.
Revised: In this paper, we extend this framework to the setting of model-predictive control, when we have expensive dynamics (i.e. transition function) which we treat as a black-box function in the BAX framework.

Original: Via this strategy, we are able to use similar techniques to develop acquisition functions for data collection in reinforcement learning.
Revised: Via this strategy, we are able to use similar techniques to develop acquisition functions for data collection in reinforcement learning.

Original: Gaussian Processes (GPs) in Reinforcement Learning There has been substantial prior work using GPs in reinforcement learning.
Revised: Gaussian Processes (GPs) in Reinforcement Learning There has been substantial prior work using GPs in reinforcement learning.

Original: Most well-known is PILCO (Deisenroth & Rasmussen, 2011), which computes approximate analytic gradients of policy parameters through the GP dynamics model while accounting for uncertainty.
Revised: Most well-known is PILCO (Deisenroth & Rasmussen, 2011), which computes approximate analytic gradients of policy parameters through the GP dynamics model while accounting for uncertainty.

Original: Most related to our eventual MPC method is (Kamthe & Deisenroth, 2018), which gives a principled probabilistic model-predictive control algorithm for GPs.
Revised: Most related to our eventual MPC method is (Kamthe & Deisenroth, 2018), which gives a principled probabilistic model-predictive control algorithm for GPs.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this work we deal with ﬁnite-horizon discrete time Markov decision processes (MDPs) which consist of a tuple (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 functionmapping T :
Revised: In this work we deal with ﬁnite-horizon discrete-time Markov decision processes (MDPs) which consist of a tuple (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 :

Original: S ×
Revised: S ×

Original: A → P ( S ) (using the convention that P ( X ) is the set of probability measures over
Revised: A → P ( S ) (using the convention that P ( X ) is the set of probability measures over

Original: X ), r : S ×
Revised: X ), r : S ×

Original: A × S → R  a reward function, p 0 ( s )  a distribution over S of start states, and H ∈ N  a horizon.
Revised: A × S → R is a reward function, p 0 ( s ) is a distribution over S of start states, and H ∈ N is a horizon.

Original: We always assume S , A , p 0 , H are known.
Revised: We always assume S , A , p 0 , H are known.

Original: 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.
Revised: 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.

Original: Our primary function of interest is the transition function T , which we learn from data.
Revised: Our primary function of interest is the transition function T , which we learn from data.

Original: Our aim is to ﬁnd a policy π : S → P ( A ) that maximizes the rewards obtained by executing the policy.
Revised: Our aim is to ﬁnd a policy π : S → P ( A ) that maximizes the objective given below.

Original: We can describe the execution of π in the MDP as a ﬁnite collection of random variables generated by s 0 ∼ p 0 and a i ∼ π ( s i ) , s i+1 ∼ T ( s i  , a i  ) for i ∈ 0 , . . .
Revised: We can describe the execution of π in the MDP as a ﬁnite collection of random variables generated by s 0 ∼ p 0 and a i ∼ π ( s i ) , s i ∼ T ( s i − i , a i − 1 ) for i ∈ 1 , . . .

Original: , H− 1 .
Revised: , H .

Original: Then this objective can be written
Revised: Then this objective can be written

Original: [Equation]
Revised: [Equation]

Revised: i =

Revised: r ( s i , a i , s i +1 ) (cid:35) .

Original: 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.
Revised: 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.

Original: We denote  π ∗ =
Revised: We denote the optimal policy as π ∗ =

Original: argmax π J T ( π ) as an optimal policy, which we can assume to be deterministic (Sutton & Barto, 1998) but not necessarily unique.
Revised: argmax π J T ( π ) , which we can assume to be deterministic (Sutton & Barto, 1998) but not necessarily unique.

Original: Finally, we assume that there is some prior P ( T ) for which the posterior
Revised: Finally, we assume that there is some prior P ( T ) for which the posterior

Original: P ( T | D ) is available for sampling given a dataset D .
Revised: P ( T | D ) is available for sampling given a dataset D .

Original: 3. T RANSITION Q UERY R EINFORCEMENT L EARNING 
Revised: 3. T RANSITION Q UERY R EINFORCEMENT L EARNING (TQRL)

Original: In the  online RL setting, one assumes data D
Revised: In the standard online RL setting, one assumes data D

Original: = { ( s i , a i , s (cid:48) i ) } i ∈
Revised: = { ( s i , a i , s (cid:48) i ) } i ∈

Original: [ n ] must be collected in length- H trajectories  where the initial state s 0 is sampled from p 0  and after an action a i is chosen the next state s (cid:48)
Revised: [ n ] must be collected in length- H trajectories ( rollouts ) where the initial state s∼ p 0 , and after an action a i is chosen, the next state s (cid:48)

Original: i =
Revised: i =

Original: s i +1 is sampled from T ( s i , a i ) up to i
Revised: s i +1 is sampled from T ( s i , a i ) up to i

Original: =
Revised: =

Original: H  at which point a new initial state s H +1 ∼ p 0 is sampled.
Revised: H , at which point the process repeats.

Original: We pay special attention to the TQRL setting, where the agent sequentially acquires datapoints x i = ( s i , a i , s (cid:48) i ) in arbitrary order by choosing ( s i , a i ) from S × A and sampling s (cid:48)
Revised: In this work, we consider the TQRL setting, where the agent sequentially acquires data ( s i , a i , s (cid:48) i ) in arbitrary order by querying a state action pair ( s i , a i ) from S ×A and recieving a sample s (cid:48)

Original: i ∼ T ( s, a )  (Kearns et al., 2002; Kakade, 2003; Azar et al., 2013).
Revised: i ∼ T ( s, a ) from the black-box transition function T (Kearns et al., 2002; Kakade, 2003; Azar et al., 2013).

Original: Performance in this setting is still evaluated by Objective 1.
Revised: The goal in both settings is to ﬁnd a policy which optimizes the objective in Equation (1).

Original: It has been shown for ﬁnite MDPs in Azar et al. that the PAC sample complexity, which is the number of samples required to identify with high probability a policy that achieves almost optimal value, of this setting is ˜ O ( |S||A| ) , ignoring the PAC factors.
Revised: It has been shown for ﬁnite MDPs in (Azar et al., 2013) that the PAC sample complexity, which is the number of samples required to identify with high probability a policy that achieves almost optimal value, of this setting is ˜ O ( |S||A| ) , ignoring the PAC factors.

Original: This is notably better than the bound of ˜
Revised: This is notably better than the bound of ˜

Original: O ( |S| 2 |A| ) in the online RL setting given in Section 8.3 of Kakade (2003).
Revised: O ( |S| 2 |A| ) in the online RL setting given in Section 8.3 of Kakade (2003).

Original: The improvement shown in ﬁnite cases suggests that there could be similar reductions available in a continuous setting.
Revised: The improvement shown in ﬁnite cases suggests that there could be similar reductions available in a continuous setting.

Original: 3. M ODEL -P REDICTIVE C ONTROL
Revised: 3. M ODEL -P REDICTIVE C ONTROL

Original: Though our acquisition function is derived for any policy search procedure including planning or model-free reinforcement learning algorithms, we focus on model-predictive control (MPC) as it is simple and requires minimal components to function.
Revised: Though our acquisition function is derived for any policy search procedure including planning or model-free reinforcement learning algorithms, we focus on model-predictive control (MPC) as it is simple and requires minimal components to function.

Original: MPC is a standard technique in model-based reinforcement learning (Chua et al., 2018; Wang & Ba, 2020).
Revised: MPC is a standard technique in model-based reinforcement learning (Chua et al., 2018; Wang & Ba, 2020).

Original: Using a probabilistic dynamics model ˆ T =
Revised: Using an estimated dynamics model ˆ T : S ×

Original: E [ T | D ] and an optimization algorithm, an MPC strategy with a planning horizon h ∈ N choosing an action at a state s 0 solves the planning problem
Revised: A → S and an optimization algorithm, an MPC strategy with a planning horizon h ∈ N choosing an action at a state s 0 solves the planning problem

Original: [Equation]
Revised: [Equation]

Original: with a stochastic optimization technique.

Original: Planning is typically redone periodically, often every timestep, as actions are executed in the real environment.
Revised: Planning is typically redone periodically, often every timestep, as actions are executed in the real environment.

Original: The cross-entropy method (CEM) ismost often used to solve the optimization problem.
Revised: The cross-entropy method (CEM) is often used to solve this optimization problem.

Original: In this work, we use an improved variant of CEM described in Pinneri et al.
Revised: In this work, we use an improved variant of CEM described in Pinneri et al.

Original: Given sufﬁcient computational resources and an accurate model, model-predictive control will typically ﬁnd very good if not optimal actions in an MDP.

Original: In this paper we will refer to π T as the stochastic policy obtained by running MPC over a dynamics function T .
Revised: Here we will refer to π T as the stochastic policy obtained by running MPC over a dynamics function T .

Original: The randomness in the policy is due to any randomness inˆ T and the randomness used by the optimizer.
Revised: The randomness in the policy is due to any randomness in T and the randomness used by the optimizer.

Original: We give details on the hyperparameters involved in this optimization problem and their values in Section B.
Revised: We give details on the hyperparameters involved in this optimization problem and their values in Section B.

Original:  A N A CQUISITION F UNCTION FOR M ODEL -B ASED RL
Revised:  A N A CQUISITION F UNCTION FOR M ODEL -B ASED RL

Original: We draw inspiration from BO and BOED in constructing an acquisition function suitable for control applications.
Revised: We draw inspiration from BO and BOED in constructing an acquisition function suitable for control applications.

Original: For our purposes an acquisition function is a computationally tractable function
Revised: For our purposes, an acquisition function is a computationally tractable function

Original: S ×
Revised: S ×

Original: A → R that describes the marginal contribution of collecting a transition sample at a state-action pair inimproving the performance of the policy on the MDP conditioned on all previously observed data.
Revised: A → R that describes the marginal improvement in the performance of the policy on the MDP (conditioned on all previously observed data) when observing one additional state-action pair ( s, a ) ∈ S × A .

Original: As acquisition functions are greedy, they aren’t necessarily optimal data-selection strategies given a budget in the sense that solving a Bayes-adaptive MDP is optimal.
Revised: As acquisition functions are greedy, they aren’t necessarily optimal data-selection strategies given a ﬁxed budget, compared to non-myopic strategies such as solving the Bayes-adaptive MDP.

Original: However, they are tractable and allow for good one-step choices to be made that account for all information currently available.
Revised: However, greedy strategies using mutual information are tractable and are often effective due to the submodularity of the expected information gain.

Original: In this work we focus on acquisition functions written as an Expected Information Gain (EIG). These functions output the mutual information  between a query of our transition model and a target random variable.
Revised: More speciﬁcally, our acquisition function is an expected information gain (EIG) or equivalently the mutual information (MI) between a query of our transition model and a representation of the optimal policy, as we elaborate in this section.

Original: The target variable should contain the minimal information required for an optimal solution of our MDP with as low a total entropy as possible. The typical framing of reinforcement learning is as a search for an optimal policy π ∗ .
Revised: A typical approach in a Bayesian setting might be to gather data such that the entropy H [ π ∗ ] of the belief of the optimal policy π ∗ is minimized.

Original: However,fully specifying such a  π ∗ would require giving a map for states which are never visited  in rollouts of π ∗ with start states sampled from p 0 .
Revised: However, a full speciﬁcation of π ∗ includes the behavior of the policy in all parts of the state space including states that are not visited at all, or visited less often in rollouts of π ∗ when the start state is sampled from p 0 .

Original: Thus, optimizing the EIG about π ∗in the TQRL setting would result in queries across areas of  the state space which wouldn’t ever be visited by a good policy well after they can be determined to be irrelevant.
Revised: As a result, not all points in the state space are equally important when learning an optimal policy aimed at maximizing the expected reward. Optimizing the entropy of π ∗ would lead to a uniform treatment of all the points in the state space, and hence would be far from optimal for the standard goal in RL.

Original: Rather than the policy π ∗ deﬁned over S , we consider as our variable ofinterest the optimal trajectory 
Revised: We instead propose to minimize the entropy of the optimal trajectory τ ∗ =

Original: O = { s i } Hi =1 given by a random vector of states generated by ﬁrst sampling s 0 ∼ p 0 then sampling a i = π ∗ ( s i ) , s i +1 ∼ T ( s i , a i ) for H timesteps.
Revised: { s i } Hi =1 deﬁned as a random vector of states generated by ﬁrst sampling s∼ p 0 then sampling a i = π ∗ ( s i ) , s i +1 ∼ T ( s i , a i ) for H timesteps.

Original: Clearly, O is determined up to the randomness of the MDP byspecifying the optimal policy π ∗ .
Revised: The optimal trajectory is completely speciﬁed by π ∗ and the randomness arising from the MDP.

Original: O also contains the  information needed about  T to solve the MDP, since any state that could ever be visited by π ∗ is in the support of O .
Revised: Furthermore, τ ∗ contains the necessary information needed about the transition function T to solve the MDP, since any state that could ever be visited by π ∗ is in the support of τ ∗ .

Original: Wedon’t claim O is the minimal-entropy choice of target random variable but it empirically seems toacquire the necessary data in an efﬁcient manner.
Revised: We empirically observe that this leads to an efﬁcient strategy for active RL.

Original: There are three sources of uncertainty about O : the randomness in  the start state distribution p 0 , therandomness in the dynamics T , which is commonly called aleatoric uncertainty , and the uncertainty in our estimate of the model T due to our limited experience which is usually discussed as epistemic uncertainty .
Revised: The randomness in τ ∗ arises from three sources: the start state distribution p 0 , the dynamics T constituting the aleatoric uncertainty , and the uncertainty in our estimate of the model T due to our limited experience which constitutes the epistemic uncertainty .

Original: As it is impossible to reduce the ﬁrst two components of uncertainty, any appropriate information gain quantity will reduce epistemic uncertainty about O as desired.
Revised: The ﬁrst two sources of uncertainty being aleatoric in nature cannot be reduced by experience. Our proposed acquisition function based on information gain naturally leads to reduction in the epistemic uncertainty about τ ∗ as desired.

Original: In this paper, we use H to denote the Shannon entropy of a random variable. Finally, we can give
Revised: Finally, our acquisition function for a given state-action pair ( s, a ) is given as

Original: [Equation]
Revised: [Equation]

Original: as our acquisition function.

Original: Here we assume a posterior model of the dynamics T(cid:48) ( s, a | D ) for a dataset D we havealready observed.
Revised: Here we assume a posterior model of the dynamics T ( s, a | D ) for a dataset D we have observed.

Original: The second equality is true because s 0 ⊥ s (cid:48) | s, a .
Revised: The second equality is true because s 0 ⊥ s (cid:48) | s, a .

Original: In our work, we assume the MPC policy for the ground truth model π T ≈ π ∗ , though in principle π ∗ could be approximated using any method.
Revised: In this paper, we assume the MPC policy using the ground truth transition function is approximately optimal, i.e. π T ≈ π ∗ , though in principle π ∗ could be approximated using any method.

Original: Of course, our method never actually has access to π T or π ∗ .
Revised: Of course, our method never actually has access to π T or π ∗ .

Original: 4. E STIMATING EIG O VIA P OSTERIOR F UNCTION S AMPLING
Revised: 4. E STIMATING EIG τ ∗ VIA P OSTERIOR F UNCTION S AMPLING

Original: For EIG O to be of practical beneﬁt in selecting data, we must be able to  approximate it in a tractable manner.
Revised: For EIG τ ∗ to be of practical beneﬁt, we must be able to tractably approximate it.

Original: Here we show how to attain such an approximation.
Revised: Here we show how to obtain such an approximation.

Original: By the symmetry of mutual information, we can rewrite Equation 4 as[Equation]
Revised: By the symmetry of MI, we can rewrite Equation (3) as

Original: [Equation]
Revised: [Equation]

Original: Since H [ s (cid:48) | s, a, D, s 0 ]
Revised: Since H [ s (cid:48) | s, a, D, s 0 ]

Original: = H [ s (cid:48) | s, a, D ] doesn’t depend on O or s 0 , we can simply compute it as the entropy of the posterior predictive distribution P ( s (cid:48) | s, a, D ) given by our posterior over 
Revised: = H [ s (cid:48) | s, a, D ] doesn’t depend on τ ∗ or s 0 , we can simply compute it as the entropy of the posterior predictive distribution P ( s (cid:48) | s, a, D ) given by our posterior over the transition function P ( T | D ) .

Original: Agent
Revised: In order to compute the other term, we must take samples τ ∗ ij ∼ P ( τ ∗ | D ) .

Original: Draw samples from posterior of Run RL on samples{+p 
 (we always set m = 1 in [Figure]+}
Revised: To do this, we ﬁrst sample m start states s i 0 from {+p 
 (we always set m = 1 in [Figure]+}

Original: MaximizeEIG
Revised: Maximize

Original: Environment Query Observe (a) −−  20 x −−−−     y (b) the transition function P ( T | D ) .
Revised: [Figure]

Original: [-In order to compute the other term, we must take samples O ij ∼ P ( O | D ) . To do this, we ﬁrst sample m start states s i 0 from p 
-]
Revised: [-In order to compute the other term, we must take samples O ij ∼ P ( O | D ) . To do this, we ﬁrst sample m start states s i 0 from p 
-]Query

Original: (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.
Revised: Observe (a) −−  20 x −−−−     y (b) 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.

Original: Then we can run the
Revised: We then run the

Original: MPC procedure on each of the posterior functions from s i 0 using T (cid:48) ij for T and π T (cid:48) ij for π ∗ (using our assumption that π ∗ ≈ π T ), giving our sampled O ij .
Revised: MPC procedure on each of the posterior functions from s i 0 using T (cid:48) ij for T and π T (cid:48) ij for π ∗ (using our assumption that π ∗ ≈ π T ), giving our sampled τ ∗ ij .

Original: This is an expression of the generative process for O as described in the previous section that accounts for the uncertainty in T .
Revised: This is an expression of the generative process for τ ∗ as described in the previous section that accounts for the uncertainty in T .

Original: Formally, we can approximate EIG O via Monte-Carlo as
Revised: Formally, we can approximate EIG τ ∗ via Monte-Carlo as

Original: [Equation]
Revised: [Equation]

Original: Finally, we must calculate the entropy H
Revised: Finally, we must calculate the entropy H

Original: [ s (cid:48) | s, a, O i , D ] .
Revised: [ s (cid:48) | s, a, τ ∗ i , D ] .

Original: For this, we follow a similar strategy asdone in Neiswanger et al.
Revised: For this, we follow a similar strategy as Neiswanger et al.

Original: In particular, since O i is a set of states output from the transition model, we can treat them as additional noiseless datapoints for our dynamics model and condition on them.
Revised: In particular, since τ ∗ i is a set of states output from the transition model, we can treat them as additional noiseless datapoints for our dynamics model and condition on them.

Original: In the following section we describe our instantiation of this EIG estimate, and how we can use it in reinforcement learning procedures.
Revised: In the following section we describe our instantiation of this EIG estimate, and how we can use it in reinforcement learning procedures.

Original: Though inspired by the work cited above, we modify the  acquisition function to factor p 0 as an irreducible source of uncertainty.
Revised: Though inspired by the work cited here, we modify the computation of the acquisition function to factor p 0 as an irreducible source of uncertainty.

Original: We also extend the function being queried to be vector-valued.
Revised: We also extend the function being queried to be vector-valued.

Original:  B AYESIAN A CTIVE R EINFORCEMENT L EARNING
Revised:  B AYESIAN A CTIVE R EINFORCEMENT L EARNING

Original: In this work, we take the simplest possible approach for nonlinear control in continuous spaces and assume a Gaussian process  prior  T (cid:48) to model the dynamics.
Revised: In this work, we take a simple approach for nonlinear control in continuous spaces and assume a Gaussian process (GP) prior P ( T ) to model the dynamics.

Original: Though computationally expensive, this choice ensures that we have a straightforward approximation for all necessary quantities.
Revised: Though computationally expensive, this choice ensures that we can easily approximate all necessary quantities.

Original: We assert however that the development of the acquisition function is general and anyapproximation of a Bayesian model could be used.
Revised: However, we note that the development of the acquisition function is general and any Bayesian model could be used in principle.

Original: In particular, we follow the approach of Wilson et al.

Original: (2020), based on sparse-GPs and random fourier approximations of kernels (Rahimi et al., 2007), allowing us to approximately but efﬁciently sample from the GP posterior conditioned on the observations.

Original: We model the transition function T : S×A → p ( S ) (the problem dynamics) using Gaussian processes due to their non-parametric nature andinherent ability to model uncertainties in the modeled function.
Revised: The transition function T : S × A → p ( S ) (dynamics) can be modeled with a GP due to its non- parametric nature and ability to capture uncertainties in T .

Original: The transition function takes a state action pair ( s, a ) ∈ R d + n as input, and produces a d -dimensional output denoting the next state.
Revised: The transition function takes a state action pair ( s, a ) ∈ R d + n as input, and produces a d -dimensional output denoting the next state.

Original: We model each of the d dimensions of the output using individual GPs.
Revised: We model each of the d dimensions of the output as independent GPs.

Original: More speciﬁcally we model the change in state ∆( s, a ) =
Revised: More speciﬁcally, we model the change in state ∆( s, a ) =

Original: T ( s, a ) − s rather thanmodelling the ﬁnal state T ( s, a ) directly.
Revised: T ( s, a ) − s rather than the ﬁnal state T ( s, a ) directly.

Original: This is helpful for continuous control problems since the state does not change by  a huge magnitude.
Revised: This is helpful for continuous control problems since the state often changes by only a small magnitude.

Original: Given observations D
Revised: Given observations D

Original: = { ( s i , a i , s (cid:48) i ) } , our approach requires a posterior sample of the transition function conditioned on D .
Revised: = { ( s i , a i , s (cid:48) i ) } , our approach requires a posterior sample of the transition function conditioned on D .

Original: We follow the approach of Wilson et al.
Revised: We follow the approach of Wilson et al.

Original: (2020) for fast approximate posterior sampling of GPs. Assuming access to a generative model and an initial dataset D (which we in practice use one randomly sampled datapoint ( s, a, s (cid:48) ) for), we compute EIGO for D by running MPC on posterior function samples and approximate argmax s ∈S ,a ∈A
Revised: (2020), based on sparse-GPs and random fourier approximations of kernels (Rahimi et al., 2007), allowing us to approximately but efﬁciently sample from the GP posterior conditioned on the observations. Assuming access to a generative model and an initial dataset D (for which, in practice, we use one randomly sampled datapoint ( s, a, s (cid:48) ) ), we compute EIG

Original: EIG O
Revised: τ

Original:  ( s, a ) by zeroth order approximation.
Revised: ∗ for D by running MPC on posterior function samples and approximate argmax s ∈S ,a ∈A EIG τ ∗ ( s, a ) by zeroth order approximation.

Original: Then we query s (cid:48) ∼ T ( s, a ) and add the subsequent triple to the dataset D and repeat the process.
Revised: Then we query s (cid:48) ∼ T ( s, a ) and add the subsequent triple to the dataset D and repeat the process.

Original: To evaluate, we simply perform the MPC procedure in Equation 2 and execute π E [ T(cid:48) | D ] on the real environment.
Revised: To evaluate, we simply perform the MPC procedure in Equation (2) and execute π E [ T | D ] on the real environment.

Original: We refer to this procedure as Bayesian Active Reinforcement Learning (BARL).
Revised: We refer to this procedure as Bayesian active reinforcement learning (BARL).

Original: Details are given in Algorithm 1 (here, U denotes the uniform distribution) and a schematic diagram in
Revised: Details are given in Algorithm 1 (here, U denotes the uniform distribution) and a schematic diagram in

Original: Figure 1a.
Revised: Figure 1a.

Original: We discuss details of training hyperparameters and the GP model in Appendix A.
Revised: We discuss details of training hyperparameters and the GP model in Appendix A.

Original: Inputs: budget b , # of points for optimization k , # of posterior function samples n .

Original: Initialize ( s 0 , a 0 ) ∼ U ( S × A[-) , s (cid:48) 
-]
Revised: Initialize ( x 0 , y 0 ) ∼ U ( S × A[-) , s (cid:48) 
-]

Original:  ∼ T ( s 0 , a 0 ) , and D ← { ( s 0 , a 0 , s (cid:48) 0 ) } .
Revised: , x (cid:48) 0 ∼ T ( x y 0 ) , and D ← { ( x 0 , y 0 , x (cid:48) 0 ) } .

Original: for i ∈ [ b ] do Sample posterior functions { T (cid:48) (cid:96) } n(cid:96)
Revised: for i ∈ [ b ] do Sample posterior functions { T (cid:48) (cid:96) } n(cid:96)

Original: =1 ∼ P ( T (cid:48) | D ) .
Revised: =1 ∼ P ( T (cid:48) | D ) .

Original: Sample start state s∼ p 0 .
Revised: Sample start state s∼ p 0 .

Original: Compute { O (cid:96) } n(cid:96)
Revised: Compute { executing starting argmax

Original: =1 by executing MPC policy π T (cid:48) (cid:96) on the dynamics of T (cid:48) (cid:96) starting from s 0 . ( x 1 , y 1 ) , . . . , ( x k , y k ) ∼ U ( S × A ) ( x ∗ , y ∗ ) ← argmax { ( x i ,y i ) } ki =1 EIG O
Revised: D

Original: ( x i , y i ) . D← D ∪ { ( x ∗ , y ∗ , x (cid:48)∗ ) } where x (cid:48)∗ ∼ T ( x ∗ , y ∗ ) .
Revised: D

Original: end for return π ˆ T where ˆ T ( s, a ) =
Revised: end for return π ˆ T where ˆ T ( s, a ) =

Original: E T ∼ P (T | D)
Revised: E T ∼ P ( | D

Original: [ T ( s, a )]
Revised: [ T ( s, a )]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: As discussed in the introduction, the aim of our study of acquisition functions for RL is to reduce the sample complexity of learning good policies in continuous spaces.
Revised: The aim of our study of acquisition functions for RL is to reduce the sample complexity of learning good policies in continuous spaces, under expensive dynamics.

Original: In this work, we demonstrate the effectiveness of  EIG O in allowing us to leverage transition queries by comparing the algorithm developed above to a variety of state-of-the-art RL algorithms.
Revised: Here, we demonstrate the effectiveness of using EIG τ ∗ to leverage transition queries by comparing against a variety of state-of-the-art RL algorithms.

Original: In particular, we compare the average return across ﬁve evaluation episodes across ﬁve runs with differing random seeds of each algorithm on ﬁve continuous control problems as data is collected.
Revised: In particular, we compare the average return across ﬁve evaluation episodes across ﬁve runs with differing random seeds of each algorithm on ﬁve continuous control problems as data is collected.

Original: We also assess the amount of data taken by each algorithm to ‘solve’ the problem, which is taken to mean performing as well as our MPC procedure using the ground truth dynamics.
Revised: We also assess the amount of data taken by each algorithm to ‘solve’ the problem, which is taken to mean performing as well as our MPC procedure using the ground truth dynamics.

Original: Our method greatly outperforms other methods across the board.
Revised: Our proposed method, BARL, greatly outperforms other methods across the board.

Original: In particular we use 5 – 1 , 000 × less data to solve problems than state-of-the-art model-based  algorithms and 10 3 – 10 5 × less  than model-free ones.
Revised: In particular, BARL uses 5 – 1 , 000 × less data to solve problems than state-of-the-art model-based RL algorithms and 10 3 – 10 5 × less data than model-free RL algorithms.

Original: Though in this section we  focus on the performance of the controller,  we  discuss the runtime of the algorithm in section A.1.
Revised: In this section we primarily focus on the performance of the controller, and in section A.1 we also discuss the runtime of the algorithm.

Original: Comparison Methods We use as our model-based comparison methods in this work PETS (Chua et al., 2018) as implemented by Pineda et al.
Revised: Comparison Methods. We use as our model-based comparison methods in this work PETS (Chua et al., 2018) as implemented by Pineda et al.

Original: (2021), which does MPC using a probabilistic ensemble of neural networks and particle sampling for stochastic dynamics and a similar MPC method using the mean of the same GP model we use for BARL to execute π ˆ T  as in the standard RL setting.
Revised: (2021), which does MPC using a probabilistic ensemble of neural networks and particle sampling for stochastic dynamics and a similar MPC method using the mean of the same GP model we use for BARL to execute π ˆ T to collect data as in the standard RL setting.

Original: We also compare against PILCO (Deisenroth & Rasmussen, 2011), which also leverages a GP to directly optimize a policy that maximizes an uncertainty-aware long term reward.
Revised: We also compare against PILCO (Deisenroth & Rasmussen, 2011), which also leverages a GP to directly optimize a policy that maximizes an uncertainty-aware long term reward.

Original: For model-free methods, we use Soft Actor-Critic (SAC) (Haarnoja et al., 2018), which is an actor-critic method that uses an entropy bonus for the policy to encourage evaluation, TD3 (Fujimoto et al., 2018) which addresses the stability questions of actor-critic methods by including twin networks for value and several other modiﬁcations, and Proximal Policy Optimization (PPO) (Schulman et al., 2017), which addresses stability by forcing the policy to change slowly in KL so that the critic remains accurate.
Revised: For model-free methods, we use Soft Actor-Critic (SAC) (Haarnoja et al., 2018), which is an actor-critic method that uses an entropy bonus for the policy to encourage evaluation, TD3 (Fujimoto et al., 2018) which addresses the stability questions of actor-critic methods by including twin networks for value and several other modiﬁcations, and Proximal Policy Optimization (PPO) (Schulman et al., 2017), which addresses stability by forcing the policy to change slowly in KL so that the critic remains accurate.

Original: As a baseline TQRL method and to better understand the GP performance, we use a method we denote EIG T , which chooses points that maximize the predictive entropy of the transition model to collect data.
Revised: As a baseline TQRL method and to better understand the GP performance, we use a method we denote EIG T , which chooses points which maximize the predictive entropy of the transition model to collect data.

Original: We believe that when given access to transition queries many unsupervised exploration methods like Pathak et al.
Revised: We believe that when given access to transition queries many unsupervised exploration methods like Pathak et al.

Original: ; Shyam et al.
Revised: ; Shyam et al.

Original: (2019) or methods which value information gain over the transtion function (Nikolov et al., 2019) would default to this behavior.
Revised: (2019) or methods which value information gain over the transtion function (Nikolov et al., 2019) would default to this behavior.

Original: Control Problems.
Revised: Control Problems.

Original: We tackle ﬁve control problems: the standard underactuated pendulum swingup problem (Pendulum-v0 from Brockman et al.
Revised: We tackle ﬁve control problems: the standard underactuated pendulum swing-up problem (Pendulum-v0 from Brockman et al.

Original: (2016)), a cartpole swing-up problem, a 2D lava path navigation problem, a 2-DOF robot arm reacher problem with 8-dimensional state (Reacher-v from Brockman et al.
Revised: (2016)), a cartpole swing-up problem, a 2D lava path navigation problem, a 2-DOF robot arm reacher problem with 8-dimensional state (Reacher-v2 from Brockman et al.

Original: (2016)), and a simpliﬁed beta tracking problem from plasma control  where the controller must maintain a ﬁxed normalized plasma pressure using as GT dynamics a model learned similarly to Abbate et al.
Revised: (2016)), and a simpliﬁed beta tracking problem from plasma control (Char et al., 2019; Mehta et al., 2020) where the controller must maintain a ﬁxed normalized plasma pressure using as GT dynamics a model learned similarly to Abbate et al.

Original: Though this is very much a simpliﬁcation of the problem, we intend to extend and apply BARL to more realistic settings in the immediate future.
Revised: The lava path is intended to test stability and exploration of algorithms. The goal is to reach a ﬁxed goal state from a narrow uniform distribution over start states.

Original: The lava
Revised: As shown in Figure 1b, the state space contains a ‘lava’ region which gives large negative rewards for every timestep. When not in lava, the reward is simply the negative squared distance to the goal, forcing the agent to navigate to the goal as quickly as possible.

Original: [Figure]
Revised: [Figure]

Original: ˙θ
Revised: ˙

Original: [Figure]
Revised: [Figure]

Original: Environment
Revised: Environment

Original: BARL
Revised: BARL

Original: MPC
Revised: MPC

Original: EIG T
Revised: EIG T

Original: PETS
Revised: PETS

Original: SAC
Revised: SAC

Original: TD
Revised: TD

Original: PPO
Revised: PPO

Original: Lava Path n/a n/a n/a
Revised: Lava Path N/A N/A N/A

Original: Pendulum5, 6, 57, 13,
Revised: Pendulum

Original: Cartpole 1, 31, 18, n/a
Revised: Cartpole N/A

Original: Beta Tracking n/a 9, 6, 16,
Revised: Beta Tracking N/A

Original: Reacher n/a 23, 13, n/a
Revised: Reacher N/A

Original: ‘Solved’ is determined by operating an MPC policy similar to the one used for evaluation in many methods using the ground truth dynamics to predict actions.
Revised: N/A used for evaluation) on the ground truth dynamics to predict actions.

Original: ‘n/a’ is recorded when the median run didn’t solve the problem by the end ofour training.
Revised: We record ‘N/A’ when the median run is unable to solve the problem by the end of training.

Original: path is intended to test stability and exploration of algorithms.

Original: The goal is to reach a ﬁxed goal state from a relatively narrow uniform distribution over start states.

Original: As shown in Figure 1b, the state space contains a ‘lava’ region which gives very large negative rewards for every timestep.

Original: Other than when in lava, the reward is simply the negative squared distance to the goal, forcing the agent to navigate to the goal as quickly as possible.

Original: Since the lava has a narrow path through, the actor is forced to explore a policy which will realize that it is safe and efﬁcient to cross.
Revised: Since there is a narrow path through the lava, we want to explore a policy which crosses efﬁciently and safely.

Original: Agents who fail to ﬁnd this solution will be forced to go around, incurring penalties.
Revised: Agents who fail to ﬁnd this solution will be forced to go around, incurring penalties.

Original: We see in both the sample complexity ﬁgures in Table 1, the learning curves in Figure 3, and visually in Figure 2, that BARL leverages EIG O to greatly reduce the data requirements of learning controllers on the problems presented.
Revised: We see in both the sample complexity ﬁgures in Table 1, the learning curves in Figure 3, and visually in Figure 2 that BARL leverages EIG τ ∗ to signiﬁcantly reduce the data requirements of learning controllers on the problems presented.

Original: We’d like to additionally point out several failure cases of related algorithms that BARL avoids.
Revised: We’d like to additionally point out several failure cases of related algorithms that BARL avoids.

Original: Though it performs well on the simplest environments (pendulum and cartpole), EIG T suffers from an inability to focus on acquiring data relevant to the control problem and not just learning dynamics as the state space becomes higher-dimensional in the reacher problem, or less smooth as in the beta tracking problem.
Revised: Though it performs well on the simplest environments (pendulum and cartpole), EIG T suffers from an inability to focus on acquiring data relevant to the control problem and not just learning dynamics as the state space becomes higherdimensional in the reacher problem, or less smooth as in the beta tracking problem.

Original: The MPC method performs  well across the board and is competitive with BARL on the plasma problem but suffers from many morereqired samples inthe smaller environments where the model uncertainty can point to meaningfully underexplored areas.
Revised: The MPC method performs reasonably well across the board and is competitive with BARL on the plasma problem but requires relatively more samples in smaller environments where the model uncertainty can point to meaningfully underexplored areas.

Original: PETS is strong across the board but suffers from more required samples due to both its neural network dynamics model and its inability to make transition queries.
Revised: PETS is strong across the board but suffers from more required samples due to both its neural network dynamics model and its inability to make transition queries.

Original: All algorithms besides BARL suffer substantial instability on the lava path problem, which is designed to be challenging to explore in a sequential fashion and require a precise understanding of which areas are safe to enter.
Revised: All algorithms besides BARL suffer substantial instability on the lava path problem, which is designed to be challenging to explore in a sequential fashion and require a precise understanding of which areas are safe to enter.

Original: BARL manages to learn where it is safe to operate in a handful of queries, which is an exciting result and will bear further investigation.
Revised: BARL manages to learn where it is safe to operate in a handful of queries, which is an exciting result and will bear further investigation.

Original: Figure 1b gives some intuition as to why: points are initally queried close to the start and as those dynamics are understood they are subsequently queried farther and farther along the execution paths.
Revised: Figure 1b gives some intuition as to why: points are initally queried close to the start and as those dynamics are understood they are subsequently queried farther and farther along the execution paths.

Original: This allows BARL to use transition queries to avoid traversing well-understood areas of state space to reach the areas which are worth learning.
Revised: This allows BARL to use transition queries to avoid traversing well-understood areas of state space to reach the areas which are worth learning.

Original: We see a speedup in sample complexity reminiscent of a move from quadratic to linear, which mirrors some of the theoretical improvements given in the prior work discussed on tabular methods.
Revised: We see a speedup in sample complexity reminiscent of a move from quadratic to linear, which mirrors some of the theoretical improvements given in the prior work discussed on tabular methods.

Original: We further support our assertion that BARL is picking ‘meaningful’ points to the control problem by the evidence in Figure 4.
Revised: We further support our assertion that BARL is picking ‘meaningful’ points to the control problem by the evidence in Figure 4.

Original: Here, BARL is able to solve the reacher problem while EIG T is not.
Revised: Here, BARL is able to solve the reacher problem while EIG T is not.

Original: [Figure]

Original: [Figure] [Figure] uniformly the However, it has much worse model predictions onthe random datapoints than EIG T while doing a much better job modelling the datapoints used by the MPC procedure.
Revised: However, BARL has much worse model predictions on random data than EIG T while doing a much better job modeling data used by the MPC procedure.

Original: Clearly, the EIG O acquisition function captures in some way which data would be valuable to acquire to not just learn about the transition function but actually solve the control problem.
Revised: Clearly, the EIG τ ∗ acquisition function captures in some way which data would be valuable to acquire to not just learn about the transition function but actually solve the control problem.

Original: We see this pattern across other tasks as well. 
Revised: We see this pattern across other tasks as well. In [Figure] [Figure]

Revised: 2011) on Pendulum.

Revised: PILCO makes assumptions about the initial state distribution and suffers from numerical instability under long control horizon so we were unable to reach representative performance on the other problems.

Original: In section B.1, we study whether the acquisition function we see here is able to work with a suboptimal controller on posterior samples of the dynamics.
Revised: [Figure] section B.1, we study whether the acquisition function we see here is able to work with a suboptimal controller on posterior samples of the dynamics.

Original: Our experiments show that EIG O seems to work well even when the policy used to generate O is subpoptimal.
Revised: Our experiments show that EIG τ ∗ seems to work well even when the policy used to generate τ ∗ is suboptimal.

Original:  D ISCUSSION AND F UTURE W ORK
Revised:  D ISCUSSION AND F UTURE W ORK

Original: In this work, we proposed an acquisition function for reinforcement learning, and were able to apply to the setting of TQRL leading to a novel algorithm for addressing the problem of efﬁcient data collection in RL.
Revised: In this work, we proposed an acquisition function for reinforcement learning, and applied it to the setting of TQRL, leading to a novel algorithm for addressing the problem of efﬁcient data collection in RL.

Original: We experimented with several control problems and demonstrated that this approach leads to substantial improvements with respect to the sample efﬁciency.
Revised: We experimented with several control problems and demonstrated that this approach leads to substantial improvements with respect to the sample efﬁciency.

Original: However, there are some drawbacks to this method as well.
Revised: However, there are some drawbacks to this method as well.

Original: Computing our proposed acquisition function relies on executing the entire control algorithm over a set of posterior function samples leading to higher computational requirements.
Revised: Computing the proposed acquisition function relies on executing the entire control algorithm over a set of posterior function samples, leading to high computational requirements.

Original: In our current implementation we use GPs to model the transition function, which are computationally expensive and do not scale well to higher dimensions and large datasets.
Revised: In our current implementation we use Gaussian processes to model the transition function, which are computationally expensive and do not scale well to higher dimensions and large datasets.

Original: In the future, we plan to extend this idea to use other types of Bayesian models such as Bayesian neural networks and take advantage of GPU compute for better scalability.
Revised: In the future, we plan to extend this idea to use other types of Bayesian models such as Bayesian neural networks and take advantage of GPU compute for better scalability.

Original: A T RAINING D ETAILS
Revised: A CKNOWLEDGMENTS

Revised: We would like to acknowledge anonymous reviewers for valuable feedback.

Revised: VM acknowledges

Revised: Swapnil Pande for providing the idea for and implementation of the Lava Path environment and Ian

Revised: Char for providing the trained model for the Beta Tracking environment.

Revised: This work was funded in part by DOE grant number DE-SC0021414.

Revised: WN acknowledges the helpful feedback from members of the

Revised: Ermon Group studying RL.

Revised: WN was supported in part by NSF (#1651565), ONR (N000141912145),

Revised: AFOSR (FA95501910024), ARO (W911NF-21-1-0125), DOE (DE-AC02-76SF00515) and Sloan

Revised: Fellowship.

Revised: A T RAINING D ETAILS

Original: We vary our budget based on our understanding of how much data would be required to solve the problem.
Revised: We vary our budget based on our understanding of how much data would be required to solve the problem.

Original: The other hyperparameters of the BARL algorithm are constant but listed for completeness in Table2.
Revised: The other hyperparameters of the BARL algorithm are constant but listed for completeness in Table

Original: For all of our experiments

Original: [Table]
Revised: [Table]

Original:  we use a squared exponential kernel with automatic relevance determination (MacKay et al., 1994; Neal, 1995).
Revised: For all of our experiments, we use a squared exponential kernel with automatic relevance determination (MacKay et al., 1994; Neal, 1995).

Original: The parameters of the kernel were estimated by maximizing the likelihood of the parameters after marginalizing over the posterior GP (Williams & Rasmussen, 1996).
Revised: The parameters of the kernel were estimated by maximizing the likelihood of the parameters after marginalizing over the posterior GP (Williams & Rasmussen, 1996).

Original: To optimize the transition function, we simply sampled a set of points from the domain, evaluated the acquisition function
Revised: 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.

Original: [Table] and chose the maximum of the set. This set was chosed uniformly for every problem but Reacher, for which we chose a random subset of ∪ i ∪ j O ij  since the space of samples is 10-dimensional and uniform random sampling will not get good coverage of interesting regions of the state space.
Revised: This set was chosed uniformly for every problem but Reacher, for which we chose a random subset of ∪ i ∪ j τ ∗ ij (the posterior 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.

Original: A. R UNTIME D ETAILS
Revised: A. R UNTIME D ETAILS

Original: Based on these choices and the MPC hyperparameters below in Section B, each of these problems results in a varying runtime for the BARL algorithm.
Revised: Based on these choices and the MPC hyperparameters below in Section B, each of these problems results in a varying runtime for the BARL algorithm.

Original: In Table 3, we report the time taken for an iteration of BARL and how it breaks down by step.
Revised: In Table 3, we report the time taken for an iteration of BARL and how it breaks down by step.

Original: We give these as ranges, as the computational time requires increases as the learning process continues since GP computational costs scale with the size of the dataset.
Revised: We give these as ranges, as the computational time requires increases as the learning process continues since GP computational costs scale with the size of the dataset.

Original: We also include for completeness the time taken to execute the MPC policy on the ground truth problem, which is not strictly part of the BARL algorithm
Revised: We also include for completeness the time taken to execute the MPC policy on the ground truth problem, which is not strictly part of the BARL algorithm

Original: but still relevant to practitioners.
Revised: but still relevant to practitioners.

Original: Clearly, BARL is a relatively slow algorithm computationally.
Revised: Clearly, BARL is a relatively slow algorithm computationally.

Original: But in settings where samples are scarce, BARL is much cheaper than alternative methods which might use less compute for the RL algorthms but require many more samples.
Revised: But in settings where samples are scarce, BARL is much cheaper than alternative methods which might use less compute for the RL algorthms but require many more samples.

Original: When compared to the costs of running an hour-long simulation or running a costly experiments, spending a few minutes computing the acquisition function seems like a good use of resources.
Revised: When compared to the costs of running an hour-long simulation or running a costly experiments, spending a few minutes computing the acquisition function seems like a good use of resources.

Original: B MPC D ETAILS
Revised: B MPC D ETAILS

Original: As we’ve discussed, we use model-predictive control in this work to choose actions which maximize future reward given a model of the dynamics.
Revised: As we’ve discussed, we use model-predictive control in this work to choose actions which maximize future reward given a model of the dynamics.

Original: In particular we use the improved Cross-Entropy
Revised: In particular we use the improved Cross-Entropy

Original: Method from Pinneri et al.
Revised: Method from Pinneri et al.

Original: (2020) to solve the optimization problem in Equation 2, which uses several 
Revised: (2020) to solve the optimization problem in Equation 2, which uses several tricks including colored noise samples and caching to reduce the number of queries to the planning

Original: [Table]
Revised: [Table]

Original: tricks including colored noise samples and caching to reduce the number of queries to the planning model.
Revised: model.

Original: There is a natural trade-off in any search method between computational cost and quality of actions found in terms of predicted reward.
Revised: There is a natural trade-off in any search method between computational cost and quality of actions found in terms of predicted reward.

Original: In this work, we chose hyperparameters for each task that were as computationally light as possible which attained a similar reward to larger hyperparameters when executing MPC using the ground-truth model ( π T , in our terms).
Revised: In this work, we chose hyperparameters for each task that were as computationally light as possible which attained a similar reward to larger hyperparameters when executing MPC using the ground-truth model ( π T , in our terms).

Original: 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 .
Revised: 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 .

Original: 3 for the amount of caching.
Revised: 3 for the amount of caching.

Original: We manually tuned the base number of samples, planning horizon, number of elites to take from the sampled action sequences, number of iterations of planning, and the replanning period of the model.
Revised: We manually tuned the base number of samples, planning horizon, number of elites to take from the sampled action sequences, number of iterations of planning, and the replanning period of the model.

Original: Here we give the ultimate values for those parameters, which were used for all ablations using our GP model and MPC.
Revised: Here we give the ultimate values for those parameters, which were used for all ablations using our GP model and MPC.

Original: The values we used across all experiments for each problem are given in Table 4.
Revised: The values we used across all experiments for each problem are given in Table 4.

Original: B. R OBUSTNESS OF EIG O TO A SUBOPTIMAL CONTROLLER
Revised: B. R OBUSTNESS OF EIG τ ∗ TO A SUBOPTIMAL CONTROLLER

Original: In order to compute EIG O in this work, we perform model-predictive control on posterior transition function samples (execute π T (cid:48) (cid:96) on T (cid:48) (cid:96) , in our notation).
Revised: In order to compute EIG τ ∗ in this work, we perform model-predictive control on posterior transition function samples (execute π T (cid:48) (cid:96) on T (cid:48) (cid:96) , in our notation).

Original: We assume that π T (cid:48) (cid:96) is close to the optimal policy for the MDP with transition function T (cid:48) (cid:96) .
Revised: We assume that π T (cid:48) (cid:96) is close to the optimal policy for the MDP with transition function T (cid:48) (cid:96) .

Original: However, this assumption could lead to pathologies in the method if it doesn’t hold in practice.
Revised: However, this assumption could lead to pathologies in the method if it doesn’t hold in practice.

Original: In this section, we emprically investigate the consequences of using a suboptimal controller when ﬁnding samples of O on posterior samples of the transition function.
Revised: In this section, we emprically investigate the consequences of using a suboptimal controller when ﬁnding samples of τ ∗ on posterior samples of the transition function.

Original: In order to understand the sensitivity of EIG O to the MPC policy executed on posterior samples, we ran experiments where we reduced the planning budget or horizon for the posterior function policy in order to see whether the acquisition function fails.
Revised: In order to understand the sensitivity of EIG τ ∗ to the MPC policy executed on posterior samples, we ran experiments where we reduced the planning budget or horizon for the posterior function policy in order to see whether the acquisition function fails.

Original: In particular, we ran the reacher and cartpole experiments from the main paper with varying MPC budgets for the posterior function policy π T [Table] and a ﬁxed MPC budget at test time.
Revised: In particular, we ran the reacher and cartpole experiments from the main paper with varying MPC budgets for the posterior function policy π T (cid:48) (cid:96) and a ﬁxed MPC budget at test time.

Original: This allows us to isolate the effect of a subpoptimal policy generating samples of O .
Revised: This allows us to isolate the effect of a suboptimal policy generating samples of τ ∗ .

Original: On the reacher experiment, we vary the number of CEM iterations ranging from 1 to 5.
Revised: On the reacher experiment, we vary the number of CEM iterations ranging from 1 to 5.

Original: This is straightforwardly linked to the amount of search the policy conducts before executing an action.
Revised: This is straightforwardly linked to the amount of search the policy conducts before executing an action.

Original: On the cartpole experiment, we varied the length of the planning horizon, which affects how far in the future the policy will consider actions as it is deciding what is the good immediate next action.
Revised: On the cartpole experiment, we varied the length of the planning horizon, which affects how far in the future the policy will consider actions as it is deciding what is the good immediate next action.

Original: In both cases, we see in Figure 5 that performance hardly changes as the budget for MPC is reduced.
Revised: In both cases, we see in Figure 5 that performance hardly changes as the budget for MPC is reduced.

Original: Only on the reacher problem when the number of CEM iterations is reduced to 1 (effectively reducing
Revised: Only on the reacher problem when the number of CEM iterations is reduced to 1 (effectively reducing

Original: CEM to simple random search) do we see a signiﬁcant drop in performance.
Revised: CEM to simple random search) do we see a signiﬁcant drop in performance.

Original: This supports the notion that the quality of the approximation of the optimal policy in EIG O is not critical to the performance of the acquisition function as a data selection strategy.
Revised: This supports the notion that the quality of the approximation of the optimal policy in EIG O is not critical to the performance of the acquisition function as a data selection strategy.

Original: We also plot in the ﬁgure the performance of an MPC controller on the ground truth dynamics with these reduced MPC budgets.
Revised: We also plot in the ﬁgure the performance of an MPC controller on the ground truth dynamics with these reduced MPC budgets.

Original: It is clear that if we were to execute these degraded policies at test time, they would be much worse.
Revised: It is clear that if we were to execute these degraded policies at test time, they would be much worse.

Original: We ﬁnd it interesting and intend to further study the robustness of using cheaper policies to decide where to acquire data.
Revised: We ﬁnd it interesting and intend to further study the robustness of using cheaper policies to decide where to acquire data.

Original: C D ESCRIPTION OF C ONTINUOUS C ONTROL P ROBLEMS
Revised: C D ESCRIPTION OF C ONTINUOUS C ONTROL P ROBLEMS

Original: Lava Path.
Revised: Lava Path.

Original: The lava path has 4-dimensional state (position and velocity) and 2-dimensional action (an applied force in the plane).
Revised: The lava path has 4-dimensional state (position and velocity) and 2-dimensional action (an applied force in the plane).

Original: The goal is to reach a ﬁxed goal state from a relatively narrow uniform 
Revised: The goal is to reach a ﬁxed goal state from a relatively narrow uniform distribution over start states.

Revised: As shown in Figure 1b, the state space contains a ‘lava’ region which

Original: [Figure]
Revised: [Figure]

Original: distribution over start states. As shown in Figure 1b, the state space contains a ‘lava’ region which gives very large negative rewards for every timestep.
Revised: gives very large negative rewards for every timestep.

Original: Other than when in lava, the reward is simply the negative squared distance to the goal, forcing the agent to navigate to the goal as quickly as possible.
Revised: Other than when in lava, the reward is simply the negative squared distance to the goal, forcing the agent to navigate to the goal as quickly as possible.

Original: Since the lava has a narrow path through, the actor is forced to explore a policy which will realize that it is safe and efﬁcient to cross.
Revised: Since the lava has a narrow path through, the actor is forced to explore a policy which will realize that it is safe and efﬁcient to cross.

Original: Agents who fail to ﬁnd this solution will be forced to go around, incurring penalties.
Revised: Agents who fail to ﬁnd this solution will be forced to go around, incurring penalties.

Original: Pendulum.
Revised: Pendulum.

Original: The pendulum swing-up problem is the standard one found in the OpenAI gym (Brockman et al., 2016).
Revised: The pendulum swing-up problem is the standard one found in the OpenAI gym (Brockman et al., 2016).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The reward function here penalizes deviation from an upright pole and squared torque.
Revised: The reward function here penalizes deviation from an upright pole and squared torque.

Original: Cartpole.
Revised: Cartpole.

Original: 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).
Revised: 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).

Original: Here, the difﬁculty lies in translating the horizontal motion of the cart into effective torque on the pole.
Revised: Here, the difﬁculty lies in translating the horizontal motion of the cart into effective torque on the pole.

Original: The reward function here is a negative sigmoid function penalizing the distance betweent the tip of the pole and a centered upright goal position.
Revised: The reward function here is a negative sigmoid function penalizing the distance betweent the tip of the pole and a centered upright goal position.

Original: Reacher.
Revised: Reacher.

Original: The reacher problem simulates a 2-DOF robot arm aiming to move the end effector to a randomly resampled target provided.
Revised: The reacher problem simulates a 2-DOF robot arm aiming to move the end effector to a randomly resampled target provided.

Original: The problem requires joint angles and velocities as well as an indication of the direction of the goal, giving an 8-dimensional state space with the mentioned 2-D control.
Revised: The problem requires joint angles and velocities as well as an indication of the direction of the goal, giving an 8-dimensional state space with the mentioned 2-D control.

Original: Our results on this problem are particularly encouraging as they show that BARL can scale to some problems with higher dimensionalities.
Revised: Our results on this problem are particularly encouraging as they show that BARL can scale to some problems with higher dimensionalities.

Original: Beta Tracking (Nuclear Fusion).
Revised: Beta Tracking (Nuclear Fusion).

Original: Finally, the beta tracking problem has 4-dimensional state consisting of the current normalized plasma performance β N in the DIII-D tokamak.
Revised: Finally, the beta tracking problem has 4-dimensional state consisting of the current normalized plasma performance β N in the DIII-D tokamak.

Original: β N is given by an appropriately normalized ratio between the plasma pressure and the magnetic pressure and is a common ﬁgure of merit in fusion energy research.
Revised: β N is given by an appropriately normalized ratio between the plasma pressure and the magnetic pressure and is a common ﬁgure of merit in fusion energy research.

Original: In addition to β n the state space contains its most recent change as well as the current power injection level and its most recent change.
Revised: In addition to β n the state space contains its most recent change as well as the current power injection level and its most recent change.

Original: The action is the next change in the power injection level.
Revised: The action is the next change in the power injection level.

Original: The “ground-truth” dynamics for this problem are given by a neural network model learned from data processed as in Abbate et al.
Revised: The “ground-truth” dynamics for this problem are given by a neural network model learned from data processed as in Abbate et al.

Original: Control is done at a timestep of 200ms and the reward function is the negative absolute deviation from β n = 2 .
Revised: Control is done at a timestep of 200ms and the reward function is the negative absolute deviation from β n = 2 .

Original: Reliably controlling plasmas to sustain high performance is a major goal of research efforts for fusion energy, and though this is very much a simpliﬁcation of the problem, we intend to extend and apply BARL to more realistic settings in the immediate future.
Revised: Reliably controlling plasmas to sustain high performance is a major goal of research efforts for fusion energy, and though this is very much a simpliﬁcation of the problem, we intend to extend and apply BARL to more realistic settings in the immediate future.