Excluding the Irrelevant: Focusing Reinforcement Learning through Continuous Action Masking

Authors: roland  stolz
hanna  krasowski
jakob  thumm
michael  eichelbeck
philipp  gassert
matthias  althoff
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 15:38:04

Abstract: Continuous action spaces in reinforcement learning (RL) are commonly defined as multidimensional intervals. While intervals usually reflect the action boundaries for tasks well, they can be challenging for learning because the typically large global action space leads to frequent exploration of irrelevant actions. Yet, little task knowledge can be sufficient to identify significantly smaller state-specific sets of relevant actions. Focusing learning on these relevant actions can significantly improve training efficiency and effectiveness. In this paper, we propose to focus learning on the set of relevant actions and introduce three continuous action masking methods for exactly mapping the action space to the state-dependent set of relevant actions. Thus, our methods ensure that only relevant actions are executed, enhancing the predictability of the RL agent and enabling its use in safety-critical applications. We further derive the implications of the proposed methods on the policy gradient. Using proximal policy optimization ( PPO), we evaluate our methods on four control tasks, where the relevant action set is computed based on the system dynamics and a relevant state set. Our experiments show that the three action masking methods achieve higher final rewards and converge faster than the baseline without action masking.

Main paper: 
 introduction: 
 Reinforcement learning (RL) can solve complex tasks in areas such as robotics [13] , games [35] , and large language models [26] . Yet, training RL agents is often sample-inefficient due to frequent exploration of actions, which are irrelevant to learning a good policy. Irrelevant actions are actions that are either physically impossible, forbidden due to some formal specification, or evidently counterproductive for solving the task. Since the global action space is typically large in relation to the relevant actions in each state, exploring these actions frequently can introduce unnecessary costs, lead to slow convergence, or even prevent the agent from learning a suitable policy.Action masking mitigates this problem by constraining the exploration to the set of relevant statespecific actions, which can be obtained based on task knowledge. For example, when there is no opponent within reach in video games, attack actions are masked from the action space [43] . Leveraging task knowledge through action masking usually leads to faster convergence and also improves the predictability of the RL agent, especially when the set of relevant actions has a specific notion, such as being the set of safe actions. For instance, if the set of relevant actions is a set of verified safe actions, action masking can be used to provide safety guarantees [10, 19] .When the relevant action set is easy to compute, action masking usually benefits RL by improving the sample efficiency and is the quasi-standard for discrete action spaces [34, 46] , e.g., in motion planning [8, 18, 23, 30, 42] , games [14, 15, 43] , and power systems [21, 38] . However, real-world systems operate in continuous space and discretizing it might prevent learning optimal policies. Furthermore, simulation of real-world systems is computationally expensive, and developing RL agents for them often requires additional real-world training [45] . Thus, sample efficiency is particularly valuable for these applications.In this work, we propose three action masking methods for continuous action spaces. They can employ convex set representations, e.g., polytopes or zonotopes, for the relevant action set. Our action masking methods generalize previous work in [19] , which is constrained to intervals as relevant action sets. This extends the applicability of continuous action masking to expressive convex relevant action sets, which is especially useful when action dimensions are coupled, e.g., for a thrust-controlled quadrotor. To integrate the relevant action set into RL, we introduce three methods: the generator mask, which exploits the generator representation of a zonotope, the ray mask, which projects an action into the relevant action set based on radial directions, and the distributional mask, which truncates the policy distribution to the relevant action set. In summary, our main contributions are:• We introduce continuous action masking based on convex sets representing the statedependent relevant action sets;• We present three methods to utilize the relevant action sets and derive their integration in the backward pass of RL with stochastic policies;• We evaluate our approach on four benchmark environments that demonstrate the applicability of our continuous action masking approaches. 
related literature: 
 Action masking has been mainly applied to discrete action spaces [8, 10, 14, 15, 19, 18, 21, 23, 30, 38, 41, 42, 43, 46] . Huang et al. [15] derive implications on policy gradient RL and evaluate their theoretical findings on real-time strategy games. They show that masking actions leads to higher training efficiency and scales better with an increasing number of actions than penalizing the agent for selecting irrelevant actions. Huo et al. [14] have extended [15] to off-policy RL and have observed similar empirical results.Action masking for discrete action spaces can be categorized by the purpose of the state-dependent relevant action set. Often, the set is obtained by removing irrelevant actions based on task knowledge [8, 14, 15, 21, 30, 42] , e.g., executing a harvesting action before goods are produced [15] . While the relevant action sets are usually manually engineered, a recent study [46] takes a data-driven approach and identifies redundant actions based on similarity metrics. Another common interpretation of the relevant action set is that it only includes safe actions [10, 18, 23, 38, 41] . These works typically use the system dynamics to verify the safety of actions. A safe action avoids defined unsafe areas or complies with logic formulas. In our experiments, the relevant action sets are either state-dependent safe action sets or a global relevant action set modeling power supply constraints.For continuous action spaces, there is work on utilizing action masking with multidimensional intervals (hereafter only referred to as intervals) as relevant action sets [19] . In particular, the proposed continuous action masking represents relevant action sets by intervals that reflect safety constraints and employs straightforward re-normalization to map the action space to the relevant action set. In this paper, we generalize to more expressive relevant action sets and demonstrate the applicability to four benchmark environments. 
preliminaries: 
 As the basis for our derivations of the three masking methods, we provide a concise overview of RL with policy gradients. Further, we define the considered system dynamics as well as the set representations used in this work.A Markov decision process is a tuple (S, A, T, r, γ) consisting of the following elements: the observable and continuous state set S ⊂ R n S , the action set A ⊂ R N , the state-transition distribution T (s ′ |a, s), which is stationary and describes the transition probability to the next state s ′ ∈ S ⊂ R n S from the current state s ∈ S when executing action a ∈ A, the reward r : S × A → R, and the discount factor γ for future rewards [37] . The goal of RL is to learn a parameterized policy π θ (a|s) that maximizes the expected reward max θ E π θ ∞ t=0 γ t r(s t , a t ). For policy gradient algorithms, learning the optimal policy π * θ (a|s) is achieved by updating its parameters θ using the policy gradient [36, Thm. 2 ]∇J(π θ ) = E π θ [∇ θ log π θ (a|s)A π θ (a, s)] ,(1) where A π θ (a, s) is the advantage function, which represents the expected improvement in reward by taking action a in state s compared to the average action taken in that state according to the policy π θ (a|s). An estimation of the advantage A π θ (a, s) is usually provided by a neural network.We consider general continuous-time systems of the formEQUATIONwhere w ∈ W ⊂ R n W denotes a disturbance. The input is piece-wise constant with a sampling interval ∆t. We assume S, A, and W to be convex.Zonotopes are a convex set representation well-suited for representing the relevant action set, due to the efficiency of computing their Minkowski sums and linear maps. A zonotope Z ⊂ R N with center c ∈ R N , generator matrix G ∈ R N ×P , and scaling factors β ∈ R P is defined asEQUATIONAdditionally, let us denote that G (•,i) returns the i-th column vector of the generator matrix G. The Minkowski addition Z 1 ⊕ Z 2 of two zonotopes Z 1 , Z 2 and the linear map M Z 1 of a zonotope Z 1 are given by [2, Eq. 2.1]EQUATIONEQUATION)Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: For all three methods, we assume convex relevant action sets and a stochastic policy. Although, as mentioned in Section 5.3, the ray mask and generator mask could be adapted to non-stochastic policies. For the generator mask, we additionally assume the set representation to be a zonotope. We propose the adapted policy gradient for each method in Propositions 1, 3, and 4, and provide proofs for all three.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
