Paper: Discovering Temporally-Aware Reinforcement Learning Algorithms (ICLR-2024)

Reviewer comment (Experiments):
Hyperparameter: Setting the entropy coefficient of PPO to 0.0 in the Brax environment appears to significantly limit exploration in my view and poses questions about the experimental design. To address this and strengthen the validity of your results, I recommend conducting a systematic hyperparameter search on the entropy coefficient of PPO within the Brax environment.

Author response:
## Hyperparameters
This is a good discussion point. Our reasoning for setting entropy=$0.0$ is based on [OpenAI's original PPO implementation](https://github.com/openai/baselines/blob/master/baselines/ppo2/defaults.py), which uses that setting for continuous control. [Extremely thorough follow-up work](https://arxiv.org/abs/2006.05990) (See Decision C32, Figure 77) [1] has found no evidence that entropy significantly improves performance on continuous control tasks. While we can tune this coefficient per-environment, we believe that this is against the spirit of the benchmark and is not commonly done.

For a more in-depth discussion of this, we recommend [this blog post](https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/) [2].

Nonetheless, we have run the requested set of experiments and have done a hyperparameter sweep over the entropy coefficient in similar settings to [1] with 5 seeds over 5 values and have attached the results to the supplementary material in Figure 9. In general, our findings align with [1] and we find that an entropy coefficient of 0.0 performs comparably to the others. While an entropy coefficient of 0.005 performs significantly better in Walker2d, it performs significantly worse than 0.0 in Ant and Humanoid. If the reviewer believes this would significantly improve the analysis of our paper, we can provide updated figures that tune the coefficient per-environment and per-algorithm.

[1] Andrychowicz, Marcin, et al. "What matters for on-policy deep actor-critic methods? a large-scale study." International conference on learning representations. 2020.

[2] Huang, et al., "The 37 Implementation Details of Proximal Policy Optimization", ICLR Blog Track, 2022.

Reviewer comment (Novelty):
The idea of incorporating time-step information is not novel in RL, such as this [work](https://proceedings.mlr.press/v80/pardo18a.html), which should be included in the related work.

Author response:
We would like to thank the reviewer for pointing out this line of related work. We have updated the manuscript to include the suggested paper. However,  there is an important difference in the use of the term “time-step” between these works. In our work, “time-step” refers to the total training budget (i.e. a “global” time-step). For the mentioned work, “time-step” refers to the time-step *within* an episode. We’ve further updated the manuscript to clarify this.

Even with this distinction, we agree that incorporating time step information is not novel in RL. Any learning rate schedule incorporates global time-step information. However, incorporating time-step information in RL algorithm discovery methods  (e.g. LPG and LPO) is novel, and leads to the significant improvements in performance demonstrated in our work.

Reviewer comment (Experiments):
Lack of ablation study: In Section 4.1 Equation 9, $n/N$ and $\log(N)$ are included as part of the input to the agent. It is not clear or argued how good this option is compared with other options, such as $n$ and $N$, $N-n$, or $n/N$.

Author response:
Thank you for raising this question, we have included a discussion of alternative representations in the revised paper and are running the ablations outlined below.

For LPG, our choice of representation was fairly arbitrary. We decided to include $\log(N)$ as a measure of total training budget - applying a logarithmic transformation in order to handle large values - and $n/N$ as a measure of training progress that would be invariant to $N$. We did not condition on $N$ or $n$ directly, since these numbers can be very large (e.g. $1e7$ for Min-Atar) and lead to dramatically out-of-distribution values when transferring to much longer or shorter horizons. As long as we have a measure of both total and relative train step, and apply suitable transformations, we are confident that the representation of total and relative train step will have a negligible impact on performance.

To verify this, we are rerunning the experiments and will provide results with alternative parameterisations for training progress - $\log(N-n)$ and $(\log(N), \log(n))$ - in addition to another experiment directly conditioning on the total train steps, $N$.

Reviewer comment (Experiments):
Section 4.1, after Equation 10, "Since TA-LPO is only meta-trained on a single environment and horizon, we do not include the $\log(N)$ term since it does not vary across updates." Would including the $\log(N)$ term decrease the performance?

Author response:
> Would adding $\log(N)$ to TA-LPO decrease performance?

Since TA-LPO is trained on a single task with a constant $N$, $\log(N)$ would simply be a constant bias term. Therefore, adding $\log(N)$ to TA-LPO would be unlikely to impact in-distribution performance. However, including $\log(N)$ could affect transfer to out-of-distribution tasks with different $N$, since the input would then be out of distribution.

Reviewer comment (Writing):
Section 5.2, Paragraph "Lifetime conditioning improves performance on out-of-distribution environments." Figure Figure 2: typo.

Author response:
> Figure 2 typo

Thank you! We’ve fixed this in our updated manuscript.

Reviewer comment (Theory):
Although the experimental results support the authors' hypothesis that conditioning on the lifetime is helpful in meta-learning, there are no theoretical justifications.

Author response:
## On theoretical justifications 
 
We’ve included a theoretical justification and proof in Appendix E. The high-level justification is that without temporal awareness, the algorithm cannot effectively trade off exploration and exploitation. For example, in the case of multi-armed bandits, the agent should converge to a single lever at the end of training, but perform exploration during training.

Reviewer comment (Experiments):
The way to incorporate the lifetime depends on the algorithm, and the experimenter carefully has to design augmentation.

Author response:
##  On “careful” augmentations

For LPG, our choice of augmentation was arbitrary. We decided to include $\log(N)$ as a measure of total training budget - applying a logarithmic transformation in order to handle large values - and $n/N$ as a measure of training progress that would be invariant to $N$. As long as we have a measure of each of these quantities and are scaled to a reasonable range, we expect the representation of total and relative train step to have a negligible impact on performance.

To verify this, we are rerunning the experiments and will provide results with alternative parameterisations for training progress - $\log(N-n)$ and $(\log(n))$ - in addition to another experiment directly conditioning on the total train steps, $N$.

We have updated our revision to include new experiments with alternative temporal representations (Appendix F), which can also be found with other new results in the supplementary material. As expected, these results verify our hypothesis that TA-LPG is robust to the representation of temporal information, under the condition that it is suitably scaled or bounded.

We hope these results address the reviewer's concern about the careful design of the temporal representation, and our new theoretical justification (Appendix E) addresses their concern over the lack of theory.

Reviewer comment (Writing):
The function $U_\phi (x_t \mid x_{t+1}, \ldots, x_T)$ is defined in Section 3.2, but it is unclear because $T$ is not explained. Is it the total number of interactions, denoted by $N$ later?

Author response:
> $T$ is not explained

Thank you for catching this. $T$ is the trajectory/rollout length that the actor collects each update step, not the total number of environment interactions $N$. We have updated the manuscript to include this!

Reviewer comment (Writing):
Two variables $n/N$ and $\log N$ are augmented in TA-LPG. Is $n/N = 0$ when $N$ is unbounded? In my understanding, $N$ is determined and fixed before training. Would you explain $N$ and provide an example task?

Author response:
> Is $n/N$ unbounded when $N$ is unbounded? Would you explain $N$ and provide an example task?

$N$ is the total number of environment interactions we provide to the agent to learn. For example, in MinAtar-Breakout, the agent interacts with the environment for $1e7$ timesteps, so $N=1e7$. From this, we calculate $n/N$, which measures the proportion through training, and $\log(N)$, which provides the agent with the total number of training steps.

When $N$ is unbounded, $n/N$ is a constant $0$. In practice, the total training budget $N$ is usually known, as this allows RL algorithms to trade off exploration over training.

Reviewer comment (Theory):
In TA-LPO, $ \frac{n}{N} x_{r, A}$ is augmented. It is not straightforward to me because it is proportional to $x_{r, A}$, which means linearly-dependent. Would you explain the problem if $n/N$ and $\log N$ are added as the authors did in TA-LPG?

Author response:
> Would you explain the problem if $n/N$  and $\log(N)$ are added [to TA-LPO]?

For TA-LPO it is *slightly* more involved. Given a quick understanding of the theoretical framework of LPO, TA-LPO is a simple augmentation.

The theoretical framework behind LPO is described in Section 3.3. The key takeaway from is that the network needs to output $0$ when the likelihood ratio $p$ is equal to $1$. The network used for LPO has no biases and uses a tanh activation, so if the input is $0$, then the output is $0$. Thus, the input to the LPO network, $x_{r,A}$, is designed to be $0$ when $p=1$ (Equation 7).

If we simply append $n/N$ as an input to TA-LPO, then the requirement doesn’t hold since it is not $0$ when $p=1$. Thus, we multiply $n/N$ by $x_{r,A}$ (which is $0$ when $p=1$).

We do not include any variant of $\log(N)$ because LPO is only trained with a single $N$, which would simply reduce that term to a bias term. 

We believe this is an insignificant challenge for future work and does not reduce the strength of our contribution.

Reviewer comment (Theory):
In the final paragraph of Section 4.2, the authors mentioned, "We found this stabilized training and led to higher performance..." Would you explain "this stabilized training" in detail? If it means the rank transformation, it is used in Salimans et al. (2017). Discussing the relation between two equations would be better if the stabilized training means Eq. (12) rather than (11).

Author response:
> Would you explain "this stabilized training" in detail? 

Thank you for pointing this out, we have revised the paper to make this clearer. The update we propose in (12), “antithetic task sampling”, is designed to reduce the variance from applying ES in the multi-task setting. Using the update from Salimans et al. (2017), each candidate is evaluated on a randomly sampled task, before a rank transformation is applied to their fitness. This can lead to instability in the update, when the fitness across tasks has varying scales. In (12), we propose evaluating each antithetic candidate pair *on the same task*, before applying a rank transformation over the pair (equivalent to selecting the higher-performing candidate). 

This allowed us to normalize fitness across tasks, which led to faster meta-training convergence and improved final performance. We do not examine this rigorously in the paper since it is tangential from the focus of the work. However, we include it since it is a novel and impactful implementation detail, which we had to apply consistently throughout the LPG experiments.

Reviewer comment (Writing):
I do not fully understand the final paragraph of Section 3.1. In the manuscript, LPG and LPO are selected as the base algorithms. However, the authors mentioned as follows: In our work, we focus on instances of meta-RL that parameterize surrogate loss functions with $\phi$ and apply gradient-based updates to $\pi_\theta$ (Houthooft et al., 2018; Kirsch et al., 2019; Bechtle et al., 2021). Does it mean that three algorithms are implemented somewhere?

Author response:
> Does it mean that three algorithms are implemented somewhere?

Thank you for pointing this out! We have updated the manuscript to clarify this. We implemented LPG and LPO, which are successors to those papers.

Reviewer comment (Writing):
Equation (10): Is $x_{r, A}$ a typo of $x_{p, A}$?

Author response:
Thank you for catching the typo in equation 10 and the incorrect year on the reference! We’ve fixed these in the updated manuscript.

Reviewer comment (Writing):
Please update the reference Kirsch et al. (2019) to Kirsch et al. (2020). Please see https://openreview.net/forum?id=S1evHerYPr

Author response:
Thank you for catching the typo in equation 10 and the incorrect year on the reference! We’ve fixed these in the updated manuscript.