Train Once, Get a Family: State-Adaptive Balances for Offline-to-Online Reinforcement Learning
Offline-to-online reinforcement learning (RL) is a training paradigm that combines pre-training on a pre-collected dataset with fine-tuning in an online environment. However, the incorporation of online fine-tuning can intensify the well-known distributional shift problem. Existing solutions tackle this problem by imposing a policy constraint on the policy improvement objective in both offline and online learning. They typically advocate a single balance between policy improvement and constraints across diverse data collections. This one-size-fits-all manner may not optimally leverage each collected sample due to the significant variation in data quality across different states. To this end, we introduce Family Offline-to-Online RL (FamO2O), a simple yet effective framework that empowers existing algorithms to determine state-adaptive improvement-constraint balances. FamO2O utilizes a universal model to train a family of policies with different improvement/constraint intensities, and a balance model to select a suitable policy for each state. Theoretically, we prove that state-adaptive balances are necessary for achieving a higher policy performance upper bound. Empirically, extensive experiments show that FamO2O offers a statistically significant improvement over various existing methods, achieving state-of-the-art performance on the D4RL benchmark. Codes are available at https://github.com/LeapLabTHU/FamO2O.
Paper
Similar papers
Peer review
Summary
Authors propose a framework for offline-to-online tuning of offline RL algorithms tunning. The idea is to train an additional network which desides helps to keep improvement-constraint balance during finetuning.
Strengths
Approach improves all of the checked algorithms performance and can be applied to different offline RL algorithms. Good range of benchmarking tasks and algorithms.
Weaknesses
-
Questions
How does the modification affect the compute time required to train algorithms?
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
3 good
Presentation
3 good
Contribution
3 good
Limitations
Adding another network requires additional choice of hyperparameters which might be hard to find
Summary
This paper approaches offline-to-online RL from the intuition that at a particular state, if the dataset already contains good actions, then the subsequent online tuning should be more conservative to retain the good actions in the dataset; but if the dataset's actions is poor, then more radical policy improvement is needed. To this end, this paper introduces a framework, Family Offline-to-Online RL (FamO2O), which aims at a state-adaptive improvement-constraint balance for each state. Specifically, from the collected dataset, the authors train a diverse policy family ranging from conservative to radical and use the environmental feedback to select an appropriate policy from this family at each state. Practically, this is achieved by a universal model, which determines the degree of policy conservatism; and a balance model, which learns the balance coefficients at each state. Experimental results show that FamO2O improves on various offline-to-online RL methods and achieves competitive performance.
Strengths
1. The paper is well-written and generally easy-to-follow. 2. The proposed method is well theoretically justified. 3. The empirical discussion and ablation study are thorough.
Weaknesses
1. The proposed method seems require abundant diverse data, which may not be feasible on harder settings, e.g., the Adroit domain in the D4RL benchmark where the data is limited and the data distribution is narrow and lack of diversity. 2. Maybe I misunderstand something. I think the purpose of the $\pi_b(s)$ model is to find the $\beta_s$ that corresponds to the optimal sequence of constraints $\\{\epsilon_s, s \in \mathcal{S}\\}$ in Eqn. (6), which relates to $\beta_s$ by the Language multipliers $\mu(s)$ (Line 537, Appendix C.2). In this regard, the paper's story of balancing policy improvement and constraint can be slightly over-complicated and somewhat confusing. 3. The proposed method is in spirit similar to decision-transformer-style methods, and is therefore less a surprise.
Questions
1. Could you explain more on how the updating rule Eqn. (10) train the balance model $\pi_b$ to control policy improvement and constraint? Why does maximizing the $Q^k$ values achieve this balance? 2. The significant variation in data quality across different states depicted by Figure 2 may come from the nature of the medium-expert datasets, which by construction contains medium and expert trajectories. Does this variation in data quality exist in other type of datasets, say expert, medium, or medium replay? More generally, is this phenomenon of significant data-quality variation ubiquitous or specific? 3. Could you explain more on the benefit of "state-adaptive improvement-constraint balances"? Even if we do not make such a balance and simply just online finetune an offline-pretrained policy, the environmental feedback should be able to tell us which actions are good and should be retained and which actions are poor and should be improved. 4. Could you explain more why you consider the term $\log \pi (a|s)$ in Eqn. (1) as a policy constraint? What is the target of this constraint, especially when $(s,a)$ is an online interaction samples? 5. Could you explain more on how you get Eqn. (12) and (13) from Eqn. (5)? It would be better if you can expand Appendix C.1 to include more details and explanations. The current version is a bit hard to follow. 6. nit: in Eqn. (6), are you missing a "$\forall \epsilon > 0$" before $\exists \\{...\\}$? 7. [L161] Could you explain more on the cooperation between $\pi_u$ and $\pi_b$? Why couldn't we still randomly sample $\beta_s$ during online fine-tuning? 8. [L159] How many $\beta_s$ vectors are required to learn the universal model? How does this number scale with the number of states in the dataset? And how to select/design the balance coefficient space $\mathcal{B}$?
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
The paper proposes a new algorithm to perform offline-to-online reinforcement learning. The core idea is to consider a state-adaptive balance parameter, which aims to encourage imitation of dataset behavior only if the corresponding advantages / values are high, while prior works have mostly assumed fixed balances. The authors provide detailed experimental evaluation of their approach and show superiority over relevant baselines.
Strengths
I think offline-to-online RL is a very relevant and promising direction for research since if reliable it would enable much more practical applications of RL in real-world tasks. The core idea of the paper, training a collection of policies that adapts on a per state level, appears logical and powerful (however please see [2,3,4], which I think put forward similar ideas) and will in the future probably significantly influence the way offline-to-online RL is performed and thought about. The paper is overall very well written and understandable & offers a strong statistical analysis of the proposed algorithms performance compared to relevant baselines.
Weaknesses
The term data quality is frequently used but not really introduced (could mean accuracy / truth of the information in the data, but I think it means something like return). Not all D4RL locomotion datasets were used - its unclear whether that's an issue since the selection is not justified. The considered baselines appear relevant (TD3+BC & CQL less so since its already shown in the AWAC paper that using the same conservative formulations in the online setting does not work well), however I think the closest existing methods are missing: E.g. Confidence-conditioned value functions [1] automatically adjusts its policy confidence level (as far as I understand on a per state level since every state goes into the considered history). Also methods like RvS [4], which condition on return, could easily be extended to this setting (i.e. always try to maximise conditioned return during online data collection). Generally I have the feeling that some very relevant related work regarding offline adaptive policies was not considered, i.e. also: [2] User-Interactive Offline RL (which considers policies conditioned on a balance between conservative and very liberal) [3] offline policies should be trained to be adaptive (which adapts policies based on the online collected history) I'm not sure the state-adaptive balance between conservative and radical idea for offline-to-online is thus entirely novel. I see some issues regarding clarity that could easily be fixed (see questions)
Questions
Could you please define how you use the term data quality (I'm assuming it has something to do with the data's return, but I'm not sure)? Please elaborate why the random datasets in the D4RL datasets were not considered. There exist offline-to-online methods that already consider something like state-adaptive policy changes, like confidence-conditioned value functions [1], as well as other adaptive offline-to-online policy concepts [2,3] that could easily be extended to your case - why were they not considered? (I realise that some are rather recent, but I think they deserve at least brief discussion) I do not understand section 6.1 / figure 8 at all - what is meant by guidance? what is meant by high / low quality data? Starting positions are denoted as triangles, suggesting a moving direction but I don't think that's intended. The text says high balance coefficients are found where data quality is also high (quality=reward / return?) & that higher quality data is found at the lower crossing point in the 5th row, however from the color encoding it seems that the balance coefficient there is actually the lowest... I think the ratio between offline and online training steps as well as collected interaction steps is a crucial parameter for reproducibility / future comparisons, however that information is only conveyed in the figure 10 axis & the appendix - I think it should be explicitly stated somewhere in the experimental section text. When I look at figure 10, it seems that after all online training steps, the attained performance is almost identical to that at the end of offline training - isn't that a very disappointing result and couldn't you just throw away the whole online part altogether then? Since FamO2O outperforms prior offline-to-online baselines, does that mean that these baselines only get worse with online training? [1] Hong, J., Kumar, A., & Levine, S. (2022). Confidence-Conditioned Value Functions for Offline Reinforcement Learning. ICLR 2023 [2] Swazinna, P., Udluft, S., & Runkler, T. (2022). User-Interactive Offline Reinforcement Learning. ICLR 2023 [3] Ghosh, D., Ajay, A., Agrawal, P., & Levine, S. (2022, June). Offline rl policies should be trained to be adaptive. ICML 2022 [4] Emmons, S., Eysenbach, B., Kostrikov, I., & Levine, S. (2021). RvS: What is Essential for Offline RL via Supervised Learning?. arXiv preprint arXiv:2112.10751.
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
4 excellent
Presentation
4 excellent
Contribution
3 good
Limitations
no, limitations are not discussed - perhaps looking at the not improved performance after online fine-tuning in figure 10, as well as the often only small improvements over standard IQL, a brief discussion would be good.
Summary
The paper introduces a new method to mitigate the distribution shift problem in the offline to online RL problem. The paper states the intuition that the policy should behave differently on states with different values, that is, the policy should be more conservative on high return states and exploratory on the low return states. With this intuition, the paper proposed to train a family of policies in the offline to online setting, specifically, train another "policy" to parameterize the rollout policy via the state-adaptive balance coefficient. The experiments show that FamO2O outperforms previous O2O baselines, and show via a toy experiment that FamO2O indeed leans the state-wise adaptivity, and various ablations show the importance of each design choices.
Strengths
1. The experiment result is solid as it evaluates on extensive D4RL dataset with different data quality and both locomotion and maze tasks. 2. The discussion sections validate the algorithm design choices, and section 6.1 verify that the algorithm indeed learns a state-wise adaptive policy, which support the intuition and motivation of the algorithm.
Weaknesses
1. The exploitation vs. exploration intuition is not brand new in the offline to online setting, there is also some work with similar intuition [1]. I believe proper comparison is required given the similarity of the intuition, although I believe training a family of policies (or conditionally parametrized the policy) seems like a slight generalization. 2. It is confusing that in the universal model training (eq. (9)), the exponential of the advantages are weighted by the balance parameter $\beta$, but when training the balance model, the Q-function (as in the loss) is not weighted by the balance parameter. There seems to be some consistency. I can tell that the unweighted objective (the current form of eq (10)) would be more computationally friendly, but theoretically, it seems more natural to optimize over a weighted version where eq. 10 is also weighted by $\beta, which sample from the balance model. 3. The action distance in eq. (11) may not be the best metric to measure the discrepancy between a policy and a trajectory. For example, if the policy that induces the trajectory only takes $a_1$ in $s$, and the evaluated policy takes $a_1$ with $p=0.51$ and $a_2$ with $p=0.49$, which may induce a very negative reward, or cause great trajectory derailment (which is not recorded in the offline trajectory), the proposed metric will still be 0 but in reality the evaluated policy is not that close to the trajectory.
Questions
See above.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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
3 good
Presentation
3 good
Contribution
2 fair
Limitations
N/A
Rebuttal Response
Thank you very much for the detailed responses to my questions, they have been very helpful in better understanding your work, especially Q3,4 & 6. I am however still not quite sure I understand Figure 8 (Q4): When you say guidance means "directing the agent", how exactly do you direct it? Is there an explicit reward signal given only at this point or does a separate policy take over which "knows" the way or ...? I believe illustrative examples like this one are important & I understand you have limited space, but I think a little more information is needed to make it really helpful. If the color encoding is correct & what you wrote in your response > the agent typically outputs lower balance coefficients for high-quality samples is correct, I might still misunderstand your method. In Eq(1) it seems to me that high balance coefficients would lead to the agent more likely copying the behaviour that was present in the dataset - since I would like to repeat behaviour that has yielded high return, I would expect high balance coefficients in states where you have high quality data. In the example it is however the other way around... Could you please elaborate? One more clarification regarding your answer on Q6: Does that mean the plot shows performance only on a single offline dataset? If so, which one?
Explanations of Figure 8 (Q4), Eq. (1), and Figure 10 (Q6)
Thank you for your thought-provoking questions. We've addressed your inquiries in detail below. > The meaning of "directing the agent". "Directing the agent" refers to compelling the agent to adhere to the route and direction that yield the shortest path to the goal, rather than letting the agent decide the route and direction on its own. This aligns with your perception that "a separate policy take over which 'knows' the way". Thank you for bringing this to our attention. We will incorporate the above explanation into the next version of our paper to make this point clearer. > The effect of the balance coefficient value. Thank you for your thoughtful question. You're suggesting that high balance coefficients should be used with high-quality data, which makes sense at first glance. But our method works differently, and here's how: - **For High-Quality Data:** Utilizing high balance coefficients might lead the policy to aggressively pursue actions with the highest possible advantage, $Q(\mathbf{s}, \mathbf{a})-V(\mathbf{s})$. **But since the advantages of the high-quality data are already high, trying to push for even higher advantages can easily lead to mistakes due to overestimation in Q values**. So, we use lower balance coefficients for high-quality data, making sure the policy stays safe by following the known good actions. - **For Low-Quality Data:** On the other hand, with low-quality actions, it makes sense to use higher balance coefficients. **Copying what the low-quality data does will surely end up in failure, so it's worth the risk to try for something better**. This leads to a more daring or "radical" policy that looks for higher-quality actions. By using balance coefficients this way, depending on whether the data is high or low quality, our method reduces the risks and finds a good middle ground. It doesn't chase after the highest advantages in a way that can cause mistakes, but it also doesn't just copy what's in the bad data. It's a careful balance that helps the policy make the best decisions. > One more clarification regarding your answer on Q6: Does that mean the plot shows performance only on a single offline dataset? If so, which one? Yes, Figure 10 displays the performance solely on one offline dataset, namely `antmaze-umaze-diverse`, as referenced on line L297 of our manuscript. We selected this specific dataset because, on it, IQL exhibits the most significant decline in performance when transitioning from offline pre-training to online fine-tuning. Although we've alluded to this rationale on lines L296-297 of our manuscript, we will articulate it more explicitly in the upcoming version of our paper.
Rebuttal Response 2
Thank you for the additional clarifications. From your response regarding the balance coefficients, I am however now not sure whether I understand the balance model correctly: You say that low balance coefficients should be provided by $\pi_b$ if the data quality (i.e. return) in the current state is expected to be high. However, as far as I understand the balance model training, it basically just maximises Q-value (i.e. expected return) of the resulting universal model actions. If what you say is correct, then that means that the Q-maximizing actions are provided by the universal model when the input balance coefficient is low - this is possible, but I wonder why it would be the case - is there any extra effect / regularisation that prefers low balance coefficients if quality / Q-value is high? I'll try to give a simplified example: If I optimise Eq. 9 not over the entire dataset, but only a single state, and this state has high quality data, then the value in the expectation is maximised if high balance coefficients coincide with high log likelihoods from the universal policy model. Of course the universal model could output the highest log likelihood for low balances, but why would it / what mechanism is pulling the likelihood maximising actions towards the low balance coefficients in the universal model?
Explanation of the Balance Coefficient Selection Mechanism
Thank you for your thoughtful question. It brings attention to a critical aspect of our methodology that may appear perplexing at first glance. Indeed, Equation (9) does not seem to explicitly contain any effect or regularization to guide the balance model, $\pi_b$, to choose low or high balance coefficients, $\beta$, depending on data quality. **However, there is an implicit effect -- namely, online rewards -- that directs the behavior of the balance model in the desired manner.** To elucidate this mechanism, let's focus on states with high-quality data (the process for low-quality data is the inverse). During **the initial stage of online fine-tuning**, selecting a large $\beta$ for high-quality data states often leads to an action with an overestimated Q-value. This, in turn, results in comparatively low online rewards, subsequently causing lower Q-values. Thus, $\pi_b$ learns to shy away from large $\beta$ values for high-quality data states. **As online fine-tuning proceeds**, because small $\beta$ values are selected more often for high-quality data states, the universal policy $\pi_u$ becomes increasingly trained with small $\beta$, enhancing its performance compared to large $\beta$ values. **Returning to your simplified example** involving a single state with high-quality data, the mechanism described above still holds. The online rewards will lead the balance model $\pi_b$ to favor smaller balance coefficients $\beta$ for the universal model $\pi_u$. **Therefore, even in this minimalist scenario, the interplay between the balance coefficients and online rewards enforces the selection strategy we have described.**
Rebuttal Response 3
Thank you for your insightful explanation. I had considered the setting more from an offline perspective at first, where my intuition was high quality should lead to large β. However, I hadn't thought about the fact that this leads to low returns in the online finetuning due to overestimation & that the online feedback would then turn the whole thing around since the balance model needs to counteract. Thank you for clearing up the confusion. While thinking about this, I came up with another question: How do you determine the balance coefficient ranges and how do you enforce them? I.e. it seems from appendix fig.11 & 12 that the algorithm is not particularly sensitive the chosen interval, yet different ones are used for the different environments - why is that? I assume it has something to do with the reward scale? How would I choose the interval when I would deploy FAMO2O in a new environment? Lastly, do you clip network outputs to the relevant range or how do you enforce valid β values?
Reply to Rebuttal Response 3
Thank you for your insightful inquiry. We're glad that our previous explanations were clear, and we welcome the opportunity to address your new questions. --- > Choosing the range of $\beta$ when applying FamO2O to a new algorithm In selecting the range of $\beta$ for FamO2O, we adhere to **two principles** that have guided our work: (1) The region should encompass the original algorithm's $\beta$, as this value has been carefully chosen for the algorithm, providing an indicator for the general scale of the range; (2) The upper and lower bounds of the $\beta$ range should lead to reasonable performance in the original algorithm, ensuring the reasonableness of the range. > Different $\beta$ intervals for different environments Regarding Figures 11 & 12 in the appendix and your question on why different $\beta$ intervals are used for various environments: as stated in lines L630-631, FamO2O is shown to be insensitive to the range of coefficient balances **within a reasonable scope**. This scope can be assured by the second principle mentioned above. Nevertheless, **due to differences in environment characteristics** such as reward scale, sparsity/density of reward, and dynamics, **the reasonable scope for different sets of environments (like locomotion and antmaze) may vary**. **This phenomenon aligns with most offline / offline-to-online algorithms**, which often require distinct $\beta$ values (or other hyperparameters) for different environmental sets. > Enforcing valid $\beta$ values To enforce valid $\beta$ values, we employ **a $\tanh$ activation layer** as the last layer of the balance model, **rescaling the $\tanh$ output range of $(-1, 1)$ to the desired $\beta$ range**. Theoretically, clipping could also serve this purpose and may be an avenue worth exploring in future studies. ---- We sincerely hope that these explanations address your queries comprehensively. Thank you once again for your thoughtful questions
Thank you
Thank you for your thoughtful comments and advice. We have carefully examined the EOP paper and will take it into consideration for inclusion in the next version of our work. Your insights are greatly appreciated.
Response to the authors
Dear authors, Thank you so much for the detailed response, which clears out all my questions. I will increase my rating to 7.
Thank you
Dear Reviewer urL3, Thank you for your thoughtful review and for taking the time to reconsider our work. We appreciate your positive feedback and the increased rating. Your insights have been invaluable to us, and we are pleased to have addressed your concerns. Best regards.
Response
I appreciate the authors' detailed response and my concerns (2 & 3) are addressed, and I increased my score accordingly. I also apologize for not specifying the reference, and if I recall correctly, [1] should be [1] Zhang, Haichao, We Xu, and Haonan Yu. "Policy Expansion for Bridging Offline-to-Online Reinforcement Learning." arXiv preprint arXiv:2302.00935 (2023).
Thank you
Thank you for your considerate review and for identifying an important reference. Concerning [1], which details the method of alternating between offline policy $\pi_\beta$ and online policy $\pi_\theta$ (starting from scratch), it indeed stands as a significant and relevant work to our research. We appreciate your recommendation, and we commit to including a comparison with this work in the revised version of our paper. Your insight and guidance are gratefully acknowledged, and we thank you once again for your constructive feedback.
Thank you
Thank you for your thoughtful examination of our work and for the constructive feedback. We are pleased that our clarifications have addressed your concerns, and we appreciate the increased score. Your insights have been instrumental in refining our paper.
Decision
Accept (spotlight)