continuous action masking: 
 To apply action masking, a relevant action set A r (s) ⊆ A has to be available, which constrains the action space A based on task knowledge. Let us denote the state-dependent relevant action set as A r (s) ⊆ A. From now on, we omit the dependency on the state to simplify notation. For our continuous action masking methods, we specifically require the following two assumptions: Assumption 1. The relevant action set A r is convex and its center and boundary points are computable. Assumption 2. The policy π θ : S × A → R + of the agent is represented by a parameterized probability distribution a ∼ π θ (a|s).Common convex set representations that fulfill Assumption 1 are polytopes or zonotopes. Our continuous action masking methods transform the policy π θ (a|s), for which the parameters θ usually specify a neural network, into the relevant policy π r θ : S × A r → R + through a functional h : (Π × P(A)) → Π r . Here, Π is the space of all policies, Π r is the space of all relevant policies, and P(A) is the power set of all A r . The transformation is defined as a r ∼ π r θ (a r |s) = h π θ (a|s), A r , (5) and thereby ensures that a r ∈ A r always holds. Please note that policy gradient methods only require the ability to sample from and to compute the gradient of the policy distribution. Therefore, explicit closed-form expressions for π r θ (a r |s) and h are not necessary. In the following subsections, we introduce and evaluate three masking methods: generator mask, ray mask, and distributional mask, as shown in Fig. 1 . We derive the effect of each masking approach on the gradient of the objective function for stochastic policy gradient methods. A r A G (•,3) a 3 G (•,2) a 2 G (•,1) a 1 A l c a r (b) Generator mask A r A π θ (a|s) π r θ (a r |s) a r (c) Distributional maskFigure 1 : Illustration of masking methods in action space A with a hexagon-shaped relevant action set A r . The ray mask radially maps the actions towards the center of the relevant action set. The generator mask employs the latent action space A l , which is the generator space of the zonotope modeling the relevant action set. The distributional mask augments the policy probability density function so that it is zero outside the relevant action set.The ray mask maps the action set A to A r by scaling each action a alongside a ray from the center of A r to the boundary of A, as shown in Figure 1a . Specifically, the relevant policy π r θ results from mapping the action a, sampled from π θ (a|s), using the function g R : A → A r :EQUATIONHere, λ A (a) and λ A r (a) denote the distances of the action a to the boundaries of the relevant action set and action space, respectively, measured from the center of the relevant action set c in the direction of a. Note that computing λ A r (a) for a zonotope requires solving a convex optimization problem, as specified in Appendix A.4. Yet, the ray mask is applicable for all convex sets, for which we can compute the center and boundary points. Since g R (a) is bijective (see Appendix A.1 for a detailed proof), we can apply a change of variables [5, Eq. 1.27] to compute the relevant policyEQUATIONwhere g -1 R (a r ) = a is the inverse of g R . In general, there is no closed form of the distribution π r θ available. However, for stochastic policy gradient-based RL, we only require to sample from π r θ and compute its policy gradient. Samples from π r θ (a r |s) are created by sampling from the original policy a ∼ π θ (a|s) followed by computing a r = g R (a). The policy gradient is derived next. Proposition 1. Policy gradient for the ray mask.The policy gradient of π r θ (a r |s) for the ray mask isEQUATIONwhere A π r θ (a r , s) is the advantage function associated with π r θ (a r |s).Proof. The determinant in (7) is independent of θ, i.e.,EQUATIONwhich simplifies the score function of π r θ (a r |s) to ∇ θ log π r θ (a r |s) = ∇ θ log π θ (a|s).Combining (10) and the general form of the policy gradient in (1) for π r θ (a r |s) results in (8).Zonotopes can be interpreted as the map of a hypercube in the generator space to a lower-dimensional space (see Section 3.2). The generator mask is based on exploiting this interpretation by letting the RL agent select actions in the hypercube of the generator space. Since the size of the output layer of the policy network cannot change during the learning process, we fix the dimension of the generator space. The generator mask requires the following assumption: Assumption 3. The relevant action set A r (s) is represented by a zonotope ⟨c(s), G(s)⟩ Z , with G ∈ R N ×P and c ∈ R N , and a state-invariant number of generators P .Note that in practice, Assumption 3 can often be trivially fulfilled by choosing sufficiently many generators, and the number of generators P is usually the output dimension of the parametrized policy. The domain of the policy π θ (a|s) is the hypercube A l = [-1, 1] P , which can be interpreted as a latent action space, and the domain of the relevant policy π r θ (a r |s) is a subset of the action space A r ⊆ A (see Figure 1b ). To derive the policy gradient of the generator mask, we assume: Assumption 4. π θ (a|s) is a parametrized normal distribution N (a; µ θ , Σ θ ).Proposition 2. The relevant policy of the generator mask isEQUATIONProof. The generator mask g G : A l → A r is:EQUATIONwhich is a linear function. Therefore, the proof directly follows from the linear transformation of multivariate normal distributions [Thm. 3.3.3][40] .Note that the ray mask and generator mask are mathematically equivalent to the approach in [19] if the relevant action set is constrained to intervals. Since g G (a) is not bijective in general, we cannot derive the gradient through a change of variables, as for the ray mask. Proposition 3. The policy gradient for π r θ (a r |s) as defined in (11) with respect to µ θ and Σ θ isEQUATIONEQUATION(a r -c -Gµ θ ) T (GΣ θ G T ) -1 G .Proof. The proposition is proven in Appendix A.2.Note that for the special case where G -1 exists, i.e., G is square and non-singular, the expressions in ( 13) and ( 14) simplify to ∇ µ θ log π θ (a|s), and ∇ Σ θ log π θ (a|s), respectively, when Assumption 4 holds (see Proposition 5 in Appendix A.2). While the generator matrix will commonly not be square, as usually P > N , there are cases where P = N is a valid choice, e.g., if there is a linear dependency between the action dimensions as for the 2D Quadrotor dynamics (see (40) ).The intuition behind the distributional mask comes from discrete action masking, where the probability for irrelevant actions is set to 0 [15] . For continuous action spaces, we aim to achieve the same by ensuring that actions are only sampled from the relevant action set A r , by setting the density values of the relevant policy distribution π r θ (a r |s) to zero for actions outside the relevant action set (see Figure 1c ). For the one-dimensional case, this can be expressed by the truncated distribution [7] . In higher dimensions, the resulting policy distribution isEQUATIONwhere ϕ(a, s) is the indicator functionEQUATION)Since there is no closed form of this distribution, we employ Markov chain Monte Carlo sampling to sample actions from the policy. More specifically, we utilize the random direction hit-and-run algorithm: a geometric random walk that allows sampling from a non-negative, integrable function f : R N → R + , while constraining the samples to a bounded set [44] . The algorithm iteratively chooses a random direction from the current point, computes the one-dimensional, truncated probability density of f along this direction, and samples a new point from this density. The approach is particularly effective for high-dimensional spaces where other sampling methods might struggle with convergence or efficiency. For the distributional mask, f is the policy π θ (a|s), and A r is the set. As suggested by [22] , we execute N 3 iterations before accepting the sample. To estimate the integral in (15) , we use numerical integration with cubature [11] , which is a method to approximate the definite integral of a function l : R N → R over a multidimensional geometric set. Proposition 4. The policy gradient for the distributional mask isEQUATIONProof. Equation ( 17) is obtained by calculating the gradient of the logarithm of ( 15). The indicator function ϕ(a, s) is not continuous and differentiable, which would necessitate the use of the subgradient for learning. However, since a r ∈ A r always holds, the gradient has to be computed for the continuous part of π r θ (a r |s) only, and thus ϕ(a, s) can be omitted.Since the gradient of the numeric integral A r π θ (ã|s)dã is intractable for zonotopes, we treat the integral as a constant in practice, and use ∇ θ log π r θ (a r |s) ≈ ∇ θ log π θ (a|s). We discuss potential improvements in Section 5.3.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We describe the necessary details to implement the three different masking approaches in Sections 4.1, 4.2, and 4.3. We provide information needed to re-implement the environments in Sections 5.1 and A.5. The computation of the relevant action set is specified in A.3. Further, the hyperparameters used for the training runs are reported in A.9. Any remaining ambiguity can be resolved by the code attached to this submission. 
numerical experiments: 
 We compare the three continuous action masking methods in four different environments: The simple and intuitive Seeker Reach-Avoid environment, the 2D Quadrotor environment to demonstrate the generalization from the continuous action masking approach in [19] , and the 3D Quadrotor and Mujoco Walker2D environment to show the applicability to action spaces of higher dimension. Because the derivation of the relevant action set is not trivial in practice, we selected four environments for which we could compute intuitive relevant action sets. In particular, for the Seeker and Quadrotor environments, the relevant action set is a safe action set since it is computed so that the agent does not collide (Seeker) or leave a control invariant set (2D and 3D Quadrotor). For the Walker2D environment, the relevant action set is state-independent and models a power supply constraint for the actuators.For the experiments, we extend the stable-baseline3 [29] implementation of proximal policy optimization (PPO) [33] by our masking methods. PPO is selected because it is a widely used algorithm in the field of RL and fulfills both Assumptions 2 and 4 by default. Apart from the masking agents, we also train a baseline agent with standard PPO that uses the action space A and a replacement agent, for which an action outside of A r is replaced by a uniformly sampled action from A r (see [19] for details). The replacement agent is an appropriate comparison to the masking agents since only relevant actions are executed while the replacement is implemented as part of the environment, which is usually easier than an implementation as part of the policy as for the masking methods. We conduct a hyperparameter optimization with 50 trials for each masking method and environment. The resulting hyperparameters are reported in Appendix A.9. All experiments are run on a machine with a Intel(R) Xeon(R) Platinum 8380 2.30 GHz processor and 2 TB RAM.We briefly introduce the three environments and their corresponding relevant action sets A r . Parameters and dynamics for the environments are detailed in Appendix A.5. The Seeker Reach-Avoid environment with state and action space. The agent (black) has to reach the goal (gray) while avoiding the obstacle (red). The center of the action space is illustrated by a cross and the relevant action set A r for the current state is shown in green. The state set reachable at the next time step, by the relevant action set, is S ∆t .This episodic environment features an agent navigating a 2D space, tasked with reaching a goal while avoiding a circular obstacle (see Figure 2 ). It is explicitly designed to provide an intuitive relevant action set A r . The system is defined as in Section 3.2: The dynamics for the position of the agents = [s x , s y ]T and the action a = [a x , a y ] T is ṡ = a, and there are no disturbances.The environment is characterized by the position of the agent, the goal position s * , the obstacle position o, and the obstacle radius r o . These values are pseudo-randomly sampled at the beginning of each episode, with the constraints that the goal cannot be inside the obstacle and the obstacle blocks the direct path between the initial position of the agent and the goal. The reward for each time step isEQUATIONWe compute the relevant action set A r so that all actions that cause a collision with the obstacle or the boundary are excluded (see Appendix A.3.3).The 2D Quadrotor environment models a stabilization task and employs an action space where the two action dimensions are coupled, i.e., rotational movement is originating from differences between the action values and vertical movement is proportional to the sum of the action values. The relevant action set is computed based on the system dynamics and a relevant state set (see Appendix, Eq. ( 33)). The reward function is defined asEQUATIONwhere s * = 0 is the stabilization goal state, a = [a 1 , a 2 ] is the two-dimensional action, a i,min is the lower bound for the actions in dimension i, and a i,range is the absolute difference between the lower and upper bound for the actions in dimension i.The third environment models a stabilization task for a quadrotor defined in [16] . The quadrotor has four action dimensions. We use the same reward (see (19) ) and the same calculation approach for the relevant action set (see Appendix, Eq. ( 33)) as for the 2D Quadrotor.The relevant action sets of the two Quadrotor and the Seeker environments are sets that only contain safe actions for the current state. Computing safe action sets requires considerable domain knowledge and for our case solving an optimization problems (see Appendix A.3). However, action masking is not restricted to safe relevant action sets, which we aim to demonstrate on the Mujoco Walker2D environment [39] . We extend the environment with a termination criterion, which ends an episode, when the the action violates the constraint ∥a∥ 2 ≤ α p . This can be viewed as constraining the cumulative power output on all joints to a maximum value α p . We motivate this constraint by having a battery with a maximum power output that should not be exceeded in practice. Accordingly, we define the relevant action set as the static setEQUATIONWe under-approximate this relevant action set with a zonotope that consists of 36 generators.The reported training results are based on ten random seeds for each configuration. Figure 3 shows the mean reward and the bootstrapped 95% confidence intervals [27] for the four environments and Table 1 depicts the mean and standard deviation of the episode return during deployment. First, we present the results for the Seeker and Quadrotor environments for which the relevant action set is state-dependent and only includes safe actions. Then, we detail the results for the Walker2D environment with a static relevant action set.For the Seeker and Quadotor environments, the baseline, i.e., PPO, converges significantly slower or not at all (see Figure 3 ). Additionally, the initial rewards when using action masking are significantly higher than for the baseline, indicating that the exploration is indeed constrained to relevant taskfulfilling actions. During deployment (see Table 1 ), the baseline episode returns are significantly lower than for action masking, while the three masking methods perform similarly. More specifically, for the Seeker environment, the relative volume of the relevant action set compared to the global action space is on average 70%, and all action masking methods converge significantly faster to high rewards than the baseline. The generator mask and the distributional mask achieve the highest final reward. Further, action replacement performs better than the baseline but significantly worse than masking. For the 2D Quadrotor, the relative volume is on average 28%, which is much smaller than for the Seeker environment. In the 2D Quadrotor environment, the ray mask, generator mask, and action replacement achieve the highest reward. While the baseline converges significantly slower, the final average reward is similar to the one of the distributional mask. Please note that the confidence interval for the baseline is significantly larger, because three of the ten runs do not converge and, on average, exhibit a reward of one throughout training. Additionally, we observed that if we constrain the relevant action set to intervals for this environment, our optimization problem in (33) often renders infeasible, because the maximal relevant action set cannot be well approximated by an interval. Thus, the masking method proposed in [19] is not suitable for the 2D Quadrotor task. The results on the 3D Quadrotor are similar to those on the 2D Quadrotor; again, the generator mask converges the fastest but yields a final reward similar to that of the ray mask and action replacement. The relative volume of the relevant action set to the global action space is on average 25%. Notably, in this environment, the baseline does not learn a meaningful policy. Based on these three environments with state-dependent relevant action sets, it seems that action masking performs better than action replacement when the relative volume is not too small.The training results of the Walker2D experiment are shown in Figure 3 . While the generator and ray mask both learn a performant policy, the ray mask outperforms by a significant margin. The lower performance of the generator mask is likely due to the high-dimensional generator space with 36 dimensions. This is supported by initial experiments with 12 generators (i.e., a more conservative under-approximation of the L 2 -norm) where the generator mask performed better compared to the ray mask. Replacement performs significantly worse than the two masking approaches, and the PPO baseline does not learn a meaningful policy since the environment is frequently reset due to violations of the power constraint in (20) . The frequent terminations occur since in six action dimensions the relative volume of the unit ball compared to the unit box is ≈ 8%, i.e., more than 92% of actions are violating the power constraint. To compare masking to a learning baseline, we also evaluated standard PPO without constraints, which performs slightly better than ray masking but almost always uses actions outside A r (see Appendix A.7). The deployment results in Table 1 reflect similar results as the training; the ray mask achieves better rewards than the generator mask, followed by replacement, and the baseline performs the worst. We excluded the distributional mask for the Walker2D, since its computation time is approximately 170 times slower than the baseline, compared to a 1.6 increase for the generator, 2.7 for the ray mask, and 2.5 for action replacement (see Table 3 ). The severely increased computational cost for the distributional mask arises from the geometric random walks, which scale cubically with action dimensions.Our experimental results indicate that continuous action masking with zonotopes can improve both the sample efficiency and the final policy of PPO. While the sample efficiency is higher in our experiments, computing the relevant action set adds computational load as shown by the increased computation times (see Appendix A.6). Thus, in practice, a tight relevant action set might require more computation time than the additional samples for standard RL algorithms. Yet, if the relevant action set provides guarantees, e.g., is a set of verified safe actions, this increased computation time is often acceptable. Additionally, the computational effort for the masks differs. Given a relevant action zonotope, the generator mask adds a matrix-vector multiplication, which scales quadratically with action dimensions, the ray mask is dominated by the computation of the boundary points, which scales polynomially with action dimensions [20] , and the distributional mask scales cubically with the dimension of the action space due to the mixing time of the hit-and-run algorithm [22] . Note that for the Walker2D environment with six action dimensions, the computation time for the hit-and-run algorithm is so high that the distributional mask evaluation was infeasible.The ray mask and generator mask are based on functions g R and g G that map to relevant actions. There are two different approaches of incorporating these functions into RL algorithms. One is to apply the mapping as part of the environment on an action that is sampled by a standard RL policy. The second option, which we use, is to consider the masking as part of the policy, which creates the relevant policy as in (5) . This has three main benefits over integrating masking as part of the environment. First, the actions passed to the environment are better interpretable, e.g., for the generator mask, adding the masking mapping function to the environment leads to a policy that samples actions from the generator space, which commonly does not have an intuitive interpretation. Second, more formulations of the functional h are possible, e.g., the distributional mask, and, third, the mapping function can be included in the gradient calculation. For mathematically sound masking as part of the policy, the correct gradient needs to be derived for each RL algorithm, and the standard RL algorithms need to be adapted accordingly. However, the empirical benefit could be minor. Thus, future work should investigate the significance of the correct gradient on a variety of tasks. Note that we showed in Proposition 1 that for the ray mask, the PPO gradient with respect to the original policy and relevant policy is the same. Thus, for the ray mask, it does not matter if it is viewed as part of the policy or the environment.PPO is a common RL algorithm. However, off-policy algorithms such as Twin Delayed DDPG (TD3) [9] and Soft Actor-Critic (SAC) [12] are frequently employed as well. The ray mask and generator mask are conceptually applicable for deterministic policies as used in TD3. Yet, the implications on the gradient must be derived for each RL algorithm and are subject to future work.For the distributional mask, treating the integral in (15) as constant with respect to θ is a substantial simplification, which might be an explanation for the slightly worse convergence of the distributional mask since this introduces an off-policy bias. To address this in future work, one could approximate the integral with a neural network, which has the advantage that it is easily differentiable.We focus this work on the integration of a convex relevant action set into RL and assume that an appropriate relevant action set can be obtained. While convex sets are a significant generalization from previous work [19] , they might not be sufficient for some applications, e.g., tasks where relevant action sets are disjoint. Thus, future work could include investigating hybrid RL approaches [25] to increase the applicability to multiple convex sets or non-convex sets, such as constrained polynomial zonotopes [17] . Further, obtaining the relevant action set can be a major challenge in practice, in particular when the relevant action set is different for each state. Such a high state-dependency is likely when the notion of relevance is safety, while for other definitions of action relevance, the relevant action set might be easy to pre-compute, e.g., excluding high steering angles at high velocities.Additionally, there might be an optimal precision of the relevant action set due to two opposing mechanisms. On the one hand, the larger the relevant action set is with respect to the action space, the smaller the sample efficiency gain from action masking might get. On the other hand, a tight relevant action set might require significant computation time to obtain. Thus, future work should investigate efficient methods to obtain sufficiently tight relevant action sets. 
