Original: Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning?
Revised: Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning?

Original: Abstract
Revised: Abstract

Original: Causal confusion is a phenomenon where an agent learns a policy that reflects imperfect spurious correlations in the data.
Revised: Causal confusion is a phenomenon where an agent learns a policy that reflects imperfect spurious correlations in the data.

Original: Such a policy may falsely appear to be optimal during training if most of the training data contains such spurious correlations.
Revised: Such a policy may falsely appear to be optimal during training if most of the training data contains such spurious correlations.

Original: This phenomenon is particularly pronounced in domains such as robotics, with potentially large gaps between open- and closed-loop performance of an agent.
Revised: This phenomenon is particularly pronounced in domains such as robotics, with potentially large gaps between open- and closed-loop performance of an agent.

Original: In such settings, causally confused models may appear to perform well according to open-loop metrics during training but fail catastrophically when deployed in the real world.
Revised: In such settings, causally confused models may appear to perform well according to open-loop metrics during training but fail catastrophically when deployed in the real world.

Original: In this paper, we investigate whether selectively sampling appropriate points from the dataset may enable offline RL agents to disambiguate the underlying causal mechanisms of the environment, alleviate causal confusion in offline reinforcement learning, and produce a safer model for deployment.
Revised: In this paper, we investigate whether selectively sampling appropriate points from the dataset may enable offline RL agents to disambiguate the underlying causal mechanisms of the environment, alleviate causal confusion in offline reinforcement learning, and produce a safer model for deployment.

Original: To answer this question, we consider a set of tailored offline reinforcement learning datasets that exhibit causal ambiguity and assess the ability of active sampling techniques to reduce causal confusion at evaluation.
Revised: To answer this question, we consider a set of tailored offline reinforcement learning datasets that exhibit causal ambiguity and assess the ability of active sampling techniques to reduce causal confusion at evaluation.

Original: We provide empirical evidence that uniform and active sampling techniques are able to consistently reduce causal confusion as training progresses and that active sampling is able to do so significantly more efficiently than uniform sampling.
Revised: We provide empirical evidence that uniform and active sampling techniques are able to consistently reduce causal confusion as training progresses and that active sampling is able to do so significantly more efficiently than uniform sampling.

Original:  Introduction
Revised:  Introduction

Original: Offline learning offers opportunities to scale reinforcement learning to domains where offline data is plentiful but online interaction with the environment is costly.
Revised: Offline learning offers opportunities to scale reinforcement learning to domains where offline data is plentiful but online interaction with the environment is costly.

Original: The fundamental challenge of offline reinforcement learning is to identify the cause and effect of actions from a fixed dataset, which is often intractable.
Revised: The fundamental challenge of offline reinforcement learning is to identify the cause and effect of actions from a fixed dataset, which is often intractable.

Original: In the absence of online interactions, our hope is that the dataset covers a uniform distribution of an exhaustive set of plausible scenarios.
Revised: In the absence of online interactions, our hope is that the dataset covers a uniform distribution of an exhaustive set of plausible scenarios.

Original: This is often not the case in datasets for robotic control, which are heavy-tailed and often contain only a handful of samples for rare (and informative) events.
Revised: This is often not the case in datasets for robotic control, which are heavy-tailed and often contain only a handful of samples for rare (and informative) events.

Original: Causal confusion occurs when agents misinterpret the underlying causal mechanisms of the environment and erroneously associate certain actions or states with a given reward.
Revised: Causal confusion occurs when agents misinterpret the underlying causal mechanisms of the environment and erroneously associate certain actions or states with a given reward.

Original: For example, if an agent happens to simultaneously observe independent events X and Y in its environment whenever it receives a reward R , and R causally depends on Y but not on X , the agent may attribute the reward
Revised: For example, if an agent happens to simultaneously observe independent events X and Y in its environment whenever it receives a reward R , and R causally depends on Y but not on X , the agent may attribute the reward

Original: R to X and Y occurring jointly even though R may be independent of Y .
Revised: R to X and Y occurring jointly even though R may be independent of Y .

Original: Problematically, if the spurious correlation between Y and R observed at training time ceases to hold at deployment time, a causally-confused model may show a significant deterioration in performance.
Revised: Problematically, if the spurious correlation between Y and R observed at training time ceases to hold at deployment time, a causally-confused model may show a significant deterioration in performance.

Original: Although spurious features might not perfectly explain all the input-label pairs in offline data, optimisation methods such as stochastic gradient descent benefit from correlations in the data when seeking to reduce the training loss.
Revised: Although spurious features might not perfectly explain all the input-label pairs in offline data, optimisation methods such as stochastic gradient descent benefit from correlations in the data when seeking to reduce the training loss.

Original: Therefore spurious correlations are unwittingly transferred from the data to the learned models.
Revised: Therefore spurious correlations are unwittingly transferred from the data to the learned models.

Original: This phenomenon is especially pronounced in models trained on high-dimensional visual inputs, since extracting the true causal factors of an environment (and their interplay) from an image is a particularly difficult problem.
Revised: This phenomenon is especially pronounced in models trained on high-dimensional visual inputs, since extracting the true causal factors of an environment (and their interplay) from an image is a particularly difficult problem.

Original: On the other hand, it is noticeably easy for neural networks to find shortcuts for prediction, for instance, using co-occurrence patterns between objects and backgrounds to successfully perform object detection, as has been extensively documented in the deep computer vision literature [Beery et al., 2018, Geirhos et al., 2020].
Revised: On the other hand, it is noticeably easy for neural networks to find shortcuts for prediction, for instance, using co-occurrence patterns between objects and backgrounds to successfully perform object detection, as has been extensively documented in the deep computer vision literature [Beery et al., 2018, Geirhos et al., 2020].

Original: Several works have also reported causal confusion in control policy learning and used heuristic loss reweighting or dataset-balancing schemes based on domain knowledge of the prediction task to reduce the amount of causal confusion in the learned models [Wijmans et al., 2019, Kumar et al., 2022].
Revised: Several works have also reported causal confusion in control policy learning and used heuristic loss reweighting or dataset-balancing schemes based on domain knowledge of the prediction task to reduce the amount of causal confusion in the learned models [Wijmans et al., 2019, Kumar et al., 2022].

Original: However, these heuristics require a practitioner to know the source of causal confusion a priori, since they require explicitly increasing the loss for specific spurious-correlation-breaking samples in the dataset [Wen et al., 2021, 2020].
Revised: However, these heuristics require a practitioner to know the source of causal confusion a priori, since they require explicitly increasing the loss for specific spurious-correlation-breaking samples in the dataset [Wen et al., 2021, 2020].

Original: While this may be limiting in practice, the success of these heuristics demonstrates that it is possible to recover less causally confused models from a fixed dataset, as long as we have access to an oracle sampling scheme for the training data.
Revised: While this may be limiting in practice, the success of these heuristics demonstrates that it is possible to recover less causally confused models from a fixed dataset, as long as we have access to an oracle sampling scheme for the training data.

Original: Building on recent work that highlighted active sampling as a tool for sample-efficient causal estimation, we study uncertainty-based and loss-based active sampling techniques and find that active sampling alleviates causal confusion in offline training and yields higher rewards at evaluation.
Revised: Building on recent work that highlighted active sampling as a tool for sample-efficient causal estimation, we study uncertainty-based and loss-based active sampling techniques and find that active sampling alleviates causal confusion in offline training and yields higher rewards at evaluation.

Original: We further show that active sampling is able to alleviate causal confusion at a significantly higher sample efficiency than naive approaches, such as uniform sampling, and that the usefulness of active sampling in alleviating causal confusion is highly related to the quality of predictive uncertainty estimates used in the best-performing, uncertainty-based acquisition function.
Revised: We further show that active sampling is able to alleviate causal confusion at a significantly higher sample efficiency than naive approaches, such as uniform sampling, and that the usefulness of active sampling in alleviating causal confusion is highly related to the quality of predictive uncertainty estimates used in the best-performing, uncertainty-based acquisition function.

Original: Our results demonstrate that while dominant spurious correlations in large offline datasets may lead to poor performance, and hence to harmful decision-making at deployment, active sampling techniques can significantly reduce the negative effects of spurious correlations on offline RL agents and that they are able to do so with high sample efficiency.
Revised: Our results demonstrate that while dominant spurious correlations in large offline datasets may lead to poor performance, and hence to harmful decision-making at deployment, active sampling techniques can significantly reduce the negative effects of spurious correlations on offline RL agents and that they are able to do so with high sample efficiency.

Original: The study conducted in this paper is located at the intersection between the fields of reinforcement learning, Bayesian active learning, and causal inference.
Revised: The study conducted in this paper is located at the intersection between the fields of reinforcement learning, Bayesian active learning, and causal inference.

Original: In Sections 2 and 3, we provide relevant background on and connections between these three areas.
Revised: In Sections 2 and 3, we provide relevant background on and connections between these three areas.

Original:  Preliminaries
Revised:  Preliminaries

