Original: R EINFORCEMENT L EARNING WITH B AYESIAN C LASSIFIERS : E FFICIENT S KILL L EARNING FROM O UTCOME
Revised: R EINFORCEMENT L EARNING WITH B AYESIAN C LASSIFIERS : E FFICIENT S KILL L EARNING FROM O UTCOME

Original: E XAMPLES
Revised: E XAMPLES

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: Exploration in reinforcement learning is, in general, a challenging problem.
Revised: Exploration in reinforcement learning is, in general, a challenging problem

Original: In this work, we study a more tractable class of reinforcement learning problems deﬁned by data that provides examples of successful outcome states.
Revised: In this work , we study a more tractable class of reinforcement learning problems deﬁned by data that provides examples of successful outcome states.

Original: In this case, the reward function can be obtained automatically by training a classiﬁer to classify states as successful or not.
Revised: In this case the reward function can be obtained automatically by training a classiﬁer to classify states as successful or not

Original: We argue that, with appropriate representation and regularization, such a classiﬁer can guide a reinforcement learning algorithm to an effective solution.
Revised: We argue that, with appropriate representation and regularization such a classiﬁer can guide a reinforcement learning algorithm to an effective solution.

Original: However, as we will show, this requires the classiﬁer to make uncertaintyaware predictions that are very difﬁcult with standard deep networks.
Revised: However, as we will show, this requires the classiﬁer to make uncertainty aware predictions that are very difﬁcult with standard deep networks

Original: To address this, we propose a novel mechanism for obtaining calibrated uncertainty based on an amortized technique for computing the normalized maximum likelihood distribution.
Revised: To address this we propose a novel mechanism for obtaining calibrated uncertainty based on an amortized technique for computing the normalized maximum likelihood distribution

Original: We show that the resulting algorithm has a number of intriguing connections to both count-based exploration methods and prior algorithms for learning reward functions from data, while  being able to guide algorithms towards the speciﬁed goal more effectively.
Revised: We show that the resulting algorithm has a number of intriguing connections to both count-based exploration methods and prior algorithms for learning reward functions from data, while also being able to guide algorithms towards the speciﬁed goal more effectively.

Original: We show how using amortized normalized maximum likelihood for reward inference is able to provide effective reward guidance for solving a number of challenging navigation and robotic manipulation tasks which prove difﬁcult for other algorithms.
Revised: We show how using amortized normalized maximum likelihood for reward inference is able to provide effective reward guidance for solving a number of challenging navigation and robotic manipulation tasks which prove difﬁcult for other algorithms

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: While reinforcement learning (RL) has been shown to successfully solve problems with careful reward design (Rajeswaran et al., 2018; OpenAI et al., 2019), RL in its most general form, with no assumptions on the dynamics or reward function, requires solving a challenging uninformed search problem in which rewards are sparsely observed.
Revised: While reinforcement learning (RL) has been shown to successfully solve problems with careful reward design (Rajeswaran et al., 2018; OpenAI et al., 2019), RL in its most general form, with no assumptions on the dynamics or reward function, requires solving a challenging uninformed search problem in which rewards are sparsely observed.

Original: Techniques which explicitly provide “rewardshaping” (Ng et al., 1999), or modify the reward function to guide learning, can help take some of the burden off of exploration, but shaped rewards can be difﬁcult to obtain without domain knowledge.
Revised: Techniques which explicitly provide “rewardshaping” (Ng et al., 1999), or modify the reward function to guide learning, can help take some of the burden off of exploration, but shaped rewards can be difﬁcult to obtain without domain knowledge.

Original: In this paper, we aim to reformulate the reinforcement learning problem to make it easier for the user to specify the task and to provide a tractable reinforcement learning objective.
Revised: In this paper, we aim to reformulate the reinforcement learning problem to make it easier for the user to specify the task and to provide a tractable reinforcement learning objective.

Original: Instead of requiring a reward function designed for an objective, our method instead assumes a user-provided set of successful outcome examples : states in which the desired task has been accomplished successfully.
Revised: Instead of requiring a reward function designed for an objective, our method instead assumes a user-provided set of successful outcome examples : states in which the desired task has been accomplished successfully.

Original: The algorithm aims to estimate the distribution over these states and maximize the probability of reaching states that are likely under the distribution.
Revised: The algorithm aims to estimate the distribution over these states and maximize the probability of reaching states that are likely under the distribution.

Original: Prior work on learning from success examples (Fu et al., 2018b; Zhu et al., 2020) focused primarily on alleviating the need for manual reward design.
Revised: Prior work on learning from success examples (Fu et al., 2018b; Zhu et al., 2020) focused primarily on alleviating the need for manual reward design.

Original: In our work, we focus on the potential for this mode of task speciﬁcation to produce more tractable
Revised: In our work, we focus on the potential for this mode of task speciﬁcation to produce more tractable

Original: RL problems and solve more challenging classes of tasks.
Revised: RL problems and solve more challenging classes of tasks.

Original: Intuitively, when provided with explicit examples of successful states, the RL algorithm should be able to direct its exploration, rather than simply hope to randomly chance uponpotentially high reward states.
Revised: Intuitively, when provided with explicit examples of successful states, the RL algorithm should be able to direct its exploration, rather than simply hope to randomly chance upon high reward states.

Original: The main challenge in instantiating this idea into a practical algorithm is performing appropriate uncertainty quantiﬁcation in estimating whether a given state corresponds to a successful outcome.
Revised: The main challenge in instantiating this idea into a practical algorithm is performing appropriate uncertainty quantiﬁcation in estimating whether a given state corresponds to a successful outcome.

Original: Our approach trains a classiﬁer to distinguish successful states, provided by the user, from those generated by the current policy, analogously to generative adversarial networks (Goodfellow et al., 2014) and previously proposed methods for inverse reinforcement learning (Fu et al., 2018a).
Revised: Our approach trains a classiﬁer to distinguish successful states, provided by the user, from those generated by the current policy, analogously to generative adversarial networks (Goodfellow et al., 2014) and previously proposed methods for inverse reinforcement learning (Fu et al., 2018a).

Original: In general, such a classiﬁer is not guaranteed to provide a good optimization landscape for learning the policy.
Revised: In general, such a classiﬁer is not guaranteed to provide a good optimization landscape for learning the policy.

Original: We discuss how a particular form of uncertainty quantiﬁcation based on the normalized maximum likelihood (NML) distribution produces better reward guidance for learning.
Revised: We discuss how a particular form of uncertainty quantiﬁcation based on the normalized maximum likelihood (NML) distribution produces better reward guidance for learning.

Original: We also connect our approach to count-based exploration methods, showing that a classiﬁer with suitable uncertainty estimates reduces to a count-based exploration method in the absence of any generalization across states, while also discussing how it improves over count-based exploration in the presence of good generalization.
Revised: We also connect our approach to count-based exploration methods, showing that a classiﬁer with suitable uncertainty estimates reduces to a count-based exploration method in the absence of any generalization across states, while also discussing how it improves over count-based exploration in the presence of good generalization.

Original: We then propose a practical algorithm to train success classiﬁers in a computationally efﬁcient way with NML, and show how this form of reward inference allows us to solve difﬁcult problems more efﬁciently, providing experimental results which outperform existing algorithms on a number of navigation and robotic manipulation domains.
Revised: We then propose a practical algorithm to train success classiﬁers in a computationally efﬁcient way with NML, and show how this form of reward inference allows us to solve difﬁcult problems more efﬁciently, providing experimental results which outperform existing algorithms on a number of navigation and robotic manipulation domains.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: A number of techniques have been proposed to improve exploration.
Revised: A number of techniques have been proposed to improve exploration.

Original: These techniques either add reward bonuses that encourage a policy to visit novel states in a task-agnostic manner (Wiering and
Revised: These techniques either add reward bonuses that encourage a policy to visit novel states in a task-agnostic manner (Wiering and

Original: Schmidhuber, 1998; Auer et al., 2002; Schaul et al., 2011; Houthooft et al., 2016; Pathak et al., 2017; Tang et al., 2017; Stadie et al., 2015; Bellemare et al., 2016; Burda et al., 2018a; O’Donoghue, 2018) or perform Thompson sampling or approximate Thompson sampling based on a prior over value functions (Strens, 2000; Osband et al., 2013; 2016).
Revised: Schmidhuber, 1998; Auer et al., 2002; Schaul et al., 2011; Houthooft et al., 2016; Pathak et al., 2017; Tang et al., 2017; Stadie et al., 2015; Bellemare et al., 2016; Burda et al., 2018a; O’Donoghue, 2018) or perform Thompson sampling or approximate Thompson sampling based on a prior over value functions (Strens, 2000; Osband et al., 2013; 2016).

Original: While these techniques are uninformed about the actual task, we consider a constrained set of problems where examples of successes can allow for more task directed exploration.
Revised: While these techniques are uninformed about the actual task, we consider a constrained set of problems where examples of successes can allow for more task-directed exploration.

Original: In real world problems, designing well-shaped reward functions makes exploration easier but often requires signiﬁcant domain knowledge (Andrychowicz et al., 2020), access to privileged information about the environment (Levine et al., 2016) and/or a human in the loop providing rewards (Knox and Stone, 2009; Singh et al., 2019b).
Revised: In real world problems, designing well-shaped reward functions makes exploration easier but often requires signiﬁcant domain knowledge (Andrychowicz et al., 2020), access to privileged information about the environment (Levine et al., 2016) and/or a human in the loop providing rewards (Knox and Stone, 2009; Singh et al., 2019b).

Original: Prior work has considered specifying rewards by providing example demonstrations, and inferring rewards with inverse RL (Abbeel and Ng, 2004; Ziebart et al., 2008; Ho and Ermon, 2016;
Revised: Prior work has considered specifying rewards by providing example demonstrations and inferring rewards with inverse RL (Abbeel and Ng, 2004; Ziebart et al., 2008; Ho and Ermon, 2016;

Original: Fu et al., 2018a).
Revised: Fu et al., 2018a).

Original: This requires expensive expert demonstrations to be provided to the agent.
Revised: This requires expensive expert demonstrations to be provided to the agent.

Original: In contrast, our work has the minimal requirement of simply providing successful outcome states, which can be done cheaply.
Revised: In contrast, our work has the minimal requirement of simply providing successful outcome states, which can be done cheaply and more intuitively.

Original: This subclass of problems is also related to goal conditioned RL (Kaelbling, 1993;Schaul et al., 2015;
Revised: This subclass of problems is also related to goal conditioned RL (Kaelbling, 1993;

Original:  Zhu et al., 2017; Andrychowicz et al., 2017; Nair et al., 2018; Veeriah et al., 2018; Rauber et al., 2018; Warde-Farley et al., 2018; Colas et al., 2019; Ghosh et al., 2019; Pong et al., 2020) but is more general, since it allows for a more abstract notion of task success.
Revised: Schaul et al., 2015; Zhu et al., 2017; Andrychowicz et al., 2017; Nair et al., 2018; Veeriah et al., 2018; Rauber et al., 2018; Warde-Farley et al., 2018; Colas et al., 2019; Ghosh et al., 2019; Pong et al., 2020) but is more general, since it allows for a more abstract notion of task success.

Original: A core idea behind our work is using a Bayesian classiﬁer to learn a suitable reward function.
Revised: A core idea behind our work is using a Bayesian classiﬁer to learn a suitable reward function.

