Paper: Decoupling regularization from the action space (ICLR-2024)

Reviewer comment (Experiments):
I think the key hyperparameter $\alpha$ needs more discussion. It will be helpful to analyze the (empirical) effects of different choices of the value of $\alpha$ and recommend the values or the strategy of value selection.

Author response:
We will include an ablation for $\alpha$. However, it might not be ready before the end of the discussion period. 

We added a small note in the paper that we recommend our proposed value of $\alpha$ and that it should be tuned the same way one would tune the temperature as the interpretation of $\alpha$ and $\tau$ is fairly similar, the higher $\alpha$, the higher entropy the policy will be. As for a recommended default, we chose $\alpha=0.77$ (obtained mathematically, not via hyperparameter tuning) to match SAC’s rule at the default scale (it was noted in the last paragraph of 8.2). This is why the performance of 4.a is almost identical to SAC’s.

Reviewer comment (Reproducibility):
The content in the experiment part lacks of sufficient details, which is also missing in the appendix. I recommend the authors to add the key details of experiments.

Author response:
We have included all details of the experiments, including the experiment files (see `rexp.py` in the rl codebase) in the appendix. We do not understand which details are missing.

Reviewer comment (Evaluation):
Is it possible to have some case-analysis for the results of DMC (e.g., in Figure 3,4)? For example, the proposed method achieves significantly better performance in tasks like BallInCupCatch, FingerTurnHard, HopperStand while the baseline fails totally.

Author response:
We will include an ablation for the temperature in the static temperature. However, it might not be ready before the end of the discussion period.

As for Figure 4.b, the scale was chosen such that the heuristic of SAC becomes infeasible, as noted in the last paragraph of 8.2. 

Do you have something particular in mind?

Reviewer comment (Reproducibility):
Which is the type of temperature setting adopted for the experiments in Section 8.3, i.e., fixed temperature or dynamic temperature?

Author response:
We added a remark that the temperature is static. We assumed that the fact that we called it soft Q-learning (SQL) and not SAC would convey this information as the dynamic temperature was introduced with SAC, not SQL.

Reviewer comment (Theory):
It seems that the coefficients $\gamma$ and $\tau$ are missing for the entropy term in Equation 3a.

Author response:
We added the missing $\tau$; thank you for noting this omission.

Reviewer comment (Experiments):
One possible extension to increase the contribution of the paper would be to include regularization to a base policy, as is often done in RLHF for LLMs via a KL divergence regularization term (e.g., see Bai et al., "Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback"). There have been various papers attempting to find a systematic way of setting the regularization coefficient (e.g., the above paper and Gao et al., "Scaling Laws for Reward Model Overparameterization"), but none have considered possibly setting the coefficient differently at different states. This could be an interesting direction to apply the ideas presented in this paper.

Author response:
> One possible extension to this work that could increase the contribution of the paper is to also include regularization to a base policy, as is often done in RLHF for LLMs via a KL divergence regularization term (e.g., see Bai et al., "Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback"). There have been various papers attempting to find a systematic way of setting the regularization coefficient (e.g., the above paper and Gao et al., "Scaling Laws for Reward Model Overparameterization"), but none have considered possible setting the coefficient differently at different states, so this could be an interesting direction to apply the ideas presented in this paper.


We are grateful for the suggestion, but these suggestions would make the paper extremely long and are out of the scope of this work as they have a focus somewhat orthogonal to the focus of the paper.

Reviewer comment (Novelty):
A relevant related work is "LESS is More: Rethinking Probabilistic Models of Human Behavior" by Bobu et al., which approaches the problem from the motivation of modeling human behavior. Under the "Boltzmann rational" model, this is equivalent to solving a regularized MDP. They have a different approach to solving the issue of how the number of actions or number of trajectories that can achieve a particular reward affects the optimal regularized solution, and it would be good to compare the approach in this paper with theirs.

Author response:
> Also, a relevant paper is "LESS is More: Rethinking Probabilistic Models of Human Behavior" by Bobu et al. They approach the problem from the motivation of modeling human behavior, which under the "Boltzmann rational" model is equivalent to solving a regularized MDP. They have a different approach to solving the issue of how the number of actions or number of trajectories that can achieve a particular reward affects the optimal regularized solution. It would be good to compare the approach in this paper with theirs.
> 

