Original: Toward Causal-Aware RL: State-Wise Action-Reﬁned Temporal Difference
Revised: Toward Causal-Aware RL: State-Wise Action-Reﬁned Temporal Difference

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Although it is well known that exploration plays a key role in Reinforcement
Revised: Although it is well known that exploration plays a key role in Reinforcement

Original: Learning (RL), prevailing exploration strategies for continuous control tasks in
Revised: Learning (RL), prevailing exploration strategies for continuous control tasks in

Original: RL are mainly based on naive isotropic Gaussian noise regardless of the causalityrelationship between action space and the task and consider all dimensions ofactions equally important.
Revised: RL are mainly based on naive isotropic Gaussian noise regardless of the causality relationship between action space and the task and consider all dimensions of actions equally important.

Original: In this work, we propose to conduct interventions onthe primal action space to discover the causal relationship between the actionspace and the task reward.
Revised: In this work, we propose to conduct interventions on the primal action space to discover the causal relationship between the action space and the task reward.

Original: We propose the method of State-Wise Action Reﬁned(SWAR), which addresses the issue of action space redundancy and promotecausality discovery in RL.
Revised: We propose the method of State-Wise Action Reﬁned (SWAR), which addresses the issue of action space redundancy and promote causality discovery in RL.

Original: We formulate causality discovery in RL tasks as a statedependent action space selection problem and propose two practical algorithmsas solutions.
Revised: We formulate causality discovery in RL tasks as a statedependent action space selection problem and propose two practical algorithms as solutions.

Original: The ﬁrst approach, TD-SWAR, detects task-related actions duringtemporal difference learning, while the second approach, Dyn-SWAR, revealsimportant actions through dynamic model prediction.
Revised: The ﬁrst approach, TD-SWAR, detects task-related actions during temporal difference learning, while the second approach, Dyn-SWAR, reveals important actions through dynamic model prediction.

Original: Empirically, both methodsprovide approaches to understand the decisions made by RL agents and improvelearning efﬁciency in action-redundant tasks.
Revised: Empirically, both methods provide approaches to understand the decisions made by RL agents and improve learning efﬁciency in action-redundant tasks.

Original:  Introduction
Revised:  Introduction

Original: Although model-free RL has achieved great success in various challenging tasks and outperformsexperts in most cases [21, 26, 17, 33, 4], the design of action space always requires elaboration.
Revised: Although model-free RL has achieved great success in various challenging tasks and outperforms experts in most cases [21, 26, 17, 34, 4] the design of action space always requires elaboration.

Original: Forexample, in the game StarCraftII, hundreds of units can be selected and controlled to perform variousactions.
Revised: For example in the game StarCraftII, hundreds of units can be selected and controlled to perform various actions.

Original: To tackle the difﬁculty in exploration caused by the extremely large action and state space,hierarchical action space design and imitation learning are used [27, 33] to reduce the explorationspace.
Revised: To tackle the difﬁculty in exploration caused by the extremely large action and state space, hierarchical action space design and imitation learning are used [27, 34] to reduce the exploration space.

Original: Both of those approaches require expert knowledge of the task.
Revised: Both of those approaches require expert knowledge of the task.

Original: On the other hand, even in thecontext of imitation learning where expert data is assumed to be accessible, causal confusion will stillhinder the performance of an agent [8].
Revised: On the other hand, even in the context of imitation learning where expert data is assumed to be accessible, causal confusion will still hinder the performance of an agent [8].

Original: Those defects motivate us to explore the causality-awarenessof an agent that permits an agent to discover the causal relationship for the environment and selectuseful dimensions of action space during policy learning in pursuance of improved learning efﬁciency.
Revised: Those defects motivate us to explore the causality-awareness of an agent that permits an agent to discover the causal relationship for the environment and select useful dimensions of action space during policy learning in pursuance of improved learning efﬁciency.

Original: Another motivating example is the in-hand manipulation tasks [2]: robotics equipped with touch sensors outperforms the policies learned without sensors by a clear margin in hand-in manipulation tasks [20], showing the importance of causality discovery between actions and feedbacks in RL.
Revised: Another motivating example is the in-hand manipulation tasks [2]: robotics equipped with touch sensors outperforms the policies learned without sensors by a clear margin in hand-in manipulation tasks [20], showing the importance of causality discovery between actions and feedbacks in RL.

Original: A similar example can be found in human learning: knowing nothing about how to control the ﬁnger joints ﬂexibly may not hinder a baby learns to walk, and a baby has not learned how to walk can still learn to use forks and spoons skillfully, inspiring us to believe that the challenge for exploration can be greatly eased after the causality between action space and the given task is learned
Revised: A similar example can be found in human learning: knowing nothing about how to control the ﬁnger joints ﬂexibly may not hinder a baby learns to walk, and a baby has not learned how to walk can still learn to use forks and spoons skillfully, inspiring us to believe that the challenge for exploration can be greatly eased after the causality between action space and the given task is learned.

Original: In this work, the recent advance of instance-wise feature selection technique [37] is improved to be more suitable in large-scale state-wise action selection tasks and adapted to the time-series causal discovery setting to select state-conditioned action space in RL with redundant action space
Revised: In this work, the recent advance of instance-wise feature selection technique [38] is improved to be more suitable in large-scale state-wise action selection tasks and adapted to the time-series causal discovery setting to select state-conditioned action space in RL with redundant action space.

Original: With the proposed method, the agent learns to perform intervention, discover the true structural causal model(SCM) and select task-related actions for a given task, remarkably reduces the burden of explorationand obtains on-par learning efﬁciency as well as asymptotic performance compared with agentstrained in the oracle settings where the action spaces are pruned according to given tasks manually.
Revised: With the proposed method, the agent learns to perform intervention, discover the true structural causal model (SCM) and select task-related actions for a given task, remarkably reduces the burden of exploration and obtains on-par learning efﬁciency as well as asymptotic performance compared with agents trained in the oracle settings where the action spaces are pruned according to given tasks manually.

Original:  Preliminary
Revised:  Preliminary

Original: Markov Decision Processes RL tasks can be formally deﬁned as Markov Decision Processes(MDPs), where an agent interacts with the environment and learns to make decision at every timestep.
Revised: Markov Decision Processes RL tasks can be formally deﬁned as Markov Decision Processes (MDPs), where an agent interacts with the environment and learns to make decision at every timestep.

Original: Formally, we consider the deterministic MDP with a ﬁxed horizon H 2 N + denoted by a tuple( S , A , H, r,  , T , ⇢ 0 ) , where S and A are the |S| -dimensional state and |A| -dimensional action space;45 r : S ⇥ A 7! R denotes the reward function;  2 (0 , 1] is the discount factor indicating importance46 of present returns compared with long-term returns; 
Revised: Formally, we consider the deterministic MDP with a ﬁxed horizon H ∈ N + denoted by a tuple ( S , A , H, r, γ, T , ρ 0 ) , where S and A are the |S| -dimensional state and |A| -dimensional action space; r : S × A (cid:55)→ R denotes the reward function; γ ∈ (0 , 1] is the discount factor indicating importance of present returns compared with long-term returns; T : S × A (cid:55)→ S denotes the transition dynamics; ρ 0 is the initial state distribution.

Original: T : S ⇥ A 7! S denotes the transition dynamics; 47 ⇢ 0 is the initial state distribution. We use ⇧ to represent the stationary deterministic policy class, i.e., ⇧ = { ⇡ : S 7!
Revised: We use Π to represent the stationary deterministic policy class, i.e., Π = { π : S (cid:55)→ A} .

Original: A} . The49 learning objective of an RL algorithm is to ﬁnd ⇡ ⇤ 2 ⇧ as the solution of the following optimizationproblem: max ⇡ 2 ⇧ E ⌧ ⇠ ⇢ 0 , ⇡ , T [ 
Revised: The learning objective of an RL algorithm is to ﬁnd π ∗ ∈ Π as the solution of the following optimization problem: max π ∈ Π E τ ∼ ρ 0 ,π, T [ (cid:80)

Original: P
Revised: Ht

Original: Ht =1  t r t ] where the expectation is taken over the trajectory ⌧ =( s 1 , a 1 , r 1 , . . .
Revised: =1 γ t r t ] where the expectation is taken over the trajectory τ = ( s 1 , a 1 , r 1 , . . .

Original: , s H , a H , r H ) generated by policy ⇡ under the environment T , starting from s 0 ⇠ ⇢ 0 .
Revised: , s H , a H , r H ) generated by policy π under the environment T , starting from s 0 ∼ ρ 0 .