Original: Bayesian inference with expressive models and high dimensional data can often be intractable, requiring assumptions on the form of the posterior (Hoffman et al., 2013; Blundell et al., 2015; Maddox et al., 2019).
Revised: Bayesian inference with expressive models and high dimensional data can often be intractable, requiring assumptions on the form of the posterior (Hoffman et al., 2013; Blundell et al., 2015; Maddox et al., 2019).

Original: In this work, we build on the concept of normalized maximum likelihood (Rissanen, 1996;
Revised: In this work, we build on the concept of normalized maximum likelihood (Rissanen, 1996;

Original: Shtar’kov, 1987), or NML, to learn Bayesian classiﬁers.
Revised: Shtar’kov, 1987), or NML, to learn Bayesian classiﬁers.

Original: Although NML is typically considered from the perspective of optimal coding (Grünwald, 2007; Fogel and Feder, 2018), we show how it can be used to learn success classiﬁers, and discuss its connections to exploration and reward shaping.
Revised: Although NML is typically considered from the perspective of optimal coding (Grünwald, 2007; Fogel and Feder, 2018), we show how it can be used to learn success classiﬁers, and discuss its connections to exploration and reward shaping in RL.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this paper, we study a modiﬁed reinforcement learning problem, where instead of the standard reward function, the agent is provided with successful outcome examples .
Revised: In this paper, we study a modiﬁed reinforcement learning problem, where instead of the standard reward function, the agent is provided with successful outcome examples .

Original: This reformulation not only provides a modality for task speciﬁcation that may be more natural for users to provide in some settings (Fu et al., 2018b; Zhu et al., 2020; Singh et al., 2019a), but, as we will show, can also make learning easier.
Revised: This reformulation not only provides a modality for task speciﬁcation that may be more natural for users to provide in some settings (Fu et al., 2018b; Zhu et al., 2020; Singh et al., 2019a), but, as we will show, can also make learning easier.

Original: We also derive a meta-learned variant of the conditional normalized maximum likelihood (CNML) distribution for representing our reward function, in order to make evaluation tractable.
Revised: We also derive a meta-learned variant of the conditional normalized maximum likelihood (CNML) distribution for representing our reward function, in order to make evaluation tractable.

Original: We discuss background on successful outcome examples and CNML in this section.
Revised: We discuss background on successful outcome examples and CNML in this section.

Original: 3. R EINFORCEMENT L EARNING WITH E XAMPLES OF S UCCESSFUL O UTCOMES
Revised: 3. R EINFORCEMENT L EARNING WITH E XAMPLES OF S UCCESSFUL O UTCOMES

Original: We follow the framework proposed by Fu et al.
Revised: We follow the framework proposed by Fu et al.

Original: (2018b)
Revised: (2018b)

Original: and assume that we are provided with a Markov decision process (MDP) without a reward function, given by M , where M = ( S , A , T , γ, µ 0 ) , as well as successful outcome examples S
Revised: and assume that we are provided with a Markov decision process (MDP) without a reward function, given by M , where M = ( S , A , T , γ, µ 0 ) , as well as successful outcome examples S

Original: + =
Revised: + =

Original: { s k + } Kk =1 , which is a set of states in which the desired task has been accomplished.
Revised: { s k + } Kk =1 , which is a set of states in which the desired task has been accomplished.

Original: This formalism is easiest to describe in terms of the control as inference framework (Levine, 2018).
Revised: This formalism is easiest to describe in terms of the control as inference framework (Levine, 2018).

Original: The relevant graphical model in Figure 9 consists of states and actions, as well as binary success variables e t which represent the occurrence of a particular event.
Revised: The relevant graphical model in Figure 9 consists of states and actions, as well as binary success variables e t which represent the occurrence of a particular event.

Original: The agent’s objective is to cause this event to occur (e.g., a robot that is cleaning the ﬂoor must cause the “ﬂoor is clean” event to occur).
Revised: The agent’s objective is to cause this event to occur (e.g., a robot that is cleaning the ﬂoor must cause the “ﬂoor is clean” event to occur).

Original: Formally, we assume that the states in S + are sampled from the distribution p ( s t | e t
Revised: Formally, we assume that the states in S + are sampled from the distribution p ( s t | e t

Original: =
Revised: =

Original: True ) – that is, states where the desired event has taken place.
Revised: True ) – that is, states where the desired event has taken place.

Original: In this work, we focus on efﬁcient methods for solving this reformulation of the RL problem, by utilizing a novel uncertainty quantiﬁcation method to represent the distribution p ( e t | s t ) .
Revised: In this work, we focus on efﬁcient methods for solving this reformulation of the RL problem, by utilizing a novel uncertainty quantiﬁcation method to represent the distribution p ( e t | s t ) .

Original: In practice, prior methods that build on this and similar reformulations of the RL problem (Fu et al., 2018b) derive an algorithm where the reward function in RL is produced by a classiﬁer that estimates p ( e t =
Revised: In practice, prior methods that build on this and similar reformulations of the RL problem (Fu et al., 2018b) derive an algorithm where the reward function in RL is produced by a classiﬁer that estimates p ( e t =

Original: True | s t ) .
Revised: True | s t ) .

Original: Following the adversarial inverse reinforcement learning (AIRL) derivation (Fu et al., 2018a; Finn et al., 2016), it is possible to show that the correct source of negative examples for training this classiﬁer is the state distribution of the policy itself, π ( s ) .
Revised: Following the adversarial inverse reinforcement learning (AIRL) derivation (Fu et al., 2018a; Finn et al., 2016), it is possible to show that the correct source of negative examples for training this classiﬁer is the state distribution of the policy itself, π ( s ) .

Original: This insight results in a simple algorithm: at each iteration of the algorithm, the policy is updated to maximize the current reward, given by log p ( e t =
Revised: This insight results in a simple algorithm: at each iteration of the algorithm, the policy is updated to maximize the current reward, given by log p ( e t =

Original: True | s t ) , then samples from the policy are added to the set of negative examples
Revised: True | s t ) , then samples from the policy are added to the set of negative examples

Original: S − , and the classiﬁer is retrained on the original positive set S
Revised: S − , and the classiﬁer is retrained on the original positive set S

Original: + and the updated negative set S − .
Revised: + and the updated negative set S − .

Original: 3. C ONDITIONAL N ORMALIZED M AXIMUM L IKELIHOOD
Revised: 3. C ONDITIONAL N ORMALIZED M AXIMUM L IKELIHOOD

Original: Our method builds on the principle of conditional normalized maximum likelihood (NML) (Rissanen and Roos, 2007; Grünwald, 2007; Fogel and Feder, 2018), which we review brieﬂy.
Revised: Our method builds on the principle of conditional normalized maximum likelihood (NML) (Rissanen and Roos, 2007; Grünwald, 2007; Fogel and Feder, 2018), which we review brieﬂy.

Original: CNML is a method for performing k -way classiﬁcation, given a model class Θ and a dataset
Revised: CNML is a method for performing k -way classiﬁcation, given a model class Θ and a dataset

Original: D = { ( x 0 , y 0 ) , ( x 1 , y 1 ) , ..., ( x n , y n ) } , and has been shown to provide better calibrated predictions and uncertainty estimates with minimax regret guarantees (Bibas et al., 2019).
Revised: D = { ( x 0 , y 0 ) , ( x 1 , y 1 ) , ..., ( x n , y n ) } , and has been shown to provide better calibrated predictions and uncertainty estimates with minimax regret guarantees (Bibas et al., 2019).

Original: To predict the class of a query point x q , CNML constructs k augmented datasets by adding x q with a different label in each datasets, which we write as D ∪ ( x q , y = i ) , i ∈ (1 , 2 , ..., k ) .
Revised: To predict the class of a query point x q , CNML constructs k augmented datasets by adding x q with a different label in each datasets, which we write as D ∪ ( x q , y = i ) , i ∈ (1 , 2 , ..., k ) .

Original: CNML then deﬁnes the class distribution by solving the maximum likelihood estimation problem at query time for each of these augmented datasets to convergence, and normalize the likelihoods as follows:
Revised: CNML then deﬁnes the class distribution by solving the maximum likelihood estimation problem at query time for each of these augmented datasets to convergence, and normalize the likelihoods as follows:

Original: [Equation]
Revised: [Equation]

Original: Intuitively, if x q is close to other datapoints in D , then the model will struggle to assign a high likelihood to labels that differ substantially from other nearby points.
Revised: Intuitively, if x q is close to other datapoints in D , then the model will struggle to assign a high likelihood to labels that differ substantially from other nearby points.

Original: However, if x q is far from all datapoints in D , then the different augmented MLE problems can easily classify x q as an arbitrary class, providing us with a likelihood closer to uniform.
Revised: However, if x q is far from all datapoints in D , then the different augmented MLE problems can easily classify x q as an arbitrary class, providing us with a likelihood closer to uniform.

Original: We refer readers to Grünwald (2007) for an in-depth discussion.
Revised: We refer readers to Grünwald (2007) for an in-depth discussion.

Original: A major limitation of CNML is that it requires training an entire neural network to convergence on the entire augmented dataset every time we want to evaluate a test point’s class probabilities.
Revised: A major limitation of CNML is that it requires training an entire neural network to convergence on the entire augmented dataset every time we want to evaluate a test point’s class probabilities.

Original: We will address this issue in Section 5.
Revised: We will address this issue in Section 5.

Original:  B AYESIAN S UCCESS C LASSIFIERS FOR R EWARD I NFERENCE
Revised:  B AYESIAN S UCCESS C LASSIFIERS FOR R EWARD I NFERENCE

Original:  Ideally, training a classiﬁer with the policy samples as negative examples as described in Section 3.1 should yield a smooth decision boundary between the well-separated negative and positive examples.
Revised: [Figure] Ideally, training a classiﬁer with the policy samples as negative examples as described in Section 3.1 should yield a smooth decision boundary between the well-separated negative and positive examples.

Original: For example, Figure 2 depicts a simple 1-D example, where the agent starts at the left ( s 0 ) and the positive outcomes are at the right ( s + ) side of the environment.
Revised: For example, Figure 2 depicts a simple 1-D scenario, where the agent starts at the left ( s 0 ) and the positive outcomes are at the right ( s + ) side of the environment.

Original: Since the positives are on the right and the negatives are on the left, one might expect a classiﬁer to gradually increase its prediction of a success as we move to the right (Figure 2a), which would provide a dense reward signal for the policy to move to the right.
Revised: Since the positives are on the right and the negatives are on the left, one might expect a classiﬁer to gradually increase its prediction of a success as we move to the right (Figure 2a), which would provide a dense reward signal for the policy to move to the right.

Original: However, this idealized scenario rarely happens in practice.
Revised: However, this idealized scenario rarely happens in practice.

Original: Without suitable regularization, the decision boundary between the positive and negative examples may
Revised: Without suitable regularization, the decision boundary between the positive and negative examples may

Original: [Figure]
Revised: [Figure]

Original: not be smooth.
Revised: not be smooth.

Original: In fact, the decision boundary of an optimal classiﬁer may take on the form of a sharp boundary anywhere between the positive and negative examples in the early stages of training (Figure 2b).
Revised: In fact, the decision boundary of an optimal classiﬁer may take on the form of a sharp boundary anywhere between the positive and negative examples in the early stages of training (Figure 2b).