conclusion: 
 We propose action masking methods for continuous action spaces that focus the exploration on the relevant part of the action set. In particular, we extend previous work on continuous action masking from using intervals as relevant action sets to using convex sets. To this end, we have introduced three masking methods and have derived their implications on the gradient of PPO. We empirically evaluated our methods on four benchmarks and observed that the generator mask and ray mask perform best. If the relevant action set can be described by a zonotope with fixed generator dimensions and the policy follows a normal distribution, the generator mask is straightforward to implement. If the assumptions for the generator mask cannot be fulfilled, the ray mask is recommended based on our experiments. Because of subpar performance and longer computation time, the distributional mask needs to be further improved. Future work should also investigate a broad range of benchmarks to identify the applicability and limits of continuous action masking with convex sets more clearly. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification: The paper introduces a method to exploit task knowledge via relevant action sets to achieve more efficient and effective reinforcement learning. Given the limited scope of our experiments, a detailed discussion on broader societal impact would be very speculative. However, we believe our approach has the potential to expand the use of reinforcement learning in applications where training opportunities are limited but task knowledge is easily transferable to relevant action sets and where enforcing hard constraints is important.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not work with, or release any datasets, and our models have no risk for potential misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Proof of bijectivity of the ray mapping function.Lemma 1. The mapping function g : A → A r of the ray mask g(a) = c + λ A r (a) λ A (a) (a -c) is bijective.Proof. We show that g(a) is bijective by showing that the function is both injective and surjective. For any convex set A r with center c, we can construct a ray from c through any a r ∈ A r as r γ (t) = c + γt, where γ = a-c ∥a-c∥2 , and t ∈ [0, t max ]. Since A r ⊆ A, this also holds ∀a ∈ A. By construction, any two distinct rays only intersect in c. For all points on a given ray, the scaling factors λ A and λ A r are constants, allowing us to rewrite g(a) for all a = r γ (t) as the linear function:g (r γ (t)) = c + λ A r λ A (r γ (t) -c) .To show that g(a) is injective, i.e. ∀a 1 , a 2 ∈ A, if a 1 ̸ = a 2 =⇒ g(a 1 ) ̸ = g(a 2 ), consider the following two cases for a 1 ̸ = a 2 .Case 1: If a 1 and a 2 are on the same ray, then, a 1 = r γ (t 1 ) and a 2 = r γ (t 2 ) with t 1 ̸ = t 2 . Since g (r γ (t)) is linear and thereby monotonic, it follows that g (r γ (t 1 )) ̸ = g (r γ (t 2 )) and consequently g(a 1 ) ̸ = g(a 2 ). Case 2: Otherwise, a 1 and a 2 are on different rays and a 1 ̸ = a 2 ̸ = c, so g(a 1 ) ̸ = g(a 2 ) follows directly as the rays only intersect in c.Thus, g(a) is injective.For showing g(a) to be surjective, it is sufficient to show that ∀a r ∈ A r , there exists an a ∈ A, for which g(a) = a r .Consider the ray r γ (t) passing through a r . We know that g (r γ (0)) = g(c) = c, and g (r γ (t max )) is the boundary point of A r from c in the direction d. Moreover, a r lies on the line segment between c and g (r γ (t max )). Since g (r γ (t)) is linear and continuous, according to the intermediate value theorem, ∃t * ∈ [0, t max ], for which g (r γ (t * )) = a r for any a r ∈ A r along the ray r γ (t). As we can construct such a ray through any point in A r , we have shown that for every a r ∈ A r , there exists an a ∈ A such that g(a) = a r , thus proving surjectivity.The log probability density function of the relevant policy for the generator mask isEQUATIONWe can derive the gradient w.r.t µ θ asEQUATIONand w.r.t Σ θ as∇ Σ θ log π r θ (a r |s) = = - 1 2 ∇ Σ θ log |GΣ θ G T | - 1 2 ∇ Σ θ (a r -c -Gµ θ ) T (GΣ θ G T ) -1 (a r -c -Gµ θ ) = - 1 2 (GΣ θ G T ) -1 ∇ Σ θ GΣ θ G T (23) + 1 2 (GΣ θ G T ) -1 (a r -c -Gµ θ )(a r -c -Gµ θ ) T (GΣ θ G T ) -1 ∇ Σ θ GΣ θ G T = - 1 2 G T (GΣ θ G T ) -1 G -G T (GΣ θ G T ) -1 (a r -c -Gµ θ )(a r -c -Gµ) T (GΣ θ G T ) -1 GWe can state the following for the special case when the inverse of G exists.Proposition 5. If G is invertible, ∇ θ log π rGiven an initial state s 0 , an input trajectory a (•) , and a disturbance trajectory w (•) , we denote the solution of (2) at time t as ξ t (s 0 , a (•) , w (•) ). Assuming a sampled controller with piecewise-constant input a, we define the reachable set of (2) after one time step ∆t given a set of initial states S 0 and a set of possible inputs Ã ⊆ A asEQUATIONThe reachable set over the time interval [0, ∆t] is defined asEQUATION)For many system classes it is impossible to compute the reachable set exactly [28] , which is why we generally compute overapproximations R(•) ⊇ R e (•).To guarantee constraint satisfaction over an infinite time horizon, we use a relevant state set S r . In this work, we choose S r to be a robust control invariant set in the sense that we can always find an input which guarantees that the reachable set at the next time step is contained in S r and that all state constraints are satisfied in the time interval in between [32] :EQUATION)We compute the relevant action set as the largest set of inputs that allows us to keep the system in the relevant state set in the next time step. We define a parameterized action set A r (p), where p ∈ R n p is a parameter vector. The relevant action set is then computed with the optimal programEQUATIONwhere Vol(•) is a proxy function for the volume in the sense that a maximization of Vol(•) is suboptimally maximizing the volume. In the following, we provide a detailed formulation of ( 29) as an exponential cone program.We consider the discrete-time linearization of our systemEQUATIONwhere w ′ k ∈ W ′ (s k ) additionally contains linearization errors and the enclosure of all possible trajectory curvatures between the two discrete time steps [2] . Furthermore, we assume S, A, W, and S r to be zonotopes. With regard to the input set parameterization, we consider a template zonotope with a predefined template generator matrix G. We use the vector p ∈ R P >0 of generator scaling factors to scale the generator matrix. The parameterized template zonotope is then given byEQUATIONwhere p = [c p] ⊤ . In the following, we denote the center and generator matrix of a specific zonotope Z by c Z and G Z , respectively. A zonotope Z 1 is contained in a zonotope Z 2 if there existEQUATION)Based on the linearized system dynamics in (30) , the definitions from (4), and using R as a shorthand for R ∆t (•), we have G R = AG S0 BG A r G W and c R = Ac S0 Bc A r c W . Since we already consider trajectory curvature in the disturbance, we only need to guarantee that the relevant input set is contained in the feasible input set and that the reachable set of the next time step is contained in the relevant state set. Using the containment condition from (32), we define the auxiliary variables Γ R , β R , Γ A r , and β A r and solveEQUATIONwhere 0 are zero matrices of appropriate dimensions. By vectorizing Γ R and Γ A r and resolving the infinity norm constraints [4, Sec. 1.3], we can obtain a formulation with purely linear constraints. Since computing the exact volume of a zonotope is computationally expensive, we approximate it by computing the geometric mean of the parameterization vector:EQUATION)With the cost function from (34) , the problem in ( 33) is an exponential cone program, which is convex and can be efficiently computed [6, Sec. 2.5].Because of the dynamics of the Seeker Reach-Avoid environment, we can simplify the relevant action set computation to the following optimization problem:EQUATIONwhere S ∆t is the reachable set of the agent in the next time step, the set O represents the obstacle, and the box [-10, 10] 2 is the outer boundary of the state space. The constraints represent three intuitive geometric constraints. First, containment of the relevant action set in the action set (usually, the interval [-1, 1] 2 , second, the containment of the relevant state set of the next time step in the environment boundaries, and the last constraint enforces that there is no collision with the obstacle possible. To enforce these constraints for zonotopes, we use the support function [3, Lemma 1] EQUATION)where l is the vector in the direction of interest. For the first constraint, we need to ensure that the support function for the basis vectors e 1 , and e 2 (and there negative versions) are less than or equal to 1, i.e.,EQUATION)From ( 36) it is apparent that using the negative basis vectors -e 1 and -e 2 , just flips the sign of c A r in (37) . The constraints for S r ∆t ⊆ [-10, 10] 2 can be constructed similarly. For simplicity, we express the element-wise inequality in (37) using the L ∞ -norm, and write the optimization problem asEQUATIONThe last constraint represents an under-approximation of S r ∆t ∩ O = ∅ enforcing the containment of A r in the halfspace {x|n T x ≤ b} through the support function ρ A r (n). The directional vector n = o-s ∥o-s∥ where o is the center of the obstacle, and the offset b = n T (o -n)r where r is the radius of the obstacle. This ensures that the intersection between the reachable set S r ∆t and the obstacle O will be empty since the halfspace is constructed tangential to the obstacle at the intersection point of the line from the agent to the center of the obstacle.The boundary point p ∈ R N on a zonotope ⟨c,EQUATIONand computing p = x + αd [20] .We provide the action space bounds, the state space bounds, and the generator template matrix G for the Seeker and Quadrotor environments in Table 2 . Template matrix G 1 1 1 0 1 -1 0 1 1 1 1 1 -1 0 I 4The system dynamics of the 2D Quadrotor is:EQUATIONwhere the state is s= x, z, ẋ, ż, θ, θ T , the action is a = [a 1 , a 2 ] T , and the disturbance is w = [w 1 , w 2 ]T . The dynamics are adapted from [24, Eq. 43] so that the actions are two independent thrusts. Additionally, the parameters used in our experiments are g =9.81 m s-2 , k =1 1/kg, d 0 =70, d 1 =17, n 0 =55, and W = [-0.08, -0.08] × [0.08, 0.08].The 3D Quadrotor is modeled in a twelve-dimensional state space with state s = [x, y, z, ẋ, ẏ, ż, θ, ϕ, ψ, θ, φ, ψ] T and four-dimensional action space with action a = [a 1 , a 2 , a 3 , a 4 ] T . The system dynamics are ṡ = [ ẋ, ẏ, ż, -9.81ϕ, 9.81θ, a 1 , θ, φ, ψ, a 2 , a 3 , a 4 ] T [16] .For the Walker2D environment, we use the standard parameters of the gymnasium implementation2 . The relevant action set is a zonotope under-approximation of the relevant action set stated in (20) with 36 generators and α P = 1.Note that the increased computation time of the masking approaches compared to the baseline is mainly due to the computation of the relevant action sets. Since the relevant action set is a zonotope, the generator mask does not require expensive additional computations. The increased runtime for the ray mask mainly originates from the computation of the boundary points (see Appendix A.4). For the distributional mask, the increased runtime is mostly caused by the Markov chain Monte Carlo sampling of the action. A.7 Quantitative results for Walker2D environment compared to standard PPO If we remove the power constraint that resets the environment whenever an action outside of A r is selected, the standard PPO is able to learn a policy that performs slightly better than the ray mask towards the end of training while converging slightly slower (see Figure 4 ). However, during deployment, only 0.05% of the actions from this policy are from within A r . Yet, the reward that standard PPO achieves during deployment is on average 2177.57 with standard deviation 1238.42. This is slightly higher than the deployment results for the ray mask (see Table 1 ).A.8 Qualitative results for the Seeker environment Figure 5 presents ten example trajectories for a trained agent using each of the five approaches under consideration. Notably, the continuous action masking agents solve the task by safely reaching the goal while the PPO baseline still collides with the obstacle. The replacement approach is ineffective, with the agent frequently failing to reach the goal. In fact, for the ten displayed rollouts, no replacement agent reaches the goal. This discrepancy with respect to the training performance where the replacement agent reaches the goal frequently is due to using the policy in a deterministic setting for deployment. The simple dynamics of the Seeker environment make it possible to directly visualize the relevant action set A r at each time step along the trajectory, which is depicted in the lower half of Figure 5 . For the generator mask, the least amount of A r is plotted, which indicates that the generator mask agent reaches the goal most efficiently, i.e., the lowest amount of steps are required.We specify the hyperparameters for the three masking approaches, replacement and baseline in the Seeker (Table 4 ), the 2D Quadrotor (Table 5 ), and the 3D Quadrotor environment (Table 6 ). These were obtained through hyperparameter optimization with 50 trials using the tree-structured Parzen estimator of Optuna with the default parameters [1] .Generator Distributional Replacement • The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Justification: We train each configuration with ten random seeds and report the training results with 95% bootstrapped confidence intervals in transparent color in Figure 3 . For the reported deployment results and computational runtime (see Table 1 and 3 , respectively), we report the mean and 1-sigma standard deviation.Guidelines:• The answer NA means that the paper does not include experiments. Justification: We reported the compute resources utilized in Section 5. The average computation times are reported in Section A.6. Note that the compute resource used for the experiments is a shared resource where also other researcher ran their experiments at the same time. Thus, runtime evaluations will only be partially reproducible.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).Justification: Our experiments do not involve human subjects or datasets. Regarding safety and security, our approaches can be directly used to prevent agents from performing unsafe actions. We cannot identify any potentially harmful consequences of our algorithms and work in general.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).https://gymnasium.farama.org/environments/mujoco/walker2d/ 