Original: INVASE INVASE is proposed by [37] to perform instance-wise feature selection to reduce overﬁtting in predictive models.
Revised: INVASE INVASE is proposed by [38] to perform instance-wise feature selection to reduce overﬁtting in predictive models.

Original: The learning objective is to minimize the KL-Divergence of the fullconditional distribution and the minimal-selected-features-only conditional distribution of the outcome, i.e., min F L , with
Revised: The learning objective is to minimize the KL-Divergence of the fullconditional distribution and the minimal-selected-features-only conditional distribution of the outcome, i.e., min F L , with

Original: [Equation]
Revised: [Equation]

Original: where F : X ! { 0 , 1 } d is a feature selection function and | F ( x ) | 0 denotes the cardinality ( l 0norm) of selected features, i.e., the number of 1 ’s in F ( x ) .
Revised: where F : X → { 0 , 1 } d is a feature selection function and | F ( x ) | 0 denotes the cardinality ( lnorm) of selected features, i.e., the number of 1 ’s in F ( x ) .

Original: 1 d is the dimension of input features.
Revised: 1 d is the dimension of input features.

Original: x ( F ( x ))
Revised: x ( F ( x ))

Original: =
Revised: =

Original: F ( x )  x denotes the element-wise product of x and generated mask m = F ( x ) .
Revised: F ( x ) (cid:12) x denotes the element-wise product of x and generated mask m = F ( x ) .

Original: Ideally, the optimal selection function F should be able to minimize the two terms in Equation (1)simultaneously.
Revised: Ideally, the optimal selection function F should be able to minimize the two terms in Equation (1) simultaneously.

Original: INVASE applies the Actor-Critic framework in the optimization of F through sampling, wheref ✓ ( ·| x ) , parameterized by a neural network ✓ 2 , is used as a stochastic actor.
Revised: INVASE applies the Actor-Critic framework in the optimization of F through sampling, where f θ ( ·| x ) , parameterized by a neural network θ 2 , is used as a stochastic actor.

Original: Two predictive networks
Revised: Two predictive networks

Original: C  ( · ) , B  ( · ) are considered as the critic and the baseline network used for variance reduction [35]and trained with the Cross-Entropy loss to produce return signal L , based on which f ✓ ( ·| x ) can beoptimized through policy gradient:
Revised: C φ ( · ) , B ψ ( · ) are considered as the critic and the baseline network used for variance reduction [36] and trained with the Cross-Entropy loss to produce return signal L , based on which f θ ( ·| x ) can be optimized through policy gradient:

Original: [Equation]
Revised: [Equation]

Original: Finally, F ( x ) = ( F 1 ( x ) , ..., F d ( x )) can be get by sampling from f ( ·| x ) =
Revised: Finally, F ( x ) = ( F 1 ( x ) , ..., F d ( x )) can be get by sampling from f ( ·| x ) =

Original: ( f 1 ( x ) , ..., f d ( x )) , with 
Revised: ( f 1 ( x ) , ..., f d ( x )) , with F i ( x

Original: [Equation]
Revised: [Equation]

Original: i
Revised: = (cid:26) [Equation]

Original: [Equation]
Revised: [Equation]

Original:  Proposed Method
Revised:  Proposed Method

Original: The objective of this work is to carry out state-wise action selection in RL through intervention,and thereby enhance the learning efﬁciency with a pruned task-related action space after ﬁnding thecorrect causal model.
Revised: The objective of this work is to carry out state-wise action selection in RL through intervention, and thereby enhance the learning efﬁciency with a pruned task-related action space after ﬁnding the correct causal model.

Original: Section 3.1 starts with the formalization of the action space reﬁnery objectivein RL tasks under the framework of causal discovery.
Revised: Section 3.1 starts with the formalization of the action space reﬁnery objective in RL tasks under the framework of causal discovery.

Original: Section 3.2 introduces SWAR, which improvesthe scalability of INVASE in high dimensional variable selection tasks.
Revised: Section 3.2 introduces SWAR, which improves the scalability of INVASE in high dimensional variable selection tasks.

Original: We integrate SWAR withdeterministic policy gradient methods [25] in Section 3.3 to perform state-wise action space pruning,resulting in two practical causality-aware RL algorithms.
Revised: We integrate SWAR with deterministic policy gradient methods [25] in Section 3.3 to perform state-wise action space pruning, resulting in two practical causality-aware RL algorithms.

Original: 3.1 Temporal Difference Objective with Structural Causal Models
Revised: 3.1 Temporal Difference Objective with Structural Causal Models

Original: In modern RL algorithms, the most general approach is based on the Actor-Critic framework [15],where the critic Q w ( s, a ) approximates the return of given state-action pair ( s, a ) and guides the
Revised: In modern RL algorithms, the most general approach is based on the Actor-Critic framework [15], where the critic Q w ( s, a ) approximates the return of given state-action pair ( s, a ) and guides the

Original: Actor to maximize the approximated return at state s .
Revised: Actor to maximize the approximated return at state s .

Original: The Critic is optimized to reduce Temporal
Revised: The Critic is optimized to reduce Temporal

Original: Difference (TD) error [28], deﬁned as
Revised: Difference (TD) error [29], deﬁned as

Original: [Equation]
Revised: [Equation]

Original: ) where B = ( s i , a i , r i , s 0 i ) i
Revised: where B = ( s i , a i , r i , s (cid:48) i ) i

Original: =1 , 2 ,... is the replay buffer used for off-policy learning [17, 10, 12], anda 0 i = ⇡
Revised: =1 , 2 ,... is the replay buffer used for off-policy learning [17, 10, 12, 28], and a (cid:48)

Original:  ( s 0 i ) is the predicted action for state s 0 i .
Revised: i = π ( s (cid:48) i ) is the predicted action for state s (cid:48) i .

Original: In practice, the calculations of Q w ( s [-
-] i , a 0 i ) are usuallybased on another set of slowly updated target networks for stability [10, 12].
Revised: In practice, the calculations of Q w ( s [-
-](cid:48) i , a (cid:48) i ) are usually based on another set of slowly updated target networks for stability [10, 12].

Original: Henceforth, TD-learningcan be roughly simpliﬁed as regression with notion y i =
Revised: Henceforth, TD-learning can be roughly simpliﬁed as regression with notion y i =

Original: r i +  Q w ( s [-
-]
Revised: r i + γQ w ( s [-
-](cid:48)

Original: i , a 0 i ) :
Revised: i , a (cid:48) i ) :

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

Original: Assume there are only M < L actions are related to a speciﬁc task among the L -dimensionalactions a i = a (1) i , ..., a ( L ) i , i.e., Q w ( · , · ) is function of s i , a (1) i , ..., a ( M ) i .
Revised: Assume there are only M < L actions are related to a speciﬁc task among the L -dimensional actions a i = a (1) i , ..., a ( L ) i , i.e., Q w ( · , · ) is function of s i , a (1) i , ..., a ( M ) i .

Original: Learning with the primalredundant action space will lead to around L + |S| M + |S| times sample complexity [9, 38].
Revised: Learning with the primal redundant action space will lead to around L + |S| M + |S| times sample complexity [9, 39].

Original: Therefore, we aremotivated to improve the learning efﬁciency of Q by pruning those task-irrelevant action dimensionsa ( M +1) i , ..., a ( L ) i by ﬁnding an action selection function G , satisfying
Revised: Therefore, we are motivated to improve the learning efﬁciency of Q by pruning those task-irrelevant action dimensions a ( M +1) i , ..., a ( L ) i by ﬁnding an action selection function G , satisfying

Original: [Equation]
Revised: [Equation]

Original: where y [-
-] i =
Revised: where y [-
-](cid:48) i =

Original: r i +  Q w ( s [-
-]
Revised: r i + γQ w ( s [-
-](cid:48)

Original: i , a 0 G ( a 0 i | s i ) i ) .
Revised: i , a (cid:48) G ( a (cid:48) i | s i ) i ) .

Original: Such a problem can be addressed from the perspective of causal discovery.
Revised: Such a problem can be addressed from the perspective of causal discovery.