Original: As a result, the classiﬁer might provide little to no reward signal for the policy, since it can assign arbitrarily small probabilities to the states sampled from the policy.
Revised: As a result, the classiﬁer might provide little to no reward signal for the policy, since it can assign arbitrarily small probabilities to the states sampled from the policy.

Original: We note that this issue is not pathological: our experiments in Section 6 show that this poor reward signal issue happens in practice and can greatly hinder learning.
Revised: We note that this issue is not pathological: our experiments in Section 6 show that this poor reward signal issue happens in practice and can greatly hinder learning.

Original: In this section, we will discuss how an appropriate classiﬁer training method can avoid these uninformative rewards.
Revised: In this section, we will discuss how an appropriate classiﬁer training method can avoid these uninformative rewards.

Original: 4. R EGULARIZED S UCCESS C LASSIFIERS VIA N ORMALIZED M AXIMUM L IKELIHOOD
Revised: 4. R EGULARIZED S UCCESS C LASSIFIERS VIA N ORMALIZED M AXIMUM L IKELIHOOD

Revised: 1: User provides success examples S + 2: Initialize policy π , replay buffer S − , and reward classiﬁer parameters θ R 3: for iteration i = 1 , 2 , ... do 4:

Revised: Add on-policy examples to S − by executing π .

Revised: Sample n test points from S + (label 1) and n test points from S − (label 0) to construct a dataset D 6:

Original:  To create effective shaping, we need to impose a prior on our classiﬁer so that it provides a more informative reward when evaluated at rarely visited states that lie on the path to successful outcomes.
Revised: Assign state p = 1 | s, D ) 7: Train π with RL algorithm To create effective shaping, we would like our classiﬁer to provide a more informative reward when evaluated at rarely visited states that lie on the path to successful outcomes.

Original: A more informative reward function is one that assigns higher rewards to the fringe of the states visited by the policy, because thisreward function will encourage the policy to explore and move towards the desired states.
Revised: A more informative reward function is one that assigns higher rewards to the fringe of the states visited by the policy, because this will encourage the policy to explore and move towards the desired states.

Original: We can construct such a reward function by imposing the prior that novel states have a non-negligible chance of being a success state.
Revised: We can construct such a reward function by imposing the prior that novel states have a non-negligible chance of being a success state.

Original: To do so, we train a Bayesian classiﬁer using conditional normalized maximum likelihood (CNML) (Shtar’kov, 1987), as we described in Section 3, which is essentially imposing a uniform prior on the space of outcomes.
Revised: To do so, we train a Bayesian classiﬁer using conditional normalized maximum likelihood (CNML) (Shtar’kov, 1987), as we described in Section 3, which corresponds to imposing a uniform prior on the output class probabilities.

Original: To use CNML for reward inference, the procedure is similar to the one described in Section 3.
Revised: To use CNML for reward inference, the procedure is similar to the one described in Section 3.

Original: We construct a dataset using the provided successful outcomes as positives and the on-policy samples as negatives.
Revised: We construct a dataset using the provided successful outcomes as positives and the on-policy samples as negatives.

Original: However, the label probabilities for RL are then produced by the CNML procedure described in Equation 1 to obtain the rewards r ( s ) =
Revised: However, the label probabilities for RL are then produced by the CNML procedure described in Equation 1 to obtain the rewards r ( s ) =

Original: p CNML ( e = 1 | s ) .
Revised: p CNML ( e = 1 | s ) .

Original: How does this affect reward assignment during learning? To illustrate the impact of using CNML, we visualize a potential assignment of rewards with a CNML-based classiﬁer on the problem described earlier.
Revised: To illustrate how this affects reward assignment during learning, we visualize a potential assignment of rewards with a CNMLbased classiﬁer on the problem described earlier.

Original: When the success classiﬁer is trained with CNML instead of standard maximum likelihood, intermediate unseen states would receive non-zero rewards rather than simply having vanishing likelihoods like in Figure 2b.
Revised: When the success classiﬁer is trained with CNML instead of standard maximum likelihood, intermediate unseen states would receive non-zero rewards rather than simply having vanishing likelihoods like in Figure 2b.

Original: The didactic illustrations in Fig 2c and Fig 2d show how the rewards obtained via NML might incentivize exploration.
Revised: The didactic illustrations in Fig 2c and Fig 2d show how the rewards obtained via NML might incentivize exploration.

Original: When there is no generalization across states, the CNML likelihood will correspond directly to count based exploration (as we show below). However, when there is actually some degree of generalization across states, the classiﬁer can provide more directed shaping towards the goal.
Revised: In fact, the CNML likelihood corresponds to a form of count-based exploration (as we show below), while also providing more directed shaping towards the goal when generalization exists across states.

Original: 4. R ELATIONSHIP TO C OUNT -B ASED E XPLORATION
Revised: 4. R ELATIONSHIP TO C OUNT -B ASED E XPLORATION

Original: In this section we relate the success likelihoods obtained via CNML to commonly used exploration methods based on counts.
Revised: In this section we relate the success likelihoods obtained via CNML to commonly used exploration methods based on counts.

Original: Formally, we prove that the success classiﬁer trained with CNML is equivalent to a version of count-based exploration  in the absence of any generalization across states (i.e., a fully tabular setting).
Revised: Formally, we prove that the success classiﬁer trained with CNML is equivalent to a version of count-based exploration with a sparse reward function in the absence of any generalization across states (i.e., a fully tabular setting).

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Suppose we are estimating success probabilities p ( e = 1 | s ) in the tabular setting, where we have a separate parameterindependently for each state.
Revised: Suppose we are estimating success probabilities p ( e = 1 | s ) in the tabular setting, where we have an independent parameter for each state.

Original: Let N ( s ) denote the number of times state s has been visited by the policy, and let G ( s ) be the number of occurrences of state s in the successful outcome examples.
Revised: Let N ( s ) denote the number of times state s has been visited by the policy, and let G ( s ) be the number of occurrences of state s in the set of goal examples.

Original: Then the CNML  probability p CNML ( e = 1 | s ) is equal to
Revised: Then the CNML success probability p CNML ( e = 1 | s ) is equal to

Original: G ( s )+ [Equation] N ( s )+ G ( s )+2 .
Revised: G ( s )+1 N ( s )+ G ( s )+2 .

Original: For states that are never observed to be successful, we then recover inverse counts 1 N ( s )+2 .
Revised: For states that are not represented in the goal examples, i.e. G(s) = 0, we then recover inverse counts 1 N ( s )+2 .

Original: Refer to Appendix A.7 for a full proof.
Revised: Refer to Appendix A.7 for a full proof.

Original: 4. R EWARD S HAPING WITH B AYESIAN S UCCESS C LASSIFIERS
Revised: 4. R EWARD S HAPING WITH B AYESIAN S UCCESS C LASSIFIERS

Original: While the analysis above suggests that a CNML classiﬁer would give us something akin to  an exploration bonus, the structure of the problem and the state space actually provides us more information to guide us towards the goal.
Revised: While the analysis above suggests that a CNML classiﬁer would give us something akin to a sparse reward plus an exploration bonus, the structure of the problem and the state space actually provides us more information to guide us towards the goal.

Original: In most environments (Brockman et al., 2016;
Revised: In most environments (Brockman et al., 2016;

Original: Yu et al., 2019)
Revised: Yu et al., 2019)

Original: the state space does not consist of independent and uncorrelated categorical variables, but is  provided in a representation that relates at least roughly to the dynamics structure in the environment.
Revised: the state space does not consist of independent and uncorrelated categorical variables, and is instead provided in a representation that relates at least roughly to the dynamics structure in the environment.

Original: For instance, states close to the goal dynamically are also typically close to the goal in the metric space deﬁned by the states.
Revised: For instance, states close to the goal dynamically are also typically close to the goal in the metric space deﬁned by the states.

Original: Indeed, this observation is the basis of many commonly used heuristic reward shaping methods, such as rewards given by Euclidean distance to target states.
Revised: Indeed, this observation is the basis of many commonly used heuristic reward shaping methods, such as rewards given by Euclidean distance to target states.

Original: In this case, the task speciﬁcation can actually provide more information than simply performing uninformed count-based exploration.
Revised: In this case, the task speciﬁcation can actually provide more information than simply performing uninformed count-based exploration.

Original: Since the uncertainty aware classiﬁer described in Section 4. is built on top of features that are correlated with environment dynamics,  it is able to incentivize taskaware directed exploration.
Revised: Since the uncertainty-aware classiﬁer described in Section 4.1 is built on top of features that are correlated with environment dynamics, and is trained with knowledge of the desired outcomes, it is able to incentivize task-aware directed exploration.

Original: This phenomenon is illustrated intuitively in Fig 2d, and demonstrated empirically in our experimental analysis in Section 6.

Original: As compared to CNML without generalization in Fig 2c, we expect the intermediate rewards to provide more shaping towards the goal.
Revised: As compared to CNML without generalization in Fig 2c, we expect the intermediate rewards to provide more shaping towards the goal.

Original: This  is evident from our empirical results in Section 6 where BayCRL is able to signiﬁcantly outperform methods for task agnostic exploration.
Revised: This phenomenon is illustrated intuitively in Fig 2d, and visualized and demonstrated empirically in our experimental analysis in Section 6, where BayCRL is able to signiﬁcantly outperform methods for task-agnostic exploration.

Original: 4. O VERVIEW
Revised: 4. O VERVIEW

Original: In this section, we introduced the idea of Bayesian classiﬁers trained via CNML as a means to provide rewards for RL problems speciﬁed via examples of successful outcomes.
Revised: In this section, we introduced the idea of Bayesian classiﬁers trained via CNML as a means to provide rewards for RL problems speciﬁed by examples of successful outcomes.

Original: Standard MLE classiﬁers can provide exploration  and reward shaping, whereas properly regularized classiﬁers can directly provide beneﬁts for exploration(as shown in Theorem 4.1) and reward shaping, as discussed in Section 4.3.
Revised: Concretely, a CNML-based scheme has the following advantages: • Natural exploration behavior due to accurate uncertainty estimation in the output success probabilities. This is explained by the connection between CNML and count-based exploration in the discrete case, and beneﬁts from additional generalization in practical environments, as we will see in Section 6. • Better reward shaping by utilizing goal examples to guide the agent more quickly and accurately towards the goal.

Original: We describe an efﬁcient technique to train these types of Bayesian classiﬁers via meta-learning in the following section.
Revised: We have established this beneﬁt intuitively, and will validate it empirically through extensive visualizations and experiments in Section 6.

Original:  B AY CRL: T RAINING B AYESIAN S UCCESS C LASSIFIERS FOR O UTCOME D RIVEN RL VIA M ETA -L EARNING AND CNML
Revised:  B AY CRL: T RAINING B AYESIAN S UCCESS C LASSIFIERS FOR O UTCOME D RIVEN RL VIA M ETA -L EARNING AND CNML

Original: In Section 4, we discussed how Bayesian success classiﬁers can incentivize exploration and provide reward shaping to guide RL.
Revised: In Section 4, we discussed how Bayesian success classiﬁers can incentivize exploration and provide reward shaping to guide RL.

Original: However, the reward inference technique via CNML described in
Revised: However, the reward inference technique via CNML described in

Original: Section 4.1 is computationally intractable, as it requires optimizing maximum likelihood estimation problems to convergence on every data point we want to query.
Revised: Section 4.1 is computationally intractable, as it requires optimizing maximum likelihood estimation problems to convergence on every data point we want to query.