Decision: The paper proposes three continuous action masking methods to focus reinforcement learning on state-specific relevant actions, enhancing training efficiency, effectiveness, and safety. Although there is consensus about the overall positive evaluation of this paper, the reviewers suggest that the paper would benefit from (a) additional experiments supporting the need for the gradient derivations. (b) improved clarity of presentation. Apart from this, the reviews indicate that evaluations on more complex domains would strengthen the work.

Reviews: review: paper_summary: paper_summary: This paper proposes mathematical formulations for continuous action masking in reinforcement learning, to incorporate domain-knowledge in the form of state-specific sets of relevant actions. It introduces 3 functional forms to extract relevant actions from the original action space, and consider its effect on the policy gradient. The policy gradient does not change much, and the paper shows that the forms compare similarly and better than learning an agent without any knowledge of the continuous action mask at all.
strength_weakness: strength_weakness: ## 1. General applicability of this paper's ideas
Obtaining a state-specific relevant action set can be really hard. The paper, however, makes contradictory statements about this:
- L5: "little task knowledge can be sufficient to identify significantly smaller state-specific sets of relevant actions."
- L284-285: "assume that an appropriate relevant action set can be obtained. Yet, obtaining this can be a major challenge in practice."

From the experiments on the 3 environments, it already seems like defining the relevant action set requires a lot of domain knowledge about the state space features and the dynamics function.

As of now, there does not seem to be any way the ideas in this paper could be useful for any practical domain.
- Can the authors provide some concrete examples of how one can obtain such relevant action sets for problems of practical interest and scale?
- Can the authors provide any results on a commonly used continuous action space RL benchmark?

## 2. Gains not coming from the policy gradient, but only because of constraining the action space
The paper's proposed formulation is interesting because it uses continuous action masking as part of the learning policy and informs the policy gradient update about the continuous action mask. However, when we look at the resultant policy gradients for each mask in Eq. 10, Line173, and Line199, it seems that the policy gradient simply reduces to $\lambda_\theta log \pi_\theta(a | s)$ for all cases.

So, the effective change in implementation is just how the action to take in the environment is model: $a^r = g(a)$. But, this **doesn't utilize continuous action masking to improve the policy learning objective** in any meaningful way. Is my understanding correct in this?

Another observation that validates the claim that policy learning is not influenced much is seen from the results and L248-249. The initial rewards themselves are significantly higher, which means that the action mask just reduces the space of exploration of the agent so much that, as long as it takes a valid action, it would get a high reward.

## 3. Simpler baselines for continuous action masking
Continuing from the above point, if all that needs to be done is to compare different formulations of g(a), there is a much simpler alternative perspective:
- Action-Masking as part of environment: Simply, apply the action mask as part of the environment, without changing the PPO objective at all. So, there is an action-masking filter before executing an agent's action in the environment, and ignore if the action is invalid.
- Sampling-augmented action-masking: Keep sampling actions from $\pi_\theta$ until you find a valid action that can pass through the known continuous action-masking map.