Original: Formally, we can usethe Structural Causal Models (SCMs) to represent the underlying causal structure of a sequentialdecision making process, as shown in Figure 2.
Revised: Formally, we can use the Structural Causal Models (SCMs) to represent the underlying causal structure of a sequential decision making process, as shown in Figure 2.

Original: Under this language, we use the notion of causalactions to denote a (1 ,...,M ) i , and nuisance actions for other dimension of actions.
Revised: Under this language, we use the notion of causal actions to denote a (1 ,...,M ) i , and nuisance actions for other dimension of actions.

Original: In our work, we use
Revised: In our work, we use

Original: IC-INVASE for causal discovery.
Revised: IC-INVASE for causal discovery.

Original: Ideally, the action selection function G should be able to distinguishbetween nuisance action dimensions and the causal ones that has causal relation with either dynamicsor reward mechanism.
Revised: Ideally, the action selection function G should be able to distinguish between nuisance action dimensions and the causal ones that has causal relation with either dynamics or reward mechanism.

Original: We present in the next section our causal discovery algorithms.
Revised: We present in the next section our causal discovery algorithms.

Original: Figure temporal difference

Original: Among all executable actions, there can be only a subset have effect on the dynamical changes or the reward mechanism.
Revised: Among all executable actions, there can be only a subset have effect on the dynamical changes or the reward mechanism.

Original: In our work, we use
Revised: In our work, we use

Original: IC-INVASE as a causal discovery tool to distinguish the causal irrelevant actions and hence improve learning efﬁciency.
Revised: IC-INVASE as a causal discovery tool to distinguish the causal irrelevant actions and hence improve learning efﬁciency.

Original: 3.2 Iterative Curriculum INVASE (IC-INVASE)
Revised: 3.2 Iterative Curriculum INVASE (IC-INVASE)

Original: Instead of directly applying INVASE to solve Equation (6).
Revised: Instead of directly applying INVASE to solve Equation (6).

Original: We ﬁrst propose two improvementsto make the vanilla INVASE more suitable for large-scale variable selection tasks as the actiondimension in RL might be extremely large [33].
Revised: We ﬁrst propose two improvements to make the vanilla INVASE more suitable for large-scale variable selection tasks as the action dimension in RL might be extremely large [34].

Original: Speciﬁcally, the ﬁrst improvement, based oncurriculum learning, is introduced to tackle the exploration difﬁculty when  in Equation (1) is large,where INVASE tends to converge to poor sub-optimal solutions and prune all variables including theuseful ones [37].
Revised: Speciﬁcally, the ﬁrst improvement, based on curriculum learning, is introduced to tackle the exploration difﬁculty when λ in Equation (1) is large, where INVASE tends to converge to poor sub-optimal solutions and prune all variables including the useful ones [38].

Original: The second improvement is based on the iterative structure of variable selectiontasks: the feature selection operator G can be applied multiple times to conduct hierarchical featureselection without introducing extra computation expenses.
Revised: The second improvement is based on the iterative structure of variable selection tasks: the feature selection operator G can be applied multiple times to conduct hierarchical feature selection without introducing extra computation expenses.

Original: 3.2.1 Curriculum Learning For High Dimensional Variable Selection
Revised: 3.2.1 Curriculum Learning For High Dimensional Variable Selection

Original: The work of [3] ﬁrst introduces Curriculum Learning to mimic human learning by gradually learnmore complex concepts or handle more difﬁcult tasks.
Revised: The work of [3] ﬁrst introduces Curriculum Learning to mimic human learning by gradually learn more complex concepts or handle more difﬁcult tasks.

Original: Effectiveness of the method has beendemonstrated in various set-ups [3, 19, 7, 34, 36].
Revised: Effectiveness of the method has been demonstrated in various set-ups [3, 19, 7, 35, 37].

Original: In general, it should be easier to select M usefulvariables out of L input variables when M is larger.
Revised: In general, it should be easier to select M useful variables out of L input variables when M is larger.

Original: The most trivial case is to select all L variables,with an identical mapping x
Revised: The most trivial case is to select all L variables, with an identical mapping x

Original: ( G ( x ))
Revised: ( G ( x ))

Original: = G ( x )  x = x .
Revised: = G ( x ) (cid:12) x = x .

Original: Formally, we have
Revised: Formally, we have

Original: Proposition 1 (Curriculum Property in Variable Selection) .
Revised: Proposition 1 (Curriculum Property in Variable Selection) .

Original: Assume M out of L variables areoutcome-related, let M  N 1 < N 2  L , G N 1 ( x ) minimizes D KL ( p ( Y | X = x ) || p
Revised: Assume M out of L variables are outcome-related, let M ≤ N 1 < N 2 ≤ L , G N 1 ( x ) minimizes D KL ( p ( Y | X = x ) || p

Original: ( Y | X ( G ( x )) =
Revised: ( Y | X ( G ( x )) =

Original: x ( G ( x )) ))
Revised: x ( G ( x )) ))

Original: +  || G ( x ) | [-
 -] N 1{+
+} | .
Revised: + λ || G ( x ) | [-
 -]0 − N {+
+} | .

Original: Then
Revised: Then

