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
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: 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) 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 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).

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 onon 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 a function for which only a black-box model is available and 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: 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.

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 data to collect in continuous spaces as otherwise queries would be wasted learning the optimal policy in irrelevant parts of the state space.
Revised: This is a critical distinction when evaluating data to collect in continuous spaces as otherwise queries would be wasted learning the 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
Revised: 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

Original: MDP literature.
Revised: 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 the full Bayes-Adaptive MDP and unlike the VPI heuristic, this function takes into account the current estimates of which states the optimal policy will actually visit and values potential queries accordingly.
Revised: However, like the full Bayes-Adaptive MDP and unlike the VPI heuristic, this function takes into account the current estimates of which states the optimal policy will actually visit and values potential queries accordingly.

Original: It is also scalable enough to apply to multidimensional continuous control problems.
Revised: It is also 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 stateaction 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 function is able to select relevant datapoints for control purposes leading to much-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 widely varying 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 transitions in 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 sampling transitions in 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 .

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

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 TQRL setting, 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, though every standard RL algorithm can be applied in this setting by continuing to collect data in an online manner.

Original: In this work, we give an algorithm speciﬁcally designed for this setting in continuous MDPs.
Revised: In this work, we give an algorithm speciﬁcally designed for this setting in continuous MDPs.

Original: Our work empirically shows sample complexity beneﬁts reminiscent of those theoretically shown in the tabular setting.
Revised: Our work empirically 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 (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, and an entropy bonus all try to explore 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: Though UCB, Thompson Sampling, and an entropy bonus all try to explore 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.

Original: As we mentioned, an ideal method of exploration would be to solve the intractable bayes-adaptive
Revised: As we mentioned, an ideal method of exploration would be to solve the intractable Bayes-adaptive

Original: MDP (Ross et al., 2007), giving an optimal tradeoff between exploration and exploitation.
Revised: MDP (Ross et al., 2007), giving an optimal tradeoff between exploration and exploitation.

Original: Kolter &
Revised: Kolter &

Original: Ng (2009); Guez et al.
Revised: 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 PACMDP 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: 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.

Original: These methods were further extended with the development of  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 (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.

Original: A tutorial on Bayesian methods can be found in Ghavamzadeh et al.
Revised: A tutorial on Bayesian 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 , 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).

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), or a density estimate (Bellemare et al., 2016).

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 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.

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.

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

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.

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,
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,

Original: A is the action space, T is the transition function mapping T :
Revised: A is the action space, T is the transition function mapping 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 a reward function, p 0 ( s ) a distribution over S of start states, and H ∈ N 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 objective given below.
Revised: Our aim is to ﬁnd a policy π : S → P ( A ) that maximizes the rewards obtained by executing the policy.

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  ∼ T ( s i− i , a i− 1 ) for i ∈ 1 , . . .
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 +1 ∼ T ( s i , a i ) for i ∈ 0 , . . .

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

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

Original: [Equation]
Revised: [Equation]

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 π ∗ =

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 ( π ) as an optimal policy, 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

Original: In the online RL setting, one assumes data D
Revised: In the 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 thethe next state s (cid:48)
Revised: [ 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)

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 a new initial state s H +1 ∼ p 0 is sampled.

Original: We pay special attention to the TQRL setting, where the agent sequentially acquires datapoints x i
Revised: We pay special attention to the TQRL setting, where the agent sequentially acquires datapoints x i

Original: =
Revised: =

Original: ( 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: ( 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)

Original: i ∼ T ( s, a ) (Kearns et al., 2002; Kakade, 2003; Azar et al., 2013).
Revised: i ∼ T ( s, a ) (Kearns et al., 2002; Kakade, 2003; Azar et al., 2013).

Original: Performance in this setting is still evaluated by
Revised: Performance in this setting is still evaluated by

Original: Objective 1.
Revised: Objective 1.

Original: It has been shown for ﬁnite MDPs in Azar et al.
Revised: It has been shown for ﬁnite MDPs in Azar et al.

Original: 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: 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 a probabilistic dynamics model ˆ T =

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: 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

Original: [Equation]
Revised: [Equation]

Original: with a stochastic optimization technique.
Revised: 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) is most often used to solve the optimization problem.
Revised: The cross-entropy method (CEM) is most often used to solve the 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.
Revised: 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: In this paper 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 in improving the performance of the policy on the MDP conditioned on all previously observed data.
Revised: A → R that describes the marginal contribution of collecting a transition sample at a state-action pair in improving the performance of the policy on the MDP conditioned on all previously observed data.

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 budget in the sense that solving a Bayes-adaptive MDP is optimal.

Original: However, they are tractable and allow for good one-step choices to be made that account for all information currently available.
Revised: However, they are tractable and allow for good one-step choices to be made that account for all information currently available.

Original: In this work we focus on acquisition functions written as an Expected Information Gain (EIG).
Revised: In this work we focus on acquisition functions written as an Expected Information Gain (EIG).