Original: We consider an environment formulated as a Markov Decision Process (MDP) M defined by a tuple ( S , A , P , r, d 0 , γ ) , where S is the state space, A is the action space, P ( s ′ | s , a ) is the transition probability distribution, r : S ×
Revised: We consider an environment formulated as a Markov Decision Process (MDP) M defined by a tuple ( S , A , P , r, d 0 , γ ) , where S is the state space, A is the action space, P ( s ′ | s , a ) is the transition probability distribution, r : S ×

Original: A → R is the reward function, d 0 is the initial state distribution, and γ ∈ (0 , 1] is the discount factor.
Revised: A → R is the reward function, d 0 is the initial state distribution, and γ ∈ (0 , 1] is the discount factor.

Original: The goal of reinforcement learning (RL) is to find an optimal policy π ( a | s ) that maximizes the cumulative discounted reward E s t , a t [ (cid:80) ∞ t =

Revised: The goal of reinforcement learning (RL) is to find an optimal policy π ( a | s ) that maximizes the cumulative discounted reward E s t , a t [ (cid:80) ∞ t =

Original: γ t r ( s t , a t )] , where s 0 ∼ d 0 ( · ) , a t ∼ π ( · | s t ) , and s t +1 ∼ P ( · | s t , a t ) .
Revised: γ t r ( s t , a t )] , where s 0 ∼ d 0 ( · ) , a t ∼ π ( · | s t ) , and s t +1 ∼ P ( · | s t , a t ) .

Original: Q -learning-based RL algorithms learn an optimal state–action value function Q ∗ ( s, a ) , representing the expected cumulative discounted reward starting from s with action a
Revised: Q -learning-based RL algorithms learn an optimal state–action value function Q ∗ ( s, a ) , representing the expected cumulative discounted reward starting from s with action a

Original: and then acting optimally according to policy π ∗ thereafter, ie Q ∗ ( s, a ) =
Revised: and then acting optimally according to policy π ∗ thereafter, ie Q ∗ ( s, a ) =

Original: E π ∗ [ (cid:80) ∞
Revised: E π ∗ [ (cid:80) ∞

Original: t =0 γ t r ( s t , a t ) | s 0 =
Revised: t =0 γ t r ( s t , a t ) | s 0 =

Original: s, a t = a ] .
Revised: s, a t = a ] .

Original: Analogously, the value function V ( s ) =
Revised: Analogously, the value function V ( s ) =

Original: E π [ (cid:80) ∞ t =

Revised: E π [ (cid:80) ∞ t =

Original: γ t r ( s t , a t ) | s 0 = s ] represents the expected cumulative discounted reward achievable from state s when following policy π .
Revised: γ t r ( s t , a t ) | s 0 = s ] represents the expected cumulative discounted reward achievable from state s when following policy π .

Original: Q -learning is trained on the Bellman equation defined as follows with the Bellman optimal operator B defined by:
Revised: Q -learning is trained on the Bellman equation defined as follows with the Bellman optimal operator B defined by:

Original: [Equation]
Revised: [Equation]

Original: The Q -function is updated by minimizing the Bellman Squared Error E (cid:2) ( Q − B Q ) 2 (cid:3) where a frozen, periodically updated copy of the Q -network weights are used to compute the target B Q .
Revised: The Q -function is updated by minimizing the Bellman Squared Error E (cid:2) ( Q − B Q ) 2 (cid:3) where a frozen, periodically updated copy of the Q -network weights are used to compute the target B Q .

Original: Offline reinforcement learning algorithms aim to learn an optimal policy by learning estimates of the value (or Q -value) function from a static dataset of transitions D
Revised: Offline reinforcement learning algorithms aim to learn an optimal policy by learning estimates of the value (or Q -value) function from a static dataset of transitions D

Original: = { ( s , a , r, s ′ ) } collected by a behaviour policy π β .
Revised: = { ( s , a , r, s ′ ) } collected by a behaviour policy π β .

Original: Since the agent does not interact with the environment, naively using online RL techniques in the offline case leads to value function overestimation on unseen states and actions.
Revised: Since the agent does not interact with the environment, naively using online RL techniques in the offline case leads to value function overestimation on unseen states and actions [Hasselt, 2010].

Original: This happens because the max operator in the target update in Equation (1) propagates erroneously high values from the next state’s value estimate into the update.

Original: Thus, algorithms based on the pessimism principle of underestimating the Q -values to optimise the worst-case regret bound have been successful at learning policies from datasets containing either good coverage of the state–action space or high return trajectories [Kumar et al., 2020, Kumar et al., 2019, Buckman et al., 2021, Bai et al., 2022].
Revised: Thus, algorithms based on the pessimism principle of underestimating the Q -values to optimise the worst-case regret bound have been successful at learning policies from datasets containing either good coverage of the state–action space or high return trajectories [Kumar et al., 2020, Kumar et al., 2019, Buckman et al., 2021, Bai et al., 2022].

Original: Prior work categorised pessimistic offline RL algorithms into (1) proximal and (2) uncertainty-aware algorithms [Buckman et al., 2021]. The former penalises action-value estimates based on deviations from the  actions  seen in the dataset, while the latter conservatively updates the value functions, taking the uncertainty of their targets  into account.
Revised: Prior work categorised pessimistic offline RL algorithms into (1) proximal and (2) uncertainty-aware algorithms, where the former penalise actions not seen in the data and the latter construct value function updates taking the uncertainty of the targets B Q into account [Buckman et al., 2021].

Original: Proximal pessimistic algorithms like CQL, BEAR [Kumar et al., 2019, Kumar et al., 2020] are known to work well with exactly the kind of narrow and biased data distributions that are most prone to causal confusion [Fu et al., 2020].
Revised: Proximal pessimistic algorithms like CQL, BEAR [Kumar et al., 2019, Kumar et al., 2020] are known to work well with exactly the kind of narrow and biased data distributions that are most prone to causal confusion [Fu et al., 2020].

Original: We, therefore, analyse the proximal class of algorithms in this work and leave the study of causal confusion in uncertainty-aware pessimistic offline RL algorithms to future work.
Revised: We, therefore, analyse the proximal class of algorithms in this work and leave the study of causal confusion in uncertainty-aware pessimistic offline RL algorithms to future work.

Original: Conservative Q -Learning.
Revised: Conservative Q -Learning.

Original: We choose CQL [Kumar et al., 2020] as an instantiation of a proximal pessimistic offline RL algorithm in our experiments, due to its simplicity and competitive performance on offline RL benchmarks.
Revised: We choose CQL [Kumar et al., 2020] as an instantiation of a proximal pessimistic offline RL algorithm in our experiments, due to its simplicity and competitive performance on offline RL benchmarks.

Original: The CQL objective, which combines the standard TD-error of Q -learning with a penalty constraining deviations from the behaviour policy, is defined as:
Revised: The CQL objective, which combines the standard TD-error of Q -learning with a penalty constraining deviations from the behaviour policy, is defined as:

Original: [Equation]
Revised: [Equation]

Original: where α 0 controls the degree of conservatism and θ is the parameterisation of the learned Q -function.
Revised: where α 0 controls the degree of conservatism and θ is the parameterisation of the learned Q -function.

Original:  Alleviating Causal Confusion in Offline RL via Active Sampling
Revised:  Alleviating Causal Confusion in Offline RL via Active Sampling

Original: In this section, we briefly describe the active sampling schemes we use in this work in conjunction with offline RL.
Revised: In this section, we briefly describe the active sampling schemes we use in this work in conjunction with offline RL.

Original: In Appendix C, we define our assumptions on long-tailed datasets that may lead to causally confused models, and in Appendix D, we discuss how active sampling can be seen as facilitating causal discovery in offline (observational) data.
Revised: In Appendix C, we define our assumptions on long-tailed datasets that may lead to causally confused models, and in Appendix D, we discuss how active sampling can be seen as facilitating causal discovery in offline (observational) data.

Original: 3.1 Active Sampling
Revised: 3.1 Active Sampling

Original: Active sampling techniques are used to selectively sample from a given dataset during training to enable sample-efficient learning or to improve learning from noisy data [Loshchilov and Hutter, 2015, Mirzasoleiman et al., 2019].
Revised: Active sampling techniques are used to selectively sample from a given dataset during training to enable sample-efficient learning or to improve learning from noisy data [Loshchilov and Hutter, 2015, Mirzasoleiman et al., 2019].

Original: To perform active sampling, data points are scored based on a given information-theoretic acquisition function, and a small number of points are sampled based on these scores using a pre-specified weighting scheme.
Revised: To perform active sampling, data points are scored based on a given information-theoretic acquisition function, and a small number of points are sampled based on these scores using a pre-specified weighting scheme.

Original: In this work, we build on Jesson et al.
Revised: In this work, we build on Jesson et al.

Original: [2021] and study whether active sampling can enable sample-efficient resolution of causal confusion in models trained on long-tailed datasets  —without necessitating any modifications to the learning objective.
Revised: [2021] and study whether active sampling can enable sample-efficient resolution of causal confusion in models trained on long-tailed datasets . In particular, the focus of this work is to investigate whether active sampling techniques can be used to address the challenges of causal ambiguity in state–action trajectories used in offline RL —without necessitating any modifications to the learning objective.