Original: G N 2 ( x ) minimizes D KL ( p ( Y | X = x ) || p 
Revised: G N 2 ( x ) minimizes D KL ( p ( Y | X = x ) || p ( Y | X G ( x ) = x G ( x ) ))

Original: [Equation]
Revised: +

Original: Y | [Equation] G [Equation]
Revised: λ

Original:  G
Revised: || G

Original: [Equation]
Revised: [Equation]

Original: G
Revised: x

Original: [Equation]
Revised: [Equation]

Original:  0  N 2 | can be get through:116 G
Revised: | 0 − N 2 | can be get through:

Original: [Equation]
Revised: [Equation]

Original: where [ · ] 1 N 2  N 1 means keep N 2  N 1 none-zero elements unchanged while replacing other elements118 by 0 .
Revised: where [ · ] 1 N 2 − N 1 means keep N 2 − N 1 none-zero elements unchanged while replacing other elements by 0 .

Original: Proof.
Revised: Proof.

Original: By the deﬁnition of the [ · ] 1 N 2  N 1 operator, || G ( x ) | [-
 -] N 2 | = 0 is minimized.
Revised: By the deﬁnition of the [ · ] 1 N 2 − N 1 operator, || G ( x ) | [-
 -]0 − N 2 | = 0 is minimized.

Original: On the otherhand, starting from N 1 = M , minimizing D KL ( p ( Y | X = x ) || p
Revised: On the other hand, starting from N 1 = M , minimizing D KL ( p ( Y | X = x ) || p

Original: ( Y | X ( G ( x )) =
Revised: ( Y | X ( G ( x )) =

Original: x ( G ( x )) )) requiresall the M outcome-related variables being selected by G N 1 .
Revised: x ( G ( x )) )) requires all the M outcome-related variables being selected by G N 1 .

Original: Therefore, G N 2 also minimizes the
Revised: Therefore, G N 2 also minimizes the

Original: KL-divergence by the independent assumption of the other L  M variables with the outcomes.
Revised: KL-divergence by the independent assumption of the other L − M variables with the outcomes.

Original: The proposition indicates the difﬁculty of selecting N useful out of L variables decreases monotonically as N  M increase from M, M + 1 , ..., L .
Revised: The proposition indicates the difﬁculty of selecting N useful out of L variables decreases monotonically as N ≥ M increase from M, M + 1 , ..., L .

Original: In this work, two classes of practical curriculumare designed: 1.
Revised: In this work, two classes of practical curriculum are designed: 1.

Original: curriculum on the l 0 penalty coefﬁcient, and 2.
Revised: curriculum on the l 0 penalty coefﬁcient, and 2.

Original: curriculum on the proportion ofvariables to be selected.
Revised: curriculum on the proportion of variables to be selected.

Original: Curriculum on l 0 Penalty Coefﬁcient In this curriculum design, the penalty coefﬁcient  in
Revised: Curriculum on l 0 Penalty Coefﬁcient In this curriculum design, the penalty coefﬁcient λ in

Original: Equation (1) is increased from 0 to a pre-determined number (e.g., 1 . 0 ).
Revised: Equation (1) is increased from 0 to a pre-determined number (e.g., 1 . 0 ).

Original: Increasing the value of will lead to a larger penalty on the number of variables selected by the feature selector.
Revised: Increasing the value of λ will lead to a larger penalty on the number of variables selected by the feature selector.

Original: Experimentsin [37] has shown a large  always lead to a trivial selector that does not select any variable.
Revised: Experiments in [38] has shown a large λ always lead to a trivial selector that does not select any variable.

Original: Curriculum on the Proportion of Selected Features In this curriculum design, the proportion ofvariables to be selected, denoted by p r , is adjusted from the default setting 0 to a decreasing numberfrom a pre-determined value (e.g., 0 . 5 ) to 0 .
Revised: Curriculum on the Proportion of Selected Features In this curriculum design, the proportion of variables to be selected, denoted by p r , is adjusted from the default setting 0 to a decreasing number from a pre-determined value (e.g., 0 . 5 ) to 0 .

Original: i.e., the l 0 penalty term  | G ( x ) | 0 in Equation (1) isrevised to be 
Revised: i.e., the l 0 penalty term λ | G ( x ) | 0 in Equation (1) is revised to be λ

Original: || G ( x ) | [-
 -] d · p r | , where d is the dimension of input x .
Revised: || G ( x ) | [-
 -]0 − d · p r | , where d is the dimension of input x .

Original: When the proportion is set
Revised: When the proportion is set

Original: Initialize critic networks C  1 , C  2 , baseline networks B  1 , B  2 and actor network ⇡ ⌫ , IC-INVASE selector network G ✓
Revised: Initialize critic networks C φ 1 , C φ 2 , baseline networks B ψ 1 , B ψ 2 and actor network π ν , IC-INVASE selector network G θ

Original: Initialize target networks  0 1 ,  0 2 , 01 , 02 , ⌫⌫
Revised: Initialize target networks φ (cid:48) 1 ← φ 1 , φ (cid:48) 2 ← φ 2 , ψ (cid:48) 1 ← ψ 1 , ψ (cid:48) 2 ← ψ 2 , ν (cid:48) ← ν

Original: Initialize replay buffer B for t = 1 , H do
Revised: Initialize replay buffer B for t = 1 , H do

Original: Interact with environment and store transition tuple ( s, a, r, s 0 ) in B
Revised: Interact with environment and store transition tuple ( s, a, r, s (cid:48) ) in B

Original: Sample mini-batch of transitions { ( s, a, r, s 0 ) } from B Calculate perturbed next action by ˜ a ⇡ ⌫ 0 ( s 0 )
Revised: Sample mini-batch of transitions { ( s, a, r, s (cid:48) ) } from B Calculate perturbed next action by ˜ a ← π ν (cid:48) ( s (cid:48) )

Original: + ✏ , ✏ is sampled from a clipped Gaussian.
Revised: + (cid:15) , (cid:15) is sampled from a clipped Gaussian.

Original: Select actions with target selector network ˜ 
Revised: Select actions with target selector network ˜ a ( G (˜ a | s (cid:48) ))

Original: [Equation] G [Equation] ˜
Revised: ← G θ (cid:48) (˜ a | s (cid:48) )

Original: [Equation]
Revised: (cid:12) ˜ a

Original:  ˜ a Calculate target critic value y c and baseline value y b : y c  r +  min i =1 , 2 C [- 
-] i ( s 0 ,˜ [Equation]
Revised: Calculate target critic value y c and baseline value y b : y c ← r + γ min i =1 , 2 C [- 
-]φ (cid:48) i ( s (cid:48) ,

Original: [Equation] (˜ [Equation] | s 0 )) )
Revised: ˜ a ( G (˜ a | s (cid:48) )) )

Original: y b  r +  min i =1 , 2 B 0 i ( s 0 , ˜ a )
Revised: y b ← r + γ min i =1 , 2 B ψ (cid:48) i ( s (cid:48) , ˜ a )

Original: Update critics and baselines with selected actions: a( G ( a | s ))
Revised: Update critics and baselines with selected actions: a

Original: [Equation]
Revised: [Equation]

Original:  [Equation] Update IC-
Revised: a φ i ← arg min φ i MSE ( y c ,

Original: [Table] network by [Table]
Revised: C φ i ( s, a ( G ( a | s )) ))