Thank you for the suggestion. However, we do not think they are directly comparable. First and foremost, LESS is for the inverse problem, not the forward one. Indeed, in the inverse problem (learning the reward function given observations, not a simulator), papers like [8] already estimate the optimal temperature per state. 

The second issue is that we do not assume any particular structure, while LESS assumes a reward that is linear in features. This is not learnable in the forward setting, as the linear feature may end up being the reward itself, defeating the whole purpose of LESS.

A third issue is that it is unclear how LESS would be extensible to stochastic environments.

While using trajectory similarity is interesting, LESS does not include a Bellman equation, and writing it seems nontrivial.

Reviewer comment (Writing):
In Section 2, it is not specified that the MDPs considered have discrete actions, although it later appears this is an unspoken assumption. Then later, the paper switches to continuous MDPs without any explicit distinction.

Author response:
> In Section 2, it is not specified that the MDPs considered have discrete actions, although it later appears this is an unspoken assumption. Then later, the paper switches to continuous MDPs without any explicit distinction.
> 

Our theory applies to both cases. In Section 4, we mention that “We look at differential entropy and continuous actions in Section 7.” We moved the phrase up so it’s more visible.

Reviewer comment (Theory):
In Section 2, the first time regularization is introduced it uses $\tau \Omega(\pi)$ subtracted from the value function; then immediately below it is instead added to the value function.

Author response:
> In Section 2, the first time regularization is introduced it's with subtracted from the value function; then immediately below it's instead added to the value function.
> 

We fixed this typo; thank you for finding it.

Reviewer comment (Theory):
Proof of Lemma 1 appears insufficient; it is unclear why exactly convexity implies that any other point must be smaller.

Author response:
> Proof of Lemma 1: this proof seems insufficient; why exactly does convexity imply any other point must be smaller?
> 

We added a clarification that we are looking at points between 0 and 1 as we are working with discrete distributions, which follows the notice in Section 4. Hence, this follows directly from strict convexity: $f(p)<pf(1) + (1 - p) f(0)$ for $0\leq p \leq 1$(see Section 4 of [3] in particular Theorem 4.1). Since we assumed that $f(1)=f(0)=0$ (see Assumption 1), $f(p)<p0+(1-p)0=0$.

Reviewer comment (Theory):
In SAC, the entropy is actually calculated before applying $\tanh$ to squash the actions into the action space. This means that the entropy is unbounded, and thus the idea in this paper does not apply directly.

Author response:
> In SAC, the entropy is actually calculated before applying $tanh$ to squash the actions into the action space. This means that the entropy is unbounded, and thus the idea in this paper does not apply directly.

Please see our answer to the first question.

Reviewer comment (Theory):
Why should the ideas in this paper apply to SAC when the entropy as measured in SAC is actually unbounded?

Author response:
> **Questions:**
> 
> - Why should the ideas in this paper apply to SAC when the entropy as measured in SAC is actually unbounded?

Below, we provide ample evidence that not only the log prob is calculated after the tanh which makes the entropy bounded, but the entropy before the tanh is almost always bounded from above and below as the standard deviation is often clipped. If the entropy is measured before the tanh, then the action of selecting the highest movement can be represented by a very large average and high standard deviation which is odd. Lastly, the clipping of the temperature is the only component allowing SAC with the $-n$ rule to work in the experiment with scaled actions to [-0.1,0.1] as it would otherwise drift to infinity.



[7, equation 20 in the supplementary] defines the policy as a Gaussian distribution with squashing.