Original: In this section, we describe a novel approximation that allows us to instantiate this method in practice.
Revised: In this section, we describe a novel approximation that allows us to instantiate this method in practice.

Original: 5. M ETA -L EARNING FOR CNML
Revised: 5. M ETA -L EARNING FOR CNML

Original: We adopt ideas from meta-learning to amortize the cost of obtaining the CNML distribution.
Revised: We adopt ideas from meta-learning to amortize the cost of obtaining the CNML distribution.

Original: As noted in Section 4.1, the computation of the CNML distribution involves repeatedly solving maximum likelihood problems.
Revised: As noted in Section 4.1, the computation of the CNML distribution involves repeatedly solving maximum likelihood problems.

Original: While computationally daunting, these problems share a signiﬁcant amount of common structure, which we should exploit to quickly obtain CNML estimates without redoing the slow training process repeatedly.
Revised: While computationally daunting, these problems share a signiﬁcant amount of common structure, which we can exploit to quickly obtain CNML estimates.

Original: One set of techniques that are directly applicable is meta-learning for few shot classiﬁcation.
Revised: One set of techniques that are directly applicable is meta-learning for few shot classiﬁcation.

Original: Meta-learning uses a distribution of training problems to explicitly learn models that can quickly adapt to new problems.
Revised: Meta-learning uses a distribution of training problems to explicitly learn models that can quickly adapt to new problems.

Original: To apply meta-learning to the CNML problem, we can formulate each of the maximum likelihood problems described in Equation 1 as a separate task for meta-learning, and apply any standard meta-learning technique to obtain a model capable of few-shot adaptation to the MLE problems required for CNML.
Revised: To apply meta-learning to the CNML problem, we can formulate each of the maximum likelihood problems described in Equation 1 as a separate task for meta-learning, and apply any standard meta-learning technique to obtain a model capable of few-shot adaptation to the MLE problems required for CNML.

Original: While any meta-learning algorithm is applicable, we found model agnostic meta-learning (MAML) to be an effective choice of algorithm.
Revised: While any meta-learning algorithm is applicable, we found model agnostic meta-learning (MAML)(Finn et al. to be an effective choice of algorithm.

Original: In short, MAML tries to learn a model that can quickly adapt to new tasks via a few steps of gradient descent.
Revised: In short, MAML tries to learn a model that can quickly adapt to new tasks via a few steps of gradient descent.

Original: We refer readers to Finn et al.

Original: (2017) for a detailed overview.

Original: This procedure is illustrated in Fig 10, and can be described as follows: given a dataset D
Revised: This procedure is illustrated in Fig 10, and can be described as follows: given a dataset D

Original: = { ( x 0 , y 0 ) , ( x 1 , y 1 ) , ..., ( x n , y n ) } , 2 n different tasks τ
Revised: = { ( x 0 , y 0 ) , ( x 1 , y 1 ) , ..., ( x n , y n ) } , 2 n different tasks τ

Original: i can be constructed, each corresponding to performing maximum likelihood estimation on the dataset with a certain proposed label for x i : max θ E( x,y )
Revised: i can be constructed, each corresponding to performing maximum likelihood estimation on the dataset with a certain proposed label for x i : max [Equation] E

Original: ∼D∪ ( x i ,y =0) log p ( y | x, θ ) or max θ E ( x,y ) ∼D∪ ( x i ,y =1) log p( y | x, θ ) .
Revised: [Equation] x,y ) ∼D∪ [Equation] x i ,y =0) [log p

Revised: [Equation]

Revised: y | x, θ )] or

Revised: [Equation]

Original:  Given these constructed tasks S ( τ ) , meta-training can be performed by performing the bi-level optimization described in Finn et al.
Revised: E [Equation] x,y ) ∼D∪ ( x i ,y =1) [log p ( y | x, θ )] . Given these constructed tasks S ( τ ) , meta-training as described in Finn et al.

Original: [Equation]
Revised: [Equation]

Original: This training procedure gives us parameters θ that can then be quickly adapted to provide the CNML distribution simply by performing a step of gradient descent.
Revised: This training procedure gives us parameters θ that can then be quickly adapted to provide the CNML distribution simply by performing a step of gradient descent.

Original: The model can be queried for the CNML distribution by starting from θ and taking one step of gradient descent for the query point augmented dataset, each with a different potential label.
Revised: The model can be queried for the CNML distribution by starting from θ and taking one step of gradient descent for the query point augmented dataset, each with a different potential label.

Original: These likelihoods are then normalized to provide the
Revised: These likelihoods are then normalized to provide the

Original: CNML distribution as follows:
Revised: CNML distribution as follows:

Original: [Equation]
Revised: [Equation]

Original: This algorithm, which we call meta-NML , allows us to obtain normalized likelihood estimates without having to retrain maximum likelihood to convergence at every single query point, since the model can now solve maximum likelihood problems of this form very quickly.
Revised: This algorithm, which we call meta-NML , allows us to obtain normalized likelihood estimates without having to retrain maximum likelihood to convergence at every single query point, since the model can now solve maximum likelihood problems of this form very quickly.

Original: A complete detailed description and pseudocode of this algorithm are provided in Appendix A.2.
Revised: A complete detailed description and pseudocode of this algorithm are provided in Appendix A.2.

Original: [Table]
Revised: [Table]

Original: [Table] [Table] Crucially, we ﬁnd that meta-NML is able to approximate the true NML outputs with just one or a few gradient steps.
Revised: Crucially, we ﬁnd that metaNML is able to approximate the true NML outputs with just one or a few gradient steps.

Original: This makes it several orders of magnitude faster than naive CNML, which would normally require multiple passes through the entire dataset on each input point in order to train to convergence.
Revised: This makes it several orders of magnitude faster than naive CNML, which would normally require multiple passes through the entire dataset on each input point in order to train to convergence . 5. A PPLYING M ETA -NML TO S UCCESS C LASSIFICATION

Original: Table 2 shows a comparison of runtimes for naive CNML and meta-NML on a simple two-dimensional dataset
Revised: Algorithm 1: User provides success examples S + 2: Initialize policy π , replay buffer S − , and reward classiﬁer parameters θ R 3: for iteration i = 1 , 2 , ... do 4:

Original: . 5. A PPLYING M ETA -NML TO S UCCESS C LASSIFICATION{+− by executing π . 5: if iteration i mod k == 
 then 6:+}
Revised: Collect on-policy examples to add to S {+− by executing π . 5: if iteration i mod k == 
 then 6:+}

Original: We apply the meta-NML algorithm described above to learning Bayesian success classiﬁers for providing rewards for reinforcement learning, in our proposed algorithm, which we term BayCRL— Bayesian classiﬁers for reinforcement learning.
Revised: Sample n train states from S − to create 2 n train metatraining tasks 7: Sample n test total test points equally from S + (label 1) and S − (label 0) 8: Meta-train θ R via meta-NML using Equation9: Assign state rewards via Equation10:

Revised: Train π with RL algorithm

Revised: We apply the meta-NML algorithm described above to learning Bayesian success classiﬁers for providing rewards for reinforcement learning, in our proposed algorithm, which we term BayCRL— Bayesian classiﬁers for reinforcement learning.

Original: Similarly to Fu et al.
Revised: Similarly to Fu et al.

Original: (2018b), we can train our Bayesian classiﬁer by ﬁrst constructing a dataset D for binary classiﬁcation.
Revised: (2018b), we can train our Bayesian classiﬁer by ﬁrst constructing a dataset D for binary classiﬁcation.

Original: This is done by using the provided examples of successful outcomes as positives, and on-policy examples collected by the policy as negatives, balancing the number of sampled positives and negatives in the dataset.
Revised: This is done by using the provided examples of successful outcomes as positives, and on-policy examples collected by the policy as negatives, balancing the number of sam- pled positives and negatives in the dataset.

Original: Given this dataset, the Bayesian classiﬁer parameters θ R can be trained via meta-NML as described in Equation 2.
Revised: Given this dataset, the Bayesian classiﬁer parameters θ R can be trained via meta-NML as described in Equation 2.

Original: The classiﬁer can then be used to directly and quickly assign rewards to a state s according to its probabilities r ( s ) =
Revised: The classiﬁer can then be used to directly and quickly assign rewards to a state s according to its probabilities r ( s ) =

Original: p meta-NML ( e = 1 | s ) (via a step of gradient descent, as described in Equation 4), and perform standard reinforcement learning.
Revised: p meta-NML ( e = 1 | s ) (via a step of gradient descent, as described in Equation 4), and perform standard reinforcement learning.

Original: [Equation]
Revised: [Equation]

Original: An overview of this algorithm is provided in Algorithm 1, and full details are in Appendix A.2.
Revised: An overview of this algorithm is provided in Algorithm 2, and full details are in Appendix A.2.

Original: The rewards start off at an uninformative value of 0 .
Revised: The rewards start off at an uninformative value of 0 .

Original: 5 for all unvisited states at the beginning, and close to 1 for successful outcomes.
Revised: 5 for all unvisited states at the beginning, and close to 1 for successful outcomes.

Original: As training progresses, more states are visited, added to the buffer and
Revised: As training progresses, more states are visited, added to the buffer and

Original: BayCRL starts to assign them progressively lower reward as they get visited more and more, thereby encouraging visiting of under-visited states.
Revised: BayCRL starts to assign them progressively lower reward as they get visited more and more, thereby encouraging visiting of under-visited states.

Original: At convergence, all the non successful states will have a reward of close to 0 and states at the goal will have a reward of 0 .
Revised: At convergence, all the non successful states will have a reward of close to 0 and states at the goal will have a reward of 0 .

Original: 5 , since the numbers of positive and negative labels for successful outcomes will be balanced as described above.
Revised: 5 , since the numbers of positive and negative labels for successful outcomes will be balanced as described above.

Original:  E XPERIMENTAL E VALUATION
Revised:  E XPERIMENTAL E VALUATION

Original: In our experimental evaluation we aim to answer the following questions: (1)
Revised: In our experimental evaluation we aim to answer the following questions: (1)

Original: Do the learning dynamics of prior classiﬁer-based reward learning methods provide informative rewards for
Revised: Do the learning dynamics of prior classiﬁer-based reward learning methods provide informative rewards for

Original: RL? (2) Does using BayCRL help address the exploration challenge when solving RL problems speciﬁed by successful outcomes?
Revised: RL? (2) Does using BayCRL help address the exploration challenge when solving RL problems speciﬁed by successful outcomes?

Original: (3) Does using BayCRL help provide better reward shaping than simply performing naïvely uninformed exploration?
Revised: (3) Does using BayCRL help provide better reward shaping than simply performing naïvely uninformed exploration?

Original: To evaluate these questions, we evaluate our proposed algorithm BayCRL with the following setup.
Revised: To evaluate these questions, we evaluate our proposed algorithm BayCRL with the following setup.

Original: Further details and videos can be found at https://sites.google.com/view/baycrl/home
Revised: Further details and videos can be found at https://sites.google.com/view/baycrl/home

Original: 6. E XPERIMENTAL S ETUP
Revised: 6. E XPERIMENTAL S ETUP

Original: [Figure]
Revised: [Figure]

Original: We start off by understanding the algorithm behavior by evaluating it on maze navigation problems, which require avoiding several local optima before truly reaching the goal.
Revised: We start off by understanding the algorithm behavior by evaluating it on maze navigation problems, which require avoiding several local optima before truly reaching the goal.