Original: [Table] 1 : ✓ [Table] ( l
Revised: ψ i ← arg min ψ i MSE ( y b , B ψ i ( s, a ))

Original: [Table]
Revised: Update IC-INVASE selector network by the policy gradient, with learning rate η 1 : θ ← θ − η 1 ( l b − l c ) ∇ θ log G θ ( a | s ) , l b , l c are MSE losses in the previous step.

Original: l
Revised: ν

Original: [Table] log [Table] [Table]
Revised: ← ν − η 2 ∇ a C φ 1 ( s, a ) | a = π ν ( s ) ∇ ν π ν ( s ) φ (cid:48)

Original: [Table] are [Table]
Revised: i ← τφ i + (1 − τ ) φ (cid:48)

Original: step. Update ⌫ by the deterministic policy gradient, with learning rate ⌘ 2 : ⌫ ⌫  ⌘ 2 r a C  1 ( s, a ) | a = ⇡ ⌫ ( s ) r ⌫ ⇡ ⌫ ( s ) Update target networks, with ⌧ 2 (0 , 1) :  0 i ⌧ i + (1  ⌧ )  0 i 0 i [-⌧ i + (1  ⌧ ) 
-]
Revised: i [-⌧ i + (1  ⌧ ) 
-]ψ (cid:48)

Original: i ⌫ ⌧⌫ + (1  [Equation]
Revised: i ← τψ i + (1 − τ ) ψ (cid:48)

Original: ⌫ [Equation] end for to be p r = 0 .
Revised: i ν (cid:48) ← τν + (1 − τ ) ν (cid:48) end for to be p r = 0 .

Original: 5 , the selector will be penalized whenever less or more than half of all variables areselected.
Revised: 5 , the selector will be penalized whenever less or more than half of all variables are selected.

Original: Such a curriculum design forces the feature selector to learn to select less but increasinglymore important variables gradually.
Revised: Such a curriculum design forces the feature selector to learn to select less but increasingly more important variables gradually.

Original: Thus, we get the learning objective of curriculum-INVASE:
Revised: Thus, we get the learning objective of curriculum-INVASE:

Original: [Equation]
Revised: [Equation]

Original: where  increases from 0 to some value and p r decreases from a value in [0 , 1) to 0 .
Revised: where λ increases from 0 to some value and p r decreases from a value in [0 , 1) to 0 .

Original: 3.2.2 Iterative Variable Selection
Revised: 3.2.2 Iterative Variable Selection

Original: The second improvement proposed in this work is based on the iterative structure of variable selectiontasks.
Revised: The second improvement proposed in this work is based on the iterative structure of variable selection tasks.

Original: Speciﬁcally, the G ( x ) mapping x 2 X to { 0 , 1 } d is an iterative operator, which can be appliedfor multiple times to perform coarse-to-ﬁne variable selection.
Revised: Speciﬁcally, the G ( x ) mapping x ∈ X to { 0 , 1 } d is an iterative operator, which can be applied for multiple times to perform coarse-to-ﬁne variable selection.

Original: Although in practice we follow [37]to apply an element-wise product in producing x ( G ( x )) :
Revised: Although in practice we follow [38] to apply an element-wise product in producing x ( G ( x )) :

Original: x ( G ( x )) =
Revised: x ( G ( x )) =

Original: G ( x )  x 2 X .
Revised: G ( x ) (cid:12) x ∈ X .

Original: In more generalcases, the i-th element of x ( G ( x ))
Revised: In more general cases, the i-th element of x ( G ( x ))

Original: i is
Revised: i is

Original: [Equation]
Revised: [Equation]

Original: where ⇤ can be an arbitrary identiﬁable indicator that represents the variable is not selected.
Revised: where ∗ can be an arbitrary identiﬁable indicator that represents the variable is not selected.

Original: On the other hand, once the outputs G ( x ) of the selector have been recorded, ⇤ can be replaced byany label-independent variable G ( x )  z , where z ⇠ pz ( · ) is outcome-independent.
Revised: On the other hand, once the outputs G ( x ) of the selector have been recorded, ∗ can be replaced by any label-independent variable G ( x ) (cid:12) z , where z ∼ p

Original:  Then x ( G ( x ))can be regarded as a new sample and be fed into the variable selector, resulting in a hierarchicalvariable selection process:
Revised: z ( · ) is outcome-independent. Then x ( G ( x )) can be regarded as a new sample and be fed into the variable selector, resulting in a hierarchical variable selection process:

Original: [Equation]
Revised: [Equation]

Original: where z ⇠ p z ( · ) , and  is the element-wise sum operator.
Revised: where z ∼ p z ( · ) , and ⊕ is the element-wise sum operator.

Original: Moreover, if the distribution of irrelevantvariable p x ( · ) is known, applying the variable selection operator obtained from Equation (7) formultiple times with p z ( · ) d = p x ( · ) has the meaning of hierarchical variable selection: after eachoperation, the most obvious 1  p r irrelevant variables are discarded.
Revised: Moreover, if the distribution of irrelevant variable p x ( · ) is known, applying the variable selection operator obtained from Equation (7) for multiple times with p z ( · ) d = p x ( · ) has the meaning of hierarchical variable selection: after each operation, the most obvious 1 − p r irrelevant variables are discarded.

Original: e.g., when p r = 0 . 5 , ideally top50% , 25% , 12 .
Revised: e.g., when p r = 0 . 5 , ideally top50% , 25% , 12 .

Original: 5% most important variables will be selected after the ﬁrst three selection operations.
Revised: 5% most important variables will be selected after the ﬁrst three selection operations.

Original: In this work, a coarse approximation is utilized by selecting z to be z = 0 for simplicity.
Revised: In this work, a coarse approximation is utilized by selecting z to be z = 0 for simplicity.

Original: Combining those two improvements lead to an Iterative Curriculum version of INVASE (IC-INVASE)that addresses the exploration difﬁculty in high-dimensional variable selection tasks.
Revised: Combining those two improvements lead to an Iterative Curriculum version of INVASE (IC-INVASE) that addresses the exploration difﬁculty in high-dimensional variable selection tasks.

Original: Curriculumlearning helps IC-INVASE to achieve better asymptotic performance, i.e., achieve higher True Positive
Revised: Curriculum learning helps IC-INVASE to achieve better asymptotic performance, i.e., achieve higher True Positive

Original: Rate (TPR) and lower False Discovery Rate (FDR), while iterative application of the selection operatorcontributes to higher learning efﬁciency: selectors models with different level of TPR/FDR can begenerated on-the-ﬂy.
Revised: Rate (TPR) and lower False Discovery Rate (FDR), while iterative application of the selection operator contributes to higher learning efﬁciency: selectors models with different level of TPR/FDR can be generated on-the-ﬂy.

Original: 3.3 State-Wise Action Reﬁnery with IC-INVASE3.3.1 Temporal Difference State-Wise Action Reﬁnery
Revised: 3.3 State-Wise Action Reﬁnery with IC-INVASE 3.3.1 Temporal Difference State-Wise Action Reﬁnery

Original: With the techniques introduced in the previous section, higher dimensional variable selection taskscan be better solved, therefore we are ready to use IC-INVASE to solve Equation (6).
Revised: With the techniques introduced in the previous section, higher dimensional variable selection tasks can be better solved, therefore we are ready to use IC-INVASE to solve Equation (6).

Original: The resultingalgorithm is called Temporal Difference State-Wise Action Reﬁnery (TD-SWAR).
Revised: The resulting algorithm is called Temporal Difference State-Wise Action Reﬁnery (TD-SWAR).

Original: In this work, TD3 [10] is used as the basic algorithm we build TD-SWAR up on.
Revised: In this work, TD3 [10] is used as the basic algorithm we build TD-SWAR up on.

Original: In addition to thepolicy network ⇡ ⌫ , double critic networks C  1 , C  2 and their corresponding target networks usedin vanilla TD3, TD-SWAR includes an action selector model G ✓ and two baseline networks B  1 ,
Revised: In addition to the policy network π ν , double critic networks C φ 1 , C φ 2 and their corresponding target networks used in vanilla TD3, TD-SWAR includes an action selector model G θ and two baseline networks B ψ 1 ,

Original: B  2 following [37] to reduce the variance in policy gradient learning.
Revised: B ψ 2 following [38] to reduce the variance in policy gradient learning.

Original: Pseudo-code for the proposedalgorithm is shown in Algorithm 1.
Revised: Pseudo-code for the proposed algorithm is shown in Algorithm 1.

Original: And the block diagram in Figure 1 illustrates how differentmodules in TD-SWAR updates their parameters.
Revised: And the block diagram in Figure 1 illustrates how different modules in TD-SWAR updates their parameters.

Original: 3.3.2 Static Approximation: Model-Based Action Selection
Revised: 3.3.2 Static Approximation: Model-Based Action Selection

Original: While IC-INVASE can be formally integrated with temporal difference learning, the learning stabilityis not guaranteed.
Revised: While IC-INVASE can be formally integrated with temporal difference learning, the learning stability is not guaranteed.

Original: Different from general regression tasks where the label for every instance is ﬁxedacross training, in temporal difference learning, the regression target is closely related to the presentcritic function C  , the policy ⇡ ⌫ that generates the transition tuples used for training, and the selectormodel of IC-INVASE itself.
Revised: Different from general regression tasks where the label for every instance is ﬁxed across training, in temporal difference learning, the regression target is closely related to the present critic function C φ , the policy π ν that generates the transition tuples used for training, and the selector model of IC-INVASE itself.

Original: In this section, a static approach is proposed to approximately solve thechallenge of instability in TD-SWAR 4 .
Revised: In this section, a static approach is proposed to approximately solve the challenge of instability in TD-SWAR 4 .

Original: Other than applying the IC-INVASE algorithm to solve Equation (6), another way of leveraging
Revised: Other than applying the IC-INVASE algorithm to solve Equation (6), another way of leveraging

Original: IC-INVASE in action space pruning is to combine it with the model-based methods [11, 16, 13, 14],where a dynamic model P : S ⇥ A 7! S is learned through regression:
Revised: IC-INVASE in action space pruning is to combine it with the model-based methods [11, 16, 13, 14], where a dynamic model P : S × A (cid:55)→ S is learned through regression:

Original: [Equation]
Revised: [Equation]

Original: Although the task of precise model-based prediction is in general challenging [24], in this work, weonly adopt model-based prediction in action selection, and the target is action discovery other thanprecise prediction.
Revised: Although the task of precise model-based prediction is in general challenging [24], in this work, we only adopt model-based prediction in action selection, and the target is action discovery other than precise prediction.

Original: As the dynamic models are always static across learning, such an approach can bemuch more stable than TD-SWAR.
Revised: As the dynamic models are always static across learning, such an approach can be much more stable than TD-SWAR.

Original: We name this method as Dyn-SWAR and present the pseudo-codein Algorithm 2, where we infuse IC-INVASE to Equation (10) and get the learning objective:
Revised: We name this method as Dyn-SWAR and present the pseudo-code in Algorithm 2, where we infuse IC-INVASE to Equation (10) and get the learning objective:

Original: [Equation]
Revised: [Equation]

Original:  Experiment
Revised:  Experiment

Original: In this section, we demonstrate our proposed methods in ﬁve continuous control RL tasks withredundant action space where our proposed methods can perform causality-aware RL.
Revised: In this section, we demonstrate our proposed methods in ﬁve continuous control RL tasks with redundant action space where our proposed methods can perform causality-aware RL.

Original: We providequantitatively comparison between IC-INVASE and the vanilla INVASE on synthetic datasets toshow its improved scalability in Appendix B.
Revised: We provide quantitatively comparison between IC-INVASE and the vanilla INVASE on synthetic datasets to show its improved scalability in Appendix B.

Original:  Initialize critic networks Q w 1 , Q w 2 , Dynamics critic model C  , dynamic baseline model B  , actor network ⇡ , and IC-INVASE selector network G ✓
Revised: Algorithm Initialize critic networks Q w 1 , Q w 2 , Dynamics critic model C φ , dynamic baseline model B ψ , actor network π ν , and IC-INVASE selector network G θ

Original: Initialize target networks w [-
 [Equation]-]
Revised: Initialize target networks w [-
 [Equation]-](cid:48) 1 ← w 1 , w (cid:48) 2 ← w 2 , ν (cid:48) ← ν

Original: Initialize replay buffer B for t = 1 , H do
Revised: Initialize replay buffer B for t = 1 , H do

Original: Interact with environment and store transition tuple ( s, a, r, s 0 ) in B
Revised: Interact with environment and store transition tuple ( s, a, r, s (cid:48) ) in B

Original: Sample mini-batch of transitions { ( s, a, r, s 0 ) } from B Update dynamic critics and dynamic baselines with equation (10): 
Revised: Sample mini-batch of transitions { ( s, a, r, s (cid:48) ) } from B Update dynamic critics and dynamic baselines with equation (10): φ ← arg min φ MSE ( s (cid:48) , C φ ( s, a ( G ( a | s )) ))

Original: [Equation]
Revised: ψ ← arg min ψ MSE ( s (cid:48) , B ψ ( s, a ))

Original: Update IC-INVASE selector network by the policy gradient, with learning rate ⌘ 1 : ✓ [Equation] [Table] [Equation]
Revised: Update IC-INVASE selector network by the policy gradient, with learning rate η 1 : θ ← θ − η 1 ( l b − l c ) ∇ θ log G θ ( a | s ) , l b , l c are MSE losses in the previous step.

Original: [Table] [Equation]
Revised: Calculate perturbed next action by ˜ a ← π ν (cid:48)

Original: [Table]
Revised: ( s (cid:48) )

Original: are MSE losses in the previous step. Calculate perturbed next action by ˜ a ⇡ ⌫ 0 ( s 0 ) + ✏ , ✏ is sampled from a clipped Gaussian.
Revised: + (cid:15) , (cid:15) is sampled from a clipped Gaussian.

Original: Select actions with selector network ˜ 
Revised: Select actions with selector network ˜ a ( G (˜ a | s (cid:48) ))

Revised: ← G θ (cid:48) (˜ a | s (cid:48) )

Original: [Equation]
Revised: (cid:12) ˜ a

Original: Calculate target critic value y and update critic networks: 
Revised: Calculate target critic value y and update critic networks: y ← r + γ min i =1 , 2 Q w (cid:48) i ( s (cid:48) , ˜ a ( G (˜ a | s (cid:48) )) )

Original: [Equation]
Revised: w i ← arg min w i MSE ( y, Q w i ( s, a ( G ( a | s )) ))

Original: Update ⌫ by the deterministic policy gradient, with learning rate ⌘ 2 : 
Revised: Update ν by the deterministic policy gradient, with learning rate η 2 : ν ← ν − η 2 ∇ a Q w 1 ( s, a ) | a = π ν ( s ) ∇ ν π

Original: [Equation]
Revised: [Figure] s )

Original: Update target networks, with ⌧ 2 (0 , 1) : 
Revised: Update target networks, with τ ∈ (0 , 1) : w (cid:48)

Original: [Equation]
Revised: i ← τw i + (1 − τ ) w (cid:48) i ν (cid:48) ← τν + (1 − τ ) ν (cid:48) end for +

Original: end for + 
Revised: + + + (

Original: [Table]
Revised: [Figure]

Original: (a) 4Rew.-Maze(b) Pendulum (c)
Revised: ) 4Rew.-Maze

Original: [Figure]
Revised: [Figure]

Original: For this set of experiments, we use ﬁve RL environments (Figure 5) that are listed in Table 1 5 .
Revised: For this set of experiments, we use ﬁve RL environments (Figure 5) that are listed in Table 1 5 .

Original: |S|means the dimension of state space in each task, |A| represents the dimension of task-related action196 space, and |A red.
Revised: |S| means the dimension of state space in each task, |A| represents the dimension of task-related action space, and |A red.

Original: | indicates the dimension of redundant action space that is injected to each task.
Revised: | indicates the dimension of redundant action space that is injected to each task.

Original: Those redundant dimensions of actions will not affect the state transitions or reward calculation, butan agent needs to learn to identify those redundant dimensions to perform efﬁcient learning.
Revised: Those redundant dimensions of actions will not affect the state transitions or reward calculation, but an agent needs to learn to identify those redundant dimensions to perform efﬁcient learning.

Original: We evaluate both TD-SWAR that integrate IC-INVASE with temporal difference learning and itsstatic variant Dyn-SWAR that applies IC-INVASE in dynamics prediction.
Revised: We evaluate both TD-SWAR that integrate IC-INVASE with temporal difference learning and its static variant Dyn-SWAR that applies IC-INVASE in dynamics prediction.

Original: The results are comparedwith two baselines: the Oracle : redundant action dimensions are eliminated manually; and TD3 : thevanilla TD3 algorithm without explicit action redundancy reduction.
Revised: The results are compared with two baselines: the Oracle : redundant action dimensions are eliminated manually; and TD3 : the vanilla TD3 algorithm without explicit action redundancy reduction.

Original: In experiments, we ﬁnd the implementation of Dyn-SWAR can be much more efﬁcient in termsof both sample complexity and computational expense: while the TD-SWAR need to continuouslyupdate all parameters for the IC-INVASE selector to keep consistent with the real-time policy andvalue networks as the regression label varies along time, the Dyn-SWAR selector can be trained withfar less amount of data.
Revised: In experiments, we ﬁnd the implementation of Dyn-SWAR can be much more efﬁcient in terms of both sample complexity and computational expense: while the TD-SWAR need to continuously update all parameters for the IC-INVASE selector to keep consistent with the real-time policy and value networks as the regression label varies along time, the Dyn-SWAR selector can be trained with far less amount of data.

Original: Say, 10 , 000 to 25 , 000 timesteps of interactions with the environment.
Revised: Say, 10 , 000 to 25 , 000 timesteps of interactions with the environment.

Original: Sucha property can be naturally combined with the warm-up trick used in TD3 [10], i.e., the Dyn-SWARselector can be trained with warm-up transition tuples collected in the random exploration phase
Revised: Such a property can be naturally combined with the warm-up trick used in TD3 [10], i.e., the Dyn-SWAR selector can be trained with warm-up transition tuples collected in the random exploration phase

Original: [Figure]
Revised: [Figure]

Original: and then be ﬁxed in the later learning process.
Revised: and then be ﬁxed in the later learning process.

Original: Compared with normal RL settings where millionsof interactions with the environment are always needed, the training of Dyn-SWAR only increasesnegligible computational expense.
Revised: Compared with normal RL settings where millions of interactions with the environment are always needed, the training of Dyn-SWAR only increases negligible computational expense.

Original: The results are shown in Figure 4.
Revised: The results are shown in Figure 4.

Original: In all environments, agent learning with IC-INVASE in bothmanner (TD-/Dyn-) outperforms the vanilla TD3 baseline.
Revised: In all environments, agent learning with IC-INVASE in both manner (TD-/Dyn-) outperforms the vanilla TD3 baseline.

Original: The Dyn-SWAR achieves high learningefﬁciency that is comparable to the oracle benchmarks.
Revised: The Dyn-SWAR achieves high learning efﬁciency that is comparable to the oracle benchmarks.

Original: However, the performance of TD-SWAR inhigher dimensional tasks (Walker2d-v2 and BipedalWalker-v3) still has a lot of room for improvement.
Revised: However, the performance of TD-SWAR in higher dimensional tasks (Walker2d-v2 and BipedalWalker-v3) still has a lot of room for improvement.

Original: Improving the stability of and scalability of instance-wise variable selection in temporal differencelearning thereby should be addressed in future work.
Revised: Improving the stability of and scalability of instance-wise variable selection in temporal difference learning thereby should be addressed in future work.

Original:  Related Work
Revised:  Related Work

Original: Instance-Wise Feature Selection While traditional feature selection method like LASSO [30]aims at ﬁnding globally important features across the whole dataset, instance-wise feature selectiontry to discover the feature-label dependency on a case-by-case basis.
Revised: Instance-Wise Feature Selection While traditional feature selection method like LASSO [31] aims at ﬁnding globally important features across the whole dataset, instance-wise feature selection try to discover the feature-label dependency on a case-by-case basis.

Original: L2X [5] performs instance-wisefeature selection through mutual information maximization with the technique of Gumbel softmax.
Revised: L2X [5] performs instance-wise feature selection through mutual information maximization with the technique of Gumbel softmax.

Original: L2X requires pre-determined hyper-parameter k to indicate how many features should be selected foreach instance, which limits its performance while the number of label-relevant features varies acrossinstances.
Revised: L2X requires pre-determined hyper-parameter k to indicate how many features should be selected for each instance, which limits its performance while the number of label-relevant features varies across instances.

Original: In this work, we build our instance-wise action selection model on top of INVASE [37], where policygradient is applied to replace the Gumbel softmax trick and the size of chosen features per instance ismore ﬂexible.
Revised: In this work, we build our instance-wise action selection model on top of INVASE [38], where policy gradient is applied to replace the Gumbel softmax trick and the size of chosen features per instance is more ﬂexible.

Original: [31] considers instance-wise feature selection problems in time-series setting, andbuild generative models to capture counterfactual effects in time series data.
Revised: [32] considers instance-wise feature selection problems in time-series setting, and build generative models to capture counterfactual effects in time series data.

Original: Their work enablesevaluation of the importance of features over time, which is crucial in the context of healthcare.
Revised: Their work enables evaluation of the importance of features over time, which is crucial in the context of healthcare.

Original: [18]formally deﬁnes different types of feature redundancy and leverages mutual information maximizationin instance-wise feature group discovery and introduces theoretical guidance to ﬁnd the optimalnumber of different groups.
Revised: [18] formally deﬁnes different types of feature redundancy and leverages mutual information maximization in instance-wise feature group discovery and introduces theoretical guidance to ﬁnd the optimal number of different groups.

Original: Our work is distinguished from previous works for instance-wise feature selection in two aspects.
Revised: Our work is distinguished from previous works for instance-wise feature selection in two aspects.

Original: First, while previous works focus on static scenarios like classiﬁcation and regression, this work focuson temporal difference learning where there is no static label.
Revised: First, while previous works focus on static scenarios like classiﬁcation and regression, this work focus on temporal difference learning where there is no static label.

Original: Second, the scalability of previousmethods in variable selection is challenged as there might exist hundreds of redundant actions in thecontext of RL.
Revised: Second, the scalability of previous methods in variable selection is challenged as there might exist hundreds of redundant actions in the context of RL.

Original: Dimension Reduction in RL In the context of RL, attention models [32] have been applied tointerpret the behaviors of learned policies.
Revised: Dimension Reduction in RL In the context of RL, attention models [33] have been applied to interpret the behaviors of learned policies.

Original: [29] proposes to perceive the state information through a self-attention bottleneck in vision-based RL tasks, which concentrates on the state space redundancyreduction with image inputs.
Revised: [30] proposes to perceive the state information through a self-attention bottleneck in vision-based RL tasks, which concentrates on the state space redundancy reduction with image inputs.

Original: The work of [22] also applies the attention mechanism to learn taskrelevant information.
Revised: The work of [22] also applies the attention mechanism to learn taskrelevant information.

Original: The proposed method achieves state-of-the-art performance on Atari gameswith image input while being more understandable with top-down attention models.
Revised: The proposed method achieves state-of-the-art performance on Atari games with image input while being more understandable with top-down attention models.

Original: Different from those papers, this work considers relatively tight state representations (vector input),and focuses on the task-irrelevant action reduction.
Revised: Different from those papers, this work considers relatively tight state representations (vector input), and focuses on the task-irrelevant action reduction.

Original: We aim at ﬁnding the task-related actions andimproving the learning efﬁciency without wasting samples in learning the task-irrelevant dimensionsof actions.
Revised: We aim at ﬁnding the task-related actions and improving the learning efﬁciency without wasting samples in learning the task-irrelevant dimensions of actions.

Original: Our work is most closely related to AE-DQN [38] in that we both consider the problem ofredundant action elimination.
Revised: Our work is most closely related to AE-DQN [39] in that we both consider the problem of redundant action elimination.

Original: AE-DQN tackles action space redundancy with an action-eliminationnetwork that eliminates sub-optimal actions.
Revised: AE-DQN tackles action space redundancy with an action-elimination network that eliminates sub-optimal actions.

Original: Yet its discussion is limited in the discrete settings.
Revised: Yet its discussion is limited in the discrete settings.

Original: Incontrast, our work focuses on action elimination in continuous control tasks.
Revised: In contrast, our work focuses on action elimination in continuous control tasks.

Original:  Conclusion and Future Work
Revised:  Conclusion and Future Work

Original: In this work, we tackle the challenge of action space pruning in action redundant RL tasks.
Revised: In this work, we tackle the challenge of action space pruning in action redundant RL tasks.

Original: Recentadvance on instance-wise feature selection technique (INVASE) is exploited after curriculum learningand iterative operation are integrated for the pursuance of scalability and efﬁciency.
Revised: Recent advance on instance-wise feature selection technique (INVASE) is exploited after curriculum learning and iterative operation are integrated for the pursuance of scalability and efﬁciency.

Original: The resultingmethod, termed IC-INVASE, is then generalized to the RL setting where two different algorithms areproposed, TD-SWAR and Dyn-SWAR, to conduct causality-aware RL.
Revised: The resulting method, termed IC-INVASE, is then generalized to the RL setting where two different algorithms are proposed, TD-SWAR and Dyn-SWAR, to conduct causality-aware RL.

Original: While the former algorithmaddresses the action redundant issue directly in temporal difference learning, the latter algorithmcaptures dynamical causality with model-based prediction.
Revised: While the former algorithm addresses the action redundant issue directly in temporal difference learning, the latter algorithm captures dynamical causality with model-based prediction.

Original: Experiments on various tasks demonstratethe causality-awareness is crucial for RL agents to perform efﬁcient learning in action-redundantenvironments.
Revised: Experiments on various tasks demonstrate the causality-awareness is crucial for RL agents to perform efﬁcient learning in action-redundant environments.

Original: In future work, the iterative property can be further explored to perform ensemble methods in variableselection.
Revised: In future work, the iterative property can be further explored to perform ensemble methods in variable selection.

Original: And a more proper curriculum might be designed to better fuse multiple curricula together.
Revised: And a more proper curriculum might be designed to better fuse multiple curricula together.

Original: On the RL side, the stability of TD-SWAR might be further improved for better sample efﬁciency. 
Revised: On the RL side, the stability of TD-SWAR might be further improved for better sample efﬁciency. A On the Dynamic Model Approximation

Revised: We provide analysis on the approximation in this section based on the deterministic MDP model in ﬁnite action space where the problem degenerates to Q -Learning.

Revised: Similar results can be get to prove the Policy Evaluation Lemma, combined with Policy Improvement Lemma (given proper function approximation of the arg max operator) and result in Policy Iteration Theorem.

Revised: In deterministic MDPs with s t +1 =

Revised: T ( s t , a t ) , r t = r ( s t , a t ) , the value function of a state is deﬁned as

Revised: [Equation]

Original: {+12) given s 
+}
Revised: {+12) given s 
+}