Original: These functions output the mutual information between a query of our transition model and a target random variable.
Revised: These functions output the mutual information between a query of our transition model and a target random variable.

Original: The target  the minimal information required for an optimal solution of our MDP with as low a total entropy as possible.
Revised: The target variable should contain the minimal information required for an optimal solution of our MDP with as low a total entropy as possible.

Original: The typical framing of reinforcement learning is as a search for an optimal policy π ∗ .
Revised: The typical framing of reinforcement learning is as a search for an optimal policy π ∗ .

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, 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 .

Original: So, optimizing the expected information gain 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: 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.

Original: Rather than the policy π ∗ deﬁned over S , we consider as our variable of interest the optimal trajectory
Revised: Rather than the policy π ∗ deﬁned over S , we consider as our variable of interest 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: 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.

Original: Clearly, O is determined up to the randomness of the MDP by specifying the optimal policy π ∗ .
Revised: Clearly, O is determined up to the randomness of the MDP by specifying the optimal policy π ∗ .

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

Original: We don’t claim
Revised: We don’t claim

Original: O is the minimal-entropy choice of target random variable but it empirically seems to acquire the necessary data in an efﬁcient manner.
Revised: O is the minimal-entropy choice of target random variable but it empirically seems to acquire the necessary data in an efﬁcient manner.

Original: There are three sources of uncertainty about O : the randomness in the start state distribution p 0 , the randomness 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: There are three sources of uncertainty about O : the randomness in the start state distribution p 0 , the randomness 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 .

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: 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: In this paper, we use H to denote the Shannon entropy of a random variable.

Original: Finally, we can give
Revised: Finally, we can give

Original: [Equation]
Revised: [Equation]

Original: as our acquisition function.
Revised: as our acquisition function.

Original: Here we assume a posterior model of the dynamics T (cid:48)
Revised: Here we assume a posterior model of the dynamics T (cid:48)

Original: ( s, a | D ) for a dataset D we have already observed.
Revised: ( s, a | D ) for a dataset D we have already 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 our work, we assume the MPC policy for the ground truth model π 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 O 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 O to be of practical beneﬁt in selecting data, we must be able to approximate it in a tractable manner.

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

Original: By the symmetry of mutual information, we can rewrite Equation 4 as
Revised: By the symmetry of mutual information, we can rewrite Equation 4 as