Original: Then, to evaluate our method in more complex domains, we consider three robotic manipulation tasks that were previously covered in Singh et al. (2019a) with a Sawyer robot arm: door opening, tabletop object pushing, and 3D object picking.
Revised: Then, to evaluate our method in more complex domains, we consider three robotic manipulation tasks that were previously covered in Singh et al. (2019a) with a Sawyer robot arm: door opening, tabletop object pushing, and 3D object picking.

Original: As we show in our results, exploration in these environments is challenging and using naively chosen reward shaping often does not solve the problem at hand.
Revised: As we show in our results, exploration in these environments is challenging and using naively chosen reward shaping often does not solve the problem at hand.

Original: More details on each environment and their associated challenges are available in Appendix A.4.1.
Revised: More details on each environment and their associated challenges are available in Appendix A.4.1.

Original: We compare with a number of prior algorithms and ablations.
Revised: We compare with a number of prior algorithms and ablations.

Original: To provide a comparison with a standard previous method which uses success classiﬁers trained with an IRL-based adversarial method, we include the VICE algorithm (Fu et al., 2018b).
Revised: To provide a comparison with a standard previous method which uses success classiﬁers trained with an IRL-based adversarial method, we include the VICE algorithm (Fu et al., 2018b).

Original: Note that this algorithm is quite related to BayCRL, but it uses a standard maximum likelihood classiﬁer rather than a Bayesian classiﬁer trained with CNML and meta-learning.
Revised: Note that this algorithm is quite related to BayCRL, but it uses a standard maximum likelihood classiﬁer rather than a Bayesian classiﬁer trained with CNML and meta-learning.

Original: We also include a comparison with DDL, a recently proposed technique for learning dynamical distances (Hartikainen et al., 2019).
Revised: We also include a comparison with DDL, a recently proposed technique for learning dynamical distances (Hartikainen et al., 2019).

Original: We additionally include comparisons to algorithms for uninformed exploration to show that BayCRL does a more directed form of exploration and reward shaping.
Revised: We additionally include comparisons to algorithms for uninformed exploration to show that BayCRL does a more directed form of exploration and reward shaping.

Original: To provide an apples-to-apples comparison, we use the same VICE method for training classiﬁers, but combine it with novelty-based exploration based on random network distillation (Burda et al., 2018b) for the robotic manipulation tasks, and oracle inverse count bonuses for the maze navigation tasks.
Revised: To provide an apples-to-apples comparison, we use the same VICE method for training classiﬁers, but combine it with novelty-based exploration based on random network distillation (Burda et al., 2018b) for the robotic manipulation tasks, and oracle inverse count bonuses for the maze navigation tasks.

Original: Finally, to demonstrate the importance of well-shaped rewards, we compare to running Soft Actor-Critic (Haarnoja et al., 2018), a standard RL algorithm for continuous domains, with two naive reward functions: a sparse reward at the goal, and a heuristically shaped reward which uses L2 distance to the goal state.
Revised: Finally, to demonstrate the importance of well-shaped rewards, we compare to running Soft Actor-Critic (Haarnoja et al., 2018), a standard RL algorithm for continuous domains, with two naive reward functions: a sparse reward at the goal, and a heuristically shaped reward which uses L2 distance to the goal state.

Original: More details on each algorithm and the hyperparameters used are included in Appendix A.6.
Revised: More details on each algorithm and the hyperparameters used are included in Appendix A.6.

Original: 6. C OMPARISONS WITH P RIOR A LGORITHMS
Revised: 6. C OMPARISONS WITH P RIOR A LGORITHMS

Original: We compare with prior algorithms on the domains described above.
Revised: We compare with prior algorithms on the domains described above.

Original: As we can see in Fig 5, BayCRL is able to very quickly learn how to solve these challenging exploration tasks, often reaching better asymptotic performance than most prior methods, and doing so more efﬁciently than VICE (Fu et al., 2018b) or DDL (Hartikainen et al., 2019).
Revised: As we can see in Fig 5, BayCRL is able to very quickly learn how to solve these challenging exploration tasks, often reaching better asymptotic performance than most prior methods, and doing so more efﬁciently than VICE (Fu et al., 2018b) or DDL (Hartikainen et al., 2019).

Original: This suggests that BayCRL is able to provide directed reward shaping and exploration that is substantially better than standard classiﬁer-based methods (e.g., VICE).
Revised: This suggests that BayCRL is able to provide directed reward shaping and exploration that is substantially better than standard classiﬁer-based methods (e.g., VICE).

Original: To isolate whether the beneﬁts purely come from exploration or also from task-aware reward shaping, we compare with methods that only perform uninformed, task-agnostic exploration.
Revised: To isolate whether the beneﬁts purely come from exploration or also from task-aware reward shaping, we compare with methods that only perform uninformed, task-agnostic exploration.

Original: On the maze environments, where we can discretize the state space, we compute ground truth count-based bonuses for exploration.
Revised: On the maze environments, where we can discretize the state space, we compute ground truth count-based bonuses for exploration.

Original: For the higher dimensional robotics tasks, we use RND (Burda et al., 2018b).
Revised: For the higher dimensional robotics tasks, we use RND (Burda et al., 2018b).

Original: From these comparisons, shown in Fig 5, it is clear that
Revised: From these comparisons, shown in Fig 5, it is clear that

Original: BayCRL signiﬁcantly outperforms methods that use novelty-seeking exploration, but do not otherwise provide effective reward shaping.
Revised: BayCRL signiﬁcantly outperforms methods that use novelty-seeking exploration, but do not otherwise provide effective reward shaping.

Original: In combination with our visualizations in Section 6.4, this suggests that BayCRL is providing useful task-aware reward shaping more effectively than uniformed exploration methods.
Revised: In combination with our visualizations in Section 6.4, this suggests that BayCRL is providing useful task-aware reward shaping more effectively than uniformed exploration methods.

Original: We also compare
Revised: We also compare

Original: BayCRL to a manually heuristically-designed shaped reward function, based on Euclidean distance.
Revised: BayCRL to a manually heuristically-designed shaped reward function, based on Euclidean distance.

Original: As shown in Fig 5, BayCRL generally outperforms simple manual shaping in terms of sample complexity and asymptotic performance, indicating that the learned shaping is non-trivial and adapted to the task.
Revised: As shown in Fig 5, BayCRL generally outperforms simple manual shaping in terms of sample complexity and asymptotic performance, indicating that the learned shaping is non-trivial and adapted to the task.

Original: 6. A BLATIONS
Revised: 6. A BLATIONS

Original: We ﬁrst evaluate the importance of meta-learning for estimating the NML distribution.
Revised: We ﬁrst evaluate the importance of meta-learning for estimating the NML distribution.

Original: In Figure 6, we see that naively estimating the NML distribution by taking a single gradient step and following the same process as evaluating meta-NML, but without any meta-training, results in much worse performance.
Revised: In Figure 6, we see that naively estimating the NML distribution by taking a single gradient step and following the same process as evaluating meta-NML, but without any meta-training, results in much worse performance.

Original: Second, we analyze the importance of making the BayCRL classiﬁer aware of the task being solved, to understand whether BayCRL is informed by the success examples or simply approximates count-based exploration.
Revised: Second, we analyze the importance of making the BayCRL classiﬁer aware of the task being solved, to understand whether BayCRL is informed by the success examples or simply approximates count-based exploration.

Original: To that end, we modify the training procedure so that the dataset D consists of only the on-policy negatives, and add the inferred reward from the Bayesian classiﬁer to the reward obtained by a standard MLE classiﬁer (similarly to the VICE+RND baseline).
Revised: To that end, we modify the training procedure so that the dataset D consists of only the on-policy negatives, and add the inferred reward from the Bayesian classiﬁer to the reward obtained by a standard MLE classiﬁer (similarly to the VICE+RND baseline).

Original: We see that this performs poorly, showing that the BayCRL classiﬁer is doing more than just performing count-based exploration, and beneﬁts from better reward shaping due to the provided goal examples.
Revised: We see that this performs poorly, showing that the BayCRL classiﬁer is doing more than just performing count-based exploration, and beneﬁts from better reward shaping due to the provided goal examples.

Original: Further ablations are available in Appendix A.5.
Revised: Further ablations are available in Appendix A.5.

Original: 6. A NALYSIS OF B AY CRL
Revised: 6. A NALYSIS OF B AY CRL

Original: BayCRL and Reward Shaping.
Revised: BayCRL and Reward Shaping.

Original: To better understand how BayCRL provides reward shaping, we visualize the rewards for various slices along the z axis on the Sawyer Pick task, an environment which presents a signiﬁcant exploration challenge.
Revised: To better understand how BayCRL provides reward shaping, we visualize the rewards for various slices along the z axis on the Sawyer Pick task, an environment which presents a signiﬁcant exploration challenge.

Original: In Fig 7 we see that the BayCRL rewards clearly correlate with the distance to the object’s goal position, shown as a white star, thus guiding the robot to raise the ball to the desired location even if it has never reached the goal before.
Revised: In Fig 7 we see that the BayCRL rewards clearly correlate with the distance to the object’s goal position, shown as a white star, thus guiding the robot to raise the ball to the desired location even if it has never reached the goal before.

Original: In contrast, the
Revised: In contrast, the

Original: MLE classiﬁer has a sharp, poorly-shaped decision boundary.
Revised: MLE classiﬁer has a sharp, poorly-shaped decision boundary.

Original: BayCRL and Exploration.
Revised: BayCRL and Exploration.

Original: Next, to illustrate the connection between BayCRL and exploration, we compare the states visited by BayCRL (which uses a meta-NML classiﬁer) and by VICE (which uses a standard L2-regularized classiﬁer) in Figure 8.
Revised: Next, to illustrate the connection between BayCRL and exploration, we compare the states visited by BayCRL (which uses a meta-NML classiﬁer) and by VICE (which uses a standard L2-regularized classiﬁer) in Figure 8.

Original: We see that BayCRL naturally incentivizes the agent to visit novel states, allowing it to navigate around local minima and reach the true goal.
Revised: We see that BayCRL naturally incentivizes the agent to visit novel states, allowing it to navigate around local minima and reach the true goal.

Original: In contrast, VICE learns a misleading reward function that prioritizes closeness to the goal in xy space, causing the agent to stay on the wrong side of the wall.
Revised: In contrast, VICE learns a misleading reward function that prioritizes closeness to the goal in xy space, causing the agent to stay on the wrong side of the wall.

Original: Interestingly, despite incentivizing exploration, BayCRL does not simply visit all possible states; at convergence, it has only covered around 70% of the state space.
Revised: Interestingly, despite incentivizing exploration, BayCRL does not simply visit all possible states; at convergence, it has only covered around 70% of the state space.

Original: In fact, we see in the scatterplots in Figure 8 that BayCRL prioritizes states that bring it closer to the goal and ignores ones that don’t, thus making use of the goal examples provided to it.
Revised: In fact, we see in the scatterplots in Figure 8 that BayCRL prioritizes states that bring it closer to the goal and ignores ones that don’t, thus making use of the goal examples provided to it.

Original: This suggests that BayCRL beneﬁts from a combination of novelty-seeking behavior and effective reward shaping, allowing it to choose new states strategically.
Revised: This suggests that BayCRL beneﬁts from a combination of novelty-seeking behavior and effective reward shaping, allowing it to choose new states strategically.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this work, we consider a subclass of reinforcement learning problems where examples of successful outcomes specify the task.
Revised: In this work, we consider a subclass of reinforcement learning problems where examples of successful outcomes specify the task.