Revised: = s is the initial state and a t =

Revised: π ( s t ) comes from the deterministic policy π .

Revised: The learning objective is to ﬁnd an optimal policy π , such that an optimal state value can be achieved:

Revised: [Equation]

Revised: The state-action value function ( Q -function) is then deﬁned as

Revised: [Equation]

Revised: Formally, the objective of action space pruning in action-redundant MDPs is to ﬁnd an optimal policy π ( G ) =

Revised: G ( π ( s t ) | s t ) (cid:12) π ( s t ) with an action selector G :

Revised: S × A

Revised: (cid:55)→ { 0 ,

Revised: [Equation]

Revised: d ,

Revised: [Equation]

Revised: 15) with minimal number of actions selected, i.e., | G | 0 is minimized.

Revised: The sufﬁcient and necessary condition for Equation (15) to hold is r ( s t , π ( s t ))

Revised: =

Revised: r ( s t , π ( G ) ( s t )) and T ( s t , π ( s t ))

Revised: =

Revised: T ( s t , π ( G ) ( s t )) .

Revised: In general, the reward function r and transition dynamics T may depend on different subsets of actions and the optimal, i.e., r ( s t , a t ) =

Revised: r ( s t , a ( G 1 ) t ) , while T ( s t , a t ) =

Revised: T ( s t , a ( G 2 ) t ) , where G 1 , G 2 select different subset of given actions by a ( G 1 ) t = G 1 ( a t | s t ) (cid:12) a t , a ( G 2 ) t