Original: Algorithm 1 in the Appendix describes the CQL algorithm with active sampling of transitions, where the modifications from uniform sampling are highlighted in blue.
Revised: Algorithm 1 in the Appendix describes the CQL algorithm with active sampling of transitions, where the modifications from uniform sampling are highlighted in blue.

Original: We study the following uncertainty-based and loss-based data acquisition functions:
Revised: We study the following uncertainty-based and loss-based data acquisition functions:

Original: Uncertainty-based Sampling.
Revised: Uncertainty-based Sampling.

Original: We are interested in sampling state transitions for which the learned
Revised: We are interested in sampling state transitions for which the learned

Original: Q -network is uncertain about its predictions over the action seen in the data ( a β = π β ( s ) ) or its own greedy action ( a ∗ = arg max Q ( s ) ).
Revised: Q -network is uncertain about its predictions over the action seen in the data ( a β = π β ( s ) ) or its own greedy action ( a ∗ = arg max Q ( s ) ).

Original: We model the epistemic uncertainty in the learned Q -function, by creating an ensemble model of Q -functions { Q θ i } i and training on identical transitions across the ensemble members, with their own corresponding targets, { ¯ θ i } i , as proposed in Ghasemipour et al.
Revised: We model the epistemic uncertainty in the learned Q -function, by creating an ensemble model of Q -functions { Q θ i } i and training on identical transitions across the ensemble members, with their own corresponding targets, { ¯ θ i } i , as proposed in Ghasemipour et al.

Original: We can then use the variance of an estimate across the ensemble as the measure of the epistemic uncertainty about the estimate.
Revised: We can then use the variance of an estimate across the ensemble as the measure of the epistemic uncertainty about the estimate.

Original: However, the Q -values of different ensemble members could have arbitrary numerical offsets but still be equivalent, since they’re trained by bootstrapping value estimates.
Revised: However, the Q -values of different ensemble members could have arbitrary numerical offsets but still be equivalent, since they’re trained by bootstrapping value estimates.

Original: Instead, we estimate the uncertainty about an action by computing the variance of its advantage estimates over the ensemble, where the advantage of action a i for a Q -learner is given by
Revised: Instead, we estimate the uncertainty about an action by computing the variance of its advantage estimates over the ensemble, where the advantage of action a i for a Q -learner is given by

Original: [Equation]
Revised: [Equation]

Original: The advantage function represents a causal quantity assessing the relative effect of action a i on the outcome Q for a given state s .
Revised: The advantage function represents a causal quantity assessing the relative effect of action a i on the outcome Q for a given state s .

Original: The acquisition scores in this case can then be computed by the expression:
Revised: The acquisition scores in this case can then be computed by the expression:

Original: Var ( A π ( s, a i )) .
Revised: Var ( A π ( s, a i )) .

Original: We consider two variants, Variance-greedy and Variance-data , referring to whether the advantages for the greedy ( a i = a ∗ ) or the dataset actions ( a i = a β ) are computed.
Revised: We consider two variants, Variance-greedy and Variance-data , referring to whether the advantages for the greedy ( a i = a ∗ ) or the dataset actions ( a i = a β ) are computed.

Original: Loss-based Sampling.
Revised: Loss-based Sampling.

Original: We sample transitions based on their Temporal Difference error similar to
Revised: We sample transitions based on their Temporal Difference error similar to