Original: We analyze how solutions via standard success classiﬁers suffer from shortcomings, and training Bayesian classiﬁers allows for better exploration to solve challenging problems.
Revised: We analyze how solutions via standard success classiﬁers suffer from shortcomings, and training Bayesian classiﬁers allows for better exploration to solve challenging problems.

Original: We discuss how the NML distribution can provide us a way to train such Bayesian classiﬁers, providing beneﬁts of exploration and reward shaping.
Revised: We discuss how the NML distribution can provide us a way to train such Bayesian classiﬁers, providing beneﬁts of exploration and reward shaping.

Original: To make learning tractable, we propose a novel meta-learning approach to amortize the NML process.
Revised: To make learning tractable, we propose a novel meta-learning approach to amortize the NML process.

Original: While this work has shown the effectiveness of Bayesian classiﬁers for reward inference for tasks in simulation, it would be interesting to scale this solution to real world problems.
Revised: While this work has shown the effectiveness of Bayesian classiﬁers for reward inference for tasks in simulation, it would be interesting to scale this solution to real world problems.

Original: Additionally, obtaining a theoretical understanding of how reward shaping interacts with learning dynamics would be illuminating in designing reward schemes.
Revised: Additionally, obtaining a theoretical understanding of how reward shaping interacts with learning dynamics would be illuminating in designing reward schemes.

Original: A. D ETAILED D ESCRIPTION OF M ETA -NML
Revised: A. D ETAILED D ESCRIPTION OF M ETA -NML

Original: We provide a detailed description of the meta-NML algorithm described in Section 5, and the details of the practical algorithm.
Revised: We provide a detailed description of the meta-NML algorithm described in Section 5, and the details of the practical algorithm.

Original: Given a dataset D = { ( x 0 , y 0 ) , ( x 1 , y 1 ) , .., ( x n , y n ) } , the meta-NML procedure proceeds by ﬁrst constructing k ∗ n tasks from these data points, for a k shot classiﬁcation problem.
Revised: Given a dataset D = { ( x 0 , y 0 ) , ( x 1 , y 1 ) , .., ( x n , y n ) } , the meta-NML procedure proceeds by ﬁrst constructing k ∗ n tasks from these data points, for a k shot classiﬁcation problem.

Original: We will keep k = 2 for simplicity in this description, in accordance with the setup of binary success classiﬁers in RL.
Revised: We will keep k = 2 for simplicity in this description, in accordance with the setup of binary success classiﬁers in RL.

Original: Each task τ
Revised: Each task τ

Original: i is constructed by augmenting the dataset with a negative label D ∪ ( x i , y = 0) or a positive label D ∪ ( x i , y = 1) .
Revised: i is constructed by augmenting the dataset with a negative label D ∪ ( x i , y = 0) or a positive label D ∪ ( x i , y = 1) .

Original: Now that each task consists of solving the maximum likelihood problem for its augmented dataset, we can directly apply standard meta-learning algorithms to this setting.
Revised: Now that each task consists of solving the maximum likelihood problem for its augmented dataset, we can directly apply standard meta-learning algorithms to this setting.

Original: Building off the ideas in MAML (Finn et al., 2017), we can then train a set of model parameters θ such that after a single step of gradient descent it can quickly adapt to the optimal solution for the MLE problem on any of the augmented datasets.
Revised: Building off the ideas in MAML (Finn et al., 2017), we can then train a set of model parameters θ such that after a single step of gradient descent it can quickly adapt to the optimal solution for the MLE problem on any of the augmented datasets.

Original: This is more formally written as
Revised: This is more formally written as

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ) where L represents a standard classiﬁcation loss function, α is the learning rate, and the distribution of tasks p ( τ ) is constructed as described above.
Revised: ) where L represents a standard classiﬁcation loss function, α is the learning rate, and the distribution of tasks p ( τ ) is constructed as described above.

Original: For a new query point x , these initial parameters can then quickly be adapted to provide the CNML distribution by taking a gradient step on each augmented dataset to obtain the approximately optimal MLE solution, and normalizing these as follows:
Revised: For a new query point x , these initial parameters can then quickly be adapted to provide the CNML distribution by taking a gradient step on each augmented dataset to obtain the approximately optimal MLE solution, and normalizing these as follows:

Original: [Equation]
Revised: [Equation]

Original: This algorithm in principle can be optimized using any standard stochastic optimization method such as SGD, as described in Finn et al.
Revised: This algorithm in principle can be optimized using any standard stochastic optimization method such as SGD, as described in Finn et al.

Original: (2017), backpropagating through the inner loop gradient update.
Revised: (2017), backpropagating through the inner loop gradient update.

Original: For the speciﬁc problem setting that we consider, we have to employ some optimization tricks in order to enable learning:
Revised: For the speciﬁc problem setting that we consider, we have to employ some optimization tricks in order to enable learning:

Original: A.2. I MPORTANCE W EIGHTING ON Q UERY P OINT
Revised: A.2. I MPORTANCE W EIGHTING ON Q UERY P OINT

Original: Since only one datapoint is augmented to the training set at query time for CNML, it can get challenging for stochastic gradient descent to pay attention to this datapoint with increasing dataset sizes.
Revised: Since only one datapoint is augmented to the training set at query time for CNML, it can get challenging for stochastic gradient descent to pay attention to this datapoint with increasing dataset sizes.

Original: For example, if we train on an augmented dataset of size 2048 by cycling through it in batch sizes of 32, then only 1 in 64 batches would include the query point itself and allow the model to adapt to the proposed label, while the others would lead to noise in the optimization process, potentially worsening the model’s prediction on the query point.
Revised: For example, if we train on an augmented dataset of size 2048 by cycling through it in batch sizes of 32, then only 1 in 64 batches would include the query point itself and allow the model to adapt to the proposed label, while the others would lead to noise in the optimization process, potentially worsening the model’s prediction on the query point.

Original: In order to make sure the optimization considers the query point, we include the query point and proposed label ( x q , y ) in every minibatch that is sampled, but downweight the loss computed on that point such that the overall objective remains unbiased.
Revised: In order to make sure the optimization considers the query point, we include the query point and proposed label ( x q , y ) in every minibatch that is sampled, but downweight the loss computed on that point such that the overall objective remains unbiased.

Original: This is simply doing importance weighting, with the query point downweighted by a factor of (cid:100) b − 1 N (cid:101) where b is the desired batch size and N is the total number of points in the original dataset.
Revised: This is simply doing importance weighting, with the query point downweighted by a factor of (cid:100) b − 1 N (cid:101) where b is the desired batch size and N is the total number of points in the original dataset.

Original: To see why the optimization objective remains the same, we can consider the overall loss over the dataset.
Revised: To see why the optimization objective remains the same, we can consider the overall loss over the dataset.

Original: Let f θ be our classiﬁer, L be our loss function, D (cid:48)
Revised: Let f θ be our classiﬁer, L be our loss function, D (cid:48)

Original: = { ( x i , y i ) }
Revised: = { ( x i , y i ) }

Original: Ni =1 ∪ ( x q , y ) be our augmented dataset, and B k be the k th batch seen during training.
Revised: Ni =1 ∪ ( x q , y ) be our augmented dataset, and B k be the k th batch seen during training.

Original: Using standard SGD training that cycles through batches in the dataset, the overall loss on the augmented dataset would be:
Revised: Using standard SGD training that cycles through batches in the dataset, the overall loss on the augmented dataset would be:

Original: [Equation]
Revised: [Equation]

Original: If we instead included the downweighted query point in every batch, the overall loss would be:
Revised: If we instead included the downweighted query point in every batch, the overall loss would be:

Original: [Equation]
Revised: [Equation]

Original: which is the same objective as before.
Revised: which is the same objective as before.

Original: This trick has the effect of still optimizing the same max likelihood problem required by CNML, but signiﬁcantly reducing the variance of the query point predictions as we take additional gradient steps at query time.
Revised: This trick has the effect of still optimizing the same max likelihood problem required by CNML, but signiﬁcantly reducing the variance of the query point predictions as we take additional gradient steps at query time.

Original: As a concrete example, consider querying a meta-CNML classiﬁer on the input shown in Figure 11.
Revised: As a concrete example, consider querying a meta-CNML classiﬁer on the input shown in Figure 11.

Original: If we adapt to the augmented dataset without including the query point in every batch (i.e. without importance weighting), we see that the query point loss is signiﬁcantly more unstable, requiring us to take
Revised: If we adapt to the augmented dataset without including the query point in every batch (i.e. without importance weighting), we see that the query point loss is signiﬁcantly more unstable, requiring us to take

Original: A.2. K ERNEL W EIGHTED T RAINING L OSS
Revised: A.2. K ERNEL W EIGHTED T RAINING L OSS

Original: The augmented dataset consists of points from the original dataset D and one augmented point ( x q , y ) .
Revised: The augmented dataset consists of points from the original dataset D and one augmented point ( x q , y ) .

Original: Given that we mostly care about having the proper likelihood on the query point, with an imperfect optimization process, the meta-training can yield solutions that are not very accurately representing true likelihoods on the query point.
Revised: Given that we mostly care about having the proper likelihood on the query point, with an imperfect optimization process, the meta-training can yield solutions that are not very accurately representing true likelihoods on the query point.

Original: To counter this, we introduce a kernel weighting into the loss function in Equation 6 during meta-training and subsequently meta-testing.
Revised: To counter this, we introduce a kernel weighting into the loss function in Equation 6 during meta-training and subsequently meta-testing.

Original: The kernel weighting modiﬁes the training loss function as:
Revised: The kernel weighting modiﬁes the training loss function as:

Original: [Equation]
Revised: [Equation]

Original: where x τ is the query point for task τ and K is a choice of kernel.
Revised: where x τ is the query point for task τ and K is a choice of kernel.

Original: We typically choose exponential kernels centered around x τ .
Revised: We typically choose exponential kernels centered around x τ .

Original: Intuitively, this allows the meta-optimization to mainly consider the datapoints that are copies of the query point in the dataset, or are similar to the query point, and ensures that they have the correct likelihoods, instead of receiving interfering gradient signals from the many other points in the dataset.
Revised: Intuitively, this allows the meta-optimization to mainly consider the datapoints that are copies of the query point in the dataset, or are similar to the query point, and ensures that they have the correct likelihoods, instead of receiving interfering gradient signals from the many other points in the dataset.

Original: To make hyperparameter selection intuitive, we designate the strength of the exponential kernel by a parameter λ dist , which is the Euclidean distance away from the query point at which the weight becomes 0.1.
Revised: To make hyperparameter selection intuitive, we designate the strength of the exponential kernel by a parameter λ dist , which is the Euclidean distance away from the query point at which the weight becomes 0.1.

Original: Formally, the weight of a point x in the loss function for query point x τ is computed as:
Revised: Formally, the weight of a point x in the loss function for query point x τ is computed as:

Original: [Equation]
Revised: [Equation]

Original: A.2. M ETA -T RAINING AT F IXED I NTERVALS
Revised: A.2. M ETA -T RAINING AT F IXED I NTERVALS

Original: While in principle meta-NML would retrain with every new datapoint, in practice we retrain metaNML
Revised: While in principle meta-NML would retrain with every new datapoint, in practice we retrain metaNML

Original: once every k epochs.
Revised: once every k epochs.