Revised: = G 2 ( a t | s t ) (cid:12) a t but a ( G 1 ) t (cid:54) = a ( G 2 ) t .

Revised: The ﬁnal action selector G should be generated according to G ( a | s ) =

Revised: G 1 ( a | s ) ∨ G 2 ( a | s ) , where ∨ is the element-wise OR operation.

Revised: Therefore, in our approximation of Dyn-SWAR, we assume G ( a | s ) =

Revised: G 2 ( a | s ) as an approximation for G ( a | s ) =

Revised: G 1 ( a | s ) ∨ G 2 ( a | s ) .

Revised: Future work may include another predictive model for the reward function and take the element-wise OR operation to get G .

Revised: B Additional Experiments B.1 Synthetic Data Experiment

Revised: The synthetic datasets are generated in the same way as [5, 38].

Revised: Speciﬁcally, there are 6 synthetic datasets that have inputs generated from an 11 -dim Gaussian distribution without correlations across features.

Original: {+The label Y for each dataset is generated by a Bernoulli random variable with P ( Y = 
+}
Revised: {+The label Y for each dataset is generated by a Bernoulli random variable with P ( Y = 
+}

Revised: | X ) = 1 1+ logit ( X ) .

Revised: In different tasks, logit ( X ) takes the value of:

Original: {+• Syn1 : exp( X 1 X 2 ) • Syn2 : exp( (cid:80) 6 i =
+}
Revised: {+• Syn1 : exp( X 1 X 2 ) • Syn2 : exp( (cid:80) 6 i =
+}

Revised: X 2 i − 4) • Syn3 : − 10 × sin 2 X 7 + 2 | X 8 |

Revised: + X 9 + exp( − X 10 ) • Syn4 : if X 11 < 0 , logit follows Syn1 , otherwise, logit follows Syn • Syn5 : if X 11 < 0 , logit follows Syn1 , otherwise, logit follows Syn • Syn6 : if X 11 < 0 , logit follows Syn

Revised: , otherwise, logit follows Syn

Revised: In the ﬁrst three synthetic datasets, the label Y depends on the same feature across each dataset, while in the last three datasets, the subsets of features that label Y depends on are determined by the values of X 11 .

Revised: [Table]

Revised: For each dataset, 20 , 000 samples are generated and be separated into a training set and a testing set.

Revised: In this work, we focus on ﬁnding outcome-relevant features (e.g., ﬁnding task-relevant actions in the context of RL), thus the true positive rate (TPR) and false discovery rate (FDR) are used as performance metrics.

Revised: 11-dim Feature Selection Table 2 shows the quantitative results of the proposed method, ICINVASE on the 11-dim feature selection tasks.

Revised: To accelerate training and facilitate the usage of dynamical computational graphs in curriculum learning and RL settings, the vanilla INVASE is re-implemented with PyTorch [23].

Revised: In general, the PyTorch implementation is 4 to 5 times faster than the previous Keras [1, 6] implementation, with on-par performance on the 11 -dim feature selection tasks.

Revised: In the comparison, both the reported results in [38] (denoted by INVASE (REP.) ) and our experimental results on INVASE (denoted by INVASE (EXP.) ) are presented.

Revised: The p r curriculum for IC-INVASE in all experiments are set to decrease from 0 .

Original: {+5 to 0 . 
+}
Revised: {+5 to 0 . 
+}

Revised: except in ablation studies.

Revised: Results of two different choices of the λ curriculum are reported and denoted by IC-INVASE ( λ ↑ · ) , e.g., λ ↑ 0 .

Original: {+3 means λ increases from 0 . 
+}
Revised: {+3 means λ increases from 0 . 
+}

Revised: to 0 .

Revised: 3 in the experiment.

Original: {+We omit the results on the ﬁrst three datasets ( Syn1 , Syn2 , Syn3 ) where both IC-INVASE and INVASE achieve 100 . 
+}
Revised: {+We omit the results on the ﬁrst three datasets ( Syn1 , Syn2 , Syn3 ) where both IC-INVASE and INVASE achieve 100 . 
+}

Original: {+TPR and 0 . 
+}
Revised: {+TPR and 0 . 
+}

Revised: FDR.

Revised: Iteration 1 to Iteration 4 in the table shows the results after applying the selection operator for different number of iterations.

Revised: In all experiments, IC-INVASE achieves better performance (i.e., larger TPR and lower FDR) than the vanilla INVASE with Keras and PyTorch implementation.

Revised: Iterative applying the feature selection operator can reduce the FDR with a slight cost of TPR decay.

Revised: 100-dim Feature Selection We then increase the total number of feature dimensions to 100 to demonstrate how IC-INVASE improves the vanilla INVASE in larege-scale variable selection settings.

Revised: In this experiment.

Revised: The features are generated with 100 -dim Gaussian without correlations and the rules for label generation are still the same as the 11 -dim settings.

Revised: (i.e., 89 additional labelindependent noisy dimensions of input is concatenated to the 11 -dim inputs.)

Revised: The results are shown in Table 3.

Revised: IC-INVASE achieves much better performance in all datasets, i.e., higher TPR and lower FDR.

Revised: The ablation studies on different curriculum show both an increasing λ and a decreasing p r can beneﬁt discovery of label-dependent features.

Revised: As the hyper-parameters for curriculum are not elaborated in our experiments, direct combining the two curriculum may hinder the performance.

Revised: The design for curriculum fusion is left to the future work.

Revised: C Environment Details FourRewardMaze

Revised: The FourRewardMaze is a 2-D navigation task where an agent need to ﬁnd all four solutions to achieve better performance.

Revised: The state space is 2 -D continuous vector indicating the position of the agent, while the action space is a 2 -D continuous value indicating the direction and step length of the agent, which is limited to [ − 1 , 1] .

Revised: The initial location of the agent is randomly

Revised: [Table]

Revised: 2.3 IC-INVASE W/O λ ↑ 99.

Revised: 7.0 IC-INVASE 98.

Revised: + + +

Revised: [Figure]

Revised: [Figure]

Revised: selected for each game, and each episode has the length of 32 , which is the timesteps needed to collect all four rewards from any starting position.

Revised: Pendulum-v0 The Pendulum-v0 environment is a classic problem in the control literature.

Revised: In the

Revised: Pendulum-v0 of OpenAI Gym.

Revised: The task has 3 -D state space and 1 -D action space.

Revised: In every episode the pendulum starts in a random position, and the learning objective is to swing the pendulum up and keep it staying upright.

Revised: Walker2d-v2 The Walker2d-v2 environment is a locomotion task where the learning objective is to make a two-dimensional bipedal robot walk forward as fast as possible.

Revised: The task has 17 -D state space and 6 -D action space.

Original: {+LunarLanderContinuous-v
+}
Revised: {+LunarLanderContinuous-v
+}

Revised: In the tasks of LunarLanderContinuous-v2, the agent is asked to control a lander to move from the top of the screen to a landing pad located at coordinate (0 , 0) .

Revised: The fuel is inﬁnite, so an agent can learn to ﬂy and then land on its ﬁrst attempt.

Revised: The state is as 8 -D real-valued vector and action is 2 -D vector in the range of [ − 1 , 1] , where the ﬁrst dimension controls main engine, [ − 1 , 0] off, [0 ., 1] throttle from 50% to 100% power and the second value in [ − 1 , − 0 .

Revised: 5] will ﬁre left engine, while a value in [0 . 5 , 1 . 0] ﬁres right engine, otherwise the engine is off.

Revised: BipedalWalker-v3 The BipedalWalker-v3 is a locomotion task where the state space is 24 -D and the action space is 4 -D.

Revised: The agent needs to walk as far as possible in each episode where a total timestep of 1000 are given and total 300 points might be collected up to the far end.

Revised: If the robot falls, it gets − 100 points.

Revised: Applying motor torque costs a small amount of points, more optimal agent will get better score.

Revised: D Reproduction Checklist D.1 Neural Network Structure

Revised: In all experiments, we use the same neural network structure: in TD3, we follow the vanilla implementation to use 3 -layer fully connected neural networks where 256 hidden units are used.

Revised: In the selector networks of the INVASE module, we follow the vanilla implementation to use 3 -layer fully connected neural networks where 100 , 200 hidden units are used.

Revised: D.2 Hyper-Parameters

Revised: In both TD-SWAR and the Dyn-SWAR, we apply IC-INVASE with p r reducing from 0 .

Revised: 5 to 0 .and

Revised: λ increasing from 0 .

Revised: to 0 .

Revised: While our experiments have already shown the effectiveness and robustness of those hyper-parameters, performing grid search on those hyper-parameters may lead to further performance improvement.