[The original SAC repository](https://github.com/haarnoja/sac) does the `tanh` correction for the log prob ([see here](https://github.com/haarnoja/sac/blob/8258e33633c7e37833cc39315891e77adfbe14b2/sac/policies/gaussian_policy.py#L74)) and also clamps the standard deviation ([see here](https://github.com/haarnoja/sac/blob/8258e33633c7e37833cc39315891e77adfbe14b2/sac/distributions/normal.py#L55)). Choosing a maximum standard deviation means that before the `tanh`, the maximum entropy is known by (see [“Specified mean and variance: the normal distribution”](https://en.wikipedia.org/wiki/Maximum_entropy_probability_distribution))

[The updated version from RAIL](https://github.com/rail-berkeley/softlearning) calculates the entropy after the `tanh` see `action_distribution` [here](https://github.com/rail-berkeley/softlearning/blob/13cf187cc93d90f7c217ea2845067491c3c65464/softlearning/policies/gaussian_policy.py#L74) which is the bijective transformation defined [here](https://github.com/rail-berkeley/softlearning/blob/13cf187cc93d90f7c217ea2845067491c3c65464/softlearning/policies/base_policy.py#L267).

[Spinning up](https://spinningup.openai.com/en/latest/), the RL tutorial from OpenAI uses the `tanh` log probs (see [here](https://github.com/openai/spinningup/blob/038665d62d569055401d91856abb287263096178/spinup/algos/pytorch/sac/core.py#L60) and the comment above it for an explanation) and clamps the standard deviation. (see [here](https://github.com/openai/spinningup/blob/038665d62d569055401d91856abb287263096178/spinup/algos/pytorch/sac/core.py#L42))

[Cleanrl](https://github.com/vwxyzjn/cleanrl) [4] also includes both the [tanh log prob](https://github.com/vwxyzjn/cleanrl/blob/2d660b6d3053ea9037c746b4c9f3a6faa1f20c44/cleanrl/sac_continuous_action.py#L140) and the [clamped standard deviation](https://www.notion.so/Sobhan-ICLR-2024-Rebuttal-be69138c020f4243bdf65a7615dd7d0e?pvs=21).

[Tianshou](https://github.com/thu-ml/tianshou) [5] also applies the l[og prob to the tanh](https://github.com/thu-ml/tianshou/blob/962c6d1e11d033b6faaa2da37a89794563baebf6/tianshou/policy/modelfree/sac.py#L168). 

Lastly, [6, equation 25] defines the policy with the action bound.

Reviewer comment (Novelty):
The method proposed in the paper is very simple, and involves normalizing the standard temperature by the range the regularization objective can take (e.g. minimum possible entropy). While I do not think that the simplicity of the approach should detract from the novelty of the paper, I am not convinced that the better experimental results are actually due to the proposed range-normalization, and not simply that the regularization can now be state-dependent. To my knowledge, state-dependent temperatures is not itself a novel concept, as other works have tried to dynamically set temperature based on state [1, 2]. The results would be more convincing if the authors showed that prior efforts in state-dependent regularization fail whereas theirs succeeds.

Author response:
We want first to note that we are not learning a per-state temperature; in all experiments (including DMC), the model's temperature is a single scalar, and the per-state temperature is that scalar divided by the range of the regularizer at the state. This means that in the DMC benchmark, the temperature is the same for all states for each environment. While our approach is more flexible as it effectively uses a different temperature for states with different action spaces, it differs from efforts in learning state-dependent gamma, lambda, and tau (e.g., [4]). While using a parametrized network that chooses the temperature is an interesting research venue, it is out of the scope of this work and not what we propose here.

We want to separate the approaches of [1] and [2]. 

[2] anneal the temperature as the state is visited more and more often. They effectively use temperature to both promote exploration and embed uncertainty. Indeed, the added temperature increases the likelihood of visiting that state as the value function increases, and the higher temperature increases the probability of taking nonoptimal actions. 

In many normal RL tasks, low temperature (almost zero) is desirable as we do not measure performance in a way where the added entropy is helpful. Indeed, the added entropy can be seen as a form of robustness in an adversarial environment (e.g., [3]). Yet, the test environment is often the same as the training and static, so the added entropy can hurt performance, making very small temperatures desirable. However, this insight does not transfer to the drug design MDP since our performance measure is not the best terminal state (drug) found. Instead, we measure the number of terminal states that are not similar to previously found terminal states; [6] refer to this as the number of modes. Indeed, setting the temperature too low hurts the performance. Of course, the approach of [2] is amendable as it can converge to a higher than zero temperature with minor modifications; it could even converge to our proposed temperature. In any case, [2] proposes something that does not change the final model but helps the agent learn it, making it orthogonal to our work. 

We are aware of [1] and cite it. We originally included a comparison [1] (referred to as mellowmax or MM) in an early draft in a similar but different setup. We decided not to include it in our comparisons as it worked too poorly; sampling with [1] requires solving a 1D optimization problem with bisection in the discrete case, and we are unaware of how [1]'s method can be extended to continuous control tasks. We also note that [1] only uses the dynamic temperature for the execution, not the value function evaluation. We believe this also sets apart our proposed solution with the work of [1]. Indeed, mellowmax could be used in parallel with what we propose here, and our method improves the performance of [1]. We are working on a section in the appendix comparing the performance of [1] and our method, but it could potentially not be ready before the end of the rebuttal period.

Reviewer comment (Experiments):
The algorithm does not actually circumvent the extensive hyperparameter tuning required by existing RL algorithms. It is unclear how sensitive the proposed approach is to various settings of temperature and \alpha.

Author response:
The sensitivity of the approach proposed here is not different from other methods save for one major difference: our proposed heuristic never becomes infeasible. In the state-dependent action setting, the maximum entropy a policy can have depends on the states it visits, not just the final policy, so while a target could be reasonable for the final policy, it could be infeasible for the initial policy. In the state-independent action setting, the target has to be carefully chosen to be below the maximum entropy, which is one of the things we propose. If the SAC temperature program becomes infeasible, the temperature will diverge to infinity, leading to the agent doing a random walk, which can be problematic.

Our approach shines where one temperature has to be chosen for a wide range of tasks, such as benchmarks or in the state-dependent action setting. In both cases, the number of temperatures that need choosing, if choosing one temperature per action space or environment, can be too high to be searchable or not finite. 

Lastly, we’d like to refer to the relaxed sufficient conditions we propose for undiscounted entropy regularized MDPs. This cannot be achieved with hyperparameter tuning unless one temperature is chosen per action space.

Reviewer comment (Experiments):
The empirical results swap between dynamic and static learning of the temperature parameter. Is there a reason why one of the method works better for a particular domain. Specifically, will dynamic temperature learning at least reproduce the results of using a static temperature in the drug discovery experiment?

Author response:
We want to start by noting that we used the standard setup, and the jump you mentioned is pure coincidence. We used dynamic temperature for DMC because this is what we found to be the baseline. We use the static temperature for DMC before the dynamic part because it is less common. Lastly, we use static temperature when comparing the GFNs because we want to keep the implementation as similar as possible to GFNs.

The approaches are compatible in that for temperature $\tau$ and a policy $\pi$ being the solution of the soft Bellman equation with temperature $\tau/H(U_s)$ at state $s$, there exists heuristic parameter $\alpha$ such that the dynamic temperature program of SAC does not change the temperature. 

In short, we look at the gradient of
$$
\mathbb{E}_s\left[\frac{\tau}{H(U_s)}\left(H(\pi(\cdot|a)) - \alpha H(U_s)\right)\right]
$$
with respect to $\tau$ and show that there exists a value of $\alpha$ such that the gradient is zero. The aforementioned equation is the objective we get for the temperature program of SAC when combining the two approaches we propose: a temperature that depends on the state, not with a neural network, but by scaling using the range of the regularizer and a target that is a ratio of the maximum regularizer.

The gradient is 
$$
\mathbb{E}_s\left[\frac{H(\pi(\cdot|a))}{H(U_s)} - \alpha \right]
$$

Which can be zero if we set $\alpha$ to $\mathbb{E}_s[H(\pi)/H(U_s)]$ as $\mathbb{E}_s[H(\pi)/H(U_s)]$ exists and is always non-negative and less than or equal to one as $H(U_s)$ is an upper bound for $H(\pi(\cdot|a))$.

While there is an equivalent constant target equal to $\alpha\mathbb{E}_s[H(U_s)]$ under the final policy, it does **not** need to be feasible while training the model.