The current PPO baseline is very weak, and does not utilize action-masking at all. It seems most of the learning in PPO is going into the effort of learning the action mask. To really justify this paper's proposed action masking schemes are useful, they must compare against other forms of naive action masking, including the two listed above. This perspective on considering the action mask as part of the environment is also much more generally applicable and does not require any change to the policy gradient update.
questions: questions: Listed in weaknesses.
limitations: limitations: The authors discuss several limitations of the work, but the important ones of applicability and baselines need to be addressed.
paper_summary: paper_summary: This paper discusses methods for action masking in continuous action spaces to improve convergence stability and sample efficiency in reinforcement learning. The paper introduces three methods for action masking with convex relevant action sets, proves their convergence, and experimentally verifies their effects.
strength_weakness: strength_weakness: This paper is pretty solid overall, and I have few major complaints.

The one significant issue I see is that I think the distributional mask algorithm is off-policy by nature, meaning it's use with on-policy methods like PPO is biased and will cause performance loss or divergence. This may explain the observed underperformance of this masking method in two of the three experimental tasks, and while the algorithm can clearly converge in some cases it seems like a major issue with that particular mask in the context of PPO (off-policy algorithms could use it without issue, but those are left to future work here) that should be noted, or the mask omitted from this paper and left to future off-policy methods.

Beyond that, the experimental evaluation is relatively simple (though I think it is sufficient to validate these algorithms), and more challenging tasks would be useful to demonstrate the limitations of these masking methods. That said, the paper makes it clear that defining a suitable convex relevant action set is a manual process and can be challenging (this is okay as a limitation), so it is understandable why such stress tests are not performed. If there was a way to increase difficulty without major manual action set definition work it would strengthen the evaluation to include it.

I have a few other minor issues and questions noted below, but overall this is a paper that is clear in its goals and describes methods that achieve them, validated to a reasonable standard of theoretical and experimental evidence. There's more that could be done on this topic, but the contribution of this paper is significant on its own, so I'm inclined to recommend acceptance (particularly if something is done to address my concern about the distributional mask above).
questions: questions: -Is assumption 1 (relevant action set is convex) reasonable in most cases? I can imagine disjoint action subsets being relevant in many cases- for example, a self-driving car that needs to steer either left or right to avoid collision, but not go straight ahead or backwards.

-I'm not sure it's actually necessary to compute the policy gradient across the action mask (with the exception of the distributional mask). Once an action is sampled from the policy, the mapping to the relevant action set can simply be treated as part of the environment/transition function which the policy can learn to manipulate without gradient flow. Does this simplify things or am I missing something? This would also permit arbitrary nonconvex relevant sets, I believe.

-For the gradient of the distributional mask in proposition 4, isn't this affected by off-policyness due to the constrained sampling of actions from the policy distribution? For example, if most of the policy distribution probability mass lies outside the relevant set (e.g. in the event of transfer learning to a different task with a new relevant set) the actions sampled will not follow the policy distribution closely and thus \pi_{\theta}(a|s) will not be an accurate probability of sampling action a at state s. As noted above, this seems like a big issue that should be noted or addressed, unless I'm missing something that corrects for the off-policyness.

-Small quibble: The goal in figure 2 looks black to me on two different monitors, perhaps using a lighter gray would make it more distinct from the agent position?

-It would probably be reasonable to move the environment definitions for the two quadrotor tasks to the appendix to save space in the main paper, FWIW. I'm not sure the abbreviated version that's present provides all that much context over the text descriptions of the tasks.

-It's not critical to have, and I realize it's a difficult thing to derive a relevant set for by nature, but having an experiment on an environment with 10+ action dimensions would be a nice addition to demonstrate that these masking approach can scale to higher dimensional action spaces tractably. I'd also appreciate some comments on compute cost scaling with action dimension count in the limitations or conclusions sections, if possible, since it seems like compute cost is likely to increase with the dimensionality of the action space.
limitations: limitations: This paper does an excellent job of making clear its limitations and scope.

I don't see any potential negative societal impacts from this work.
paper_summary: paper_summary: The paper addresses challenges in RL with continuous action spaces, typically defined as interval sets. These spaces often lead to inefficient exploration due to irrelevant actions. The authors propose three continuous action masking methods to focus learning on relevant actions based on current state, improving predictability and suitability for safety-critical applications. They analyze the implications on policy gradients and evaluate performance using PPO across three control tasks. Results show higher final rewards and faster convergence compared to baseline methods without action masking.
strength_weakness: strength_weakness: **Reinforcement Learning with Policy Gradients (Section 2.1)**
- L84: "r →: S × A" appears incorrect.
  
**Continuous Action Masking (Section 3)**
- Assumption 1: Clarify the definition of action relevance.
  
**Ray Mask (Section 3.1)**
- L131: Need proof that g(a) is bijective.


**Generator Mask (Section 3.2)**
- Why is A(s) suddenly state-dependent? Provide motivation and further description.
- In Proposition 2's proof, the matrix multiplication seems infeasible due to mismatched dimensions (C is N x 1 and Ga results in P x 1).


**Experiment (Section 4)**
- Justify the rationale behind the design choices for action relevance in each environment.
- Compare the chosen action relevance approach to other relevant action settings.
  
**Results (Section 4.2)**
- Why compare to a standard PPO baseline and not to prior relevant works?
- Include qualitative results to validate the proposed methods.
questions: questions: See weaknesses.
limitations: limitations: n/a
paper_summary: paper_summary: Learn a state specific mask for actions. Rather than simply a state specific interval, extend the action mask to different convex set representations. Then, derive a policy gradient for each of these masking schemes.  The masking schemes are ray masks, hypercube transform mask and distributional masks. Applies action masking to seeker and quadrotor tasks and shows that this action masking improves performance.
strength_weakness: strength_weakness: It is not clear how easy it is to recover the action masking criteria, especially under the more complex generator or distributional schemes, and it seems like this would be rare

The experiments are not particularly convincing because they all follow similar control tasks, but it also seems like these are the only tasks for which the action mask could be easily defined.
questions: questions: It is not clear why related work is not in a separate section, rather than subsection. There does not appear to be a special connection to the introduction.

It isn't obvious that if G is square and non-singular, that this does not restrict the space of possible relevant action sets, since this would ensure that the hypercube space had an invertible, i.e. one to one, mapping between itself and the action distribution. It seems like many to one would be preferred if the space of the zonotope's hypercube was higher dimension than the action set.
limitations: limitations: While noted in the limitations, deriving a method for identifying not only the policy gradient, but a policy from a learned value function is highly relevant for RL, and it is not clear how these restricted action spaces can be applied to Q value calculations, though it seems reasonable to assume it is possible.

Some suggestions in the main work of how the distributional or generator action space restriction could be defined as a function of the dynamics could be relevant since it seems like these functions have to be hand-engineered, and it is not obvious how to do that in domains where the dynamics are less well defined.

Comments: title: Answer to Response
Response to Response
Anser to Response
Response
Answer to "Thoughts regarding action masking as part of the policy versus the environment"
Answer to Response
Response to Reviewer
Response to Authors
Answer to Response
Response to Rebuttal
Thoughts regarding action masking as part of the policy versus the environment
Response
Answer to followup (2/2)
Answer to followup (1/2)
Followup
Physical implications of the gradient for the distributional mask (2/2)
Thought experiments for gradient flow and non-bijectivity (1/2)
Supplementary figures
Response
Trying to understand why action masking should be considered part of policy and not environment
comment: Thank you for your reply and clarifications.

> Regarding gradient approximation/offpolicyness issues, I'll simply say that while yes an approximation can be useful, it needs to be an unbiased approximation, or if biased that bias needs to be empirically trivial. Given that the distribution mask underperforms the others in some test cases, I'm not sure that's good evidence for the bias being trivial in practice.

We agree that the approximation introduces bias and will highlight this more prominently in the limitations. Yet, we believe that the theoretic introduction and empirical evaluation of the distributional mask add value to our paper. It shows that naively translating discrete action masking to continuous action masking creates a challenge for calculating the gradient and is not the best-performing method for our experiments. The empirical impact of the bias introduced by the approximation of the gradient was a priori unclear and could have been compensated by maintaining the shape of the policy distribution within the relevant action set. This certainly merits further methodological and empirical investigation. 

> Regarding environment masking supporting non-convex sets, removing the need to compute a gradient means non-differentiable functions can be used (for example, mapping a continuous action space from -1 to 1 to two disjoint sets, -1 to -0.5 and 0.5 to 1). The shape of the masking function is then only limited by the ability of the designer to implement it depending on the task since the policy can learn to interact with it regardless.

Thank you for clarifying. If the relevant action set is not state-dependent, it might be relatively easy to design such a mapping of disjoint sets in the environment. Yet, when the relevant action set is state-dependent and computed online, it is likely to be impossible to foresee all kinds of relevant action sets and, thus, also difficult to design a mapping function. Therefore, we believe that in such cases action replacement would be used, for which only a function that replaces irrelevant with relevant actions needs to be engineered. To reiterate, our framework also allows for non-differentiable functions, if the function is bijective (such as the function in your example).
Regarding gradient approximation/offpolicyness issues, I'll simply say that while yes an approximation can be useful, it needs to be an unbiased approximation, or if biased that bias needs to be empirically trivial. Given that the distribution mask underperforms the others in some test cases, I'm not sure that's good evidence for the bias being trivial in practice.

Regarding environment masking supporting non-convex sets, removing the need to compute a gradient means non-differentiable functions can be used (for example, mapping a continuous action space from -1 to 1 to two disjoint sets, -1 to -0.5 and 0.5 to 1). The shape of the masking function is then only limited by the ability of the designer to implement it depending on the task, since the policy can learn to interact with it regardless.
Thank you for the clarification. The relevant action set is state-dependent and, for our experiments, computed with the optimization problems specified in Eq. 35 and 38 in the Appendix. Generally, we assume that $\mathcal{A}(s)$ is computed or derived by another algorithm, which is independent of the RL agent policy (see summary rebuttal **A2. Computing relevant action sets** for more details). For our experiments, the relevant action sets become smaller the closer the agent navigates to the obstacle (Seeker) or boundary of the relevant state set $\mathcal{S}^r$ (Quadrotors). For the Quadrotor environments, the task is to stabilize at the center of $\mathcal{S}^r$. Thus, with the policy improving, the agent chooses fewer actions that lead to moving toward the boundary of the relevant state set $\mathcal{S}^r$. For the Seeker environment, the goal is to reach a position while the direct path is blocked by an obstacle. Therefore, navigating close to the obstacle is optimal, and the agent actively navigates to states with smaller relevant action sets. More specifically, using the zonotope with 4 generators (as depicted in Fig 2), the relevant action set covers approximately 87.5% of the action space when the obstacle is too far away to collide in the next time step. When the agent is at the border of the obstacle it reduces to about 40%. To put these numbers in perspective, the relevant action set coverage for the example state in Fig. 2 is 44%.