Original: (In all of our experiments we set k = 1 , but we could optionally increase k if we do not expect the meta-task distribution to change much between epochs.)
Revised: (In all of our experiments we set k = 1 , but we could optionally increase k if we do not expect the meta-task distribution to change much between epochs.)

Original: We warm-start the meta-learner parameters from the previous iteration of meta-learning, so every instance of meta-training only requires a few steps.
Revised: We warm-start the meta-learner parameters from the previous iteration of meta-learning, so every instance of meta-training only requires a few steps.

Original: We ﬁnd that this periodic training is a reasonable enough approximation, as evidenced by the strong performance of BayCRL in our experimental results in
Revised: We ﬁnd that this periodic training is a reasonable enough approximation, as evidenced by the strong performance of BayCRL in our experimental results in

Original: Section 6.
Revised: Section 6.

Original: A. M ETA -NML V ISUALIZATIONS A.3. M ETA -NML WITH A DDITIONAL G RADIENT S TEPS
Revised: A. M ETA -NML V ISUALIZATIONS A.3. M ETA -NML WITH A DDITIONAL G RADIENT S TEPS

Original: Below, we show a more detailed visualization of meta-NML outputs on data from the Zigzag Maze task, and how these outputs change with additional gradient steps.
Revised: Below, we show a more detailed visualization of meta-NML outputs on data from the Zigzag Maze task, and how these outputs change with additional gradient steps.

Original: For comparison, we also include the idealized NML rewards, which come from a discrete count-based classiﬁer.
Revised: For comparison, we also include the idealized NML rewards, which come from a discrete count-based classiﬁer.

Original: Meta-NML is able to resemble the ideal NML rewards fairly well with just 1 gradient step, providing both an approximation of a count-based exploration bonus and better shaping towards the goal due to generalization.
Revised: Meta-NML is able to resemble the ideal NML rewards fairly well with just 1 gradient step, providing both an approximation of a count-based exploration bonus and better shaping towards the goal due to generalization.

Original: By taking additional gradient steps, meta-NML can get arbitrarily close to the true
Revised: By taking additional gradient steps, meta-NML can get arbitrarily close to the true

Original: NML outputs, which themselves correspond to inverse counts of 1 n +2 as explained in Theorem 4.1.
Revised: NML outputs, which themselves correspond to inverse counts of 1 n +2 as explained in Theorem 4.1.

Original: While this would give us more accurate NML estimates, in practice we found that taking one gradient step was sufﬁcient to achieve good performance on our RL tasks.
Revised: While this would give us more accurate NML estimates, in practice we found that taking one gradient step was sufﬁcient to achieve good performance on our RL tasks.

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

Original: A.3. R UNTIME C OMPARISONS
Revised: A.3. R UNTIME C OMPARISONS

Original: Below provide the runtimes for feedforward inference, naive CNML, and meta-NML on each of our evaluation domains.
Revised: Below provide the runtimes for feedforward inference, naive CNML, and meta-NML on each of our evaluation domains.

Original: We list both the runtimes for evaluating a single input, and for completing a full epoch of training during RL.
Revised: We list both the runtimes for evaluating a single input, and for completing a full epoch of training during RL.

Original: These benchmarks were performed on an NVIDIA Titan X Pascal GPU.
Revised: These benchmarks were performed on an NVIDIA Titan X Pascal GPU.

Original: Per-input runtimes are averaged across 100 samples, and per-epoch runtimes are averaged across 20 epochs.
Revised: Per-input runtimes are averaged across 100 samples, and per-epoch runtimes are averaged across 20 epochs.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Mazes (zigzag, spiral)
Revised: Mazes (zigzag, spiral)

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: A. E XPERIMENTAL D ETAILS A.4. E NVIRONMENTS
Revised: A. E XPERIMENTAL D ETAILS A.4. E NVIRONMENTS

Original: Zigzag Maze and Spiral Maze: These two navigation tasks require moving through long corridors and avoiding several local optima in order to reach the goal.
Revised: Zigzag Maze and Spiral Maze: These two navigation tasks require moving through long corridors and avoiding several local optima in order to reach the goal.

Original: For example, on Spiral Maze, the agent must not get stuck on the other side of the inner wall, even though that position would be close in L  distance to the desired goal.
Revised: For example, on Spiral Maze, the agent must not get stuck on the other side of the inner wall, even though that position would be close in L  distance to the desired goal.

Original: On these tasks, a sparse reward is not informative enough for learning, while ordinary classiﬁer methods get stuck in local optima due to poor shaping near the goal.
Revised: On these tasks, a sparse reward is not informative enough for learning, while ordinary classiﬁer methods get stuck in local optima due to poor shaping near the goal.

Original: Both of these environments have a continuous state space consisting of the ( x, y ) coordinates of the agent, ranging from ( − 4 , − 4) to (4 , 4) inclusive.
Revised: Both of these environments have a continuous state space consisting of the ( x, y ) coordinates of the agent, ranging from ( − 4 , − 4) to (4 , 4) inclusive.

Original: The action space is the desired velocity in the x and y directions, each ranging from − 1 to 1 inclusive.
Revised: The action space is the desired velocity in the x and y directions, each ranging from − 1 to 1 inclusive.

Original: Sawyer 2D Pusher:
Revised: Sawyer 2D Pusher:

Original: This task involves using a Sawyer arm, constrained to move only in the xy plane, to push a randomly initialized puck to a ﬁxed location on a table.
Revised: This task involves using a Sawyer arm, constrained to move only in the xy plane, to push a randomly initialized puck to a ﬁxed location on a table.

Original: The state space consists of the ( x, y, z ) coordinates of the robot end effector and the ( x, y ) coordinates of the puck.
Revised: The state space consists of the ( x, y, z ) coordinates of the robot end effector and the ( x, y ) coordinates of the puck.

Original: The action space is the desired x and y velocities of the arm.
Revised: The action space is the desired x and y velocities of the arm.

Original: Sawyer Door Opening:
Revised: Sawyer Door Opening:

Original: In this task, the Sawyer arm is attached to a hook, which it must use to open a door to a desired angle of 45 degrees.
Revised: In this task, the Sawyer arm is attached to a hook, which it must use to open a door to a desired angle of 45 degrees.

Original: The door is randomly initialized each time to be at a starting angle of between 0 and 15 degrees.
Revised: The door is randomly initialized each time to be at a starting angle of between 0 and 15 degrees.

Original: The state space consists of the ( x, y, z ) coordinates of the end effector and the door angle (in radians); the action space consists of ( x, y, z ) velocities.
Revised: The state space consists of the ( x, y, z ) coordinates of the end effector and the door angle (in radians); the action space consists of ( x, y, z ) velocities.

Original: Sawyer 3D Pick and Place: The Sawyer robot must pick up a ball, which is randomly placed somewhere on the table each time, and raise it to a ﬁxed ( x, y, z ) location high above the table.
Revised: Sawyer 3D Pick and Place: The Sawyer robot must pick up a ball, which is randomly placed somewhere on the table each time, and raise it to a ﬁxed ( x, y, z ) location high above the table.

Original: This represents the biggest exploration challenge out of all the manipulation tasks, as the state space is large and the agent would normally not receive any learning signal unless it happened to pick up the ball and raise it, which is unlikely without careful reward shaping.
Revised: This represents the biggest exploration challenge out of all the manipulation tasks, as the state space is large and the agent would normally not receive any learning signal unless it happened to pick up the ball and raise it, which is unlikely without careful reward shaping.

Original: The state space consists of the ( x, y, z ) coordinates of the end effector, the ( x, y, z ) coordinates of the ball, and the tightness of the gripper (a continuous value between 0 and 1).
Revised: The state space consists of the ( x, y, z ) coordinates of the end effector, the ( x, y, z ) coordinates of the ball, and the tightness of the gripper (a continuous value between 0 and 1).

Original: The robot can control its ( x, y, z ) arm velocity as well as the gripper value.
Revised: The robot can control its ( x, y, z ) arm velocity as well as the gripper value.

Original: [Equation]
Revised: [Equation]

Original: In addition to the success rate plots in Figure 5, we provide plots of each algorithm’s distance to the goal over time according to environment-speciﬁc distance metrics.
Revised: In addition to the success rate plots in Figure 5, we provide plots of each algorithm’s distance to the goal over time according to environment-speciﬁc distance metrics.

Original: The distance metrics and success thresholds, which were used to compute the success rates in Figure 5, are listed in the table below.
Revised: The distance metrics and success thresholds, which were used to compute the success rates in Figure 5, are listed in the table below.

Original: Environment
Revised: Environment

Original: Distance Metric Used
Revised: Distance Metric Used

Original: Success Threshold
Revised: Success Threshold

Original: Zigzag Maze
Revised: Zigzag Maze

Original: Manhattan distance to goal 0.
Revised: Manhattan distance to goal 0.

Original: Spiral Maze
Revised: Spiral Maze

Original: Manhattan distance to goal 0.
Revised: Manhattan distance to goal 0.

Original: Sawyer 2D Pusher
Revised: Sawyer 2D Pusher

Original: Puck L2 distance to goal 0.
Revised: Puck L2 distance to goal 0.

Original: Sawyer Door Opening
Revised: Sawyer Door Opening

Original: Angle difference to goal (radians) 0.
Revised: Angle difference to goal (radians) 0.

Original: Sawyer 3D Pick-and-Place
Revised: Sawyer 3D Pick-and-Place

Original: Ball L2 distance to goal 0.
Revised: Ball L2 distance to goal 0.

Original: A. A DDITIONAL A BLATIONS A.5. L EARNING IN A D ISCRETE , R ANDOMIZED E NVIRONMENT
Revised: A. A DDITIONAL A BLATIONS A.5. L EARNING IN A D ISCRETE , R ANDOMIZED E NVIRONMENT

Original: In practice, many continuous RL environments such as the ones we consider in Section 6 have state spaces that are correlated at least roughly with the dynamics.
Revised: In practice, many continuous RL environments such as the ones we consider in Section 6 have state spaces that are correlated at least roughly with the dynamics.

Original: For instance, states that are closer together dynamically are also typically closer in the metric space deﬁned by the states.
Revised: For instance, states that are closer together dynamically are also typically closer in the metric space deﬁned by the states.

Original: This correlation does not need to be perfect, but as long as it exists, BayCRL can in principle learn a smoothly shaped reward towards the goal.
Revised: This correlation does not need to be perfect, but as long as it exists, BayCRL can in principle learn a smoothly shaped reward towards the goal.

Original: However, even in the case where states are unstructured and completely lack identity, such as in a discrete gridworld environment, the CNML classiﬁer would still reduce to providing an explorationcentric reward bonus, as indicated by Theorem 4.1, ensuring reasonable worst-case performance.
Revised: However, even in the case where states are unstructured and completely lack identity, such as in a discrete gridworld environment, the CNML classiﬁer would still reduce to providing an explorationcentric reward bonus, as indicated by Theorem 4.1, ensuring reasonable worst-case performance.

Original: To demonstrate this, we evaluate BayCRL on a variant of the Zigzag Maze task where states are ﬁrst discretized to a 16 × 16 grid, then "shufﬂed" so that the xy representation of a state does not correspond to its true coordinates and the states are not correlated dynamically.
Revised: To demonstrate this, we evaluate BayCRL on a variant of the Zigzag Maze task where states are ﬁrst discretized to a 16 × 16 grid, then "shufﬂed" so that the xy representation of a state does not correspond to its true coordinates and the states are not correlated dynamically.

