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.
Paper
Similar papers
Peer review
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.
Strengths
The proposed action masking covers a wide range of possible action mask definitions. The derived policy gradients are relatively straightforward given the definitions of the action boundaries. The derivations appear to be sound when applied empirically.
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 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
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.
Rating
6
Confidence
3
Soundness
3
Presentation
3
Contribution
2
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.
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.
Strengths
**Originality** - The paper presents a unique perspective on action spaces by utilizing the relevance of action utility in tasks to improve performance. Conventional methods are limited to discrete domains (tasks) so applying their methods to continuous environments was interesting to see. **Significance** - The proposed approach has practical implications, especially in complex environments where distinguishing between relevant and irrelevant actions is crucial. Regardless of the coverage of baseline, their methods significantly outperform it establishing the state of the art performance.
Weaknesses
**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
See weaknesses.
Rating
5
Confidence
2
Soundness
3
Presentation
3
Contribution
3
Limitations
n/a
Supplementary figures
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.
Response
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.
Answer to Response
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.
Response
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.
Anser to Response
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.
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.
Strengths
This paper is excellently written, defines a clear and well-motivated goal, and describes three intuitive and theoretically grounded methods to achieve that goal within well-defined and clearly stated limitations. To my knowledge these approaches are novel (though the distributional mask I suspect has been used as a one-off solution in prior work as it is conceptually very simple), and their definition and analysis are nontrivial.
Weaknesses
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
-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.
Rating
6
Confidence
4
Soundness
3
Presentation
4
Contribution
4
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.
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.
Strengths
- The problem of action masking in continuous action space is an underexplored one, but could have major impact in efficacy of agents and incorporating domain-specific knowledge. - The proposed continuous action masking could potentially be useful for safety demarcations. - The paper provides mathematical frameworks to formulate continuous action masking and also derive their (minimal) effect on the policy gradient. - The paper is mostly well-written and explains the mathematical derivations quite well. Quick note that Section 2.1 and 2.2 could be made more integrated, currently they seem completely disconnected.
Weaknesses
## 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
Listed in weaknesses.
Rating
6
Confidence
4
Soundness
3
Presentation
3
Contribution
2
Limitations
The authors discuss several limitations of the work, but the important ones of applicability and baselines need to be addressed.
Trying to understand why action masking should be considered part of policy and not environment
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.
Thought experiments for gradient flow and non-bijectivity (1/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)
Physical implications of the gradient for the distributional mask (2/2)
**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)
Followup
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?
Answer to followup (1/2)
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.
Answer to followup (2/2)
> 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.
Thoughts regarding action masking as part of the policy versus the environment
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.
Answer to "Thoughts regarding action masking as part of the policy versus the environment"
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.
Response
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?
Response to Rebuttal
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.
Answer to Response
> 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.).
Response to Response
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.
Answer to Response
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).
Response to Authors
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.
Response to Reviewer
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.
Decision
Accept (poster)