To make the above intuition clearer, we can add an extended version of Fig. 4 of the rebuttal PDF to the Appendix of our revised paper. The extension would be a subplot depicting the changes in the relevant action set coverage of the action space along the trajectories. Additionally, we can provide a plot that shows how the average relevant action set coverage changes over the training process in the Appendix of the revised paper.
So, is the relevant action set always provided to the agent, rather than the agent (e.g., the generator generating the latent action set) computing it on its own like the relevant action set being state dependent (A(s))? I wanted to observe how the relevant action set varies during the episode in relation to the agent's decision-making. When the agent performs poorly, this action relevance set fails to form a meaningful polygon in the entire action space.
Thank you sincerely for this comment as it identifies the core issue of the discussion: Whether to treat masking as part of the policy brings major benefit compared to treating it as part of the environment. Still we want to address some points in the following.

> ... which complicates the math (as noted by reviewer 8567) and limits to only differentiable masks.

We don't think that our approach limits to only differentiable masks. The ray mask is not always differentiable, since computing $g(a)$ for it can require solving optimization problems to compute the boundary points of certain set representations (e.g. zonotopes). We can still include it into our action masking framework, because we prove in Proposition 1 that the gradient remains unchanged compared to the original normal distribution.

> The author's arguments regarding interpretability, learning efficiency, and allowing more types of masks are valid (though I don't see theoretical or experimental evidence for the learning efficiency argument, ...

Your assessment for the lack of practical evidence is valid. Still, is the derivation of the actual gradients of the changed distributions in the paper, and the numerical example we provided in previous answers, not at least theoretical evidence for the learning efficiency argument?

> It seems to me that ultimately there's a tradeoff to be made, and it's not clear from existing evidence what the better tradeoff to make is. That's okay for this paper- this work focuses on one approach, and while that approach has advantages and disadvantages it's not necessary for one paper to compare to other, not previously described, approaches to validate the usefulness of the one the authors favor.

Thank you for this favourable comment!

> That said, I do think this concern is coming up because there's an obvious (and much simpler) alternative formulation which could be compared against. 

We agree that treating action masking as part of the environment would be simpler, but is it actually much simpler? Sure, the gradient derivation could be omitted, but the main effort lies in the implementation of the actual masking procedure, which remains in its entirety. Implementing the adapted gradient for e.g. the generator mask adds just a few additional lines of code with automatic differentiation engines like PyTorch (as we have detailed in the previous answer).

> I do think it would greatly strengthen the paper (and the contribution of this research direction as a whole) to experimentally compare policy-masking versus environment-masking to confirm that the supposed advantages of policy-masking actually occur compared to the environment-masking approach. At the least I'd encourage the authors to investigate such comparisons in future work regardless of the decision for this submission.

We fully agree that an experimental comparison would strengthen the paper. As we argued in previous answers, we do already provide this comparison with the ray mask. Since its gradient does not need to be adapted, our formulation is mathematically equivalent to viewing the ray mask as part of the environment. Also the only implementation change that would need to be made is moving the line `masked_action = ray_mask(action)` from the policy class to a few lines later in the `environment.step()` method, which does not change the inner workings of the code. As for the generator mask, we previously argued that it would be highly impractical to implement it into the environment as it leads to a change in the environment when changing the shape of the relevant action set (see answer "Physical implications of the gradient for the distributional mask (2/2)").
In summary, **we do already provide the requested comparison, and the results are not in favour of environment masking**. We just do not explicitly state it in the paper, which we can of course change. However, then we would have to introduce the concept of environment masking, which is not an established concept, and was not formulated prior to this discussion round.
> The exact gradient from Proposition 4 cannot be computed (as noted in lines 198-200), so this algorithm is fundamentally off-policy since it's impossible for it to be on-policy using the truncated distributions discussed in section 3.3.

You are correct by stating that if one cannot compute the gradient, the distributional mask becomes off-policy, since we update the parameters for a different distribution. However the fact that the integral is intractable does not imply that it cannot be approximated or derived in a different way (granted, our wording of "Since we cannot compute the gradient..." is suboptimal). However, this is certainly not trivial and likely warrants considerable future work. 

> Further, I disagree with the assumption that the difference between the truncated and original distributions will be trivial ... while others might be more permissive.

Saying that the difference is trivial in practice is certainly an understatement, but we believe that we did not make this statement anywhere. We agree that estimating the practical implications is challenging. Our suggestion that approximating the gradient might be valid in practice stems from the theoretical analysis presented in the "Physical implications of ... (2/2)" section of the general discussion. By approximating the gradient, we remove the option 2. for the gradient to increase the log-likelihood of the sample. Since this option can only favor actions close to the boundary of the relevant action set, it might not be applicable often in practice.

> In every PPO/A2C/etc implementation I've worked with the full distribution gets sampled. Can you clarify?

We acknowledge that our previous answer did not clearly state this point. In nearly all continuous reinforcement learning tasks, sampled actions are clipped to specific ranges (typically $[-1, 1]$) due to a fixed interval action space usually based on physical constraints. This clipping introduces a bias when using a Gaussian distribution as the policy, as its infinite support creates boundary effects. Chou et al. (2017) provide a thorough analysis of this issue in their paper "Improving Stochastic Policy ... " (Section 3.2). 
Essentially the clipping creates a form of truncated distribution with spikes at the truncation limits. One could argue that neglecting this effect inherently makes most stochastic policy gradient methods off-policy (note that we don't claim this is necessary the case!). PPO does not account for this effect, yet it demonstrates strong performance in practice. Based on this, we infer that our gradient approximation might also have little impact on performance.

> Given this, I'd strongly recommend removing the distributional mask method from the paper, or at least making it's off-policyness clear.

We believe it's important to include the distributional mask in the paper despite its current limitations. This method employs a novel technique for sampling from the policy distribution, which is rarely used in RL and thus represents a contribution. By publishing this work, even with its current need for gradient approximation, we open the door for other researchers to build upon and potentially solve the existing problems. We recognize that for this justification it is crucial to explicitly state the current limitations of the method, including its off-policy nature. 

> Regarding handling non-convex relevant action sets, I'm sorry but I find the assertion that a convex subset of the full non-convex relevant set is good enough to be unconvincing. 

We appreciate your concern about non-convex relevant action sets, particularly in robotics applications. We want to clarify that we didn't intend to suggest that convex sets are universally "good enough". Rather, they represent a significant improvement over the current practice of using only intervals as sets.
We fully acknowledge that for many applications, especially in robotics, restricting actions to convex sets can be an oversimplification, as Your example illustrates.
Our current work with convex sets is a step towards more flexible action spaces, but we agree that handling non-convex sets would be a valuable extension of our methods. However, this presents significant challenges that we believe warrant separate, focused research.

> It does seem worth noting as a limitation (sounds like this will be added) and is a mark in favor of environment-masks over policy-masking that it supports disjoint sets by default, however.

We don't really understand why you are certain that environment-masks support convex sets by default. None of our methods can be applied to non-convex sets as is, even if they are implemented in the environment. Can you please elaborate how this would be done?  We can imagine that action replacement is implemented with non-convex sets. Yet, this is conceptually different to masking and the newly added action replacement baseline performs worse than masking (see rebuttal for reviewer 8567, 3.).
Could you please reiterate what your remaining concerns are?
From our perspective, your detailed review helped us to alleviate the two major weaknesses:
> It is not clear how easy it is to recover the action masking criteria, especially under the more complex generator or distributional schemes, and it seems like this would be rare

In our new experiments, we now provide an example with a realistic constraint (maximal battery power) that we can very easily represent as a zonotope in the action space. We are convinced that such easy-to-implement constraints are not uncommon in continuous action spaces.

 > The experiments are not particularly convincing because they all follow similar control tasks, but it also seems like these are the only tasks for which the action mask could be easily defined.

We believe that the mujoco continuous control tasks Walker2D, Ant, etc. are the de-facto standard benchmark in continuous control. From our point of view these new experiments provide more than convincing arguments for the use of continuous action masking.   

We are happy to discuss these points and any other open concerns.
I appreciate the detailed response, especially the added experiments in walker. While the changes improve the paper, I think they do not significantly improve the score since the core questions remain.
Thank you for your reply. We are sorry that we misunderstood your review question for qualitative results. Still, we are not entirely sure how to interpret your comment. If we understand you correctly, you are interested in the impact of choosing different relevant action sets on the decision-making of the agent (i.e., the learning process)?

Let us briefly provide some intuition on this. While the relevant action sets $\mathcal{A}^r$ in our experiments are all collision-avoidance action sets, their average coverage of $\mathcal{A}$ varies between 25%, 28%, and 70% for the 3D Quadrotor, 2D Quadrotor, and Seeker, respectively. It seems that there is a sweet spot for the restrictiveness of the relevant action set due to two opposing mechanisms. Sample efficiency gains from masking increase, the smaller the relevant action set in relation to the global action set, whereas the exploration capability of the agent decreases. In other words, if the relevant action is almost the entire action space, the sample efficiency gains are likely small compared to not using action masking. If the relevant action set gets too small, the agent cannot learn much due to the very limited amount of possible actions. For example, in our experiments, the reward gain in the Seeker environment is much higher than in the Quadrotor environments, which underscores less constrained exploration. In the extreme case of $\{ a^r \} = \mathcal{A}^r$ always, there is no meaning in using reinforcement learning since the agent cannot explore. Future work will need to investigate these mechanisms in more detail for a multitude of tasks and relevant action sets to provide clearer and more nuanced insights. Yet, we are happy to extend our discussion in Sec. 4.3 (Limitations) in the last paragraph with this intuition.

If you are interested in a visual demonstration of the relevant action set during a training episode, we can refer you to Fig. 2 of our original paper. It shows the relevant action set at a certain state in the Seeker environment, which prohibits the agent from colliding with the red obstacle.
Thanks for your detailed comments.

-Regarding the off-policyness of the distributional mask, consider the case where the policy is a unit normal distribution and the mask is a truncated unit normal (which I believe is consistent with section 3.3). Actions sampled from this truncated distribution will have a different probability than they would under the original unit normal (e.g. because some actions are excluded, those that are included have a higher probability than otherwise). The exact gradient from Proposition 4 cannot be computed (as noted in lines 198-200), so this algorithm is fundamentally off-policy since it's impossible for it to be on-policy using the truncated distributions discussed in section 3.3 (maybe there's other distributions that would avoid this issue but those aren't discussed and seem like a non-trivial extension to define).

Further, I disagree with the assumption that the difference between the truncated and original distributions will be trivial in practice, both because it can be hard to estimate where policy gradients will push the distribution in complex tasks with limited samples per gradient step and because just how truncated the distribution is depends greatly on the task- some tasks could require very aggressive truncation to limit to only the relevant set, while others might be more permissive. The other masks seem fine to me, but this one seems flawed in a way that is hard to solve and will matter in practice as well as theory.

I'm not sure what "the clipping of samples from the normal distribution" refers to for PPO/etc. In every PPO/A2C/etc implementation I've worked with the full distribution gets sampled. Can you clarify?

Given this, I'd strongly recommend removing the distributional mask method from the paper, or at least making it's off-policyness clear. I'm certainly willing to accept experimental evidence that it works ok in practice, but the experiments presented here don't seem all that reassuring and don't motivate why a theoretically flawed algorithm might be superior in practice (it doesn't outperform the other masking algorithms in the best case and seems to underperform sometimes).