Original: [Equation] [Equation]
Revised: [Equation] [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 the transition function P ( T | D ) .
Revised: = 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 Agent Draw samples from posterior of Run RL on samples

Revised: Maximize EIG

Original:  In order to compute the other term, we must take samples
Revised: Environment Query Observe (a) −−  20 x −−−−     y (b) the transition function P ( T | D ) . In order to compute the other term, we must take samples

Original: O ij ∼ P ( O | D ) .
Revised: O ij ∼ P ( O | D ) .

Original: To do this, we ﬁrst sample m start states s i 0 from p 

Revised: To do this, we ﬁrst sample m start states s i 0 from p

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: (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.

Original: Then we can run the
Revised: Then we can 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 .This is an expression of the generative process [Figure] Environment
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 O ij .

Original: Query Observe (a) −−  20 x −−−−     y (b) 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 O 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 O 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, O i , D ] .

Original: For this, we follow a similar strategy as done in Neiswanger et al.
Revised: For this, we follow a similar strategy as done in 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 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.

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 thecomputation of the acquisition function to factor p 0 as an irreducible source of uncertainty.
Revised: Though inspired by the work cited above, we modify 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 the simplest possible approach for nonlinear control in continuous spaces and assume a Gaussian process prior T (cid:48) 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 have a straightforward approximation for all necessary quantities.

Original: We assert however that the development of the acquisition function is general and any approximation of a
Revised: We assert however that the development of the acquisition function is general and any approximation of a

Original: Bayesian model could be used.
Revised: Bayesian model could be used.

Original: In particular, we follow the approach of Wilson et al.
Revised: 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.
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.

Original: We model the transition function T : S×A → p ( S ) (the problem dynamics) using Gaussian processes due to their non-parametric nature and inherent ability to model uncertainties in the modeled function.
Revised: We model the transition function T : S×A → p ( S ) (the problem dynamics) using Gaussian processes due to their non-parametric nature and inherent ability to model uncertainties in the modeled function.

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
Revised: We model each of the d dimensions of the output using individual

Original: GPs.
Revised: 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 than modelling the ﬁnal state T ( s, a ) directly.
Revised: T ( s, a ) − s rather than modelling 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 does not change by a huge 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.
Revised: (2020) for fast approximate posterior sampling of GPs.

Original: 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 EIG O for D  and approximate argmax s ∈S ,a ∈A EIG O
Revised: 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 EIG O for D by running MPC on posterior function samples and approximate argmax s ∈S ,a ∈A EIG O

Original: ( s, a ) by zeroth order approximation.
Revised: ( 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 model-predictive control 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 (cid:48) | D ] on the real environment.

Original: We refer to this procedure as Bayesian Active Reinforcement Learning (BARL) and give a precise speciﬁcation in Algorithm 1 as well as a schematic diagram in Figure 1a.
Revised: We refer to this procedure as Bayesian Active Reinforcement Learning (BARL). Details are given in Algorithm 1 (here, U denotes the uniform distribution) and a schematic diagram in Figure 1a.

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

Original: Algorithm 1 Bayesian active reinforcement learning (

Original: ) using

Original: O Inputs: budget b , # of points for optimization k , # of posterior function samples n .
Revised: 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 ( s 0 , a 0 ) ∼ U ( S × A ) , s (cid:48)

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

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

Original: [Equation]
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 { O (cid:96) } n(cid:96)

Original: =1 by executing  π T (cid:48) (cid:96) on the dynamics of T (cid:48) (cid:96) starting from s 0 .
Revised: =1 by executing MPC policy π T (cid:48) (cid:96) on the dynamics of T (cid:48) (cid:96) starting from s 0 .

Original: [Equation]
Revised: ( x 1 , y 1 ) , . . . , ( x k , y k ) ∼ U ( S × A ) ( x ∗ , y ∗ ) ← argmax { ( x i ,y i ) } ki =1 EIG O ( x i , y i ) .

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

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 ( T | 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 was to reduce the sample complexity of learning good policies in continuous spaces.
Revised: 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.

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 reinforcement learning algorithms.
Revised: 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.

Original: In particular, we compare the average return across ﬁve evaluatation episdoes 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 method 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 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: Though in this section we focus on the performance of the controller, we discuss the runtime of the algorithm in section A.1.

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 π ˆ Tto collect data 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 as in the standard RL setting.

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, TD (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 which 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 that 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 swingup 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-v 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 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: 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.

Original: The lava path is intended to test stability and exploration of algorithms.
Revised: The lava 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.
Revised: 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.
Revised: As shown in Figure 1b, the state space [Figure] ˙ θ

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

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: Pendulum   5, 6, 57, 13,
Revised: Pendulum   5, 6, 57, 13,

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

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

Original: Reacher  n/a 23, 13, n/a  to the goal as quickly as possible.
Revised: Reacher  n/a 23, 13, n/a ‘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. ‘n/a’ is recorded when the median run didn’t solve the problem by the end of our training. contains a ‘lava’ region which gives very large negative rewards for every timestep.

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: 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 totake advantage of transition queries 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 O to greatly 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 higher-dimensional 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 more reqired samples in the smaller environments where the model uncertainty can point to meaningfully underexplored areas.
Revised: The MPC method performs well across the board and is competitive with BARL on the plasma problem but suffers from many more reqired samples in the 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: However, it has much worse model predictions on the random datapoints than EIG T while doing a
Revised: However, it has much worse model predictions on the random datapoints than EIG T while doing a

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Original: much better job modelling the datapoints used by the MPC procedure.
Revised: much better job modelling the datapoints 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 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.

Original: We see this pattern across other tasks as well.
Revised: We see this pattern across other tasks as well.

Revised: 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: Our experiments show that EIG O seems to work well even when the policy used to generate O is subpoptimal.

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 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.

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 our proposed acquisition function relies on executing the entire control algorithm over a set of posterior function samples leading to higher computational requirements.

Original: 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.
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 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 Table 
Revised: The other hyperparameters of the BARL algorithm are constant but listed for completeness in Table 2.

Original: For all of our experiments, we use a squared exponential kernel with automatic relevance determination (MacKay et al., 1994; Neal, 1995).
Revised: For all of our experiments [Table] 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, and chose the maximum of the set.
Revised: To optimize the transition function, we simply sampled a set of points from the domain, evaluated the acquisition function [Table] and chose the maximum of the set.

Original: 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 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. A. R UNTIME D ETAILS

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.

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

Original: B MPC D ETAILS
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. 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 but still relevant to practitioners. Clearly, BARL is a relatively slow algorithm computationally.

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.

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.

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  tricks including colored noise samples and caching to reduce the number of queries to the planning model.
Revised: (2020) to solve the optimization problem in Equation 2, which uses several [Table] tricks including colored noise samples and caching to reduce the number of queries to the planning 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 3.
Revised: The values we used across all experiments for each problem are given in Table 4.

Original: C D ESCRIPTION OF C ONTINUOUS C ONTROL P ROBLEMS
Revised: B. R OBUSTNESS OF EIG O TO A SUBOPTIMAL CONTROLLER

Revised: 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: 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: However, this assumption could lead to pathologies in the method if it doesn’t hold in practice.

Revised: 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 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 particular, we ran the reacher and cartpole experiments from the main paper with varying MPC budgets for the posterior function policy π T

Revised: [Table]

Revised: and a ﬁxed MPC budget at test time.

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

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

Revised: This is straightforwardly linked to the amount of search the policy conducts before executing an 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.

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

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

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

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.

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

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

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

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  distribution over start states.
Revised: The goal is to reach a ﬁxed goal state from a relatively narrow uniform [Figure] 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.
Revised: 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.
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.