Original: BayCRL manages to solve the task, while a standard classiﬁer method (VICE) does not.
Revised: BayCRL manages to solve the task, while a standard classiﬁer method (VICE) does not.

Original: Still, BayCRL is more effective in the original state space where generalization is possible, suggesting that both the exploration and reward shaping abilities of the CNML classiﬁer are crucial to its overall performance.
Revised: Still, BayCRL is more effective in the original state space where generalization is possible, suggesting that both the exploration and reward shaping abilities of the CNML classiﬁer are crucial to its overall performance.

Original: [Figure]
Revised: [Figure]

Original: Only the goal examples in the bot- tom left corner are provided to the algorithm.
Revised: Only the goal examples in the bot- tom left corner are provided to the algorithm.

Original: The intended setup for BayCRL (and classiﬁer-based RL algorithms in general) is to provide a set of success examples to learn from, thus removing the need for a manually speciﬁed reward function.
Revised: The intended setup for BayCRL (and classiﬁer-based RL algorithms in general) is to provide a set of success examples to learn from, thus removing the need for a manually speciﬁed reward function.

Original: However, here we instead consider the case where a ground truth reward function exists which we do not fully know, and can only query through interaction with the environment.
Revised: However, here we instead consider the case where a ground truth reward function exists which we do not fully know, and can only query through interaction with the environment.

Original: In this case, because the human expert has limited knowledge, the provided success examples may not cover all regions of the state space with high reward.
Revised: In this case, because the human expert has limited knowledge, the provided success examples may not cover all regions of the state space with high reward.

Original: An additional advantage of BayCRL is that it is still capable of ﬁnding these "unspeciﬁed" goals because of its built-in exploration behavior, whereas other classiﬁer methods would operate solely based on the goal examples provided.
Revised: An additional advantage of BayCRL is that it is still capable of ﬁnding these "unspeciﬁed" goals because of its built-in exploration behavior, whereas other classiﬁer methods would operate solely based on the goal examples provided.

Original: To see this, we evaluate our algorithm on a two-sided variant of the Zigzag Maze with multiple goals, visualized in Figure 17 to the right.
Revised: To see this, we evaluate our algorithm on a two-sided variant of the Zigzag Maze with multiple goals, visualized in Figure 17 to the right.

Original: The agent starts in the middle and is provided with 5 goal examples on the far left side of the maze; unknown to it, the right side contains 5 sparse reward regions which are actually closer from its initial position.
Revised: The agent starts in the middle and is provided with 5 goal examples on the far left side of the maze; unknown to it, the right side contains 5 sparse reward regions which are actually closer from its initial position.

Original: As shown in Figures 18 and 19, BayCRL manages to ﬁnd the sparse rewards while other methods do not.
Revised: As shown in Figures 18 and 19, BayCRL manages to ﬁnd the sparse rewards while other methods do not.

Original: BayCRL, although initially guided towards the provided goal examples on the left, continues to explore in both directions and eventually ﬁnds the "hidden" rewards on the right.
Revised: BayCRL, although initially guided towards the provided goal examples on the left, continues to explore in both directions and eventually ﬁnds the "hidden" rewards on the right.

Original: Meanwhile, VICE focuses solely on the provided goals, and gets stuck in a local optima near the bottom left corner.
Revised: Meanwhile, VICE focuses solely on the provided goals, and gets stuck in a local optima near the bottom left corner.

Original: A. H YPERPARAMETER AND I MPLEMENTATION D ETAILS
Revised: A. H YPERPARAMETER AND I MPLEMENTATION D ETAILS

Original: We describe the hyperparameter choices and implementation details for our experiments here.
Revised: We describe the hyperparameter choices and implementation details for our experiments here.

Original: We ﬁrst list the general hyperparameters that were shared across runs, then provide tables of additional hyperparameters we tuned over for each domain and algorithm.
Revised: We ﬁrst list the general hyperparameters that were shared across runs, then provide tables of additional hyperparameters we tuned over for each domain and algorithm.

Original: SAC
Revised: SAC

Original: Learning Rate 3 × 10 −
Revised: Learning Rate 3 × 10 −

Original: Discount Factor γ 0 .
Revised: Discount Factor γ 0 .

Original: Policy Type
Revised: Policy Type

Original: Gaussian
Revised: Gaussian

Original: Policy Hidden Sizes (512 , 512)
Revised: Policy Hidden Sizes (512 , 512)

Original: Policy Hidden Activation
Revised: Policy Hidden Activation

Original: ReLU
Revised: ReLU

Original: RL Batch Size
Revised: RL Batch Size

Original: Reward Scaling
Revised: Reward Scaling

Original: Replay Buffer Size 500 ,
Revised: Replay Buffer Size 500 ,

Original: Q Hidden Sizes (512 , 512)
Revised: Q Hidden Sizes (512 , 512)

Original: Q Hidden Activation
Revised: Q Hidden Activation

Original: ReLU
Revised: ReLU

Original: Q Weight Decay
Revised: Q Weight Decay

Original: Q Learning Rate 3 × 10 −
Revised: Q Learning Rate 3 × 10 −

Original: [Table]
Revised: [Table]

Original: Target Network τ 5 ×
Revised: Target Network τ 5 ×

Original: [Table]
Revised: [Table]

Original: −
Revised: −

Original: [Table]
Revised: [Table]

Original: BayCRL
Revised: BayCRL

Original: Adaptation batch size
Revised: Adaptation batch size

Original: Meta
Revised: Meta

Original: [Table]
Revised: [Table]

Original: tasks per epoch
Revised: tasks per epoch

Original: Meta-test set size
Revised: Meta-test set size

Original: VICE
Revised: VICE

Original: Classiﬁer Learning
Revised: Classiﬁer Learning

Original: Rate
Revised: Rate

Original: [Table]
Revised: [Table]

Original: Size
Revised: Size

Original: Classiﬁer Optimizer
Revised: Classiﬁer Optimizer

Original: Adam
Revised: Adam

Original: [Table]
Revised: [Table]

Original: Algorithm
Revised: Algorithm

Original: SAC
Revised: SAC

Original: [Table]
Revised: [Table]

Original: Sizes (256 , 256)
Revised: Sizes (256 , 256)

Original: [Table]
Revised: [Table]

Original: Goal Examples: For the classiﬁer-based methods in our experiments (VICE and BayCRL), we provide 150 goal examples for each environment at the start of training.
Revised: Goal Examples: For the classiﬁer-based methods in our experiments (VICE and BayCRL), we provide 150 goal examples for each environment at the start of training.

Original: These are used as the pool of positive examples when training the success classiﬁer.
Revised: These are used as the pool of positive examples when training the success classiﬁer.

Original: DDL Reward: We use the version of DDL proposed in Hartikainen et al.
Revised: DDL Reward: We use the version of DDL proposed in Hartikainen et al.

Original: (2019) where we provide the algorithm with the ground truth goal state g , then run SAC with a reward function of r ( s )
Revised: (2019) where we provide the algorithm with the ground truth goal state g , then run SAC with a reward function of r ( s )

Original: = − d π ( s , g ) , where d π is the learned dynamical distance function for the policy at the current iteration of training.
Revised: = − d π ( s , g ) , where d π is the learned dynamical distance function for the policy at the current iteration of training.

Original: A.6. Z IGZAG M AZE H YPERPARAMETERS
Revised: A.6. Z IGZAG M AZE H YPERPARAMETERS

Original: [Table]
Revised: [Table]

Original: A.6. S PIRAL M AZE H YPERPARAMETERS
Revised: A.6. S PIRAL M AZE H YPERPARAMETERS

Original: [Table]
Revised: [Table]

Original: A.6. S AWYER P USH H YPERPARAMETERS
Revised: A.6. S AWYER P USH H YPERPARAMETERS

Original: [Table]
Revised: [Table]

Original: A.6. S AWYER P ICK - AND -P LACE H YPERPARAMETERS
Revised: A.6. S AWYER P ICK - AND -P LACE H YPERPARAMETERS

Original: [Table]
Revised: [Table]

Original: A.6. S AWYER D OOR O PENING H YPERPARAMETERS
Revised: A.6. S AWYER D OOR O PENING H YPERPARAMETERS

Original: [Table]
Revised: [Table]

Original: A. P ROOF OF T HEOREM 1 CONNECTING NML AND INVERSE COUNTS
Revised: A. P ROOF OF T HEOREM 1 CONNECTING NML AND INVERSE COUNTS

Original: We provide the proof of Theorem 1 here for completeness.
Revised: We provide the proof of Theorem 1 here for completeness.

Original: Theorem A.1.
Revised: Theorem A.1.

Original: Suppose we are estimating success probabilities p ( e = 1 | s ) in the tabular setting, where we have a separate parameter independently for each state.
Revised: Suppose we are estimating success probabilities p ( e = 1 | s ) in the tabular setting, where we have a separate parameter independently for each state.

Original: Let N ( s ) denote the number of times state s has been visited by the policy, and let G ( s ) be the number of occurrences of state s in the successful outcomes.
Revised: Let N ( s ) denote the number of times state s has been visited by the policy, and let G ( s ) be the number of occurrences of state s in the successful outcomes.

Original: Then the CNML probability p CNML ( e = 1 | s ) is equal to
Revised: Then the CNML probability p CNML ( e = 1 | s ) is equal to

Original: G ( s )+1 N ( s )+ G ( s )+2 .
Revised: G ( s )+1 N ( s )+ G ( s )+2 .

Original: For states that are never observed to be successful, we then recover inverse counts 1 N ( s )+2 .
Revised: For states that are never observed to be successful, we then recover inverse counts 1 N ( s )+2 .

Original: Proof.
Revised: Proof.

Original: In the fully tabular setting, our MLE estimates for p ( O | s ) are simply given by ﬁnding the best parameter p s for each state.
Revised: In the fully tabular setting, our MLE estimates for p ( O | s ) are simply given by ﬁnding the best parameter p s for each state.

Original: The proof then proceeds by simple calculation.
Revised: The proof then proceeds by simple calculation.

Original: For a state with n = N ( s ) negative occurrences and g = G ( s ) positive occurrences, the MLE estimate is simply given by gn + g .
Revised: For a state with n = N ( s ) negative occurrences and g = G ( s ) positive occurrences, the MLE estimate is simply given by gn + g .

Original: Now for evaluating CNML, we consider appending another instance for each class.
Revised: Now for evaluating CNML, we consider appending another instance for each class.

Original: The new parameter after appending a negative example is then g n + g +1 , which then assigns probability n +1 n + g + to the negative class.
Revised: The new parameter after appending a negative example is then g n + g +1 , which then assigns probability n +1 n + g + to the negative class.

Original: Similarly, after appending a positive example, the new parameter is g +1 n + g +1 , so we try to assign probability g +1 n + g +1 to the positive class.
Revised: Similarly, after appending a positive example, the new parameter is g +1 n + g +1 , so we try to assign probability g +1 n + g +1 to the positive class.

Original: Normalizing, we have
Revised: Normalizing, we have

Original: [Equation]
Revised: [Equation]

Original: When considering states that have only been visited on-policy, and are not included in the set of successful outcomes, then the likelihood reduces to
Revised: When considering states that have only been visited on-policy, and are not included in the set of successful outcomes, then the likelihood reduces to

Original: [Equation]
Revised: [Equation]