-Regarding handling non-convex relevant action sets, I'm sorry but I find the assertion that a convex subset of the full non-convex relevant set is good enough to be unconvincing. When looking at robotics (one of, if not the, biggest application domain for continuous-action RL) one finds non-convex (particularly disjoint- go left or right, but not straight ahead, etc) relevant sets all over the place. I'd need experimental evidence to believe this isn't an issue in practice, but I'm willing to accept this as a topic for future work (not every limitation of an algorithm needs to be solved in one pass). It does seem worth noting as a limitation (sounds like this will be added) and is a mark in favor of environment-masks over policy-masking that it supports disjoint sets by default, however.

Considering the above and the discussion/comments of other reviewers, I think I shall maintain my score. I do think this paper could be made stronger, but it's also reasonable to leave these questions to future work.
It's interesting to follow along with this detailed discussion of the effects of masking in the policy, I think the core issue at hand (at least to me) is whether there's major benefits to treating action masking as part of the policy, which complicates the math (as noted by reviewer 8567) and limits to only differentiable masks. The author's arguments regarding interpretability, learning efficiency, and allowing more types of masks are valid (though I don't see theoretical or experimental evidence for the learning efficiency argument, I buy that it should intuitively provide some -possibly small- benefit), but the degree to which those factors matter is unclear, and as such there's not a decisive answer to the question of policy-masking versus environment-masking. 

It seems to me that ultimately there's a tradeoff to be made, and it's not clear from existing evidence what the better tradeoff to make is. That's okay for this paper- this work focuses on one approach, and while that approach has advantages and disadvantages it's not necessary for one paper to compare to other, not previously described, approaches to validate the usefulness of the one the authors favor.

That said, I do think this concern is coming up because there's an obvious (and much simpler) alternative formulation which could be compared against. I do think it would greatly strengthen the paper (and the contribution of this research direction as a whole) to experimentally compare policy-masking versus environment-masking to confirm that the supposed advantages of policy-masking actually occur compared to the environment-masking approach. At the least I'd encourage the authors to investigate such comparisons in future work regardless of the decision for this submission.
I understand the setup of Fig. 4. However, as a qualitative result, it lacks relevance to action relevance (as you mentioned, this is merely an alternative to the quantitative results in the table). For instance, it’s not clear how this plot demonstrates the role of action relevance in decision-making.
> Your analogy about beta distribution is not correct. If your model was to output a normal distribution $\mu_\theta$ and $\sigma_\theta$ --> then you sample from it --> then you use that action to somehow create your beta distribution; then does the gradient flowing into mu, sigma differ whether or not you consider the presence of beta distribution in the latter part?

We appreciate your concern about the beta distribution analogy, and we acknowledge that it is somewhat abstract. In the following, we want to show why we still think it is correct and also applies to our formulation. 

Let us start by a short introduction on implementations of univariate distributions. Here, the inverse transform sampling is used to sample from any univariate distribution with probability density function $\mathbb{P}(a; \theta)$ that has an invertible cumulative density function (CDF) $F(a; \theta)$. The method takes a uniform sample $a \sim U(0, 1)$, and applies the inverse cumulative density function (ICDF) $F^{-1}(a; \theta)$ to create samples from the distribution $\mathbb{P}(a; \theta)$. 

For our example, let $x,y \in \mathbb{R}^N$ be two outputs of a neural network. We treat them as the two parameters of a normal distribution by sampling $a \sim \mathbb{P}(a; x, y) = \mathcal{N}(x,y)$, i.e. $\mu = x$, $\sigma = y$. If we use the gradients $\nabla_x$ and $\nabla_y$ of the log-likelihood of the normal distribution in the backward pass, $x$ and $y$ will be shaped to represent the mean of variance of the normal distribution.
Now, instead, we apply the function $g(a) = F^{-1}\_{Beta} ( F\_{Normal} (a; x, y); x, y )$, which first maps the samples of the normal distribution to the uniform distribution $U(0, 1)$, and then transforms them to samples of the beta distribution through the ICDF. Thereby we **exactly** create samples from the beta distribution parameterized with $\alpha=x$ and $\beta=y$, i.e. $a \sim \mathbb{P}(a; x, y) = \frac{a^{x-1} (1 - a)^{y-1})}{B(x, y)}$, where $B(x,y)$ is the beta function. 
If we now still use the gradients of the log-likelihood of the normal distribution to update $x$ and $y$, we shape them to represent $\mu$ and $\sigma$ instead of $\alpha$ and $\beta$, which would result in $\mathbb{P}(a; x, y)$ representing beta distributions with likely not meaningful shapes. Therefore, we should use the gradients $\nabla_x$ and $\nabla_y$ of the log-likelihood of the beta distribution to shape $x$ and $y$ to meaningful values of $\alpha$ and $\beta$.

We want to reiterate that this example is very abstract, since there is a much more straightforward way to sample from the beta distribution. However, it indicates that it is important to use the gradients of the distribution that actually produce the samples in order to learn meaningful values for the parameters of the distribution.
Thank you for your follow-up and in-depth engagement with our work, which we really appreciate.

> Your paper formulates g as a function that applies over the sampled action, right? $a^r = g(a) = c + G a$

For the generator mask, yes, resulting in samples from $\pi_\theta^r$. 

> The question is that, in the backward pass of your update step, would the gradient flowing into $\mu_\theta$ and $\Sigma_\theta$ change at all due to the presence of G? Wouldn't the backpropagation remove the effect of G because of chain rule?

Thanks for your question. Since the loss function does not use the sampled actions directly, the effect of action masking is not simply the derivative of $G a + c$, i.e., modelled with the chain rule. Instead, the sampled actions are integrated into the objective of PPO through the policy probability distribution. 
We change $\log \pi(a|s)$ to $\log \pi^r(a^r|s)$ where $\pi^r(a^r|s)$ is defined as
$$
\pi_\theta^r (a^r|s) = \mathcal{N}(a^r; G \mu_\theta + c, G \Sigma_\theta G^T).
$$
Note that this is Eq. 11 with corrected notation, where we can imagine confusion could originate.
In Proposition 3, we derive the changes of the gradient due to the change of policy distribution to $\pi_\theta^r (a^r|s)$. Let us have a closer look at Eq. 13 when replacing $a^r$ with $Ga +c$:
$$
\nabla_{\mu_\theta} \log \pi_\theta^r(a^r | s) 
    =  G^T (G \Sigma_\theta G^T)^{-1} (a^r - c - G \mu_\theta) = G^T (G \Sigma_\theta G^T)^{-1} (c + G a - c - G \mu_\theta) =
    G^T (G \Sigma_\theta G^T)^{-1} (G a - G \mu_\theta).
$$
From the transformed Eq. 13, we can observe that the effect of $G$ does not cancel out. We would like to refer you to appendix A.2 for the detailed proof.

Let us also have another look at the numerical example to clarify it further. In this example, we calculate the first column with the loss $\log \pi(a_i | s)$ and the second with the loss $\log \pi^r(a^r | s)$. These loss functions reflect the policy gradient loss with assumed advantage $A^\pi=1$ for a single action sample, i.e., $a_1$ or $a_2$. While this is not the entire PPO loss, it reflects the only adaption we make: changing the policy distribution to $\pi^r$. Thus, the first two columns provide a simplified numerical example for the backward pass, i.e., gradient flow. For further clarity, assume that $\mu_\theta = \theta$, then for $a_1$ the parameters $\theta$ would be updated with $\lambda [0.1, 0.1, -0.4]$ and $\lambda [-0.066, -0.066, -0.066]$ when calculated without and with the correction of the policy distribution, respectively. Note that $\lambda$ is the learning rate.


> Are you claiming that if I train two policies: ... there would actually be a difference in what update happens to $\theta$?

Yes. For a numerical example, see the above answer. 


> I looked into your code and I don't see any evidence of modification of the loss function. All I see are different ways of applying the mask with your policy network. Then, how why do you expect the policy gradient flowing into $\theta$ change at all because of a presence of a post-processing function g?

Thank you for looking into our code. You are correct in observing that we do not directly modify the loss function. However, our action masking framework also does not require such a modification. The change in gradient comes from how the log-likelihood of the sampled action is computed. Standard PPO uses the normal distribution for $\pi$, whose mean is defined by the output of the neural network policy, to compute the log-likelihood. We define the modified normal distribution $\pi^r$ as in Eq. 11 (with corrected notation), which is implemented in lines 41-43 in `action_masking/rlsampling/sb3_distributions/generator_dist.py`. This modified distribution is used to compute the `log_prob` in line 535 of `action_masking/sb3_contrib/ppo/ppo.py`, which in terms is used to compute the loss function. The automatic differentiation engine of PyTorch now computes the accurate gradients $\nabla_\theta \log \pi_\theta^r(a | s)$, as derived in Proposition 3 (Eq. 13 and 14). We have implemented Eq. 13 and 14 and compared them to the PyTorch's gradients to validate their correctness.

In essence, the use of our modified distribution in the log-likelihood calculation leads to the desired change in gradients, as defined by our theoretical analysis. We see that our code is not easy to navigate with respect to this implementation change and will clarify this in the README.md.
1. Your analogy about beta distribution is not correct. If your model was to output a normal distribution $\mu_\theta$ and $\sigma_\theta$ --> then you sample from it --> then you use that action to somehow create your beta distribution; then does the gradient flowing into mu, sigma differ whether or not you consider the presence of beta distribution in the latter part?

Your paper formulates g as a function that applies over the sampled action, right? $a^r = g(a) = c + G a$

2. The real question is not whether $\nabla_{\mu_\theta} \log \pi_\theta(a_i \mid s)$ and $\nabla_{\mu_\theta} \log \pi_\theta^r(a^r \mid s)$ are different. Obviously, they will be different. The question is that, in the backward pass of your update step, would the gradient flowing into $\mu_\theta$ and $\Sigma_\theta$ change at all due to the presence of G? Wouldn't the backpropagation remove the effect of G because of chain rule? In that case, the presence of the transformation g, does not add anything to the policy gradient flowing into $\theta$, which are the parameters of the neural network.

Are you claiming that if I train two policies:
1. With $\pi_\theta$, do a transformation $a^r = g(a) = c + G a$ in the environment, get rewards and update the $\theta$,
2. With $\pi^r_\theta$ as defined in Equation (11), get rewards and update the $\theta$,
then there would actually be a difference in what update happens to $\theta$ ?