Original: Prioritised Experience Replay (PER; Schaul et al.
Revised: Prioritised Experience Replay (PER; Schaul et al.

Original: and refer to this variant as TD-Error .
Revised: and refer to this variant as TD-Error .

Original: We provide further implementation-related details about active sampling in Appendix E.1. 
Revised: We provide further implementation-related details about active sampling in Appendix E.1. [Figure]

Original:  Experiments
Revised:  Experiments

Original: In the following sections and in Appendix G, we describe our experiments  investigating the following questions: (1) Can causal confusion be consistently observed in offline RL agents when sampling transitions uniformly from a long-tailed demonstration dataset?
Revised: In the following section and in Appendix G, we describe our experiments on three benchmark environments investigating the following questions: (1) Can causal confusion be consistently observed in offline RL agents when sampling transitions uniformly from a long-tailed demonstration dataset?

Original: (2) Does active sampling based on a model’s predictive uncertainty or its loss help in alleviating causal confusion?
Revised: (2) Does active sampling based on a model’s predictive uncertainty or its loss help in alleviating causal confusion?

Original: (3) To what extent does the quality of predictive uncertainty estimates affect predictive performance and sample efficiency gains under uncertainty-based active sampling?
Revised: (3) To what extent does the quality of predictive uncertainty estimates affect predictive performance and sample efficiency gains under uncertainty-based active sampling?

Original: (Appendix G.2)
Revised: (Appendix G.3)

Original: 4.1 Illustrative Example: Traffic-World
Revised: 4.1 Illustrative Example: Traffic-World

Original: We build on the environment proposed in [Anonymous, 2021] to construct a gridworld (shown inthe left panel of Figure 1b), where an agent (red triangle) starts at the leftmost pointin a row behind leading vehicles (blue circles) and needs to cross a traffic light to reach a goallocation (green square) on the right side of the grid.
Revised: We build on the environment proposed in [Anonymous, 2021] to construct a gridworld (shown in Figure 1a), where an agent starting at the leftmost point behind leading vehicles, needs to cross a traffic light to reach the goal on the bottom right.

Original: We collect data such that the probability of the traffic light turning red becomes lower as the agent approaches it, and so the data distribution contains: (1) mostly episodes where the light is green throughout the episode, (2) some episodes where the traffic light is red and the agent has to wait behind the vehicle in front (referenced here onward as the leading vehicle) before the light turns green again, and (3) only a couple of episodes where the light turns red with the agent at the front of the traffic queue.
Revised: We collect data such that the probability of the traffic light turning red becomes lower as the agent approaches it, and so the data distribution contains: (1) mostly episodes where the light is green throughout the episode, (2) some episodes where the traffic light is red and the agent has to wait behind the vehicle in front (referenced here onward as the leading vehicle) before the light turns green again, and (3) only a couple of episodes where the light turns red with the agent at the front of the traffic queue.

Original: In this setup, the agent could just learn to follow the leading vehicle, instead of learning traffic light rules.

Original: To test causal confusion explicitly here, we introduce a related spurious correlate: a flashing yellow tile at the top left of the grid (emulating the brake lights on a leading vehicle), that is yellow whenever the leading vehicle is stopped or blocked, and grey otherwise.
Revised: To test causal confusion explicitly here, we introduce a related spurious correlate: a flashing yellow tile at the top left of the grid (emulating the brake lights on a leading vehicle), that is yellow whenever the leading vehicle is stopped or blocked, and grey otherwise.

Original: The agent could  follow this as an indicator of whether to stop or go ahead, and this policy would be correct for 98% of the data points.
Revised: The agent could simply follow this as an indicator of whether to stop or go ahead, and this policy would be correct for 98% of the data points.

Original: In this case, one example of a trajectory in the dataset from a tail-case scenario is the following: when the tile flashes yellow because the leading vehicle is static but the agent is not blocked because there is space to move ahead behind the leading vehicle. Another tail case is when the leading vehicle has crossed the traffic light which has just turned red in front of the agent, and the top-left tile is not flashing yellow because the leading vehicle is not static. However, if the agent takes the tile being off as a cue, it will decide to move forward and end up breaking the red light and fetching a negative reward.
Revised: We report the average reward over three instantiations of this environment testing different types of causal confusion.

Original: Figure 1b shows the evaluation curves of CQL agents trained with uniformly-sampled data, with and without the yellow tile present in images in the dataset.
Revised: Figure 1b shows the evaluation curves of CQL agents trained with uniformly-sampled data, with and without the yellow tile present in images in the dataset.

Original: The evaluation reward plotted here is the average reward obtained across three benchmark environments testing different types of causal confusion relating to the aforementioned tail cases.

Original: We see that the performance of the former agent degrades and it takes 4 × the number of gradient steps to converge to the solution of the latter agent which is trained without the spurious correlate present.
Revised: We see that the performance of the former agent degrades significantly compared to the latter.

Original: Also shown are the active sampling variants( TD-Error , Variance-greedy and Variance-data ) trained with the spurious yellow tile, which perform very similarly to uniform sampling without the spurious correlate present.In Figure 3 in the appendix, [Figure]
Revised: Also shown are the active sampling variants trained with the spurious yellow tile, which perform very similarly to uniform sampling without the spurious correlate present.

Original: 4.2 Assessing Generalization in Offline Reinforcement Learning: Procgen

Original: The Maze environment in Procgen [Cobbe et al., 2019] defines a navigation task where the agent starts at the bottom left in the maze and receives a reward of +10 upon successfully reaching the goal which is sampled at any valid location in the maze.

Original: Langosco et al.

Original: [2022] recently showed that an agent trained on a series of environments with the goal always at the top-right will be causally confused about the source of the reward and will still navigate to the top-right even when the goal is sampled elsewhere.

Original: We generate a skewed mixture dataset containing mostly episodes where the goal is sampled at the top-right and a few episodes where the goal is sampled randomly.

Original: Further details about the setup are described in the Appendix.
Revised: Appendix F provides further details and analysis related to this experiment.

Original: Figure 2b shows the evaluation performance of random and active sampling agents trained on the mixture dataset when goals are sampled randomly in the evaluation environment.

Original: We plot the computation time for the uniform and active sampling variants in Figurein the appendix and note that variance-based sampling reaches the highest score achieved by uniform sampling in lesser wall-clock time.

Original: Qualitative evaluations show that uniform sampling agents which achieve a lower reward still successfully navigate to the top-right corner of hard mazes, and are therefore only confused about the location of the goal.

Original:  Conclusion
Revised:  Conclusion

Original: In this paper, we studied how to alleviate causal confusion in offline RL.
Revised: In this paper, we studied how to alleviate causal confusion in offline RL.

Original: We designed uncertainty- and loss-based data sampling baselines to selectively sample transitions for training, and found evidence that active sampling can recover a less causally-confused model in significantly fewer training steps as compared to uniform sampling.
Revised: We designed uncertainty- and loss-based data sampling baselines to selectively sample transitions for training, and found evidence that active sampling can recover a less causally-confused model in significantly fewer training steps as compared to uniform sampling.

Original: In future work, we hope to scale the analysis performed in this paper to larger benchmark environments, with sources of noise in the reward and transitions, to further corroborate our findings.
Revised: In future work, we hope to scale the analysis performed in this paper to larger benchmark environments, with sources of noise in the reward and transitions, to further corroborate our findings.

Original: Such an analysis would help further distinguish the quality of solutions found through loss-based and uncertainty-based active sampling since noisy transitions can have irreducible loss.
Revised: Such an analysis would help further distinguish the quality of solutions found through loss-based and uncertainty-based active sampling since noisy transitions can have irreducible loss.

Original: Additional promising avenues for future work include extending the current analysis to settings with continuous action spaces.
Revised: Additional promising avenues for future work include extending the current analysis to settings with continuous action spaces.

Original: A Appendix B Related Work
Revised: A Appendix B Related Work

Original: Causal Confusion in Supervised Policy Learning.
Revised: Causal Confusion in Supervised Policy Learning.

Original: Several works in offline imitation learning have proposed solutions to mitigate causal confusion, which was first defined by de Haan et al.
Revised: Several works in offline imitation learning have proposed solutions to mitigate causal confusion, which was first defined by de Haan et al.

Original: The general practice across these works is to select sequential-decision-making benchmarks like simulated robotic control, augment their inputs with nuisance variables to induce spurious correlations, and show that policies trained on this data are inferior to those trained without the nuisance variables in the data.
Revised: The general practice across these works is to select sequential-decision-making benchmarks like simulated robotic control, augment their inputs with nuisance variables to induce spurious correlations, and show that policies trained on this data are inferior to those trained without the nuisance variables in the data.

Original: de Haan et al.
Revised: de Haan et al.

Original: [2019] demonstrated causal mis-identification in models trained in this manner on expert trajectories collected in the Mountain-Car and the CARLA simulators [Sutton and Barto, 2018,
Revised: [2019] demonstrated causal mis-identification in models trained in this manner on expert trajectories collected in the Mountain-Car and the CARLA simulators [Sutton and Barto, 2018,

Original: Dosovitskiy et al., 2017] where inputs were augmented with the previous control command taken by the acting agent.
Revised: Dosovitskiy et al., 2017] where inputs were augmented with the previous control command taken by the acting agent.

Original: They proposed to resolve the confusion through a scheme to query an expert or collect additional rollouts in the environment to refine a learned causal graph that conditions the learned policy.
Revised: They proposed to resolve the confusion through a scheme to query an expert or collect additional rollouts in the environment to refine a learned causal graph that conditions the learned policy.

Original: Wen et al.
Revised: Wen et al.

Original: [2020] proposes adversarial training to prune out any known sources of spurious correlations from the policy’s representation, for instance, the previous control commands given to a robot; Wen et al.
Revised: [2020] proposes adversarial training to prune out any known sources of spurious correlations from the policy’s representation, for instance, the previous control commands given to a robot; Wen et al.

Original: [2021] propose re-weighting the losses of data points based on the loss of a model trained with just the spurious correlates as the input; OREO [Park et al., 2021] regularises the model’s representation to be invariant to any individual object being dropped out in a scene.
Revised: [2021] propose re-weighting the losses of data points based on the loss of a model trained with just the spurious correlates as the input; OREO [Park et al., 2021] regularises the model’s representation to be invariant to any individual object being dropped out in a scene.

Original: Lee et al.
Revised: Lee et al.

Original: [2022] propose training a diversified policy ensemble for imitation learning in the case when perfect spurious correlations exist in the data and later select from these hypotheses based on validation data Causal Confusion has also recently been studied in reward-learning from preferences [Tien et al., 2022], where spurious correlations can be drawn between a human evaluator’s preferences and certain actions or parts of the state space, for tasks in the Assistive gym
Revised: [2022] propose training a diversified policy ensemble for imitation learning in the case when perfect spurious correlations exist in the data and later select from these hypotheses based on validation data Causal Confusion has also recently been studied in reward-learning from preferences [Tien et al., 2022], where spurious correlations can be drawn between a human evaluator’s preferences and certain actions or parts of the state space, for tasks in the Assistive gym

Original: [Erickson et al., 2019].
Revised: [Erickson et al., 2019].

Original: For instance, a reward model trained to classify or rank trajectories for a feeding task can easily learn that higher force correlates to higher reward since a certain level of force needs to be inserted in the direction of the mouth in all preferred trajectories.
Revised: For instance, a reward model trained to classify or rank trajectories for a feeding task can easily learn that higher force correlates to higher reward since a certain level of force needs to be inserted in the direction of the mouth in all preferred trajectories.

Original: Active Sampling.
Revised: Active Sampling.

Original: Our work is inspired by the work of Jesson et al.
Revised: Our work is inspired by the work of Jesson et al.

Original: [2021], which looks at the problem of treatment effect estimation in settings where we want to be sample-efficient in terms of querying for outcomes of costly experiments.
Revised: [2021], which looks at the problem of treatment effect estimation in settings where we want to be sample-efficient in terms of querying for outcomes of costly experiments.

Original: They propose several causality-inspired acquisition functions that prefer data-points which have both high variance in their estimated outcomes and correspond to covariates with considerable overlap in the dataset.
Revised: They propose several causality-inspired acquisition functions that prefer data-points which have both high variance in their estimated outcomes and correspond to covariates with considerable overlap in the dataset.

Original: In Q -learning-based RL algorithms,
Revised: In Q -learning-based RL algorithms,

Original: Prioritised replay [Schaul et al., 2015] is a non-Bayesian loss-based sampling scheme proposed for off-policy learning.
Revised: Prioritised replay [Schaul et al., 2015] is a non-Bayesian loss-based sampling scheme proposed for off-policy learning.

Original: It computes acquisition scores based on the TD-error of transitions and has not been studied in offline RL.
Revised: It computes acquisition scores based on the TD-error of transitions and has not been studied in offline RL.

Original: Ensemble Models in RL.
Revised: Ensemble Models in RL.

Original: Ensembles have been studied extensively to guide exploration in online
Revised: Ensembles have been studied extensively to guide exploration in online

Original: RL [Osband et al., 2016, Lee et al., 2021], and recently to construct adaptive pessimism constraints in offline RL, to disincentives uncertain actions from having high estimated returns.
Revised: RL [Osband et al., 2016, Lee et al., 2021], and recently to construct adaptive pessimism constraints in offline RL, to disincentives uncertain actions from having high estimated returns.

Original: It was recently also shown that significantly increasing the size and diversity of the ensembled critic in Soft-ActorCritic
Revised: It was recently also shown that significantly increasing the size and diversity of the ensembled critic in Soft-ActorCritic

Original: [Haarnoja et al., 2018] performs competitively with state-of-the-art offline RL algorithms
Revised: [Haarnoja et al., 2018] performs competitively with state-of-the-art offline RL algorithms

Original: [An et al., 2021].
Revised: [An et al., 2021].

Original: We are not aware of any prior work that has explored how the uncertainty of the value function could be used to sample transitions in RL.
Revised: We are not aware of any prior work that has explored how the uncertainty of the value function could be used to sample transitions in RL.

Original: AI Alignment.
Revised: AI Alignment.

Original: AI alignment seeks to align the behavior of agents with the intentions of their creators by investigating the incentives behind demonstrated tasks.
Revised: AI alignment seeks to align the behavior of agents with the intentions of their creators by investigating the incentives behind demonstrated tasks.

Original: Recent work on Goal Misgeneralisation [Langosco et al., 2022] explores how online RL agents in Procgen [Cobbe et al., 2019] can get confused about the goal they are pursuing since those goals co-occur with irrelevant artifacts in the environment most of the time.
Revised: Recent work on Goal Misgeneralisation [Langosco et al., 2022] explores how online RL agents in Procgen [Cobbe et al., 2019] can get confused about the goal they are pursuing since those goals co-occur with irrelevant artifacts in the environment most of the time.

Original: In this case, the specification is correct, but the agent still pursues an unintended objective (as opposed to poor reward definitions that predictably lead to reward hacking).
Revised: In this case, the specification is correct, but the agent still pursues an unintended objective (as opposed to poor reward definitions that predictably lead to reward hacking).

Original: We build upon an environment introduced in this work to collect datasets to reproduce the phenomenon of causal confusion in offline RL.
Revised: We build upon an environment introduced in this work to collect datasets to reproduce the phenomenon of causal confusion in offline RL.

Original: C Setup: Causally Ambiguous State–Action Trajectories
Revised: C Setup: Causally Ambiguous State–Action Trajectories

Original: We adopt the terminology of de Haan et al.
Revised: We adopt the terminology of de Haan et al.

Original: [2019] and partition the fully observable state space S into n factors, composed of m causal factors which are causal parents of the optimal state–action value function Q ∗ ( S, · ) , and n − m nuisance variables which do not influence Q ∗ ( S, · ) : S =
Revised: [2019] and partition the fully observable state space S into n factors, composed of m causal factors which are causal parents of the optimal state–action value function Q ∗ ( S, · ) , and n − m nuisance variables which do not influence Q ∗ ( S, · ) : S =

Original: [ S 1 , S 2 , ...,
Revised: [ S 1 , S 2 , ...,

Original: S m , S m +1 ,..., S n ] .
Revised: S m , S m +1 ,..., S n ] .

Original: Note that in this section, the superscript denotes the index of a point in the dataset ( t ) whereas the subscript denotes the index of the latent causal factor.
Revised: Note that in this section, the superscript denotes the index of a point in the dataset ( t ) whereas the subscript denotes the index of the latent causal factor.

Original: In this case, the following conditional independence relation indicating a lack of causal influence of S m +1: n on Q ( S, · ) should hold for every state s and action a :
Revised: In this case, the following conditional independence relation indicating a lack of causal influence of S m +1: n on Q ( S, · ) should hold for every state s and action a :

Original: [Equation]
Revised: [Equation]

Original: We assume that the offline data D is distributed such that S m +1: n is highly correlated with Q ( S, · ) , and so the independence relation does not hold when the function Q is statistically estimated from transitions in D .
Revised: We assume that the offline data D is distributed such that S m +1: n is highly correlated with Q ( S, · ) , and so the independence relation does not hold when the function Q is statistically estimated from transitions in D .

Original: Regardless of the full observability of all the true causal factors in the input S t , several works including that of Shah et al.
Revised: Regardless of the full observability of all the true causal factors in the input S t , several works including that of Shah et al.

Original: [2020] have shown that a simplicity bias in stochastic gradient descent makes it likely for the learned Q -function to use S m +1: n as a predictive feature.
Revised: [2020] have shown that a simplicity bias in stochastic gradient descent makes it likely for the learned Q -function to use S m +1: n as a predictive feature.

Original: This becomes more likely as the correlation between S m +1: n and Q ( S, · ) becomes stronger in the dataset.
Revised: This becomes more likely as the correlation between S m +1: n and Q ( S, · ) becomes stronger in the dataset.

Original: In our experiments, we simulate the presence of the non-causal factors S tm +1: n by augmenting the dataset with minimal nuisance correlations and verify that agents trained with these additions exhibit a significant decrease in performance at evaluation.
Revised: In our experiments, we simulate the presence of the non-causal factors S tm +1: n by augmenting the dataset with minimal nuisance correlations and verify that agents trained with these additions exhibit a significant decrease in performance at evaluation.

Original: We ascribe this deterioration in performance to the model choosing to attend to the spurious correlate and mapping two meaningfully distinct states s 1 and s 2 (where arg max( Q ∗ ( s 1 , · ))
Revised: We ascribe this deterioration in performance to the model choosing to attend to the spurious correlate and mapping two meaningfully distinct states s 1 and s 2 (where arg max( Q ∗ ( s 1 , · ))

Original: ̸ = arg max( Q ∗ ( s 2 , · )) ) to a similar distribution of Q -values.
Revised: ̸ = arg max( Q ∗ ( s 2 , · )) ) to a similar distribution of Q -values.

Original: We see this during qualitative evaluation in our benchmark environments, which we design to test the influence of the spurious correlate on the learned policy’s performance.
Revised: We see this during qualitative evaluation in our benchmark environments, which we design to test the influence of the spurious correlate on the learned policy’s performance.

Original: Notably, since D does contain (a potentially small number of) datapoints that break the dominant spurious correlations, the challenge in this setting is subtly different from the broader problem of achieving good generalisation, where the goal is to learn state abstractions that lead to accurate Q -values on new states with small unseen variations around concepts represented in the training data.
Revised: Notably, since D does contain (a potentially small number of) datapoints that break the dominant spurious correlations, the challenge in this setting is subtly different from the broader problem of achieving good generalisation, where the goal is to learn state abstractions that lead to accurate Q -values on new states with small unseen variations around concepts represented in the training data.

Original: In our setting where the long-tailed
Revised: In our setting where the long-tailed

Original: D is composed of a large number of causally-ambiguous trajectories, we want a model to learn state abstractions that at the very least explain all of the training data and not only the modes of the distribution of scenarios represented in D .
Revised: D is composed of a large number of causally-ambiguous trajectories, we want a model to learn state abstractions that at the very least explain all of the training data and not only the modes of the distribution of scenarios represented in D .

Original: In Appendix D.1, we describe the connection between
Revised: In Appendix D.1, we describe the connection between

Original: Q -values (and derived quantities like the advantage function) in RL and quantities of interest in the field of causal inference.
Revised: Q -values (and derived quantities like the advantage function) in RL and quantities of interest in the field of causal inference.

Original: This will elucidate why we chose active sampling as a tool for causal discovery in the following section.
Revised: This will elucidate why we chose active sampling as a tool for causal discovery in the following section.

Original: D Causal Inference D.1 Estimating the Conditional Average Treatment Effect from Causally Ambiguous Data
Revised: D Causal Inference D.1 Estimating the Conditional Average Treatment Effect from Causally Ambiguous Data

Original: Treatment-effect estimation, where the goal is to estimate the effect of a treatment T ∈ T on the outcome Y ∈ Y for individuals described by covariates X ∈ X , is a central problem in causal inference.
Revised: Treatment-effect estimation, where the goal is to estimate the effect of a treatment T ∈ T on the outcome Y ∈ Y for individuals described by covariates X ∈ X , is a central problem in causal inference.

Original: In particular, we may wish to estimate the expected difference in potential outcomes for individuals when subjected ( t = 1 ) or not subjected ( t = 0 ) to a treatment t , measured by the
Revised: In particular, we may wish to estimate the expected difference in potential outcomes for individuals when subjected ( t = 1 ) or not subjected ( t = 0 ) to a treatment t , measured by the

Original: Conditional Average Treatment Effect (CATE, Abrevaya et al.
Revised: Conditional Average Treatment Effect (CATE, Abrevaya et al.

Original: [2015]), defined as
Revised: [2015]), defined as

Original: [Equation]
Revised: [Equation]

Original: Here the treatment t is often considered to be a binary variable, but the definition can be extended to the multivariate case with continuous or discrete values.
Revised: Here the treatment t is often considered to be a binary variable, but the definition can be extended to the multivariate case with continuous or discrete values.

Original: Realizations of the random variables
Revised: Realizations of the random variables

Original: X, T, Y are denoted by x, t, y , respectively.
Revised: X, T, Y are denoted by x, t, y , respectively.

Original: We list the full set of assumptions needed to ensure identifiability of the CATE estimator in Appendix D.2.
Revised: We list the full set of assumptions needed to ensure identifiability of the CATE estimator in Appendix D.2.

Original: To frame the problem of disambiguating the effect of different actions (i.e., treatments t ∈ T ) in a given state (i.e., a set of covariates x ∈ X ), we frame estimation of the advantage functions
Revised: To frame the problem of disambiguating the effect of different actions (i.e., treatments t ∈ T ) in a given state (i.e., a set of covariates x ∈ X ), we frame estimation of the advantage functions

Original: A ( s, a ) =
Revised: A ( s, a ) =

Original: Q ( s, a ) − V ( s ) in reinforcement learning as CATE estimation [Pan et al., 2021].
Revised: Q ( s, a ) − V ( s ) in reinforcement learning as CATE estimation [Pan et al., 2021].

Original: In particular, the outcome Y corresponds to the Q -function (i.e., the expected return) for a given state–action pair, and we can express the corresponding CATE estimator as
Revised: In particular, the outcome Y corresponds to the Q -function (i.e., the expected return) for a given state–action pair, and we can express the corresponding CATE estimator as

Original: [Equation]
Revised: [Equation]

Original: which indicates the advantage of executing an action a at state s as opposed to any other alternative action.
Revised: which indicates the advantage of executing an action a at state s as opposed to any other alternative action.

Original: With this connection in established, we can now cast active sampling of data points used to estimate a value function as a sequential process of accurately estimating treatment effects.
Revised: With this connection in established, we can now cast active sampling of data points used to estimate a value function as a sequential process of accurately estimating treatment effects.

Original: D.2 Assumptions in Causal Inference
Revised: D.2 Assumptions in Causal Inference

Original: In this section we list out the assumptions adopted within the potential outcomes framework defined in the Rubin-Neyman causal model
Revised: In this section we list out the assumptions adopted within the potential outcomes framework defined in the Rubin-Neyman causal model

Original: [Rubin, 1974, Splawa-Neyman et al., 1990], to ensure identifiability of the CATE estimator.
Revised: [Rubin, 1974, Splawa-Neyman et al., 1990], to ensure identifiability of the CATE estimator.

Original: Assumption 2.1.
Revised: Assumption 2.1.

Original: (Consistency)
Revised: (Consistency)

Original: y
Revised: y

Original: = ty t + (1 − t )
Revised: = ty t + (1 − t )

Original: y 1 − t , i.e. an individual’s observed outcome y given assigned treatment t is identical to their potential outcome y t .
Revised: y 1 − t , i.e. an individual’s observed outcome y given assigned treatment t is identical to their potential outcome y t .

Original: Assumption 2.2.
Revised: Assumption 2.2.

Original: (Unconfoundedness) (cid:0) Y 0 , Y 1 (cid:1) ⊥ T | X .
Revised: (Unconfoundedness) (cid:0) Y 0 , Y 1 (cid:1) ⊥ T | X .

Original: Assumption 2.3.
Revised: Assumption 2.3.

Original: (Overlap) 0 < π t ( x ) < 1 : ∀ t ∈ T , where π t ( x ) ≡ P ( T = t | X = x ) is the propensity for treatment for individuals described by covariates X = x .
Revised: (Overlap) 0 < π t ( x ) < 1 : ∀ t ∈ T , where π t ( x ) ≡ P ( T = t | X = x ) is the propensity for treatment for individuals described by covariates X = x .

Original: When these assumptions are satisfied, (cid:98) τ ( x ) ≡ E [ Y | T = 1 , X =
Revised: When these assumptions are satisfied, (cid:98) τ ( x ) ≡ E [ Y | T = 1 , X =

Original: x ] − E
Revised: x ] − E

Original: [ Y | T = 0 , X = x ] is an unbiased estimator of τ ( x ) and is identifiable from observational data.
Revised: [ Y | T = 0 , X = x ] is an unbiased estimator of τ ( x ) and is identifiable from observational data.

Original: In practice, one or more of these assumptions are relaxed because observational data that is collected for safety-critical tasks has very little scope for randomisation of a behaviour policy, indicating limited overlap.
Revised: In practice, one or more of these assumptions are relaxed because observational data that is collected for safety-critical tasks has very little scope for randomisation of a behaviour policy, indicating limited overlap.

Original: 1: Initialise ensemble Q -function Q θ , n ep =epochs, d sz =dataset size, b sz =batch size, T =steps-per-epoch. 2: for epoch e in {1, ..., n ep } do 3: for step t in { 1, ..., T} do 4: compute scores acq i over D train =
Revised: [Table]

Original: [-[ s i , a i ] d sz i =
-]
Revised: [-[ s i , a i ] d sz i =
-][Table] [Table] [Table]

Original: according to the acquisition function 5: acq i = acq i (cid:80) dszj =1 acq j (normalise scores) 6: sample batch B
Revised: [Table]

Original: [ s s i r
Revised: [Table]

Original: ] b sz i
Revised: [Table]

Original: D multinomial acq
Revised: [Table]

Original: Train
Revised: [Table]

Original: D from
Revised: [Table]

Original: ) 8: end for 9: end for
Revised: [Table] [Table] [Table]

Original: E Implementation
Revised: E Implementation

Original: All our environments use a discrete action space.
Revised: All our environments use a discrete action space.

Original: Therefore we build our method on top of the double-DQN implementation similar to the original CQL paper.
Revised: Therefore we build our method on top of the double-DQN implementation similar to the original CQL paper.

Original: As stated in Section 3.1, we use ensembles of Q -networks where the ensemble members are trained with independent targets.
Revised: As stated in Section 3.1, we use ensembles of Q -networks where the ensemble members are trained with independent targets.

Original: At evaluation time, we average the Q -value outputs of the ensemble and select the action with the maximum Q -value.
Revised: At evaluation time, we average the Q -value outputs of the ensemble and select the action with the maximum Q -value.

Original: At other points during training where we need to do inference (for instance: to compute Q -values for computation of the conservative loss), we take the average across the ensemble as the Q -value outputs, similar to evaluation time.
Revised: At other points during training where we need to do inference (for instance: to compute Q -values for computation of the conservative loss), we take the average across the ensemble as the Q -value outputs, similar to evaluation time.

Original: A design choice we make is the network initialisation used when we start to do active sampling: the uncertainty encoded by a randomly initialised network at epoch 0 can be very inaccurate and biased towards some subset of the data.
Revised: A design choice we make is the network initialisation used when we start to do active sampling: the uncertainty encoded by a randomly initialised network at epoch 0 can be very inaccurate and biased towards some subset of the data.

Original: We instead train with uniform sampling for a single epoch and then start active sampling second epoch onwards.
Revised: We instead train with uniform sampling for a single epoch and then start active sampling second epoch onwards.

Original: We also observe that bigger ensembles encode better measures of uncertainty than smaller ones when the network is trained for a lesser period than one epoch with uniform sampling at the start (Figure 6).
Revised: We also observe that bigger ensembles encode better measures of uncertainty than smaller ones when the network is trained for a lesser period than one epoch with uniform sampling at the start (Figure 6).

Original: E.1 Computing Acquisition Scores
Revised: E.1 Computing Acquisition Scores

Original: Algorithm 1 (line 4) indicates that acquisition scores are computed for data points in D train before every training step.
Revised: Algorithm 1 (line 4) indicates that acquisition scores are computed for data points in D train before every training step.

Original: However, in practice, computing the scores over all the transitions in the dataset can be both expensive and redundant since high-error or high-uncertainty points will likely stay high over a short window of subsequent gradient steps.
Revised: However, in practice, computing the scores over all the transitions in the dataset can be both expensive and redundant since high-error or high-uncertainty points will likely stay high over a short window of subsequent gradient steps.

Original: Thus, we only recompute all the scores after every n gradient steps and vary n as a hyper-parameter in our experiments.
Revised: Thus, we only recompute all the scores after every n gradient steps and vary n as a hyper-parameter in our experiments.

Original: In Appendix G, we also explore a scheme where scores are recomputed only over-sampled batches, as in [Schaul et al., 2015].
Revised: In Appendix G, we also explore a scheme where scores are recomputed only over-sampled batches, as in [Schaul et al., 2015].

Original: This is an approximation where a priority queue is maintained with priorities of every data point deriving from the acquisition score computed on it.
Revised: This is an approximation where a priority queue is maintained with priorities of every data point deriving from the acquisition score computed on it.

Original: In this case, the priorities are only updated for a small subset of points at every gradient step, since the acquisition scores are only recomputed on the data points in the sampled batch, potentially leading to many points in the replay buffer with stale scores.
Revised: In this case, the priorities are only updated for a small subset of points at every gradient step, since the acquisition scores are only recomputed on the data points in the sampled batch, potentially leading to many points in the replay buffer with stale scores.

Original: We adopt the same implementation as that of PER Schaul et al.
Revised: We adopt the same implementation as that of PER Schaul et al.

Original: [2015] with the scores for the priorities coming from the TD-error and Variance estimates.
Revised: [2015] with the scores for the priorities coming from the TD-error and Variance estimates.

Original: We include further details about this procedure in the appendix.
Revised: We include further details about this procedure in the appendix.

Original: We refer to the above two cases by appending -dataset and -batch to the names of the sampling schemes to indicate that scores are recomputed for the entire data versus just the sampled batch respectively.
Revised: We refer to the above two cases by appending -dataset and -batch to the names of the sampling schemes to indicate that scores are recomputed for the entire data versus just the sampled batch respectively.

Original: F Evaluation Metrics
Revised: F Traffic-World Experiment: Additional Details

Original: We present the reward curves directly in most of our experimental reporting.
Revised: In the traffic-world environment with the spurious correlate as the yellow tile, the following are two examples of tail-case scenarios in the corresponding collected dataset: (1) When the tile flashes yellow because the leading vehicle is static but the agent is not blocked because there is space to move ahead behind the leading vehicle, and (2) When the leading vehicle has crossed the traffic light which has just turned red in front of the agent, and the top-left tile is not flashing yellow because the leading vehicle is not static. However, if the agent takes the tile being off as a cue, it will decide to move forward and end up breaking the red light and fetching a negative reward.

Original: The curves are computed by taking the inter-quartile-mean (IQM) across seeds as proposed by Agarwal et al.
Revised: In Figure 2, we plot a histogram for the computed variance scores over the data–action’s advantage estimates from the model ( Variance-data scores) after one epoch of training with uniform sampling.

Original: We report the number of seeds used for all experiments in Appendix I.
Revised: We can verify that samples from the tail-case fall in the higher percentile of variance scores.

Original: Offline RL training performance is known to be non-monotonic, unlike supervised learning where the accuracy (or loss) increases (decreases) and then saturates.
Revised: The Maze environment in Procgen [Cobbe et al., 2019] defines a navigation task where the agent starts at the bottom left in the maze and receives a reward of +10 upon successfully reaching the goal which is sampled at any valid location in the maze. Langosco et al.

Original: Often the reward curves start to decrease after a period when overfitting to the data-action values (through the conservatism penalty) starts to happen.
Revised: [2022] recently showed that an agent trained on a series of environments with the goal always at the top-right will be causally confused about the source of the reward and will still navigate to the top-right even when the goal is sampled elsewhere.

Original: Since we define deterministic benchmark environments testing all the tail scenarios in the case of Traffic-World and Maze experiments, we can consider any point on the curve where the model solves the highest number of environments (achieves the highest reward) as the point of convergence (as opposed to it being a noisy spike due to stochastic evaluation).
Revised: We generate a skewed mixture dataset containing mostly episodes where the goal is sampled at the top-right and a few episodes where the goal is sampled randomly. Further details about the setup are described in the Appendix. Figure 3b shows the evaluation performance of random and active sampling agents trained on the mixture dataset when goals are sampled randomly in the evaluation environment. We plot the computation time for the uniform and active sampling

Revised: [Figure]

Revised: similarly in the fixed goal evaluation environment

Original:  This  is similar to taking the max-reward checkpoint as done in [Agarwal et al., 2020].
Revised: , but the active sampling variants achieve higher rewards the environments randomly sampled goals. This verifies that the model is not just performing well in one of the two kinds of environments, is not constrained by capacity, and the reason behind the lower performance of uniform sampling, in this case, is simply causal confusion variants in Figure 8 in the appendix and note that variance-based sampling reaches the highest score achieved by uniform sampling in lesser wall-clock time. Qualitative evaluations show that uniform sampling agents which achieve a lower reward still successfully navigate to the top right of hard , and are therefore only confused about

Original: However, to ensure that  the solutions learned by any method are recoverable, we want to be able to get a good checkpoint from the model without needing to evaluate it too often. Thus we consider a method as having achieved higher reward than other methods at any point in time if it maintains this gap for at least two subsequent post-epoch evaluations.
Revised: It is important to note that if the evaluation environments considered here had followed a distribution of scenarios similar to the training data distribution, then we would have observed an evaluation curve similar to the weighted sum of the two curves in the left and right panels of Figure 5. The weight of the curve on the right would contribute significantly more since those scenarios (fixed top-right goals) would occur much more often in the evaluation environment.

Original: This is akin to taking a windowed-max over the reward curves.
Revised: We would conclude that performance has saturated and the model has converged somewhere between 10,000–20, gradient steps since that is when the evaluation performance on the environment with fixed goals saturates. An active sampling model that was early-stopped at this point, would then fare much better on the tail cases (and do similarly otherwise) than the uniformly sampled model. This also motivates the design of our experiment setup where we define a set of deterministic evaluation environments for each scenario separately to enable model evaluation on an equally hard set of environments across seeds and methods and capture performance on both the head and tail of the distribution of scenarios.

Original: G Additional Results G.1 Causal Confusion in the ALE Benchmark
Revised: G.2 Causal Confusion in the ALE Benchmark

Original: We are now interested in evaluating our active sampling baselines on a larger pre-existing dataset with more realistic noise and variations.
Revised: We are now interested in evaluating our active sampling baselines on a larger pre-existing dataset with more realistic noise and variations.

Original: Prior work in imitation learning [Park et al., 2021, de Haan et al., 2019] has attempted to simulate causal confusion in specific Atari game-play datasets by modifying images to display the previous action taken by the agent.
Revised: Prior work in imitation learning [Park et al., 2021, de Haan et al., 2019] has attempted to simulate causal confusion in specific Atari game-play datasets by modifying images to display the previous action taken by the agent.

Original: This kind of causal confusion is inspired by robotics datasets that have trajectories with highly correlated (and predictable) sequences of actions because embodiment dictates that an agent’s state does not change too drastically between subsequent timesteps.
Revised: This kind of causal confusion is inspired by robotics datasets that have trajectories with highly correlated (and predictable) sequences of actions because embodiment dictates that an agent’s state does not change too drastically between subsequent timesteps.

Original: However, when we trained uniform sampling CQL agents with and without the previous action on the images for 20 Atari games, only a select few games actually exhibited causal confusion of this kind.
Revised: However, when we trained uniform sampling CQL agents with and without the previous action on the images for 20 Atari games, only a select few games actually exhibited causal confusion of this kind.

Original: We pick Enduro since we can consistently observe both a convergence speed degradation and final reward degradation upon adding the previous action to the image.
Revised: We pick Enduro since we can consistently observe both a convergence speed degradation and final reward degradation upon adding the previous action to the image.

Original: Enduro is a car racing game where an agent needs to overtake cars (receiving a positive reward for each car overtaken) and drive along a winding road, with a limited number of collisions allowed.
Revised: Enduro is a car racing game where an agent needs to overtake cars (receiving a positive reward for each car overtaken) and drive along a winding road, with a limited number of collisions allowed.

Original: Figure 4 shows our proposed active sampling baselines compared to uniform sampling in the following two cases that we previously discussed in Appendix E.1: (1) -dataset case when we recompute the scores across the dataset every few gradient steps and (2) -batch case when the scores are only recomputed for the sampled batch and updated in a priority queue.
Revised: Figure 4 shows our proposed active sampling baselines compared to uniform sampling in the following two cases that we previously discussed in Appendix E.1: (1) -dataset case when we recompute the scores across the dataset every few gradient steps and (2) -batch case when the scores are only recomputed for the sampled batch and updated in a priority queue.

Original: Also shown is the uniform sampling baseline trained without the previous action displayed on the image (data without causal ambiguity).
Revised: Also shown is the uniform sampling baseline trained without the previous action displayed on the image (data without causal ambiguity).

Original: We note that all of the active sampling variants perform significantly better than their uniform sampling counterpart not just in terms of convergence to their best performance but also in terms of the highest reward achieved.
Revised: We note that all of the active sampling variants perform significantly better than their uniform sampling counterpart not just in terms of convergence to their best performance but also in terms of the highest reward achieved.

Original: G.2 Effect of Predictive Uncertainty on Active Sampling and Causal Confusion
Revised: G.3 Effect of Predictive Uncertainty on Active Sampling and Causal Confusion

Original: To better understand the improvement in convergence speed of active sampling vis-‘a-vis uniform sampling, we train CQL with a smaller ( n = 3 ) and a larger ( n = 10 ) ensemble of Q -networks, all else equal, and sample identical transitions based on the predictive uncertainty of only either the smaller or the larger ensemble.
Revised: To better understand the improvement in convergence speed of active sampling vis-‘a-vis uniform sampling, we train CQL with a smaller ( n = 3 ) and a larger ( n = 10 ) ensemble of Q -networks, all else equal, and sample identical transitions based on the predictive uncertainty of only either the smaller or the larger ensemble.

Original: This experiment setup allows us to isolate the effect of the predictive variance on active sampling and therefore on the observed reduction in causal confusion.
Revised: This experiment setup allows us to isolate the effect of the predictive variance on active sampling and therefore on the observed reduction in causal confusion.

Original: Figure 5 shows the training curves for both ensembles, and we find that when transitions are sampled according to the predictive uncertainty of the larger ensemble, even the smaller ensemble converges faster (albeit to a smaller value than that achieved by the larger ensemble).
Revised: Figure 5 shows the training curves for both ensembles, and we find that when transitions are sampled according to the predictive uncertainty of the larger ensemble, even the smaller ensemble converges faster (albeit to a smaller value than that achieved by the larger ensemble).

Original: Similarly, when points are sampled according to the predictive uncertainty of the smaller ensemble, the larger ensemble converges to its highest reward more slowly.
Revised: Similarly, when points are sampled according to the predictive uncertainty of the smaller ensemble, the larger ensemble converges to its highest reward more slowly.

Original: This result implies that improved predictive uncertainty estimation, for which we use the size of the ensemble as a proxy, improves the ability of active sampling to identify samples that break spurious correlations in the data and reduce causal confusion.
Revised: This result implies that improved predictive uncertainty estimation, for which we use the size of the ensemble as a proxy, improves the ability of active sampling to identify samples that break spurious correlations in the data and reduce causal confusion.

Original: Notably, even when training CQL using the smaller ensemble, which we expect to produce worse mean Q -function estimates than the medium-size ensemble, using the larger ensemble’s predictive uncertainty for active sampling results in a higher reward than for the medium-size ensemble, implying that the positive impact of improved predictive uncertainty estimates outweighs the adverse impact of the reduction in accuracy of the
Revised: Notably, even when training CQL using the smaller ensemble, which we expect to produce worse mean Q -function estimates than the medium-size ensemble, using the larger ensemble’s predictive uncertainty for active sampling results in a higher reward than for the medium-size ensemble, implying that the positive impact of improved predictive uncertainty estimates outweighs the adverse impact of the reduction in accuracy of the

Original: Q -function estimator.
Revised: Q -function estimator.

Original: We provide a further discussion of these observations in Appendix H.
Revised: We provide a further discussion of these observations in Appendix I. H Evaluation Metrics

Revised: We present the reward curves directly in most of our experimental reporting.

Revised: The curves are computed by taking the inter-quartile-mean (IQM) across seeds as proposed by Agarwal et al.

Revised: We report the number of seeds used for all experiments in Appendix J.

Original: G.3 Code and Data
Revised: Offline RL training performance is known to be non-monotonic, unlike supervised learning where the accuracy (or loss) increases (decreases) and then saturates. Often the reward curves start to decrease after a period when overfitting to the data-action values (through the conservatism penalty) starts to happen. Since we define deterministic benchmark environments testing all the tail scenarios in the case of Traffic-World and Maze experiments, we can consider any point on the curve where the model solves the highest number of environments (achieves the highest reward) as the point of convergence (as opposed to it being a noisy spike due to stochastic evaluation). This is similar to taking the max-reward checkpoint as done in [Agarwal et al., 2020].

Revised: However, to ensure that the solutions learned by any method are recoverable, we want to be able to get a good checkpoint from the model without needing to evaluate it too often.

Revised: Thus we consider a method as having achieved higher reward than other methods at any point in time if it maintains this gap for at least two subsequent post-epoch evaluations.

Revised: This is akin to taking a windowed-max over the reward curves.

Revised: H.1 Code and Data

Original: We will release our code, data and pretrained models once the work is uploaded online.
Revised: We will release our code, data and pretrained models once the work is uploaded online.

Original: The code repository will also contain code to reproduce all the figures in this work.
Revised: The code repository will also contain code to reproduce all the figures in this work.

Original: G.4 Data Collection
Revised: H.2 Data Collection

Original: Traffic-World: To collect data for Offline RL, we trained a PPO agent on a slightly modified version of the Traffic-world environment, with reward shaping on the environment [Anonymous, 2021], to incentivise the agent to reach the goal since this is a hard exploration environment (there is the potential to receive many negative rewards before receiving a positive reward, and without reward shaping the PPO agent just learns to toggle in-place till the episode ends to avoid negative penalties).
Revised: Traffic-World: To collect data for Offline RL, we trained a PPO agent on a slightly modified version of the Traffic-world environment, with reward shaping on the environment [Anonymous, 2021], to incentivise the agent to reach the goal since this is a hard exploration environment (there is the potential to receive many negative rewards before receiving a positive reward, and without reward shaping the PPO agent just learns to toggle in-place till the episode ends to avoid negative penalties).

Original: Maze: We use the Impala-based PPO agent trained in [Langosco et al., 2022] for 200M steps to collect the expert trajectories on 6000 episodes of episodes with randomised goals and 200 episodes of episodes with fixed goals.
Revised: Maze: We use the Impala-based PPO agent trained in [Langosco et al., 2022] for 200M steps to collect the expert trajectories on 6000 episodes of episodes with randomised goals and 200 episodes of episodes with fixed goals.

Original: H Further Analysis
Revised: I Further Analysis

Original: In addition to the figure in the main paper we plot in Figure 6 the training curves of two ensembles when they’ve been initialised with a longer period of uniform sampling at the start.
Revised: In addition to the figure in the main paper, we plot in Figure 6 the training curves of two ensembles when they’ve been initialised with a longer period of uniform sampling at the start.

Original: We see that the gap between the informative-ness of samples fetched based on their uncertainties reduces (as compared to Figure 5).
Revised: We see that the gap between the informative-ness of samples fetched based on their uncertainties reduces (as compared to Figure 5).

Original: The larger ensemble’s uncertainty still converges quicker, however, than when both ensembles are training using the smaller ensemble’s uncertainty for data sampling.
Revised: The larger ensemble’s uncertainty still converges quicker, however, than when both ensembles are training using the smaller ensemble’s uncertainty for data sampling.

Original: One observation we made was related to active sampling results in the case of benchmarks with short episodes ( n steps ∼ 20 − 30 ) and sparse rewards, in our case Traffic-World and
Revised: One observation we made was related to active sampling results in the case of benchmarks with short episodes ( n steps ∼ 20 − 30 ) and sparse rewards, in our case Traffic-World and

Original: Maze.
Revised: Maze.

Original: Q-learning is trained through bootstrapping where we minimise the TD-error which involves estimating the Q -values at successive states across transitions in a trajectory.
Revised: Q-learning is trained through bootstrapping where we minimise the TD-error which involves estimating the Q -values at successive states across transitions in a trajectory.

Original: Additionally, the CQL objective has the gap-expanding property because the conservatism penalty tries to push Q -values of different actions at a state apart (pushing up the data-action value and pushing down others to some extent).
Revised: Additionally, the CQL objective has the gap-expanding property because the conservatism penalty tries to push Q -values of different actions at a state apart (pushing up the data-action value and pushing down others to some extent).

Original: Sometimes in the case of repeated sampling of a particular transition tuple, there is a potential divergence of Q -values of the nearby (preceding and following) state-action pairs.
Revised: Sometimes in the case of repeated sampling of a particular transition tuple, there is a potential divergence of Q -values of the nearby (preceding and following) state-action pairs.

Original: This can be seen as an explosion of Q -values in the training metrics and can be partially resolved by gradient clipping.
Revised: This can be seen as an explosion of Q -values in the training metrics and can be partially resolved by gradient clipping.

Original: We only observed this in the case of TD − Error and V ariance − data variants.
Revised: We only observed this in the case of TD − Error and V ariance − data variants.

Original: One experiment we tried (with some success) was that of episodic sampling, where we use a heuristic to convert individual transition-wise acquisition scores to scores over entire episodes (for example taking the maximum acquisition score over transitions in an episode).
Revised: One experiment we tried (with some success) was that of episodic sampling, where we use a heuristic to convert individual transition-wise acquisition scores to scores over entire episodes (for example taking the maximum acquisition score over transitions in an episode).

Original: This kind of episodic sampling turns out to give much more stable training curves but involves additional hyperparameters and heuristics.
Revised: This kind of episodic sampling turns out to give much more stable training curves but involves additional hyperparameters and heuristics.

Original: The reasoning for its success is likely related to the motivation of algorithms like emphaticTD [Rupam Mahmood et al., 2015], Reverse Experience Replay (RER)
Revised: The reasoning for its success is likely related to the motivation of algorithms like emphaticTD [Rupam Mahmood et al., 2015], Reverse Experience Replay (RER)

Original: [Rotinov, 2019].
Revised: [Rotinov, 2019].

Original: In these works, they propose not just prioritising transitions with high TD-error, but also increasing the priority of transitions preceding these ones in the priority queue, since these transitions will have informative
Revised: In these works, they propose not just prioritising transitions with high TD-error, but also increasing the priority of transitions preceding these ones in the priority queue, since these transitions will have informative

Original: TD-updates in the subsequent time-steps.
Revised: TD-updates in the subsequent time-steps.

Original: We show the reward curves with episodic and without episodic active sampling of data corresponding to the experiment in Traffic-World with the yellow tile present in the data.
Revised: We show the reward curves with episodic and without episodic active sampling of data corresponding to the experiment in Traffic-World with the yellow tile present in the data.

Original: H.1 Computational Cost
Revised: I.1 Computational Cost

Original: as reward Figure 8 shows a scatterplot for the wallclock times to achieve highest reward across different active and uniform baselines (labelled as TD-Best, Variance-Best and Random-Best).
Revised: Figure 8 shows a scatterplot for the wallclock times to achieve highest reward across different active and uniform baselines (labelled as TD-Best, Variance-Best and Random-Best).

Original: It also plots the time needed for active sampling variants to achieve the best reward that uniform sampling achieves (denoted as Variancepar-Random and TD-par-Random in the plot).
Revised: It also plots the time needed for active sampling variants to achieve the best reward that uniform sampling achieves (denoted as Variance-par-Random and TD-par-Random in the plot).

Original: I Hyper-parameters
Revised: J Hyper-parameters

Original: Active Sampling.
Revised: Active Sampling.

Original: We kept all the hyperparameters the same as uniform sampling (batch size, learning rate, α ).
Revised: We kept all the hyper-parameters the same as uniform sampling (batch size, learning rate, α ).

Original: The parameters related to active sampling are : (1) n : the number of gradient steps with stale scores we take before we recompute acquisition scores on the data.
Revised: The parameters related to active sampling are : (1) n : the number of gradient steps with stale scores we take before we recompute acquisition scores on the data.

Original: (2) the ensemble size which we keep constant across the active and uniform sampling variants for a fair comparison.
Revised: (2) the ensemble size which we keep constant across the active and uniform sampling variants for a fair comparison.

Original: [Table]
Revised: [Table]

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