In the field of reinforcement learning (RL), representation learning is a proven tool for complex image-based tasks, but is often overlooked for environments with low-level states, such as physical control problems. This paper introduces SALE, a novel approach for learning embeddings that model the nuanced interaction between state and action, enabling effective representation learning from low-level states. We extensively study the design space of these embeddings and highlight important design considerations. We integrate SALE and an adaptation of checkpoints for RL into TD3 to form the TD7 algorithm, which significantly outperforms existing continuous control algorithms. On OpenAI gym benchmark tasks, TD7 has an average performance gain of 276.7% and 50.7% over TD3 at 300k and 5M time steps, respectively, and works in both the online and offline settings.
Paper
Similar papers
Peer review
Summary
This paper proposes a new RL algorithm called TD7, which is based on TD3, and adopts additional techniques including (1) learning state-action representations, (2) LAP prioritized replay, (3) a behavior cloning term in the learning objective, and (4) checkpoints, where (3) and (4) are applied exclusively to the offline RL setting. They demonstrate the superior empirical performance of TD7 on the D4RL benchmark.
Strengths
1. The paper is well written and easy to follow. 2. Extensive experiments and ablation studies to evaluate the proposed methods. 3. The proposed method is shown to have strong empirical performance on established benchmarks.
Weaknesses
1. **Limited contribution**: While I appreciate TD7 as a competent algorithm for both online and offline RL, I'm not so sure about the contribution of SALE especially regarding its applicability to other existing baselines. It doesn't seem conclusive to me that the proposed embedding is generally superior for RL agents. From the ablation studies it seems many independent design choices could significantly affect the final performance, which may also be algorithm dependent.
Questions
1. Could you provide more evidence that shows SALE's applicability to other existing RL methods?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The limitations are not explicitly discussed in the paper.
Summary
The paper introduces an approach dubbed SALE for learning embeddings that model the interaction between state and action in low-level state environments. The authors extensively study the design space of these embeddings and integrate SALE into the TD3 algorithm along with 3 other components to form a new algorithm, TD7. They perform an extensive empirical evaluation over the design space to discover the most significant contributors to final performance. The paper shows that TD7 outperforms existing continuous control algorithms on MuJoCo OpenAI gym tasks.
Strengths
**Originality** The paper's originality lies in its comprehensive study of learning embeddings that model the interaction between state and action in low-level state environments. The experiments conducted provide evidence of the effectiveness of their approach in improving the performance of TD3. **Quality** See Weaknesses Section. **Clarity** The paper is well-written and organized, making it easy for readers to follow the authors' thought process and understand the methodology and results. **Significance** The authors demonstrate that their proposed method, when integrated into TD3 to form what they name TD7, performs significantly better than TD3 for the MuJoCo environments. This work brings to light the importance of learning the interaction between state and action information to improve performance in RL.
Weaknesses
While the paper presents a novel approach to learning embeddings for state and action information for RL, there are areas where it could be improved. Firstly, the evaluation of TD7 could be expanded to include more benchmarks. While the MuJoCo environments provide valuable data, including additional benchmarks such as Procgen or Brax could offer a more comprehensive understanding of the method's performance and any potential limitations. These considerations, while not detracting from the originality and potential significance of the method, do highlight areas where further work could be beneficial.
Questions
See Weaknesses Section.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
The authors adequately addressed the limitations.
Summary
This work introduces a novel state-action representation learning framework SALE and two other techniques (e.g. checkpointing, a new type of Q value clipping) that substantially improve the data efficiency and final performance of TD3 in online and offline RL.
Strengths
1. The work studies joint state-action representation learning, which is under-studied with respect to state representation learning. 2. The method is very clear, and the authors provide extensive empirical analysis and ablation explaining their design choices. 3. Some techniques introduced in this work (checkpointing, Q clipping in Eq. 6) can be applied to most RL methods and are thus influential beyond representation learning.
Weaknesses
**Core comments:** 1. I believe the paper would read more easily after reorganization. * Section 4.2 feels like an ablation study that should follow the core empirical results. Also, Section 4.2 references TD7 results, though TD7 is not introduced until section 5. * Section 5 begins by mentioning stability and extrapolation error jumped out at me; the intro and related work emphasize the representation learning side of this work much more, and extrapolation error is not mentioned in sections 1-4. I believe the paper would read better if the authors painted a clearer picture of what the stability side of the work entails in the introduction, and included transitions between section 4 and 5 that clearly indicate “we talked about one issue, and now we’re going to talk about another.” * Transition sentences before bolded subheadings (e.g. Normalizing embeddings, Fixed embeddings) would help in Section 4 – something to the effect of “We now discuss two important aspect of SALE…” 1. The checkpoint policy $\pi_c$ denotes the policy is the largest minimum return during the assessment phase. During agent evaluation, you compute the return over N rollouts of $\pi_c$ rather than $\pi_t$, correct? If so, I’m a bit skeptical of the results; is it possible that the perceived benefits of checkpointing are simply due to the batched nature of training it requires? Suppose you run vanilla TD3 within the checkpointing framework but perform separate evaluations using the checkpoint policy *and* current policy. This experiment would be equivalent to vanilla TD3 with variable length batched updates. Does the batched updates alone improve performance? **Minor comments:** 1. Figure 2 would be more readable with error bars rather than text to indicate the 95% confidence interval.
Questions
1. In principle, any latent state/action RL framework can be modified such that $Q$ and $\pi$ take both the original state/action and the latent state/action as inputs. Would it be fair to say the ablations in Fig. 2 suggests that any state/action representation learning method should have $Q$ and $\pi$ depend on both the original and latent state/actions? 1. Fig. 2: Can the authors clarify what is meant by percent loss here? Are these quantities all with respect to default values? Where do the default values fall on these figures? 1. The authors mention that checkpointing harms performance early in training – is this because the policy changes rapidly early on? Can the authors provide more detail? 1. Line 217: How does checkpointing change the relevance of data?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
4 excellent
Presentation
3 good
Contribution
3 good
Limitations
See Weaknesses.
Summary
This paper proposes TD7, an improved version of the popular TD3 algorithm with 4 additional techniques: state-action learned embeddings (SALE, the major one among the four), using policy checkpoints for stable evaluation, an existing prioritized experience replay method called LAP, an existing offline RL algorithm called TD3+BC. The four techniques strengthen vanilla TD3 algorithm from different aspects. TD7 is evaluated on MuJoCo and D4RL benchmarks for online and offline RL settings, significantly outperforming existing continuous control algorithms. Comprehensive ablation studies are also included.
Strengths
- This paper is well written and the content is clear. - The proposed improvements (almost) make sense to me and well motivated or explained. - The part of SALE (Section 4) is interesting to me, although SALE can be viewed as an improvement over OFENet. I think the detailed studies on useful techniques for stable representation learning could be inspiring to other similar problems in RL. - This paper conduct ablation studies for most design choices to rule out the alternative choices empirically. - The code implementation is neat and I assume it is easy to reproduce the experiments.
Weaknesses
- The proposed improvements are piecemeal, each of which is incremental or existing. - A few additional hyperparameters (e.g., the dimensionality of $z^{s,a}$, the episode number to maintain the policy unchanged) are introduced, although the authors give recommended values. - Although the paper has comprehensive empirical studies, intuitive explanations are lacked.
Questions
TD7 needs to keep current policy fixed for several episodes: 1. What if the episode horizon is long? 2. Will this be inconvenient to combine high update frequency (like used in REDQ)? 3. How to consider the gap between using the checkpoint policy for evaluation (exploitation) while the comparison and selection of checkpoint policies are based on the performance of exploration policies?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
4 excellent
Contribution
3 good
Limitations
- The proposed improvements seems to be general and orthogonal to different continuous control RL algorithms. However, SALE is not combined with SAC, TQC and REDQ to evaluate the generality. - TD7 is evaluated on MuJoCo and D4RL whose backends are both OpenAI MuJoCo suite. I would be interested in seeing TD7 in other continuous control benchmarks like DMC.
Summary
This paper introduces a couple of ideas to improve the empirical performance of the TD3 algorithm on continuous-action RL problems. The core contribution is to show that learning state and action embeddings that are designed to predict themselves in successive timestep can help achieve more reward. This is pretty nice because the representation learning process is done in conjunction with reward maximization and is shown to be useful and stable enough to enable faster learning in the same task. A few other contributions are presented which I list below: 1- a normalization approach that scales the output of the embeddings, and which is compared against some of the existing normalization approaches such as batch and layer norm. 2- checkpointing the policy during online RL, akin to supervised learning, so rather than using the latest learned policy one can continue learning by using the best-performing policy. 3- clipping the value estimates during training to ensure that the estimates remain within a meaningful range. Very nice empirical results are then provided, most notably on the Mujoco benchmark in online RL, where the proposed approach named TD7 is capable of beating competitive baselines such as TD3 and SAC.
Strengths
The highlight of the paper is the impressive empirical results provided on the online RL experiments with the Mujoco baseline. I did my best to cross-check the performance of the baseline agents and it does seem like that TD7 is capable of beating TD3 and SAC (these are the two baselines I checked) on these benchmarks. I do see some discrepancy between SAC results reported here and those reported (for Hopper) in the original paper, but otherwise the results are consistent with published papers.
Weaknesses
I did enjoy reading this paper, and I applaud authors for their successful implementation and empirical results. That said, I am not sure what the core message of the paper is. I think the contributions need to be motivated, framed, and highlighted better. For instance, is this paper primarily about the advantage of learning state and action embeddings in RL? If yes, existing work has already demonstrated that, so other than the fact that the empirical results are superior, what kind of statements can we make about this new approach to learning embeddings that we did not know previously? Is it better for transfer learning? Is it faster in terms of run time? Is it better motivated theoretically? Just to name a few papers, with the danger of missing other related and interesting work: - An approach very similar to the one presented here was presented earlier by Zhang et al: "Learning Action Representations for Reinforcement Learning" - Gelada et al also propose a very similar approach in "Deepmdp: Learning continuous latent space models for representation learning", and I am not clear in what sense SALE is doing something different than them. - Chandak et al show in "Learning Action Representations for Reinforcement Learning" that one can learn a low-dimensional representation for actions. These are just a few examples to show that the idea of learning state and action representations is well-explored and to me this limits the novelty of the work. Moving to the other contributions, and in terms of the normalization approach, I am not quite positive in what way this normalization is going to hedge against the issue of collapsing all states and actions to 0. To me it seems like that if the reward is not part of the process of learning the embeddings, and we use the kind of successive predictive loss in the embedding space, then trivially the 0 solution would be optimal with or without normalization. I found the checkpointing idea creative and interesting. That said, I have a few issues with it: 1- it seems like based on Figure 7 that checkpointing is rarely effective in the online case, and as the paper states, not applicable in the offline case. 2- it seems to me that when computing the goodness of a checkpointed policy, a potential issue is that during training we add some exploration noise to the action suggested by the network, so unless we do zero exploration, we need to account for the fact that we have not exactly executed the policy and therefore some off-policy learning is needed to compute the true goodness of the checkpointed policies. Do I get this right? The clipping idea also makes sense, and it is interesting, but it is still a limited contribution.
Questions
- In what way is the new embedding-learning approach different than the previous work? What are the advantages and disadvantages, and what can we learn about designing embedding-learning algorithms in light of your experiments? - How is the normalization hedging against the 0 mode collapse? - Do we not need to do off-policy learning during check-pointing to account for exploration?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
2 fair
Contribution
2 fair
Limitations
NA
Thanks for the added experiments, it does strengthen the case for the paper. I am still lingering on the effect of the proposed normalization step. Can you clearly define the optimization problem being solved here, and demonstrate how the normalization step is hedging against the network collapsing to trivial solutions? "However, this isn’t necessarily problematic since the checkpoint is more likely to prefer policies which are robust to noise in the action space." I do understand that this empirically would not be a terrible idea, however, I meant that to estimate the utility of a checkpoint policy in an unbiased way, one needs to be doing off-policy policy evaluation.
Thank you for the quick response! **Embeddings**: In the context of training the embeddings, what we do is very similar to BYOL [1] and SPR [2], who only apply normalization in the loss function. The main difference between normalizing the embedding and the loss function is that the final embedding used by a downstream network (the value function in our case) will be normalized. During training, the process is effectively identical. BYOL has a discussion on collapse in Section 3.2 of their paper which is relevant to our approach. The main intuition is that while there do exist trivial uninformative solutions, there is no reason for the method to actually converge to them. This is because we are optimizing a loss towards a target network which changes independently of the loss function. Formally studying this optimization process is an important problem in both RL and self-supervised learning, but definitely out of the scope of this project. Empirically we haven’t found any issues with representation collapse, and this observation is consistent with work based on BYOL. We also experiment with some settings that are closer to what SPR/BYOL propose in the design choices section of our paper (Section 4.2) and find our approach is better for our use case. **Checkpoints**: Yes, you are certainly correct that to get an unbiased estimate of the checkpoint we need to do a form of unbiased off-policy evaluation. What our method does is use a proxy measure of the quality of the checkpoint policy (the online exploration policy performance). There are a lot of potential alternate approaches that could be used (for example, delving into OPE literature). Even naively there are simple adjustments that could be made to minimize bias, such as running a potential checkpoint for X episodes with no exploration noise if it surpasses a threshold or the previous exploration policy. However, in practice we found that our proxy was accurate and does not come at any additional costs in terms of algorithmic complexity, run time, or sample-efficiency. Regardless, we think further studying the use of checkpoints of RL is an exciting research direction and this is just an initial outline for a potentially very valuable strategy for RL. **References:** - [1] Grill, Jean-Bastien, et al. "Bootstrap your own latent-a new approach to self-supervised learning." Advances in neural information processing systems. 2020. - [2] Schwarzer, Max, et al. "Data-Efficient Reinforcement Learning with Self-Predictive Representations." International Conference on Learning Representations. 2020.
Normalization
I am still not seeing why this normalization operation avoids mode collapse, which is claimed on line 123. To the best of my knowledge, the paper does not show empirically that mode collapse is avoided. The paper rather shows that it works better in terms of reward when compared to other other normalization strategies presented in prior work. Yes, I do agree that the embeddings will not go to 0, but they can technically mode collapse to trivial, constant solutions. I also was hoping a rough theory explanation could be provided, but it seems that arguing it would be difficult at this point. Also, the BYOL paper is providing a different intuition about why mode collapse is not manifesting itself: "... the combination of (i) the addition of a predictor to the online network and (ii) the use of a slow-moving average of the online parameters as the target network encourages encoding more and more information within the online projection and avoids collapsed solutions" To summarize, It is warranted to claim this is an empirically superior normalization approach, but the paper has not convinced me that this normalization is hedging against mode collapse. While I appreciate that the fact that this normalization works better empirically, I think the question about why mode collapse is not happening (and how to further reduce its possibility) is an open one.
Thank you for the response! We believe our statement is true based on empirical results and well established prior work. However, we also agree with the reviewer that it is best not to make claims that we cannot theoretically defend within the scope of the paper, so will adjust the statements on line 123 to remove any comments on representation collapse: > Similar to the normalized loss functions used by SPR [Schwarzer et al., 2020] and BYOL [Grill et al., 2020], AvgL1Norm protects from monotonic growth, but also keeps the scale of the downstream input constant without relying on updating statistics (e.g. BatchNorm [Ioffe and Szegedy, 2015]).
Mode Collapse
Thanks for acknowledging this. I am going to raise my score from 4->5, but just to be clear, other than the absence of theoretical evidence, empirical evidence is only that this normalization step is working better than the alternatives, meaning that it is more performant (Figure 2 from the submission). As to why exactly this is the case we are not sure yet, right? An empirical evidence for the statement would be one that shows mode collapse is actually less frequent under the proposed normalization, which having again checked the submission and the appendix, I don't think that's provided. Hope I'm not missing anything.
Response to author rebuttal
I want to thank the authors for providing thorough responses to my comments (as well as other reviewer comments). All of my comments have been addressed. I maintain my score, and would advocate for accepting this paper. Reviewers Casd and FeJr expressed concern about the work’s novelty, though I do not share these concerns. Joint state-action representation learning is under-explored compared to just state or action representation learning, and this work provides a thorough analysis of relevant design choices missing in the existing literature. Reviewer sfFv noted that the proposed improvements – when considered individually – are incremental. I think this is a valid comment regarding target Q clipping (Eq. 6) and checkpointing (Sec. 5.2), but the core contribution of this work is the novel state-action representation learning method and its analysis. While target Q clipping and checkpointing can in principle be applied to many RL algorithms without SALE, it serves a particular purpose in the context of SALE which the authors mention in section 5.1. The policy and Q function are more prone to extrapolation error because SALE increases the dimensionality of their inputs, and these techniques help mitigate this effect.
Thank you!
Thank you for advocating for the paper. We agree with your comments and appreciate your support.
Thank you for the response. I have increased my rating as my main concern on the compatibility with existing baselines has been addressed.
The results for TD7 in comparison to TD3 are compelling and offer a clearer picture of the algorithm's performance across a diverse set of tasks. Given the updated information and the detailed results you've shared, I am pleased to increase my score.
Decision
Accept (poster)