I looked into your code and I don't see any evidence of modification of the loss function. All I see are different ways of applying the mask with your policy network. Then, how why do you expect the policy gradient flowing into $\theta$ change at all because of a presence of a post-processing function g?
**3.**

Thank you for this question; it is indeed a counterintuitive observation. You are correct that the policy gradient aims to increase the log-likelihood of high-return samples by adapting the parameters of the Gaussian policy distribution. However, the situation becomes more complex with the distributional mask, as we explain in the following. The relevant policy distribution $\pi_\theta^r(a^r| s)$ essentially models a truncated normal distribution in one dimension [Burkardt, 2023, The Truncated Normal Distribution]. Due to the normalization of the distribution with the integral over the relevant action zonotope (Eq. 15), there are now two ways to increase the log-likelihood of a sample.
1. Moving the mean closer to the sample (as reflected in the first term of Eq. 17).
2. Moving the mean farther away decreases the integral in Eq. 15, thereby increasing the log-likelihood (as reflected in the second term of Eq. 17).

For the second option, while the probability density of $\pi_\theta(\cdot| s)$ decreases over $\mathcal{A}^r$, it remains the same for $\pi_\theta^r(\cdot| s)$ due to the normalization. These effects would be easier to grasp with a visualization of the distributions, which we cannot provide due to the format restrictions of comments. So please let us know if we should elaborate more.


Finally, we are curious to understand better why you consider action masking for convex sets in the environment as trivial? While it may simplify some aspects by excluding the gradient adaption, the core effort still lies in the formulation of the three masking approaches, which we believe is a key contribution of our work.

We hope our answers resolve your skepticism regarding the relevance of adding the mapping function to the policy. Please let us know if you have additional questions. Lastly, let us take a step back from the theoretical analysis and look at the discussion from a practical point of view. As we stated in the summary rebuttal, defining action masking in the environment completely excludes the distributional mask and practically excludes the generator mask. This is because we would need to modify the environment when changing the relevant action zonotope's shape, leading to impracticalities like multiple versions of the Walker2D environment with varying numbers of generators.
In summary, formulating our masking approaches as part of the policy reflects the actual policy distribution, is more intuitive and practical, and allows for more formulations of $\mathcal{F}$.

(2/2)
Thank you for your detailed response and thoughts on the intuitive and empirical effect of the changed gradient. 

**1.**

Thank you for your illustrative example. The mapping function $\mathcal{F}$ describes the relation between the action space and the relevant action space. Thus, using  $\nabla \log \pi^r_\theta(a^r \| s)$ instead of $\nabla \log \pi_\theta(a \| s)$ in the objective of PPO encodes this mapping, i.e., is the extra information (see 2. second paragraph for a numerical example).

Additionally, we think that the example, doesn't directly apply to our evaluation of stochastic policy gradients. In this context, the neural network's output is not the action itself but rather the parameters (such as the mean) of a probability distribution from which actions are sampled. The function in question is applied to these samples, effectively modifying the probability distribution.
Since the policy gradient assumes the log-likelihood to come from the distribution the action is sampled from, we argue that it is necessary to derive this gradient $\nabla_\theta \log \pi^r_\theta (a | s)$ for the modified distribution.
To illustrate this point, consider the scenario where the policy neural network models a beta distribution instead of a normal distribution (as explored in previous work [Chou et al., 2017, Improving stochastic ...]). In this case, would you agree that it is necessary to use the gradient of the beta distribution?

We are not certain that we fully covered all your questions in 1. Could you please elaborate if that is not the case.


**2.**

Thank you for your question. Intuitively, a non-bijective mapping function in the environment means that two different actions chosen by the agent can lead to exactly the same next state. Consider this abstract example. The function $g(a) = 0$ maps the samples of the policy distribution to the zero vector, making it a non-bijective function. Essentially, no matter the parameters of the neural network, the action would always be $0$, thus making any gradient updated irrelevant. If we incorporate $g(a)$ in the environment, a gradient update would still occur since the agent receives a return and $\log \pi_\theta (a | s)$ represents the log-likelihood of the unmapped sample $a$. However, if we treat $g(a)$ as part of the policy distribution, we essentially transform $\pi_\theta(a| s)$ into a Dirac distribution at $0$. Since the Dirac distribution is independent of the parameters $\theta$, the gradient $\nabla_\theta \log \pi_\theta^r(a| s)$ would be zero, which more accurately reflects the situation, as no meaningful updates to the policy can be made.



Additionally, let us illustrate the gradient flow with a numerical example for the generator mask. Consider a relevant action zonotope with center $c = [0, 0]$ and generator matrix $G$:
$$
    G = \begin{bmatrix}
1 & 0 & 0.5\\\\
0 & 1 & 0.5
\end{bmatrix}.
$$

The policy $\pi(a|s)$ is defined by a mean vector for the multi-variate Normal distribution $\mu_\theta = [0.4, 0.4, 0.4]$ and an identity covariance matrix. The relevant policy is defined as in Eq. 11. Now consider two samples: $a_1 = [0.5, 0.5, 0]$ and $a_2 = [0, 0, 1]$. Both actions result in the same relevant action $a^r = [0.5,0.5]$. The table below provides the numerical evaluation for the gradient at $\mu_\theta$ and the log probabilities for the two actions.

|         | $\nabla_{\mu_\theta} \log \pi_\theta(a_i \| s)$ | $\nabla_{\mu_\theta} \log \pi_\theta^r(a^r\|s)$ | $\log\pi_\theta(a_i\|s)$ | $\log\pi_\theta^r(a^r\|s)$ |
|---------|------------------------------------|--------------------------------------|--------------------|--------------------|
| $a_1$ | [0.1, 0.1, -0.4]                   | [-0.066, -0.066, -0.066]             | -2.84              | -2.04              |
| $a_2$ | [-0.4, -0.4, 0.6]                  | [-0.066, -0.066, -0.066]             | -3.09              | -2.04              |

This numerical example shows that $\pi^r$ leads to the same log-likelihood, reflecting that the executed relevant actions are identical. If we look at the gradient flow at $\mu_\theta$, we see that for $\pi^r$, the gradient is the same for both $a_1$ and $a_2$, thus providing consistent feedback. In contrast, for the gradient of $\pi$, we observe in this example two contradictory gradients, although the executed relevant action is identical. Thus, potentially introducing an oscillation effect for the parameter update. 

Let us provide a physical interpretation of this numerical example as well. The relevant action set described with $G$ and $c$ has two generators that are axis aligned in $\mathcal{A}$ and one that influences both dimensions of $\mathcal{A}$. A gradient step with respect to \pi^r$ includes this state-dependent but fixed dependency between the actions. If the generator mask mapping function is part of the environment and, thus, a black box for the agent, it has to learn this dependency with potentially many training samples.

(1/2)
Thanks for your question. Fig. 4 shows qualitative deployment behavior for the different RL agents in the Seeker environment. We use the same goal state (gray circle), same obstacle (red circle), and same ten initial states (colored crosses) for the agents for better comparability. The ten trajectories have the same color as their initial state. Note that for the masking approaches the trajectories are almost identical after passing the obstacle and therefore plotted on top of each other.

The masking agents (ray, generator and distributional) reach the goal state for all ten initial states, with the generator being the quickest. In contrast, the PPO baseline agent reaches the goal with four out of the ten initial states and otherwise collides with the obstacle. The qualitative behavior is aligned with the quantitative results in Table 1 of the submitted paper meaning that the generator mask agent performs best and the PPO baseline performs worst with respect to reaching the goal efficiently. 

The goal of the visualization in Fig. 4 is to provide an intuition on the trained agents' behavior. We choose the Seeker environment since we believe it is the most intutive to understand and thus also to interpret the qualitative behavior. 

Note that there is unfortunately a typo in the subplot titles: distribution should be distributional.
Thank you for your response and for proving the bijectivity as well as running a new experiment.
Unfortunately, I can't seem to find the information about Fig. 4 so could you elaborate on Fig. 4?
Thank you for your rebuttal response. I am posting here for a potentially joint discussion with other reviewers (like PubC) who shared my question about `treating action masking as part of the environment`.

I want to understand this part better. I still believe the main benefit of formulating action masking as part of the policy comes from how it can help the policy parameters learn about what the action mask is, and thus, help it learn a better policy. For this to happen, there has to be a change in the gradient objective. I have been thinking about this in depth, and have some thoughts / questions.

1. Let's think intuitively at a high level. Is it even physically possible that gradient information can be passed into the policy? Say, I am a neural network (policy) and my outputs are transformed in a particular way, via $\mathcal{F}$ from $\pi(a|s) \to \pi^r(a^r | s)$, and then it goes through a non-differentiable operation (environment, E) to give a return R. Now, the gradients for me can still only flow through my outputs, i.e., $\pi(a|s)$. Why would separating what happens after my output be any helpful to me, i.e., how can I extract some information by separately treating $\mathcal{F}$ and E? Even if I knew what $\mathcal{F}$ is, I don't know what E is, and I can only see the final effect in the form of the return R. Is the knowledge of $\mathcal{F}$ adding any extra information above the knowledge of return R, or is my gradient conditionally independent of $\mathcal{F}$ given R?

2. For the case when $a^r = g(a)$ and g is bijective, we know that the gradient of $\pi^r$ = gradient of $\pi$. Then, why do we believe that when g is not bijective there will be any extra information added into my gradient?

It would be ideal if we could see some actual effect of gradient flow into the policy due to the knowledge of the action mask in any one of the methods and experimental results. That would resolve all my concerns. But even if that's not tractable right now, I would like to see some evidence or a thought experiment to prove that it is indeed possible!

3. I also tried to understand what the distributional mask's gradient in Eq. 17 physically implies. If we look at the terms in the gradient, the first term is the standard REINFORCE term, which has the implication that if the associated return R is high with this term, then the policy gradient increases $\log \pi_\theta(a | s)$ for that state. Similarly, it would reduce the second term here because of the negative sign, i.e., $\log \int_{A^r} \pi_\theta(a|s)da $ should go down if the return is high. But, this means that the probability density of $\pi_\theta(\cdot|s)$ should go down over $A^r$. Basically, the area under curve for the segment of $A^r$ has a lower probability density than the other areas in the action space. But, I don't see any reason why this should happen? This would increase the probability of "invalid actions" (along with the action $a$ from the first RL term). I see no reason why that would be helpful at all or why that makes sense?

To sum up, I really want to believe that there is some performance benefit to considering action masking as part of policy and not environment. But, I don't have any evidence to support that yet. I hope the authors can help me understand this.

If it turns out that there is no gradient flow, then I don't see any merit in the entire mathematical formulation of this paper, because all it says is that we should just consider action masking as part of environment, which is the trivial thing to do. I sincerely hope that is not the case, and by answering 1,2,3, we can get to that conclusion.